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

app/app_iface.go

@@ -1161,4 +1161,5 @@ type AppIface interface {
 	VerifyUserEmail(userID, email string) *model.AppError
 	ViewChannel(c request.CTX, view *model.ChannelView, userID string, currentSessionId string, collapsedThreadsSupported bool) (map[string]int64, *model.AppError)
 	WriteFile(fr io.Reader, path string) (int64, *model.AppError)
+	WriteFileContext(ctx context.Context, fr io.Reader, path string) (int64, *model.AppError)
 }

app/file.go

@@ -161,6 +161,10 @@ func (a *App) MoveFile(oldPath, newPath string) *model.AppError {
 	return nil
 }
 
+func (a *App) WriteFileContext(ctx context.Context, fr io.Reader, path string) (int64, *model.AppError) {
+	return a.Srv().writeFileContext(ctx, fr, path)
+}
+
 func (a *App) WriteFile(fr io.Reader, path string) (int64, *model.AppError) {
 	return a.Srv().writeFile(fr, path)
 }
@@ -173,6 +177,30 @@ func (s *Server) writeFile(fr io.Reader, path string) (int64, *model.AppError) {
 	return result, nil
 }
 
+func (s *Server) writeFileContext(ctx context.Context, fr io.Reader, path string) (int64, *model.AppError) {
+	type ContextWriter interface {
+		WriteFileContext(context.Context, io.Reader, string) (int64, error)
+	}
+
+	var (
+		fileBackend = s.FileBackend()
+		written     int64
+		err         error
+	)
+
+	// Check if we can provide a custom context, otherwise just use the default method.
+	if cw, ok := fileBackend.(ContextWriter); ok {
+		written, err = cw.WriteFileContext(ctx, fr, path)
+	} else {
+		written, err = fileBackend.WriteFile(fr, path)
+	}
+	if err != nil {
+		return written, model.NewAppError("WriteFile", "api.file.write_file.app_error", nil, "", http.StatusInternalServerError).Wrap(err)
+	}
+
+	return written, nil
+}
+
 func (a *App) AppendFile(fr io.Reader, path string) (int64, *model.AppError) {
 	result, nErr := a.FileBackend().AppendFile(fr, path)
 	if nErr != nil {

app/opentracing/opentracing_layer.go

@@ -18514,6 +18514,28 @@ func (a *OpenTracingAppLayer) WriteFile(fr io.Reader, path string) (int64, *mode
 	return resultVar0, resultVar1
 }
 
+func (a *OpenTracingAppLayer) WriteFileContext(ctx context.Context, fr io.Reader, path string) (int64, *model.AppError) {
+	origCtx := a.ctx
+	span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.WriteFileContext")
+
+	a.ctx = newCtx
+	a.app.Srv().Store().SetContext(newCtx)
+	defer func() {
+		a.app.Srv().Store().SetContext(origCtx)
+		a.ctx = origCtx
+	}()
+
+	defer span.Finish()
+	resultVar0, resultVar1 := a.app.WriteFileContext(ctx, fr, path)
+
+	if resultVar1 != nil {
+		span.LogFields(spanlog.Error(resultVar1))
+		ext.Error.Set(span, true)
+	}
+
+	return resultVar0, resultVar1
+}
+
 func NewOpenTracingAppLayer(childApp app.AppIface, ctx context.Context) *OpenTracingAppLayer {
 	newApp := OpenTracingAppLayer{
 		app: childApp,

jobs/export_process/worker.go

@@ -4,6 +4,7 @@
 package export_process
 
 import (
+	"context"
 	"io"
 	"path/filepath"
 
@@ -19,6 +20,7 @@ const jobName = "ExportProcess"
 type AppIface interface {
 	configservice.ConfigService
 	WriteFile(fr io.Reader, path string) (int64, *model.AppError)
+	WriteFileContext(ctx context.Context, fr io.Reader, path string) (int64, *model.AppError)
 	BulkExport(ctx request.CTX, writer io.Writer, outPath string, opts model.BulkExportOpts) *model.AppError
 	Log() *mlog.Logger
 }
@@ -45,7 +47,8 @@ func MakeWorker(jobServer *jobs.JobServer, app AppIface) model.Worker {
 		errCh := make(chan *model.AppError, 1)
 		go func() {
 			defer close(errCh)
-			_, appErr := app.WriteFile(rd, filepath.Join(outPath, exportFilename))
+			// Try to write without a timeout
+			_, appErr := app.WriteFileContext(context.Background(), rd, filepath.Join(outPath, exportFilename))
 			errCh <- appErr
 		}()
 

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)
 	}