mirrors/acp
1
0
Fork 0
mirror of https://github.com/samuelncui/acp.git synced 2026-09-03 22:57:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

add bounded streaming copy API

Browse Source
This commit is contained in:
Samuel Cui
2026-08-17 21:14:42 +08:00
parent f6dde5acb6
commit 3fe7bb3702
7 changed files with 353 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acp.go
+36 -5
View File
@@ -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})
}
cleanup.go
+11 -3
View File
@@ -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
}
}
}
index.go
+67
View File
@@ -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
job.go
+5 -4
View File
@@ -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
opt.go
+2
View File
@@ -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
stream.go
+54
View File
@@ -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
}
stream_test.go
+178
View File
@@ -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)
}
}