diff --git a/cmd/put.go b/cmd/put.go index 2c316e9..36e5743 100644 --- a/cmd/put.go +++ b/cmd/put.go @@ -16,6 +16,7 @@ package cmd import ( "bytes" + "context" "errors" "fmt" "io" @@ -58,8 +59,12 @@ type uploadChunk struct { } func uploadOneChunk(dbx filesClient, args *files.UploadSessionAppendArg, data []byte) error { - return retryWithBackoff(func() error { - err := dbx.UploadSessionAppendV2Context(currentContext(), args, bytes.NewReader(data)) + return uploadOneChunkContext(currentContext(), dbx, args, data) +} + +func uploadOneChunkContext(ctx context.Context, dbx filesClient, args *files.UploadSessionAppendArg, data []byte) error { + return retryWithBackoffContext(ctx, func() error { + err := dbx.UploadSessionAppendV2Context(ctx, args, bytes.NewReader(data)) if uploadChunkAlreadyAccepted(err, args.Cursor.Offset+uint64(len(data))) { return nil } @@ -106,14 +111,17 @@ func uploadSingleShot(dbx filesClient, r io.ReadSeeker, uploadArg *files.UploadA } func uploadChunked(dbx filesClient, r io.Reader, commitInfo *files.CommitInfo, sizeTotal int64, workers int, chunkSize int64, debug bool) (metadata *files.FileMetadata, err error) { + ctx, cancel := context.WithCancel(currentContext()) + defer cancel() + t0 := time.Now() startArgs := files.NewUploadSessionStartArg() startArgs.SessionType = &files.UploadSessionType{} startArgs.SessionType.Tag = files.UploadSessionTypeConcurrent var res *files.UploadSessionStartResult - err = retryWithBackoff(func() error { + err = retryWithBackoffContext(ctx, func() error { var e error - res, e = dbx.UploadSessionStartContext(currentContext(), startArgs, nil) + res, e = dbx.UploadSessionStartContext(ctx, startArgs, nil) return e }) if err != nil { @@ -137,8 +145,15 @@ func uploadChunked(dbx filesClient, r io.Reader, commitInfo *files.CommitInfo, s args.Close = chunk.close t0 := time.Now() - if err := uploadOneChunk(dbx, args, chunk.data); err != nil { - errCh <- err + if err := uploadOneChunkContext(ctx, dbx, args, chunk.data); err != nil { + // Only the first error is needed. A worker must not block + // reporting another error while the caller waits for it. + select { + case errCh <- err: + default: + } + cancel() + return } if debug { log.Printf("Chunk upload at offset %d took: %v\n", chunk.offset, time.Since(t0)) @@ -147,8 +162,33 @@ func uploadChunked(dbx filesClient, r io.Reader, commitInfo *files.CommitInfo, s }() } + stopWorkers := sync.OnceFunc(func() { + close(workCh) + wg.Wait() + }) + defer func() { + // Abort in-flight requests and retry waits before joining workers on + // any early return. The command context remains usable by other files. + cancel() + stopWorkers() + }() + + cancellationError := func() error { + select { + case err := <-errCh: + return err + default: + return ctx.Err() + } + } + written := int64(0) for written < sizeTotal { + // A ready queue send can win a select even after cancellation. Check + // before reading so that it cannot trigger another source read. + if ctx.Err() != nil { + return nil, cancellationError() + } data, err := io.ReadAll(&io.LimitedReader{R: r, N: chunkSize}) if err != nil { return nil, err @@ -168,16 +208,17 @@ func uploadChunked(dbx filesClient, r io.Reader, commitInfo *files.CommitInfo, s } select { - case workCh <- chunk: + case <-ctx.Done(): + return nil, cancellationError() case err := <-errCh: return nil, err + case workCh <- chunk: } written += int64(len(data)) } - close(workCh) - wg.Wait() + stopWorkers() select { case err := <-errCh: return nil, err @@ -190,9 +231,9 @@ func uploadChunked(dbx filesClient, r io.Reader, commitInfo *files.CommitInfo, s t2 := time.Now() cursor := files.NewUploadSessionCursor(res.SessionId, uint64(written)) finishArgs := files.NewUploadSessionFinishArg(cursor, commitInfo) - err = retryWithBackoff(func() error { + err = retryWithBackoffContext(ctx, func() error { var e error - metadata, e = dbx.UploadSessionFinishContext(currentContext(), finishArgs, nil) + metadata, e = dbx.UploadSessionFinishContext(ctx, finishArgs, nil) return e }) if debug { diff --git a/cmd/put_cancellation_test.go b/cmd/put_cancellation_test.go new file mode 100644 index 0000000..930fdc6 --- /dev/null +++ b/cmd/put_cancellation_test.go @@ -0,0 +1,309 @@ +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/dropbox/dropbox-sdk-go-unofficial/v6/dropbox" + "github.com/dropbox/dropbox-sdk-go-unofficial/v6/dropbox/files" +) + +func TestUploadChunkedCancelsHTTPRequests(t *testing.T) { + for _, scenario := range []string{"source error", "source truncation", "sibling error", "parent cancellation", "parent deadline"} { + t.Run(scenario, func(t *testing.T) { + parent, cancel := context.WithCancel(context.Background()) + if scenario == "parent deadline" { + cancel() + parent, cancel = context.WithTimeout(context.Background(), time.Second) + } + old := commandContext + commandContext = parent + defer func() { commandContext = old }() + started := make(chan struct{}) + canceled := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "upload_session/start"): + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"session_id":"session"}`) + case strings.HasSuffix(r.URL.Path, "upload_session/append_v2"): + _, _ = io.Copy(io.Discard, r.Body) + var arg files.UploadSessionAppendArg + if err := json.Unmarshal([]byte(r.Header.Get("Dropbox-API-Arg")), &arg); err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + if arg.Cursor.Offset != 0 { + <-started + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, "permanent sibling append error") + return + } + close(started) + <-r.Context().Done() + close(canceled) + default: + t.Errorf("unexpected request %s", r.URL.Path) + w.WriteHeader(http.StatusBadRequest) + } + })) + defer server.Close() + defer cancel() // Also releases a stalled handler when a regression fails. + dbx := files.NewContext(dropbox.Config{Client: server.Client(), URLGenerator: func(_, ns, route string) string { return server.URL + "/" + ns + "/" + route }}) + var source io.Reader = bytes.NewReader([]byte("a")) + size, workers := int64(1), 1 + sourceErr := errors.New("source read failed") + switch scenario { + case "source error": + size = 2 + source = &uploadSourceFailure{started: started, failure: sourceErr} + case "source truncation": + size = 2 + source = &uploadSourceFailure{started: started, failure: io.EOF} + case "sibling error": + size = 2 + workers = 2 + source = bytes.NewReader([]byte("ab")) + } + done := make(chan error, 1) + go func() { + _, err := uploadChunked(dbx, source, files.NewCommitInfo("/file"), size, workers, 1, false) + done <- err + }() + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("append request did not start") + } + if scenario == "parent cancellation" { + cancel() + } + select { + case err := <-done: + if err == nil { + t.Fatal("expected upload failure") + } + switch scenario { + case "source error": + if !errors.Is(err, sourceErr) { + t.Fatalf("error=%v, want original source error", err) + } + case "source truncation": + if !strings.Contains(err.Error(), "failed to read") { + t.Fatalf("error=%v, want short-read error", err) + } + case "sibling error": + if !strings.Contains(err.Error(), "permanent sibling append error") { + t.Fatalf("error=%v, want original append error", err) + } + case "parent cancellation": + if !errors.Is(err, context.Canceled) { + t.Fatalf("error=%v, want cancellation", err) + } + case "parent deadline": + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("error=%v, want deadline exceeded", err) + } + } + case <-time.After(5 * time.Second): + t.Fatal("upload did not cancel its stalled request") + } + select { + case <-canceled: + case <-time.After(time.Second): + t.Fatal("HTTP server did not observe request cancellation") + } + if scenario != "parent cancellation" && scenario != "parent deadline" && parent.Err() != nil { + t.Fatalf("upload canceled the command context: %v", parent.Err()) + } + }) + } +} + +func TestUploadChunkedCancelsRetryWaitOnSourceError(t *testing.T) { + parent, cancel := context.WithCancel(context.Background()) + defer cancel() + oldContext := commandContext + commandContext = parent + defer func() { commandContext = oldContext }() + waiting := make(chan struct{}) + oldSleep := retrySleep + retrySleep = func(ctx context.Context, _ time.Duration) error { close(waiting); <-ctx.Done(); return ctx.Err() } + defer func() { retrySleep = oldSleep }() + sourceErr := errors.New("source read failed") + mock := &mockFilesClient{ + uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + return &files.UploadSessionStartResult{SessionId: "session"}, nil + }, + uploadSessionAppendV2Fn: func(*files.UploadSessionAppendArg, io.Reader) error { return io.ErrUnexpectedEOF }, + } + done := make(chan error, 1) + go func() { + _, err := uploadChunked(mock, &uploadSourceFailure{started: waiting, failure: sourceErr}, files.NewCommitInfo("/file"), 2, 1, 1, false) + done <- err + }() + select { + case err := <-done: + if !errors.Is(err, sourceErr) { + t.Fatalf("error=%v, want original source error", err) + } + case <-time.After(5 * time.Second): + cancel() + <-done + t.Fatal("source failure did not cancel retry wait") + } + if parent.Err() != nil { + t.Fatalf("command context canceled: %v", parent.Err()) + } +} + +// Each call represents one source chunk, allowing tests to cancel at precise +// points without filesystem timing or a blocked read that needs manual release. +type uploadReadObserver struct { + reads int + beforeRead func(int) +} + +func (r *uploadReadObserver) Read(p []byte) (int, error) { + r.reads++ + if r.beforeRead != nil { + r.beforeRead(r.reads) + } + p[0] = 'a' + return 1, nil +} + +func TestUploadChunkedDoesNotReadAfterCancellation(t *testing.T) { + for _, when := range []string{"before first read", "during first read"} { + t.Run(when, func(t *testing.T) { + // Repeat the second case to exercise both simultaneously-ready select arms. + for i := 0; i < 100; i++ { + ctx, cancel := context.WithCancel(context.Background()) + old := commandContext + commandContext = ctx + source := &uploadReadObserver{} + mock := &mockFilesClient{uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + if when == "before first read" { + cancel() + } + return &files.UploadSessionStartResult{SessionId: "session"}, nil + }} + wantReads := 0 + if when == "during first read" { + wantReads = 1 + source.beforeRead = func(int) { cancel() } + } + _, err := uploadChunked(mock, source, files.NewCommitInfo("/file"), 2, 1, 1, false) + cancel() + commandContext = old + if !errors.Is(err, context.Canceled) { + t.Fatalf("error=%v, want cancellation", err) + } + if source.reads != wantReads { + t.Fatalf("trial %d: source reads=%d, want %d", i, source.reads, wantReads) + } + } + }) + } +} + +type uploadContextClient struct { + *mockFilesClient + appendContext func(context.Context, *files.UploadSessionAppendArg, io.Reader) error +} + +func (c *uploadContextClient) UploadSessionAppendV2Context(ctx context.Context, arg *files.UploadSessionAppendArg, r io.Reader) error { + return c.appendContext(ctx, arg, r) +} + +func TestUploadChunkedCancellationPreservesWorkerError(t *testing.T) { + parent := context.Background() + old := commandContext + commandContext = parent + defer func() { commandContext = old }() + workerContext := make(chan context.Context, 1) + failure := errors.New("permanent worker error") + mock := &uploadContextClient{ + mockFilesClient: &mockFilesClient{uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + return &files.UploadSessionStartResult{SessionId: "session"}, nil + }}, + appendContext: func(ctx context.Context, _ *files.UploadSessionAppendArg, _ io.Reader) error { + workerContext <- ctx + return failure + }, + } + source := &uploadReadObserver{beforeRead: func(read int) { + if read == 2 { + ctx := <-workerContext + <-ctx.Done() + } + }} + _, err := uploadChunked(mock, source, files.NewCommitInfo("/file"), 3, 1, 1, false) + if !errors.Is(err, failure) { + t.Fatalf("error=%v, want original worker failure", err) + } + if source.reads > 2 { + t.Fatalf("source reads=%d, want at most 2", source.reads) + } + if parent.Err() != nil { + t.Fatalf("parent canceled: %v", parent.Err()) + } +} + +func TestPutRecursiveContinuesAfterSessionCancellation(t *testing.T) { + dir := t.TempDir() + for _, name := range []string{"a-fail", "b-good"} { + f, err := os.Create(filepath.Join(dir, name)) + if err != nil { + t.Fatal(err) + } + if err := f.Truncate(singleShotUploadSizeCutoff + 1); err != nil { + f.Close() + t.Fatal(err) + } + if err := f.Close(); err != nil { + t.Fatal(err) + } + } + sessions := 0 + var committed []string + failure := errors.New("first file append failed") + stubFilesClient(t, &mockFilesClient{ + uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + sessions++ + return &files.UploadSessionStartResult{SessionId: strconv.Itoa(sessions)}, nil + }, + uploadSessionAppendV2Fn: func(arg *files.UploadSessionAppendArg, _ io.Reader) error { + if arg.Cursor.SessionId == "1" { + return failure + } + return nil + }, + uploadSessionFinishFn: func(arg *files.UploadSessionFinishArg, _ io.Reader) (*files.FileMetadata, error) { + committed = append(committed, arg.Commit.Path) + return &files.FileMetadata{}, nil + }, + createFolderV2Fn: func(*files.CreateFolderArg) (*files.CreateFolderResult, error) { + return &files.CreateFolderResult{}, nil + }, + }) + err := putRecursive(dir, "/target", putOptions{workers: 1, chunkSize: putChunkSizeUnit, errOut: io.Discard}) + if err == nil || !strings.Contains(err.Error(), failure.Error()) { + t.Fatalf("error=%v, want aggregate failure", err) + } + if sessions != 2 || len(committed) != 1 || committed[0] != "/target/b-good" { + t.Fatalf("sessions=%d committed=%v, want second file committed", sessions, committed) + } +} diff --git a/cmd/put_workers_test.go b/cmd/put_workers_test.go new file mode 100644 index 0000000..24c6cef --- /dev/null +++ b/cmd/put_workers_test.go @@ -0,0 +1,113 @@ +package cmd + +import ( + "bytes" + "errors" + "io" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/dropbox/dropbox-sdk-go-unofficial/v6/dropbox/files" +) + +func TestUploadChunkedReturnsConcurrentAppendErrors(t *testing.T) { + for _, workers := range []int{1, 2, 4} { + t.Run(strconv.Itoa(workers)+" workers", func(t *testing.T) { + failure := errors.New("permanent append error") + var started sync.WaitGroup + started.Add(workers) + mock := &mockFilesClient{ + uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + return &files.UploadSessionStartResult{SessionId: "session"}, nil + }, + uploadSessionAppendV2Fn: func(*files.UploadSessionAppendArg, io.Reader) error { started.Done(); started.Wait(); return failure }, + uploadSessionFinishFn: func(*files.UploadSessionFinishArg, io.Reader) (*files.FileMetadata, error) { + t.Error("finish called after append failure") + return nil, nil + }, + } + done := make(chan error, 1) + go func() { + _, err := uploadChunked(mock, bytes.NewReader(make([]byte, workers)), files.NewCommitInfo("/file"), int64(workers), workers, 1, false) + done <- err + }() + select { + case err := <-done: + if !errors.Is(err, failure) { + t.Fatalf("error=%v, want %v", err, failure) + } + case <-time.After(5 * time.Second): + t.Fatal("upload hung while workers reported append errors") + } + }) + } +} + +type uploadSourceFailure struct { + started <-chan struct{} + failure error + read bool +} + +func (r *uploadSourceFailure) Read(p []byte) (int, error) { + if !r.read { + r.read = true + p[0] = 'a' + return 1, nil + } + <-r.started + return 0, r.failure +} + +func TestUploadChunkedJoinsWorkersOnSourceFailure(t *testing.T) { + for _, tc := range []struct { + name string + failure error + want string + }{ + {"read error", errors.New("source read failed"), "source read failed"}, + {"truncation", io.EOF, "failed to read"}, + } { + t.Run(tc.name, func(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + mock := &mockFilesClient{ + uploadSessionStartFn: func(*files.UploadSessionStartArg, io.Reader) (*files.UploadSessionStartResult, error) { + return &files.UploadSessionStartResult{SessionId: "session"}, nil + }, + uploadSessionAppendV2Fn: func(*files.UploadSessionAppendArg, io.Reader) error { + close(started) + <-release + return errors.New("append failed too") + }, + uploadSessionFinishFn: func(*files.UploadSessionFinishArg, io.Reader) (*files.FileMetadata, error) { + t.Error("finish called after source failure") + return nil, nil + }, + } + done := make(chan error, 1) + go func() { + _, err := uploadChunked(mock, &uploadSourceFailure{started: started, failure: tc.failure}, files.NewCommitInfo("/file"), 2, 1, 1, false) + done <- err + }() + select { + case err := <-done: + close(release) + t.Fatalf("returned before joining in-flight worker: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(release) + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("error=%v, want %q", err, tc.want) + } + case <-time.After(5 * time.Second): + t.Fatal("upload did not return after source and append errors") + } + }) + } +}