diff --git a/acp.go b/acp.go index 14ea68f..bc97ba1 100644 --- a/acp.go +++ b/acp.go @@ -95,10 +95,11 @@ func (c *Copyer) run(ctx context.Context) error { } } - // empty pipes + // Drain remaining stages in dependency order. Prepared jobs retain open sources. for range indexed { } - for range prepared { + for job := range prepared { + job.finishSource() } for range copyed { } diff --git a/copy.go b/copy.go index b4b8db1..02da853 100644 --- a/copy.go +++ b/copy.go @@ -63,17 +63,13 @@ func (c *Copyer) copy(ctx context.Context, prepared <-chan *writeJob) <-chan *ba go wrap(ctx, func() { defer copying.Done() - for { - select { - case <-ctx.Done(): - return - case job, ok := <-prepared: - if !ok { - return - } - - wrap(ctx, func() { c.write(ctx, job, ch, cntr, noSpaceDevices) }) + for job := range prepared { + if ctx.Err() != nil { + job.finishSource() + continue } + + wrap(ctx, func() { c.write(ctx, job, ch, cntr, noSpaceDevices) }) } }) } @@ -82,18 +78,26 @@ func (c *Copyer) copy(ctx context.Context, prepared <-chan *writeJob) <-chan *ba } func (c *Copyer) write(ctx context.Context, job *writeJob, ch chan<- *baseJob, cntr *counter, noSpaceDevices mapset.Set[string]) { - job.setStatus(jobStatusCopying) - var wg sync.WaitGroup defer func() { wg.Wait() - job.done() + job.finishSource() job.setStatus(jobStatusFinishing) - ch <- job.baseJob + select { + case ch <- job.baseJob: + case <-ctx.Done(): + } }() - // shortcut - if noSpaceDevices.Contains(lo.Map(job.targets, func(target string, _ int) string { return c.getDevice(target) })...) { + job.setStatus(jobStatusCopying) + if job.size != job.stat.size { + job.fail("", fmt.Errorf("source size changed, indexed=%d current=%d", job.stat.size, job.size)) + return + } + + // Skip jobs only when every requested target device is already exhausted. + targetDevices := lo.Map(job.targets, func(target string, _ int) string { return c.getDevice(target) }) + if len(targetDevices) > 0 && noSpaceDevices.Contains(targetDevices...) { job.fail("", ErrTargetNoSpace) return } @@ -145,6 +149,8 @@ func (c *Copyer) write(ctx context.Context, job *writeJob, ch chan<- *baseJob, c } if !job.copyer.toDevice.linear && job.size > 0 { if err := truncate(file, job.size); err != nil { + _ = file.Close() + _ = os.Remove(target) job.fail(target, fmt.Errorf("truncate dst file fail, %w", err)) continue } @@ -180,7 +186,11 @@ func (c *Copyer) write(ctx context.Context, job *writeJob, ch chan<- *baseJob, c job.fail(target, fmt.Errorf("write dst file fail, %w", rerr)) }() - defer file.Close() + defer func() { + if file != nil { + _ = file.Close() + } + }() for buf := range ch { n, err := file.Write(buf) if err != nil { @@ -193,20 +203,25 @@ func (c *Copyer) write(ctx context.Context, job *writeJob, ch chan<- *baseJob, c } } - if err := file.Sync(); err != nil { - rerr = fmt.Errorf("sync dst file fail, %w", err) + // A linear target publishes its durability boundary when the caller unmounts it. + if !c.toDevice.linear { + if err := file.Sync(); err != nil { + rerr = fmt.Errorf("sync dst file fail, %w", err) + return + } + } + if err := file.Close(); err != nil { + file = nil + rerr = fmt.Errorf("close dst file fail, %w", err) return } + file = nil if readErr != nil { rerr = readErr return } }) } - if len(chans) == 0 { - return - } - if c.withHash { sha := sha256Pool.Get().(hash.Hash) sha.Reset() @@ -226,17 +241,25 @@ func (c *Copyer) write(ctx context.Context, job *writeJob, ch chan<- *baseJob, c job.setHash(sha.Sum(nil)) }) } - readErr = c.streamCopy(ctx, chans, job.reader, &cntr.bytes) + if len(chans) == 0 { + return + } + var copied int64 + copied, readErr = c.streamCopy(ctx, chans, job.reader, &cntr.bytes) + if readErr == nil && copied != job.size { + readErr = fmt.Errorf("source size changed while copying, expected=%d copied=%d", job.size, copied) + } } -func (c *Copyer) streamCopy(ctx context.Context, dsts []chan []byte, src io.ReadCloser, bytes *int64) error { - for idx := int64(0); ; idx += batchSize { +func (c *Copyer) streamCopy(ctx context.Context, dsts []chan []byte, src io.ReadCloser, bytes *int64) (int64, error) { + var copied int64 + for { buf := make([]byte, batchSize) n, err := io.ReadFull(src, buf) if err != nil { if !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) { - return fmt.Errorf("slice mmap fail, %w", err) + return copied, fmt.Errorf("slice mmap fail, %w", err) } } @@ -246,14 +269,15 @@ func (c *Copyer) streamCopy(ctx context.Context, dsts []chan []byte, src io.Read } nr := len(buf) + copied += int64(nr) atomic.AddInt64(bytes, int64(nr)) if nr < batchSize { - return nil + return copied, nil } select { case <-ctx.Done(): - return ctx.Err() + return copied, ctx.Err() default: } } diff --git a/copy_test.go b/copy_test.go index 63c5a83..c63e5ac 100644 --- a/copy_test.go +++ b/copy_test.go @@ -12,6 +12,19 @@ import ( mapset "github.com/deckarep/golang-set/v2" ) +type trackingReadCloser struct { + closed int +} + +func (*trackingReadCloser) Read([]byte) (int, error) { + return 0, io.EOF +} + +func (r *trackingReadCloser) Close() error { + r.closed++ + return nil +} + func TestCopyEmptyFile(t *testing.T) { tests := []struct { name string @@ -94,3 +107,66 @@ func TestWritePublishesFinishingJob(t *testing.T) { t.Fatalf("published status = %q, want %q", status, jobStatusFinishing) } } + +func TestWriteJobWaitConsumedReturnsOnCancellation(t *testing.T) { + // Model a linear source whose reader has already moved to the Copy stage. + job := newWriteJob(nil, new(trackingReadCloser), 0, true) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // Cancellation must release Prepare without waiting for Copy to consume the reader. + if job.waitConsumed(ctx) { + t.Fatal("waitConsumed() = true after cancellation, want false") + } +} + +func TestCopyClosesPreparedSourcesAfterCancellation(t *testing.T) { + // Queue one prefetched reader before starting an already-canceled Copy stage. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + copyer := &Copyer{option: newOption(), eventCh: make(chan Event, 1)} + copyer.toDevice.threads = 1 + reader := new(trackingReadCloser) + job := newWriteJob(nil, reader, 0, true) + prepared := make(chan *writeJob, 1) + prepared <- job + close(prepared) + + // Copy owns accepted readers and must drain and close them during cancellation. + for range copyer.copy(ctx, prepared) { + } + if reader.closed != 1 { + t.Fatalf("reader closed %d times, want 1", reader.closed) + } + if !job.waitConsumed(context.Background()) { + t.Fatal("linear source was not notified that the reader was consumed") + } +} + +func TestWriteReturnsWhenCanceledBeforePublishing(t *testing.T) { + // Use an unbuffered completion channel with no receiver to expose a blocked handoff. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + copyer := &Copyer{option: newOption(), eventCh: make(chan Event, 8)} + reader := new(trackingReadCloser) + job := newWriteJob(&baseJob{ + copyer: copyer, + src: &source{}, + stat: &stat{}, + }, reader, 0, false) + done := make(chan struct{}) + go func() { + copyer.write(ctx, job, make(chan *baseJob), new(counter), mapset.NewSet[string]()) + close(done) + }() + + // Cancellation must skip the completion handoff while retaining source cleanup. + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("write did not return after cancellation") + } + if reader.closed != 1 { + t.Fatalf("reader closed %d times, want 1", reader.closed) + } +} diff --git a/job.go b/job.go index 6c11038..ce28f23 100644 --- a/job.go +++ b/job.go @@ -1,6 +1,7 @@ package acp import ( + "context" "encoding/hex" "io" "io/fs" @@ -111,36 +112,42 @@ func (j *baseJob) report() *Job { type writeJob struct { *baseJob - reader io.ReadCloser - size int64 - ch chan struct{} + reader io.ReadCloser + size int64 + consumed chan struct{} } -func newWriteJob(job *baseJob, src io.ReadCloser, size int64, needWait bool) *writeJob { +func newWriteJob(job *baseJob, src io.ReadCloser, size int64, waitConsumed bool) *writeJob { j := &writeJob{ baseJob: job, reader: src, size: size, } - if needWait { - j.ch = make(chan struct{}) + if waitConsumed { + j.consumed = make(chan struct{}) } return j } -func (wj *writeJob) done() { - wj.reader.Close() +func (wj *writeJob) finishSource() { + _ = wj.reader.Close() - if wj.ch != nil { - close(wj.ch) + if wj.consumed != nil { + close(wj.consumed) } } -func (wj *writeJob) wait() { - if wj.ch == nil { - return +func (wj *writeJob) waitConsumed(ctx context.Context) bool { + if wj.consumed == nil { + return true + } + + select { + case <-wj.consumed: + return true + case <-ctx.Done(): + return false } - <-wj.ch } type Job struct { diff --git a/prepare.go b/prepare.go index f9c2dd7..1a3deea 100644 --- a/prepare.go +++ b/prepare.go @@ -79,8 +79,16 @@ func (c *Copyer) prepare(ctx context.Context, indexed <-chan *baseJob) <-chan *w } wj := newWriteJob(job, file, size, c.fromDevice.linear) - ch <- wj - wj.wait() + select { + case ch <- wj: + case <-ctx.Done(): + wj.finishSource() + return + } + + if !wj.waitConsumed(ctx) { + return + } } } }) diff --git a/stream_test.go b/stream_test.go index ce1e147..8fa4b88 100644 --- a/stream_test.go +++ b/stream_test.go @@ -2,6 +2,8 @@ package acp import ( "context" + "crypto/sha256" + "encoding/hex" "errors" "fmt" "io" @@ -9,6 +11,7 @@ import ( "path/filepath" "sync/atomic" "testing" + "time" ) type sliceStreamSource struct { @@ -87,6 +90,33 @@ func TestRunStreamCopiesRequestsInLinearOrder(t *testing.T) { } } +func TestRunStreamHashesWithoutTargets(t *testing.T) { + // Submit one targetless request through the same bounded stream interface. + content := []byte("hash-only fixture") + input := filepath.Join(t.TempDir(), "source.txt") + if err := os.WriteFile(input, content, 0o644); err != nil { + t.Fatal(err) + } + source := &sliceStreamSource{requests: []*StreamRequest{{ID: 1, Source: input}}} + sink := new(collectingStreamSink) + + // Verify ACP reads the source once and reports its SHA-256 without creating a target. + if err := RunStream(context.Background(), source, sink, WithHash(true)); err != nil { + t.Fatal(err) + } + if len(sink.results) != 1 { + t.Fatalf("received %d results, want 1", len(sink.results)) + } + job := sink.results[0].Job + wantHash := sha256.Sum256(content) + if job.Status != JobStatusFinished || job.SHA256 != hex.EncodeToString(wantHash[:]) { + t.Fatalf("unexpected hash-only result: %#v", job) + } + if len(job.SuccessTargets) != 0 || len(job.FailTargets) != 0 { + t.Fatalf("hash-only targets = success:%v fail:%v", job.SuccessTargets, job.FailTargets) + } +} + func TestRunStreamReturnsSourceAndSinkErrors(t *testing.T) { // Verify that source failures cross the synchronous stream interface. sourceErr := errors.New("source failed") @@ -176,3 +206,66 @@ func TestRunStreamAppliesBoundedBackpressure(t *testing.T) { t.Fatalf("maximum outstanding requests = %d, want less than %d", maximum, total/2) } } + +type untilCanceledStreamSource struct { + input string + target string + nextID int64 +} + +func (s *untilCanceledStreamSource) Next(ctx context.Context) (*StreamRequest, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + s.nextID++ + return &StreamRequest{ + ID: s.nextID, + Source: s.input, + Targets: []string{filepath.Join(s.target, fmt.Sprintf("%04d", s.nextID))}, + }, nil +} + +func TestRunStreamCancellationDrainsPrefetchedJobs(t *testing.T) { + // Feed work until Prepare starts so cancellation occurs with an active pipeline. + root := t.TempDir() + input := filepath.Join(root, "source.txt") + if err := os.WriteFile(input, []byte("fixture"), 0o644); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + source := &untilCanceledStreamSource{input: input, target: filepath.Join(root, "target")} + cancelWhenPreparing := func(event Event) { + update, ok := event.(*EventUpdateJob) + if !ok { + return + } + if update.Job.Status != JobStatusPreparing { + return + } + cancel() + } + + // The canceled pipeline must drain its queues and return the context error. + done := make(chan error, 1) + go func() { + done <- RunStream( + ctx, + source, + new(collectingStreamSink), + SetToDevice(LinearDevice(true)), + WithEventHandler(cancelWhenPreparing), + ) + }() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("RunStream() error = %v, want %v", err, context.Canceled) + } + case <-time.After(5 * time.Second): + t.Fatal("RunStream did not return after cancellation") + } +}