S3Store: Concurrently write upload parts to S3 while reading from client

pkg/s3store/s3store.go

@@ -135,6 +135,11 @@ type S3Store struct {
 	// MaxObjectSize is the maximum size an S3 Object can have according to S3
 	// API specifications. See link above.
 	MaxObjectSize int64
+	// MaxBufferedParts is the number of additional parts that can be received from
+	// the client and stored on disk while a part is being uploaded to S3. This
+	// can help improve throughput by not blocking the client while tusd is
+	// communicating with the S3 API, which can have unpredictable latency.
+	MaxBufferedParts int
 }
 
 type S3API interface {
@@ -272,6 +277,70 @@ func (upload *s3Upload) writeInfo(ctx context.Context, info handler.FileInfo) er
 	return err
 }
 
+type s3ChunkProducer struct {
+	files chan<- *os.File
+	done  chan struct{}
+	err   error
+	r     io.Reader
+}
+
+func (scp *s3ChunkProducer) produce(chunkSize int64) {
+	for {
+		file, err := scp.nextChunk(chunkSize)
+		if err != nil {
+			scp.err = err
+			close(scp.files)
+			return
+		}
+		if file == nil {
+			close(scp.files)
+			return
+		}
+		select {
+		case scp.files <- file:
+		case <-scp.done:
+			close(scp.files)
+			return
+		}
+	}
+}
+
+func (scp *s3ChunkProducer) nextChunk(size int64) (*os.File, error) {
+	// Create a temporary file to store the part
+	file, err := ioutil.TempFile("", "tusd-s3-tmp-")
+	if err != nil {
+		return nil, err
+	}
+
+	limitedReader := io.LimitReader(scp.r, size)
+	n, err := io.Copy(file, limitedReader)
+
+	// If the HTTP PATCH request gets interrupted in the middle (e.g. because
+	// the user wants to pause the upload), Go's net/http returns an io.ErrUnexpectedEOF.
+	// However, for S3Store it's not important whether the stream has ended
+	// on purpose or accidentally. Therefore, we ignore this error to not
+	// prevent the remaining chunk to be stored on S3.
+	if err == io.ErrUnexpectedEOF {
+		err = nil
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	// io.Copy returns (0, nil) when it reaches EOF
+	if n == 0 {
+		os.Remove(file.Name())
+		file.Close()
+		return nil, nil
+	}
+
+	// Seek to the beginning of the file
+	file.Seek(0, 0)
+
+	return file, nil
+}
+
 func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Reader) (int64, error) {
 	id := upload.id
 	store := upload.store
@@ -315,35 +384,23 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 		src = io.MultiReader(incompletePartFile, src)
 	}
 
