mediastore: support for lower-latency streaming via wishful thinking

Need aws/aws-sdk-go-v2#1094 to remove the wishful thinking part :)
fsouza · Apr 1, 2021 · 07fdcfb · 07fdcfb
1 parent aad8124
commit 07fdcfb
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -37,14 +37,15 @@ name on mediastore:
 s3-upload-proxy configuration's is defined using the following environment
 variables:
 
-| Variable            | Default value | Required | Description                                                                      |
-| ------------------- | ------------- | -------- | -------------------------------------------------------------------------------- |
-| UPLOAD_DRIVER       | s3            | No       | Upload driver to use (options are "mediastore" or "s3")                          |
-| BUCKET_NAME         |               | Yes      | Name of the S3 bucket or the mediastore container (depends on the upload driver) |
-| HEALTHCHECK_PATH    | /healthcheck  | No       | Path for healthcheck                                                             |
-| HTTP_PORT           | 80            | No       | Port to bind (unsigned int)                                                      |
-| LOG_LEVEL           | debug         | No       | Logging level                                                                    |
-| CACHE_CONTROL_RULES |               | No       | JSON array with cache control rules (see below)                                  |
+| Variable                    | Default value | Required | Description                                                                      |
+| --------------------------- | ------------- | -------- | -------------------------------------------------------------------------------- |
+| UPLOAD_DRIVER               | s3            | No       | Upload driver to use (options are "mediastore" or "s3")                          |
+| BUCKET_NAME                 |               | Yes      | Name of the S3 bucket or the mediastore container (depends on the upload driver) |
+| HEALTHCHECK_PATH            | /healthcheck  | No       | Path for healthcheck                                                             |
+| HTTP_PORT                   | 80            | No       | Port to bind (unsigned int)                                                      |
+| LOG_LEVEL                   | debug         | No       | Logging level                                                                    |
+| CACHE_CONTROL_RULES         |               | No       | JSON array with cache control rules (see below)                                  |
+| MEDIASTORE_CHUNKED_TRANSFER | false         | No       | Whether to enable chunked transfer with MediaStore for lower latency             |
 
 ## Defining cache-control rules
 

diff --git a/internal/uploader/mediastore/mediastore.go b/internal/uploader/mediastore/mediastore.go
@@ -10,13 +10,23 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/mediastore"
 	"github.com/aws/aws-sdk-go-v2/service/mediastoredata"
+	"github.com/aws/aws-sdk-go-v2/service/mediastoredata/types"
 	awsint "github.com/fsouza/s3-upload-proxy/internal/aws"
 	"github.com/fsouza/s3-upload-proxy/internal/uploader"
 )
 
+// DriverOptions are the set of options that can change how the mediastore
+// driver behaves.
+type DriverOptions struct {
+	ChunkedTransfer bool
+}
+
 // New returns an uploader that sends objects to Elemental MediaStore.
-func New() (uploader.Uploader, error) {
-	var u msUploader
+func New(options DriverOptions) (uploader.Uploader, error) {
+	u := msUploader{uploadAvailability: types.UploadAvailabilityStandard}
+	if options.ChunkedTransfer {
+		u.uploadAvailability = types.UploadAvailabilityStreaming
+	}
 	sess, err := awsint.Config()
 	if err != nil {
 		return nil, err
@@ -26,8 +36,9 @@ func New() (uploader.Uploader, error) {
 }
 
 type msUploader struct {
-	client     *mediastore.Client
-	containers sync.Map
+	client             *mediastore.Client
+	containers         sync.Map
+	uploadAvailability types.UploadAvailability
 }
 
 func (u *msUploader) Upload(options uploader.Options) error {
@@ -36,10 +47,11 @@ func (u *msUploader) Upload(options uploader.Options) error {
 		return err
 	}
 	input := mediastoredata.PutObjectInput{
-		Path:         aws.String(options.Path),
-		ContentType:  options.ContentType,
-		CacheControl: options.CacheControl,
-		Body:         options.Body,
+		Path:               aws.String(options.Path),
+		ContentType:        options.ContentType,
+		CacheControl:       options.CacheControl,
+		Body:               options.Body,
+		UploadAvailability: u.uploadAvailability,
 	}
 	_, err = client.PutObject(options.Context, &input)
 	return err

diff --git a/main.go b/main.go
@@ -31,6 +31,7 @@ type Config struct {
 	HTTPPort        int                `envconfig:"HTTP_PORT" default:"80"`
 	LogLevel        string             `envconfig:"LOG_LEVEL" default:"debug"`
 	CacheControl    cachecontrol.Rules `envconfig:"CACHE_CONTROL_RULES"`
+	ChunkedTransfer bool               `envconfig:"MEDIASTORE_CHUNKED_TRANSFER"`
 }
 
 func loadConfig() (Config, error) {
@@ -42,6 +43,9 @@ func loadConfig() (Config, error) {
 	if cfg.UploadDriver != "s3" && cfg.UploadDriver != "mediastore" {
 		return cfg, errors.New(`invalid UPLOAD_DRIVER, valid options are "s3" and "mediastore"`)
 	}
+	if cfg.ChunkedTransfer && cfg.UploadDriver != "mediastore" {
+		return cfg, errors.New("MEDIASTORE_CHUNKED_TRANSFERS should only be defined for the mediastore UPLOAD_DRIVER")
+	}
 	return cfg, nil
 }
 
@@ -50,7 +54,7 @@ func (c *Config) uploader() (uploader.Uploader, error) {
 		return s3.New()
 	}
 	if c.UploadDriver == "mediastore" {
-		return mediastore.New()
+		return mediastore.New(mediastore.DriverOptions{ChunkedTransfer: c.ChunkedTransfer})
 	}
 	return nil, fmt.Errorf("invalid upload driver %q", c.UploadDriver)
 }