amadan

carlos / platform Public

runs: the chunker stops a write that gives up mid-way instead of spinning on it (review round 1)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
claude 39e359f4724731db1db5397d1d8fa9cb489fc43e parent 449981a
2 files changed, +116 −3
  • internal/runs/chunker.go +10 −3
  • internal/runs/chunker_test.go +106 −0

CI passed — run details

diff --git a/internal/runs/chunker.go b/internal/runs/chunker.go
index 57b59c77d..a18d44ee5 100644
--- a/internal/runs/chunker.go
+++ b/internal/runs/chunker.go
@@ -21,9 +21,12 @@ var errBadStream = errors.New("runs: chunker stream is not out or err")
// it, so the follower interleaves exactly as the agent read. A chunk is
// written every second, or as soon as it holds 256 KB. Past 64 MB in total
// nothing more is written and Close says truncated; writes never fail, so
-// the command is never blocked by a refused pipe. Each chunk is a
-// create-only Put retried until it lands or the context ends, under the
-// lock, so seqs are written in order and a slow bucket slows the copy.
+// the command is never failed by a refused pipe. Each chunk is a
+// create-only Put retried until it lands or the context ends, and the
+// retries run under the lock. A stalled bucket therefore does block the
+// command's pipes, and the command with them once its pipe buffers fill,
+// until the Put lands or the context ends. That is deliberate: it keeps
+// seqs in order and stops output piling up in memory.
type Chunker struct {
ctx context.Context
s release.Store
@@ -129,6 +132,10 @@ func (c *Chunker) write(stream string, p []byte) {
p = p[n:]
if len(c.pending) >= ChunkBytes {
c.flushLocked()
+ if c.err != nil {
+ c.truncated = true
+ return
+ }
}
}
}
diff --git a/internal/runs/chunker_test.go b/internal/runs/chunker_test.go
index dc3297952..d7937c13e 100644
--- a/internal/runs/chunker_test.go
+++ b/internal/runs/chunker_test.go
@@ -131,6 +131,112 @@ func TestChunkerGivesUpWhenItsContextEnds(t *testing.T) {
}
}
+// dyingStore lets the first ok Puts through and fails every one after.
+type dyingStore struct {
+ release.Store
+ mu sync.Mutex
+ ok int
+}
+
+func (d *dyingStore) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
+ d.mu.Lock()
+ allowed := d.ok > 0
+ if allowed {
+ d.ok--
+ }
+ d.mu.Unlock()
+ if !allowed {
+ return "", errors.New("gone")
+ }
+ return d.Store.Put(ctx, key, body, ifMatch)
+}
+
+func TestChunkerReportsTheLastSeqThatLanded(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ s := &dyingStore{Store: release.NewDirStore(t.TempDir()), ok: 1}
+ c := newChunker(ctx, s, "bab", "titogo", tid, time.Hour, MaxOutputBytes)
+ go func() { time.Sleep(50 * time.Millisecond); cancel() }()
+ io.WriteString(c.Stream(StreamOut), strings.Repeat("x", ChunkBytes)) // chunk 1 lands
+ io.WriteString(c.Stream(StreamOut), strings.Repeat("y", ChunkBytes)) // chunk 2 gives up
+ last, truncated, err := c.Close()
+ if last != 1 || !truncated || err == nil {
+ t.Fatalf("Close = %d %v %v, want 1 true and an error", last, truncated, err)
+ }
+ if got := chunksOf(t, s.Store); len(got) != 1 {
+ t.Fatalf("chunks %d, want 1", len(got))
+ }
+}
+
+// A give-up in the middle of one large write must end the write, not spin
+// on it holding the lock.
+func TestChunkerGivingUpMidWriteDoesNotHang(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+ s := &failingStore{Store: release.NewDirStore(t.TempDir()), fail: 1 << 30}
+ c := newChunker(ctx, s, "bab", "titogo", tid, time.Hour, MaxOutputBytes)
+ done := make(chan struct{})
+ var last int64
+ var truncated bool
+ var err error
+ go func() {
+ defer close(done)
+ n, werr := c.Stream(StreamOut).Write(make([]byte, 2*ChunkBytes+1))
+ if n != 2*ChunkBytes+1 || werr != nil {
+ t.Errorf("Write = %d %v", n, werr)
+ }
+ last, truncated, err = c.Close()
+ }()
+ select {
+ case <-done:
+ case <-time.After(10 * time.Second):
+ t.Fatal("Write or Close hung after the chunker gave up")
+ }
+ if last != 0 || !truncated || err == nil {
+ t.Fatalf("Close = %d %v %v", last, truncated, err)
+ }
+}
+
+func TestChunkerRejoinsEachStreamWhenBothWriteAtOnce(t *testing.T) {
+ s := release.NewDirStore(t.TempDir())
+ c := newChunker(context.Background(), s, "bab", "titogo", tid, 2*time.Millisecond, MaxOutputBytes)
+ want := map[string]string{}
+ var wg sync.WaitGroup
+ for _, stream := range []string{StreamOut, StreamErr} {
+ var b strings.Builder
+ for i := 0; i < 2000; i++ {
+ b.WriteString(stream)
+ b.WriteString(strings.Repeat(string(rune('a'+i%26)), i%37))
+ b.WriteString("\n")
+ }
+ data := b.String()
+ want[stream] = data
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ w := c.Stream(stream)
+ for len(data) > 0 {
+ n := min(len(data), 1+len(data)%50)
+ w.Write([]byte(data[:n]))
+ data = data[n:]
+ time.Sleep(time.Microsecond)
+ }
+ }()
+ }
+ wg.Wait()
+ if _, truncated, err := c.Close(); truncated || err != nil {
+ t.Fatalf("Close: %v %v", truncated, err)
+ }
+ got := map[string]string{}
+ for _, ch := range chunksOf(t, s) {
+ got[ch[:1]] += ch[2:]
+ }
+ for _, stream := range []string{StreamOut, StreamErr} {
+ if got[stream] != want[stream] {
+ t.Fatalf("stream %s: got %d bytes, want %d, or the bytes differ", stream, len(got[stream]), len(want[stream]))
+ }
+ }
+}
+
func TestChunkerRefusesAnInvalidRunAtConstruction(t *testing.T) {
s := release.NewDirStore(t.TempDir())
for _, tc := range []struct {