-	for {
-		// Create a temporary file to store the part in it
-		file, err := ioutil.TempFile("", "tusd-s3-tmp-")
-		if err != nil {
-			return bytesUploaded, err
-		}
-		defer os.Remove(file.Name())
-		defer file.Close()
-
-		limitedReader := io.LimitReader(src, optimalPartSize)
-		n, err := io.Copy(file, limitedReader)
-
-		// If the HTTP PATCH request gets interrupted in the middle (e.g. because
-		// the user wants to pause the upload), Go's net/http returns an io.ErrUnexpectedEOF.
-		// However, for S3Store it's not important whether the stream has ended
-		// on purpose or accidentally. Therefore, we ignore this error to not
-		// prevent the remaining chunk to be stored on S3.
-		if err == io.ErrUnexpectedEOF {
-			err = nil
-		}
+	fileChan := make(chan *os.File, store.MaxBufferedParts)
+	doneChan := make(chan struct{})
+	defer close(doneChan)
 
-		// io.Copy does not return io.EOF, so we not have to handle it differently.
+	chunkProducer := s3ChunkProducer{
+		done:  doneChan,
+		files: fileChan,
+		r:     src,
+	}
+	go chunkProducer.produce(optimalPartSize)
+
+	for file := range fileChan {
+		stat, err := file.Stat()
 		if err != nil {
-			return bytesUploaded, err
-		}
-		// If io.Copy is finished reading, it will always return (0, nil).
-		if n == 0 {
-			return (bytesUploaded - incompletePartSize), nil
+			return 0, err
 		}
+		n := stat.Size()
 
 		// Seek to the beginning of the file
 		file.Seek(0, 0)
@@ -373,7 +430,16 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 		offset += n
 		bytesUploaded += n
 		nextPartNum += 1
+
+		if err := os.Remove(file.Name()); err != nil {
+			return bytesUploaded, err
+		}
+		if err := file.Close(); err != nil {
+			return bytesUploaded, err
+		}
 	}
+
+	return bytesUploaded - incompletePartSize, chunkProducer.err
 }
 
 func (upload *s3Upload) GetInfo(ctx context.Context) (info handler.FileInfo, err error) {

pkg/s3store/s3store_chunk_producer_test.go

@@ -0,0 +1,155 @@
+package s3store
+
+import (
+	"errors"
+	"os"
+	"strings"
+	"testing"
+	"time"
+)
+
+type InfiniteZeroReader struct{}
+
+func (izr InfiniteZeroReader) Read(b []byte) (int, error) {
+	b[0] = 0
+	return 1, nil
+}
+
+type ErrorReader struct{}
+
+func (ErrorReader) Read(b []byte) (int, error) {
+	return 0, errors.New("error from ErrorReader")
+}
+
+func TestChunkProducerConsumesEntireReaderWithoutError(t *testing.T) {
+	fileChan := make(chan *os.File)
+	doneChan := make(chan struct{})
+	expectedStr := "test"
+	r := strings.NewReader(expectedStr)
+	cp := s3ChunkProducer{
+		done:  doneChan,
+		files: fileChan,
+		r:     r,
+	}
+	go cp.produce(1)
+
+	actualStr := ""
+	b := make([]byte, 1)
+	for f := range fileChan {
+		n, err := f.Read(b)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		if n != 1 {
+			t.Fatalf("incorrect number of bytes read: wanted %d, got %d", 1, n)
+		}
+		actualStr += string(b)
+
+		os.Remove(f.Name())
+		f.Close()
+	}
+
+	if actualStr != expectedStr {
+		t.Errorf("incorrect string read from channel: wanted %s, got %s", expectedStr, actualStr)
+	}
+
+	if cp.err != nil {
+		t.Errorf("unexpected error from chunk producer: %s", cp.err)
+	}
+}
+
+func TestChunkProducerExitsWhenDoneChannelIsClosed(t *testing.T) {
+	fileChan := make(chan *os.File)
+	doneChan := make(chan struct{})
+	cp := s3ChunkProducer{
+		done:  doneChan,
+		files: fileChan,
+		r:     InfiniteZeroReader{},
+	}
+
+	completedChan := make(chan struct{})
+	go func() {
+		cp.produce(10)
+		completedChan <- struct{}{}
+	}()
+
+	close(doneChan)
+
+	select {
+	case <-completedChan:
+		// producer exited cleanly
+	case <-time.After(2 * time.Second):
+		t.Error("timed out waiting for producer to exit")
+	}
+
+	safelyDrainChannelOrFail(fileChan, t)
+}
+
+func TestChunkProducerExitsWhenDoneChannelIsClosedBeforeAnyChunkIsSent(t *testing.T) {
+	fileChan := make(chan *os.File)
+	doneChan := make(chan struct{})
+	cp := s3ChunkProducer{
+		done:  doneChan,
+		files: fileChan,
+		r:     InfiniteZeroReader{},
+	}
+
+	close(doneChan)
+
+	completedChan := make(chan struct{})
+	go func() {
+		cp.produce(10)
+		completedChan <- struct{}{}
+	}()
+
+	select {
+	case <-completedChan:
+		// producer exited cleanly
+	case <-time.After(2 * time.Second):
+		t.Error("timed out waiting for producer to exit")
+	}
+
+	safelyDrainChannelOrFail(fileChan, t)
+}
+
+func TestChunkProducerExitsWhenUnableToReadFromFile(t *testing.T) {
+	fileChan := make(chan *os.File)
+	doneChan := make(chan struct{})
+	cp := s3ChunkProducer{
+		done:  doneChan,
+		files: fileChan,
+		r:     ErrorReader{},
+	}
+
+	completedChan := make(chan struct{})
+	go func() {
+		cp.produce(10)
+		completedChan <- struct{}{}
+	}()
+
+	select {
+	case <-completedChan:
+		// producer exited cleanly
+	case <-time.After(2 * time.Second):
+		t.Error("timed out waiting for producer to exit")
+	}
+
+	safelyDrainChannelOrFail(fileChan, t)
+
+	if cp.err == nil {
+		t.Error("expected an error but didn't get one")
+	}
+}
+
+func safelyDrainChannelOrFail(c chan *os.File, t *testing.T) {
+	// At this point, we've signaled that the producer should exit, but it may write a few files
+	// into the channel before closing it and exiting. Make sure that we get a nil value
+	// eventually.
+	for i := 0; i < 100; i++ {
+		if f := <-c; f == nil {
+			return
+		}
+	}
+
+	t.Fatal("timed out waiting for channel to drain")
+}