diff --git a/app/app_iface.go b/app/app_iface.go index 44ea3d936a..348e10e95d 100644 --- a/app/app_iface.go +++ b/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) } diff --git a/app/file.go b/app/file.go index 06a516bc86..a47e817003 100644 --- a/app/file.go +++ b/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 { diff --git a/app/opentracing/opentracing_layer.go b/app/opentracing/opentracing_layer.go index 1cea08e051..30fc5b6d92 100644 --- a/app/opentracing/opentracing_layer.go +++ b/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, diff --git a/jobs/export_process/worker.go b/jobs/export_process/worker.go index a9deaa8634..2697b65980 100644 --- a/jobs/export_process/worker.go +++ b/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 }() diff --git a/shared/filestore/filesstore_test.go b/shared/filestore/filesstore_test.go index 9bc9e281d7..c17836558d 100644 --- a/shared/filestore/filesstore_test.go +++ b/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" diff --git a/shared/filestore/s3store.go b/shared/filestore/s3store.go index 60132cd830..3dcbbe9b1a 100644 --- a/shared/filestore/s3store.go +++ b/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) }