Allow S3 uploads without a timeout for exports (#21774)

shared/filestore/filesstore_test.go

@@ -5,9 +5,12 @@ package filestore
 
 import (
 	"bytes"
+	"context"
 	"fmt"
+	"io"
 	"math/rand"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -121,6 +124,91 @@ func (s *FileBackendTestSuite) TestReadWriteFile() {
 	s.EqualValues(readString, "test")
 }
 
+func (s *FileBackendTestSuite) TestReadWriteFileContext() {
+	type ContextWriter interface {
+		WriteFileContext(context.Context, io.Reader, string) (int64, error)
+	}
+
+	data := "test"
+
+	s.T().Run("no deadline", func(t *testing.T) {
+		var (
+			written int64
+			err     error
+		)
+
+		path := "tests/" + randomString()
+
+		ctx := context.Background()
+		if cw, ok := s.backend.(ContextWriter); ok {
+			written, err = cw.WriteFileContext(ctx, strings.NewReader(data), path)
+		} else {
+			written, err = s.backend.WriteFile(strings.NewReader(data), path)
+		}
+		s.NoError(err)
+		s.EqualValues(len(data), written, "expected given number of bytes to have been written")
+		defer s.backend.RemoveFile(path)
+
+		read, err := s.backend.ReadFile(path)
+		s.NoError(err)
+
+		readString := string(read)
+		s.Equal(readString, data)
+	})
+
+	s.T().Run("long deadline", func(t *testing.T) {
+		var (
+			written int64
+			err     error
+		)
+
+		path := "tests/" + randomString()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if cw, ok := s.backend.(ContextWriter); ok {
+			written, err = cw.WriteFileContext(ctx, strings.NewReader(data), path)
+		} else {
+			written, err = s.backend.WriteFile(strings.NewReader(data), path)
+		}
+		s.NoError(err)
+		s.EqualValues(len(data), written, "expected given number of bytes to have been written")
+		defer s.backend.RemoveFile(path)
+
+		read, err := s.backend.ReadFile(path)
+		s.NoError(err)
+
+		readString := string(read)
+		s.Equal(readString, data)
+	})
+
+	s.T().Run("missed deadline", func(t *testing.T) {
+		var (
+			written int64
+			err     error
+		)
+
+		path := "tests/" + randomString()
+
+		r, w := io.Pipe()
+		go func() {
+			// close the writer after a short time
+			time.Sleep(500 * time.Millisecond)
+			w.Close()
+		}()
+		ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+		defer cancel()
+		if cw, ok := s.backend.(ContextWriter); ok {
+			written, err = cw.WriteFileContext(ctx, r, path)
+		} else {
+			// this test works only with a context writer
+			return
+		}
+		s.Error(err)
+		s.Zero(written)
+	})
+}
+
 func (s *FileBackendTestSuite) TestReadWriteFileImage() {
 	b := []byte("testimage")
 	path := "tests/" + randomString() + ".png"

shared/filestore/s3store.go

@@ -369,6 +369,13 @@ func (b *S3FileBackend) MoveFile(oldPath, newPath string) error {
 }
 
 func (b *S3FileBackend) WriteFile(fr io.Reader, path string) (int64, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), b.timeout)
+	defer cancel()
+
+	return b.WriteFileContext(ctx, fr, path)
+}
+
+func (b *S3FileBackend) WriteFileContext(ctx context.Context, fr io.Reader, path string) (int64, error) {
 	var contentType string
 	path = filepath.Join(b.pathPrefix, path)
 	if ext := filepath.Ext(path); isFileExtImage(ext) {
@@ -377,22 +384,26 @@ func (b *S3FileBackend) WriteFile(fr io.Reader, path string) (int64, error) {
 		contentType = "binary/octet-stream"
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), b.timeout)
-	defer cancel()
 	options := s3PutOptions(b.encrypt, contentType)
 
-	objSize := -1
+	objSize := int64(-1)
 	isCloud := os.Getenv("MM_CLOUD_FILESTORE_BIFROST") != ""
 	if isCloud {
 		options.DisableContentSha256 = true
-	}
-	// We pass an object size only in situations where bifrost is not
-	// used. Bifrost needs to run in HTTPS, which is not yet deployed.
-	if buf, ok := fr.(*bytes.Buffer); ok && !isCloud {
-		objSize = buf.Len()
+	} else {
+		// We pass an object size only in situations where bifrost is not
+		// used. Bifrost needs to run in HTTPS, which is not yet deployed.
+		switch t := fr.(type) {
+		case *bytes.Buffer:
+			objSize = int64(t.Len())
+		case *os.File:
+			if s, err := t.Stat(); err == nil {
+				objSize = s.Size()
+			}
+		}
 	}
 
-	info, err := b.client.PutObject(ctx, b.bucket, path, fr, int64(objSize), options)
+	info, err := b.client.PutObject(ctx, b.bucket, path, fr, objSize, options)
 	if err != nil {
 		return info.Size, errors.Wrapf(err, "unable write the data in the file %s", path)
 	}