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

support hash-only streaming jobs

Browse Source
This commit is contained in:
Samuel Cui
2026-08-21 16:05:48 +08:00
parent 4c33c74475
commit 202c9d3041
6 changed files with 256 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acp.go
+3 -2
View File
@@ -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 {
}
copy.go
+53 -29
View File
@@ -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:
}
}
copy_test.go
+76
View File
@@ -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)
}
}
job.go
+21 -14
View File
@@ -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 {
prepare.go
+10 -2
View File
@@ -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
}
}
}
})
stream_test.go
+93
View File
@@ -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")
}
}