From 3fe7bb3702159e7ec83a6b62993d0891f4564f6d Mon Sep 17 00:00:00 2001 From: Samuel Cui Date: Mon, 17 Aug 2026 21:14:42 +0800 Subject: [PATCH] add bounded streaming copy API --- acp.go | 41 ++++++++++-- cleanup.go | 14 +++- index.go | 67 +++++++++++++++++++ job.go | 9 +-- opt.go | 2 + stream.go | 54 +++++++++++++++ stream_test.go | 178 +++++++++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 353 insertions(+), 12 deletions(-) create mode 100644 stream.go create mode 100644 stream_test.go diff --git a/acp.go b/acp.go index f27ff35..14ea68f 100644 --- a/acp.go +++ b/acp.go @@ -2,6 +2,7 @@ package acp import ( "context" + "fmt" "sync" "github.com/sirupsen/logrus" @@ -10,6 +11,8 @@ import ( type Copyer struct { *option running sync.WaitGroup + errLock sync.Mutex + err error eventCh chan Event getDevice func(in string) string getDiskUsageCache func(mountPoint string) *diskUsageCache @@ -41,7 +44,8 @@ func New(ctx context.Context, opts ...Option) (*Copyer, error) { }), } - c.running.Add(1) + // Account for both pipeline and event dispatch before either goroutine starts. + c.running.Add(2) go wrap(ctx, func() { c.run(ctx) }) return c, nil @@ -51,6 +55,25 @@ func (c *Copyer) Wait() { c.running.Wait() } +// WaitErr waits for the copy pipeline and returns its first error. +func (c *Copyer) WaitErr() error { + c.Wait() + c.errLock.Lock() + defer c.errLock.Unlock() + return c.err +} + +func (c *Copyer) setError(err error) { + if err == nil { + return + } + c.errLock.Lock() + defer c.errLock.Unlock() + if c.err == nil { + c.err = err + } +} + func (c *Copyer) run(ctx context.Context) error { defer c.running.Done() defer close(c.eventCh) @@ -59,12 +82,18 @@ func (c *Copyer) run(ctx context.Context) error { indexed, err := c.index(ctx) if err != nil { + c.setError(err) return err } prepared := c.prepare(ctx, indexed) copyed := c.copy(ctx, prepared) - c.cleanupJob(ctx, copyed) + sinkFailed := c.cleanupJob(ctx, copyed) + if c.streamSink != nil && !sinkFailed { + if err := c.streamSink.Flush(ctx); err != nil { + c.setError(fmt.Errorf("flush stream sink failed, %w", err)) + } + } // empty pipes for range indexed { @@ -78,7 +107,6 @@ func (c *Copyer) run(ctx context.Context) error { } func (c *Copyer) eventLoop(ctx context.Context) { - c.running.Add(1) defer c.running.Done() chans := make([]chan Event, len(c.eventHanders)) @@ -86,13 +114,14 @@ func (c *Copyer) eventLoop(ctx context.Context) { chans[idx] = make(chan Event, 128) } + var handlers sync.WaitGroup for idx, ch := range chans { handler := c.eventHanders[idx] events := ch - c.running.Add(1) + handlers.Add(1) go wrap(ctx, func() { - defer c.running.Done() + defer handlers.Done() for { e, ok := <-events @@ -109,6 +138,7 @@ func (c *Copyer) eventLoop(ctx context.Context) { for _, ch := range chans { close(ch) } + handlers.Wait() }() for e := range c.eventCh { for _, ch := range chans { @@ -127,6 +157,7 @@ func (c *Copyer) submit(e Event) { func (c *Copyer) reportError(src, dst string, err error) { e := &Error{Src: src, Dst: dst, Err: err} + c.setError(fmt.Errorf("copy failed, source=%q target=%q, %w", src, dst, err)) c.logf(logrus.ErrorLevel, e.Error()) c.submit(&EventReportError{Error: e}) } diff --git a/cleanup.go b/cleanup.go index 49b7484..e078c24 100644 --- a/cleanup.go +++ b/cleanup.go @@ -5,12 +5,13 @@ import ( "fmt" ) -func (c *Copyer) cleanupJob(ctx context.Context, copyed <-chan *baseJob) { +func (c *Copyer) cleanupJob(ctx context.Context, copyed <-chan *baseJob) bool { + streamSinkFailed := false for { select { case job, ok := <-copyed: if !ok { - return + return streamSinkFailed } for _, dst := range job.successTargets { @@ -20,8 +21,15 @@ func (c *Copyer) cleanupJob(ctx context.Context, copyed <-chan *baseJob) { } job.setStatus(jobStatusFinished) + if c.streamSink != nil && !streamSinkFailed { + if err := c.streamSink.Write(ctx, &StreamResult{ID: job.streamID, Job: job.report()}); err != nil { + c.setError(fmt.Errorf("write stream result failed, id=%d, %w", job.streamID, err)) + streamSinkFailed = true + } + } case <-ctx.Done(): - return + c.setError(ctx.Err()) + return streamSinkFailed } } } diff --git a/index.go b/index.go index c8662f6..04cecc7 100644 --- a/index.go +++ b/index.go @@ -3,7 +3,9 @@ package acp import ( "context" "fmt" + "io" "os" + "path/filepath" "sort" "sync" "sync/atomic" @@ -19,6 +21,10 @@ type counter struct { } func (c *Copyer) index(ctx context.Context) (<-chan *baseJob, error) { + if c.streamSource != nil { + return c.indexStream(ctx), nil + } + jobs, err := c.walk(ctx) if err != nil { return nil, err @@ -40,6 +46,67 @@ func (c *Copyer) index(ctx context.Context) (<-chan *baseJob, error) { return ch, nil } +func (c *Copyer) indexStream(ctx context.Context) <-chan *baseJob { + ch := make(chan *baseJob, 128) + go wrap(ctx, func() { + defer close(ch) + + var bytes, files int64 + defer func() { + c.submit(&EventUpdateCount{Bytes: bytes, Files: files, Finished: true}) + }() + for { + request, err := c.streamSource.Next(ctx) + if err != nil { + if err != io.EOF { + c.reportError("", "", fmt.Errorf("read stream source failed, %w", err)) + } + return + } + if request == nil { + c.reportError("", "", fmt.Errorf("read stream source failed, request is nil")) + return + } + + sourcePath := filepath.Clean(request.Source) + info, err := os.Stat(sourcePath) + if err != nil { + c.reportError(sourcePath, "", fmt.Errorf("stream job get stat failed, %w", err)) + return + } + if !info.Mode().IsRegular() { + c.reportError(sourcePath, "", fmt.Errorf("stream job source is not a regular file")) + return + } + stat, err := newStat(sourcePath, info) + if err != nil { + c.reportError(sourcePath, "", fmt.Errorf("read stream job stat failed, %w", err)) + return + } + + job := &baseJob{ + copyer: c, + src: &source{base: filepath.Dir(sourcePath), path: filepath.Base(sourcePath)}, + path: sourcePath, + stat: stat, + targets: append([]string(nil), request.Targets...), + streamID: request.ID, + } + c.submit(&EventUpdateJob{job.report()}) + bytes += stat.size + files++ + + select { + case <-ctx.Done(): + c.setError(ctx.Err()) + return + case ch <- job: + } + } + }) + return ch +} + func (c *Copyer) walk(ctx context.Context) ([]*baseJob, error) { done := make(chan struct{}) var reporting sync.WaitGroup diff --git a/job.go b/job.go index 666efa5..6c11038 100644 --- a/job.go +++ b/job.go @@ -35,10 +35,11 @@ var ( ) type baseJob struct { - copyer *Copyer - src *source - path string - stat *stat + copyer *Copyer + src *source + path string + stat *stat + streamID int64 lock sync.Mutex writeTime time.Time diff --git a/opt.go b/opt.go index 90e3b69..c3a63d0 100644 --- a/opt.go +++ b/opt.go @@ -34,6 +34,8 @@ func comparePath(a, b string) int { type option struct { accurateJobs []*accurateJob wildcardJobs []*wildcardJob + streamSource StreamSource + streamSink StreamSink fromDevice *deviceOption toDevice *deviceOption diff --git a/stream.go b/stream.go new file mode 100644 index 0000000..b08cbfa --- /dev/null +++ b/stream.go @@ -0,0 +1,54 @@ +package acp + +import ( + "context" + "fmt" +) + +// StreamRequest describes one exact source and its destinations. +type StreamRequest struct { + ID int64 + Source string + Targets []string +} + +// StreamResult identifies the final report for one streamed request. +type StreamResult struct { + ID int64 + Job *Job +} + +// StreamSource supplies requests serially. It returns io.EOF when exhausted. +type StreamSource interface { + Next(context.Context) (*StreamRequest, error) +} + +// StreamSink consumes final results serially and flushes after the pipeline drains. +type StreamSink interface { + Write(context.Context, *StreamResult) error + Flush(context.Context) error +} + +// RunStream copies a bounded stream without retaining a whole-job report. +func RunStream(ctx context.Context, source StreamSource, sink StreamSink, opts ...Option) error { + if source == nil { + return fmt.Errorf("run stream failed, source is nil") + } + if sink == nil { + return fmt.Errorf("run stream failed, sink is nil") + } + + streamOption := func(option *option) *option { + option.streamSource = source + option.streamSink = sink + return option + } + copyer, err := New(ctx, append(opts, streamOption)...) + if err != nil { + return fmt.Errorf("run stream failed, %w", err) + } + if err := copyer.WaitErr(); err != nil { + return fmt.Errorf("run stream failed, %w", err) + } + return nil +} diff --git a/stream_test.go b/stream_test.go new file mode 100644 index 0000000..ce1e147 --- /dev/null +++ b/stream_test.go @@ -0,0 +1,178 @@ +package acp + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "sync/atomic" + "testing" +) + +type sliceStreamSource struct { + requests []*StreamRequest + index int + err error +} + +func (s *sliceStreamSource) Next(context.Context) (*StreamRequest, error) { + if s.index < len(s.requests) { + request := s.requests[s.index] + s.index++ + return request, nil + } + if s.err != nil { + return nil, s.err + } + return nil, io.EOF +} + +type collectingStreamSink struct { + results []*StreamResult + writes int + flushes int + err error +} + +func (s *collectingStreamSink) Write(_ context.Context, result *StreamResult) error { + s.writes++ + if s.err != nil { + return s.err + } + s.results = append(s.results, result) + return nil +} + +func (s *collectingStreamSink) Flush(context.Context) error { + s.flushes++ + return nil +} + +func TestRunStreamCopiesRequestsInLinearOrder(t *testing.T) { + // Create an ordered source stream without building ACP options per file. + root := t.TempDir() + source := new(sliceStreamSource) + for index, content := range []string{"first", "second", "third"} { + input := filepath.Join(root, "source", string(rune('a'+index))+".txt") + output := filepath.Join(root, "target", string(rune('a'+index))+".txt") + if err := os.MkdirAll(filepath.Dir(input), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(input, []byte(content), 0o644); err != nil { + t.Fatal(err) + } + source.requests = append(source.requests, &StreamRequest{ID: int64(index + 1), Source: input, Targets: []string{output}}) + } + sink := new(collectingStreamSink) + + // Run the linear pipeline and verify every final result is emitted once. + if err := RunStream(context.Background(), source, sink, WithHash(true), SetToDevice(LinearDevice(true))); err != nil { + t.Fatal(err) + } + if len(sink.results) != len(source.requests) { + t.Fatalf("received %d results, want %d", len(sink.results), len(source.requests)) + } + for index, result := range sink.results { + if result.ID != int64(index+1) { + t.Fatalf("result ID = %d, want %d", result.ID, index+1) + } + if result.Job.Status != JobStatusFinished || len(result.Job.SuccessTargets) != 1 || result.Job.SHA256 == "" { + t.Fatalf("unexpected result: %#v", result.Job) + } + } + if sink.flushes != 1 { + t.Fatalf("sink flushed %d times, want 1", sink.flushes) + } +} + +func TestRunStreamReturnsSourceAndSinkErrors(t *testing.T) { + // Verify that source failures cross the synchronous stream interface. + sourceErr := errors.New("source failed") + err := RunStream(context.Background(), &sliceStreamSource{err: sourceErr}, new(collectingStreamSink)) + if !errors.Is(err, sourceErr) { + t.Fatalf("RunStream() error = %v, want %v", err, sourceErr) + } + + // Verify that persistence failures are returned after the pipeline drains. + root := t.TempDir() + input := filepath.Join(root, "source.txt") + if err := os.WriteFile(input, []byte("fixture"), 0o644); err != nil { + t.Fatal(err) + } + sinkErr := errors.New("sink failed") + sink := &collectingStreamSink{err: sinkErr} + err = RunStream(context.Background(), &sliceStreamSource{requests: []*StreamRequest{ + {ID: 1, Source: input, Targets: []string{filepath.Join(root, "target-1.txt")}}, + {ID: 2, Source: input, Targets: []string{filepath.Join(root, "target-2.txt")}}, + {ID: 3, Source: input, Targets: []string{filepath.Join(root, "target-3.txt")}}, + }}, sink, Overwrite(true)) + if !errors.Is(err, sinkErr) { + t.Fatalf("RunStream() error = %v, want %v", err, sinkErr) + } + if sink.writes != 1 || sink.flushes != 0 { + t.Fatalf("sink calls = writes:%d flushes:%d, want writes:1 flushes:0", sink.writes, sink.flushes) + } +} + +type boundedStreamSource struct { + input string + target string + total int64 + produced int64 + consumed *int64 + maximum int64 +} + +func (s *boundedStreamSource) Next(context.Context) (*StreamRequest, error) { + id := atomic.AddInt64(&s.produced, 1) + if id > s.total { + return nil, io.EOF + } + outstanding := id - atomic.LoadInt64(s.consumed) + for { + maximum := atomic.LoadInt64(&s.maximum) + if outstanding <= maximum || atomic.CompareAndSwapInt64(&s.maximum, maximum, outstanding) { + break + } + } + return &StreamRequest{ + ID: id, Source: s.input, Targets: []string{filepath.Join(s.target, fmt.Sprintf("%04d", id))}, + }, nil +} + +type countingStreamSink struct { + consumed int64 +} + +func (s *countingStreamSink) Write(context.Context, *StreamResult) error { + atomic.AddInt64(&s.consumed, 1) + return nil +} + +func (*countingStreamSink) Flush(context.Context) error { + return nil +} + +func TestRunStreamAppliesBoundedBackpressure(t *testing.T) { + root := t.TempDir() + input := filepath.Join(root, "source.txt") + if err := os.WriteFile(input, nil, 0o644); err != nil { + t.Fatal(err) + } + const total = int64(1024) + sink := new(countingStreamSink) + source := &boundedStreamSource{ + input: input, target: filepath.Join(root, "target"), total: total, consumed: &sink.consumed, + } + if err := RunStream(context.Background(), source, sink, SetToDevice(LinearDevice(true))); err != nil { + t.Fatal(err) + } + if consumed := atomic.LoadInt64(&sink.consumed); consumed != total { + t.Fatalf("consumed %d results, want %d", consumed, total) + } + if maximum := atomic.LoadInt64(&source.maximum); maximum >= total/2 { + t.Fatalf("maximum outstanding requests = %d, want less than %d", maximum, total/2) + } +}