tusd/pkg/s3store/s3store_upload_queue.go

package s3store

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
)

type s3UploadQueue struct {
	service              S3API
	queue                chan *s3UploadJob
	disableContentHashes bool
}

type s3UploadJob struct {
	// These fields are filled out be the job customer, i.e. the entity creating job.
	ctx             context.Context
	uploadPartInput *s3.UploadPartInput
	file            *os.File
	size            int64
	resultChannel   chan<- *s3UploadJob

	// These fields are filled out by the job worker, i.e. the entity doing the actual upload.
	etag string
	err  error
}

type s3APIForPresigning interface {
	UploadPartRequest(input *s3.UploadPartInput) (req *request.Request, output *s3.UploadPartOutput)
}

// newS3UploadQueue create a new upload queue and starts the workers in the goroutines.
func newS3UploadQueue(service S3API, concurrency int64, maxBufferedParts int64, disableContentHashes bool) *s3UploadQueue {
	s := &s3UploadQueue{
		service:              service,
		queue:                make(chan *s3UploadJob, maxBufferedParts),
		disableContentHashes: disableContentHashes,
	}

	for i := 0; i < concurrency; i++ {
		go s.uploadLoop()
	}

	return s
}

// queueLength returns the number of item waiting in the queue.
func (s s3UploadQueue) queueLength() int {
	return len(s.queue)
}

// push appends another item to the queue and returns immediately.
func (s s3UploadQueue) push(job *s3UploadJob) {
	s.queue <- job
}

// stop instructs the worker goroutines to shutdown after they complete their
// current task. The function immediately returns without waiting for the shutdown.
func (s s3UploadQueue) stop() {
	close(s.queue)
}

func (s s3UploadQueue) uploadLoop() {
	for {
		job, more := <-s.queue
		if !more {
			break
		}

		job.etag, job.err = s.putPartForUpload(job.ctx, job.uploadPartInput, job.file, job.size)
		job.resultChannel <- job
	}
}

func (s s3UploadQueue) putPartForUpload(ctx context.Context, uploadPartInput *s3.UploadPartInput, file *os.File, size int64) (etag string, err error) {
	// TODO: Move this back into s3store where the file is created
	defer cleanUpTempFile(file)

	if !s.disableContentHashes {
		// By default, use the traditional approach to upload data
		uploadPartInput.Body = file
		res, err := s.service.UploadPartWithContext(ctx, uploadPartInput)
		if res.ETag != nil {
			etag = *res.ETag
		}
		return etag, err
	} else {
		// Experimental feature to prevent the AWS SDK from calculating the SHA256 hash
		// for the parts we upload to S3.
		// We compute the presigned URL without the body attached and then send the request
		// on our own. This way, the body is not included in the SHA256 calculation.
		s3api, ok := s.service.Service.(s3APIForPresigning)
		if !ok {
			return fmt.Errorf("s3store: failed to cast S3 service for presigning")
		}

		s3Req, _ := s3api.UploadPartRequest(uploadPartInput)

		url, err := s3Req.Presign(15 * time.Minute)
		if err != nil {
			return "", err
		}

		req, err := http.NewRequest("PUT", url, file)
		if err != nil {
			return "", err
		}

		// Set the Content-Length manually to prevent the usage of Transfer-Encoding: chunked,
		// which is not supported by AWS S3.
		req.ContentLength = size

		res, err := http.DefaultClient.Do(req)
		if err != nil {
			return "", err
		}
		defer res.Body.Close()

		if res.StatusCode != 200 {
			buf := new(strings.Builder)
			io.Copy(buf, res.Body)
			return "", fmt.Errorf("s3store: unexpected response code %d for presigned upload: %s", res.StatusCode, buf.String())
		}

		return res.Header.Get("ETag"), nil
	}
}