diff --git a/internal/connector/dispatcher.go b/internal/connector/dispatcher.go index 004b9dbd4..222c2186c 100644 --- a/internal/connector/dispatcher.go +++ b/internal/connector/dispatcher.go @@ -1351,8 +1351,10 @@ type refusalRecorder struct { pending int } -// refusalWriteTimeout bounds a refusal's write, which runs on the goroutine -// reading the agent's stream. +// refusalWriteTimeout bounds a refusal's write. Which goroutine it runs on is +// the driver's: on most it is the one reading the agent's stream, and the +// codex driver writes on a goroutine of its own precisely because ten +// seconds is more than a reader can afford to spend. const refusalWriteTimeout = 10 * time.Second // RecordRefusal implements driver.RefusalRecorder. diff --git a/internal/connector/driver/codex/codex.go b/internal/connector/driver/codex/codex.go index 2c62eb86b..e307ba4ce 100644 --- a/internal/connector/driver/codex/codex.go +++ b/internal/connector/driver/codex/codex.go @@ -65,6 +65,8 @@ package codex import ( "bufio" "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "errors" "fmt" @@ -333,6 +335,11 @@ func (d *Driver) start(ctx context.Context, cfg driver.SessionConfig, resumeID s updates: make(chan driver.Update, 256), readerEnd: make(chan struct{}), } + s.scribe = newScribe(s.writeRefusal, s.emit) + // The worker's output is released when this session has read it, not on + // a clock, and the reader promises to do nothing slow between reads — + // which is the scribe's whole reason for existing. + worker.ReadingDone(s.readerEnd) go s.read() //nolint:contextcheck // the reader outlives the start's context: it runs as long as the worker does return s, nil } @@ -443,6 +450,10 @@ type session struct { updates chan driver.Update readerEnd chan struct{} + scribe *scribe + // endings counts the turn endings running off the reader. Only the + // reader adds to it, and only the reader's own ending waits on it. + endings sync.WaitGroup mu sync.Mutex id string @@ -456,6 +467,9 @@ type session struct { verifyDone chan struct{} verifyErr error closed bool + // full is a session that has refused more distinct tool calls than it + // can remember having refused. It is ending; it records no more. + full bool // updatesClosed is the reader's record that the updates channel is // closed. It is read and written under the same lock every emit takes, // so a goroutine still finishing a turn cannot send on a channel that @@ -591,7 +605,14 @@ func (s *session) Cancel(context.Context) error { if t == nil { return nil } - go s.worker.Terminate(s.grace) + // The reader is asked to stop as well, as Close does. Ending the worker + // is not the end of its output when a descendant outside the group is + // holding the pipe open: the end of file never comes, and a canceled + // turn that only the reader can finish would never be finished at all. + go func() { + s.endWorker(s.grace) + s.readerDone() + }() return nil } @@ -612,8 +633,10 @@ func (s *session) Close() error { case <-s.worker.Done(): case <-time.After(s.grace): } - s.worker.Terminate(s.grace) + s.endWorker(s.grace) s.readerDone() + // Close says the ledger has every refusal this session read. + s.scribe.close() return nil } @@ -629,7 +652,13 @@ func (s *session) readerDone() { select { case <-s.readerEnd: case <-time.After(s.grace): - s.worker.CloseStdout() + // A reader that never ends means a descendant outside the worker's + // group is holding the output open, so the end of file never comes. + // It is asked to stop, not cut off: it reads what is in the pipe + // first and ends within the drain budget. From here it must never + // be made to wait handing a refusal over, or that budget would run + // while it sat on a full queue. + s.worker.StopReading() <-s.readerEnd } } @@ -655,6 +684,19 @@ func (s *session) finish(t *turn, result driver.PromptResult, err error) { s.end(t, result, err, false) } +// settled is what every ending does before it hands a turn's result back: +// the refusals that turn made are in the ledger, not still queued. The +// writes are not on the reader any more, so a caller that gets its result +// and asks the ledger would otherwise be asking early — and the dispatcher +// settles an attempt on what the ledger holds. +// +// Endings run away from the reader, so waiting here holds up no reading. +func (s *session) settled() { + if s.scribe != nil { + s.scribe.drain() + } +} + // settle ends a turn this ending claimed, giving up the claim in the same // step so that nothing can come between the two. func (s *session) settle(t *turn, result driver.PromptResult, err error) { @@ -713,10 +755,21 @@ func (s *session) closeUpdates() { // the process closes its stdout. func (s *session) read() { defer func() { - // The updates channel closes last: finishing the turn still emits - // (a refusal read from stderr), and an update emitted after this is + // The updates channel closes last, and the scribe is drained before + // it: the dispatcher settles what the ledger would not take once the + // updates are over, so every refusal the reader read has to have + // been written by then. Finishing the turn still emits — a refusal + // read from stderr — and an update emitted after the close is // dropped rather than sent. defer s.closeUpdates() + defer s.scribe.close() + // An ending already running off the reader is the one that finishes + // its turn, and it reads the worker's last word on its way. Waiting + // for it here is what keeps this from finishing the same turn as + // ended, and what keeps the scribe open until it has handed over + // whatever it read. Every ending is bounded — the policy check's + // verdict and the worker's exit both are — so this is too. + s.endings.Wait() s.mu.Lock() s.ended = true t := s.turn @@ -732,11 +785,12 @@ func (s *session) read() { case <-time.After(s.grace): } } - s.stderrRefusals() + s.stderrRefusals(t) if t != nil { s.mu.Lock() canceled := t.canceled s.mu.Unlock() + s.settled() refusals := s.refusalsOf(t) switch { case canceled: @@ -755,6 +809,14 @@ func (s *session) read() { scanner.Buffer(make([]byte, 64<<10), 64<<20) for scanner.Scan() { s.handle(scanner.Bytes()) + if s.scribe.enough() { + // A drain that has heard as many refusals as it can hold. The + // pipe stops being read, which is what its own clock was about + // to do; the scanner's bufferful is still parsed, because + // breaking out here would drop lines that are no longer in the + // pipe for anyone else to find. + s.worker.StopReading() + } } // Drain what a scanner error left, so the process never blocks writing. _, _ = io.Copy(io.Discard, s.worker.Stdout()) @@ -849,12 +911,13 @@ func (s *session) threadStarted(id string) { // for this ending is the one thing it is for. func (s *session) unsafe(err error) { t := s.claim() - s.worker.Terminate(0) + s.endWorker(0) if t == nil { return } s.readerDone() - s.lastWord() + s.lastWord(t) + s.settled() s.settle(t, driver.PromptResult{Refusals: s.refusalsOf(t)}, err) } @@ -869,8 +932,9 @@ func (s *session) unsafe(err error) { // The policy check's own goroutine has unsafe, which waits for the reader // first, for that same reason. func (s *session) finishUnsafe(t *turn, err error) { - s.worker.Terminate(0) - s.lastWord() + s.endWorker(0) + s.lastWord(t) + s.settled() s.finish(t, driver.PromptResult{Refusals: s.refusalsOf(t)}, err) } @@ -898,14 +962,14 @@ func (s *session) failedVerification() error { // and a worker that is gone is not a stream that has been read. An ending // that is not the reader's own waits for the reader too — see unsafe — before // it asks a turn what its refusals are. -func (s *session) lastWord() { +func (s *session) lastWord(t *turn) { if s.worker != nil { select { case <-s.worker.Done(): case <-time.After(s.grace): } } - s.stderrRefusals() + s.stderrRefusals(t) } // finishCanceled ends a turn the connector canceled, after the worker's last @@ -913,7 +977,8 @@ func (s *session) lastWord() { // one still running is not waited for, because the process it would judge is // being ended by the cancel anyway. func (s *session) finishCanceled(t *turn) { - s.lastWord() + s.lastWord(t) + s.settled() // The turn's refusals are read after the worker's last word: the ones it // only logged are read by lastWord just above, and the ones it put on its // stream are already the turn's when the reader is the one ending it, @@ -992,7 +1057,7 @@ func (s *session) item(kind string, e event) { } s.emit(u) if kind == "item.completed" && it.Type == "mcp_tool_call" && it.Error != nil && refusedByApproval(it.Error.Message) { - s.refused("item:"+it.ID, it.ID, u.Tool, u.ToolKind) + s.refused(nil, "item:"+it.ID, it.ID, u.Tool, u.ToolKind) } } @@ -1021,26 +1086,72 @@ func refusedByApproval(message string) bool { // session's recorder before anything else is done with it, and only the first // time its tool call id is seen (driver's "Refusals"). A refusal Codex logs // and gives no id gets the key the caller passes. -func (s *session) refused(key, id, tool string, kind driver.ToolKind) { - refusal := driver.Refusal{ToolCallID: s.red.Sanitize(id), Tool: s.red.Sanitize(tool)} +func (s *session) refused(t *turn, key, id, tool string, kind driver.ToolKind) { + // The id is the agent's, and the redactor takes things out of a string + // without making it shorter: an agent that names a tool call a megabyte + // long would otherwise have that megabyte kept on the turn, in the + // queue, and in what this session remembers having refused. + refusal := driver.Refusal{ToolCallID: cut(s.red.Sanitize(id), maxToolCallID), Tool: s.red.Sanitize(tool)} + key = identity(key) s.mu.Lock() + // The turn the caller is settling, where it named one: an ending reads + // the worker's last word after another ending may already have taken + // the turn out of the session, and a refusal read then belongs to the + // turn it was read for rather than to nobody. + on := t + if on == nil { + on = s.turn + } + if s.full { + // The session is ending for want of memory; it cannot tell a repeat + // from a first any more, and recording one twice would be worse + // than recording neither. + s.mu.Unlock() + return + } first := !s.recorded[key] + overfull := false if first { s.recorded[key] = true - if s.turn != nil { - s.turn.refusals = append(s.turn.refusals, refusal) + if on != nil { + on.refusals = append(on.refusals, refusal) + } + // The one that fills the memory is recorded like any other. Past it + // a repeat cannot be told from a first, and recording one twice + // would be worse, so the session ends — after this refusal has been + // taken, not instead of it. + if len(s.recorded) >= maxRecorded { + s.full, overfull = true, true } } s.mu.Unlock() if !first { return } - if s.recorder != nil { - // The recorder owns what happens when the ledger refuses the write; - // the refusal happened either way. - _ = s.recorder.RecordRefusal(context.Background(), refusal) + defer func() { + if overfull { + s.endedOverfull() + } + }() + // The write and the update it precedes both belong to the scribe: the + // reader hands them over and goes back to reading, which is what lets + // the bound on its reading be a clock. See scribe.go. + s.scribe.hand(pending{ + refusal: refusal, + update: driver.Update{Kind: driver.UpdatePermission, ToolCallID: id, Tool: tool, ToolKind: kind, Allowed: false}, + }) +} + +// writeRefusal is the ledger write itself, and the one slow thing this +// driver does with a refusal. +func (s *session) writeRefusal(refusal driver.Refusal) { + if s.recorder == nil { + return } - s.emit(driver.Update{Kind: driver.UpdatePermission, ToolCallID: id, Tool: tool, ToolKind: kind, Allowed: false}) + // The recorder owns what happens when the ledger refuses the write; the + // refusal happened either way. + //nolint:contextcheck // a refusal's write is not any caller's to cancel + _ = s.recorder.RecordRefusal(context.Background(), refusal) } func (s *session) turnCompleted(e event) { @@ -1055,33 +1166,57 @@ func (s *session) turnCompleted(e event) { s.mu.Unlock() if canceled { // A cancel that won does not wait out the policy check either. - s.finishCanceled(t) + s.endTurn(func() { s.finishCanceled(t) }) return } - if err := s.verified(); err != nil { - s.finishUnsafe(t, err) - return - } - // Codex exits right after the turn it completed, and its stderr is whole - // only once it has: a refusal it logged and did not put on the stream is - // in the tail by then. - select { - case <-s.worker.Done(): - case <-time.After(s.grace): - } - s.stderrRefusals() - s.mu.Lock() - result := driver.PromptResult{Stop: driver.TurnEndTurn, Refusals: slices.Clone(t.refusals)} - if t.canceled { - // Only a cancel the connector asked for reads as canceled. - result.Stop = driver.TurnCanceled - } - s.mu.Unlock() - if e.Usage != nil { - result.Usage = driver.Usage{InputTokens: e.Usage.InputTokens, OutputTokens: e.Usage.OutputTokens} - s.emit(driver.Update{Kind: driver.UpdateUsage, Usage: &result.Usage}) - } - s.finish(t, result, nil) + s.endTurn(func() { + if err := s.verified(); err != nil { + s.finishUnsafe(t, err) + return + } + // Codex exits right after the turn it completed, and its stderr is + // whole only once it has: a refusal it logged and did not put on the + // stream is in the tail by then. + select { + case <-s.worker.Done(): + case <-time.After(s.grace): + } + s.stderrRefusals(t) + s.settled() + s.mu.Lock() + result := driver.PromptResult{Stop: driver.TurnEndTurn, Refusals: slices.Clone(t.refusals)} + if t.canceled { + // Only a cancel the connector asked for reads as canceled. + result.Stop = driver.TurnCanceled + } + s.mu.Unlock() + if e.Usage != nil { + result.Usage = driver.Usage{InputTokens: e.Usage.InputTokens, OutputTokens: e.Usage.OutputTokens} + s.emit(driver.Update{Kind: driver.UpdateUsage, Usage: &result.Usage}) + } + s.finish(t, result, nil) + }) +} + +// endTurn runs a turn's ending away from the reader, and remembers it so the +// reader's own ending waits for it. +// +// A turn ends by waiting: for the policy check's verdict, and for the worker +// to be gone so its stderr is whole. Neither is reading, and neither belongs +// on the goroutine whose job is to keep a pipe empty — a reader that waits +// is a reader that is not draining, and every bound on its reading is then a +// bound on something else. This is what lets the drain budget be a clock and +// the queue's cap be a count. +// +// The count is taken here, on the reader, before the ending starts, and +// waited for from the reader's own ending: one goroutine adds, the same one +// waits. +func (s *session) endTurn(ending func()) { + s.endings.Add(1) + go func() { + defer s.endings.Done() + ending() + }() } func (s *session) turnFailed() { @@ -1095,21 +1230,76 @@ func (s *session) turnFailed() { canceled := t.canceled s.mu.Unlock() if canceled { - s.finishCanceled(t) + s.endTurn(func() { s.finishCanceled(t) }) return } - // As after a completed turn: the stderr tail is whole once Codex exits. - select { - case <-s.worker.Done(): - case <-time.After(s.grace): - } - s.stderrRefusals() - refusals := s.refusalsOf(t) - if err := s.failedVerification(); err != nil { - s.finishUnsafe(t, err) - return + s.endTurn(func() { + // As after a completed turn: the stderr tail is whole once Codex + // exits. + select { + case <-s.worker.Done(): + case <-time.After(s.grace): + } + s.stderrRefusals(t) + s.settled() + refusals := s.refusalsOf(t) + if err := s.failedVerification(); err != nil { + s.finishUnsafe(t, err) + return + } + s.finish(t, driver.PromptResult{Refusals: refusals}, errors.New("codex: the turn failed")) + }) +} + +// maxToolCallID is how much of a tool call id a refusal keeps, and +// maxRecorded how many refusals a session remembers having recorded. The +// agent writes both, and neither is a session's to grow without end — the +// acp driver bounds the same two for the same reason. +// +// Reaching maxRecorded is not a bound that can be applied quietly: past it a +// repeat cannot be told from a first, and recording the same refusal twice +// would be worse than stopping. The session ends instead. +const ( + maxToolCallID = 256 + maxRecorded = 4096 +) + +// cut shortens a string the agent wrote to what this driver keeps of it. +func cut(s string, keep int) string { + if len(s) <= keep { + return s } - s.finish(t, driver.PromptResult{Refusals: refusals}, errors.New("codex: the turn failed")) + return s[:keep] +} + +// identity is what a session remembers a refusal by. It is bounded, because +// the agent writes it and a session's memory is not the agent's to grow; and +// it is not a prefix, because two tool call ids that begin alike are two +// refusals and cutting them to a common start would record only the first. +// A digest is bounded and keeps them apart. +func identity(key string) string { + if len(key) <= maxToolCallID { + return key + } + sum := sha256.Sum256([]byte(key)) + return hex.EncodeToString(sum[:]) +} + +// endedOverfull ends a session that has refused more distinct tool calls +// than it can remember having refused. +func (s *session) endedOverfull() { + s.endWorker(0) +} + +// endWorker ends the worker, and frees the reader from waiting for room +// before it does. Ending the worker asks its reader to stop, which starts +// the drain's clock, and a reader held at the scribe's queue is a reader +// spending that clock without reading — the pipe would be abandoned with +// the worker's output still in it. From here the queue grows rather than +// holds anyone up. +func (s *session) endWorker(grace time.Duration) { + s.scribe.noWaiting() + s.worker.Terminate(grace) } // refusalsOf is a turn's refusals so far. @@ -1122,7 +1312,7 @@ func (s *session) refusalsOf(t *turn) []driver.Refusal { // stderrRefusals counts the refusals Codex logs but does not put on its JSON // stream: an edit outside the working directory. Best effort: the stderr // kept is a tail. -func (s *session) stderrRefusals() { +func (s *session) stderrRefusals(t *turn) { if s.worker == nil { return } @@ -1143,7 +1333,7 @@ func (s *session) stderrRefusals() { // way are two, and reading the same output again — every way a turn // can end reads it — records each of them once. seen[line]++ - s.refused("stderr:"+strconv.Itoa(seen[line])+":"+line, "", tool, kind) + s.refused(t, "stderr:"+strconv.Itoa(seen[line])+":"+line, "", tool, kind) } } diff --git a/internal/connector/driver/codex/codex_test.go b/internal/connector/driver/codex/codex_test.go index d5fa4674a..e5bcf3e4b 100644 --- a/internal/connector/driver/codex/codex_test.go +++ b/internal/connector/driver/codex/codex_test.go @@ -6,6 +6,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "os" "path/filepath" "slices" @@ -791,10 +792,12 @@ func TestACompletedTurnThatWasCanceledDoesNotWaitForTheCheck(t *testing.T) { s := &session{verifyDone: make(chan struct{}), verifyAfter: time.Hour} turn := &turn{done: make(chan struct{}), canceled: true} s.turn = turn - done := make(chan struct{}) - go func() { s.turnCompleted(event{}); close(done) }() + // The turn's own end is what is waited for. A turn ends away from the + // reader now, so the call returning says nothing about whether the + // ending waited on anything. + s.turnCompleted(event{}) select { - case <-done: + case <-turn.done: case <-time.After(5 * time.Second): t.Fatal("a canceled turn waited for the policy check") } @@ -1146,6 +1149,10 @@ func TestAnUnsafeVerdictWaitsForWhatTheWorkerPutOnItsStream(t *testing.T) { got := <-answers require.ErrorIs(t, got.err, driver.ErrUnsafeMode, "an unsafe session is reported as unsafe, whoever gets to the turn first") require.ErrorContains(t, got.err, `Codex applied "on-request"`) + // The ledger is whole when the session is closed: the writes are + // not on the reader any more, and the turn's end is not where + // they are promised — the session's updates are. + require.NoError(t, s.Close()) assert.Len(t, ledger.Recorded(), 2, "both refusals reach the ledger") assert.Len(t, got.result.Refusals, 2, "the result carries what the ledger carries, and carries %d of the %d recorded", @@ -1154,6 +1161,71 @@ func TestAnUnsafeVerdictWaitsForWhatTheWorkerPutOnItsStream(t *testing.T) { } } +// A ledger that takes its time does not cost the worker its last lines. +// +// Writing a refusal is allowed ten seconds. While that ran on the goroutine +// reading the worker's output there was no bound that could be put on the +// reading: a clock runs out inside a write and takes the pipe away with the +// worker's next refusal still in it — recorded nowhere, not merely missing +// from a result — and a bound that waits for the pipe to fall idle never +// runs out at all against a descendant that keeps writing. +// +// The writes are a goroutine of their own now, so the reading is only ever +// reading. The hold here outlasts every clock in this path, with a refusal +// sitting in the pipe the whole time, and it is still read and still +// written. +func TestALedgerThatTakesItsTimeDoesNotCostTheWorkerItsLastLines(t *testing.T) { + ledger := &heldLedger{writing: make(chan struct{}, 1), release: make(chan struct{})} + denial := func(id string) string { + return `{"type":"item.completed","item":{"id":"` + id + `","type":"mcp_tool_call","server":"other","tool":"write",` + + `"error":{"message":"MCP tool call requires approval, but approval policy is never"},"status":"failed"}}` + } + unsafe := safeTurnContext() + unsafe["approval_policy"] = "on-request" + late := filepath.Join(t.TempDir(), "say-the-second") + h := newHarness(t, scenario{ + TurnContext: unsafe, + TurnContextAfterEvents: true, + Events: []string{`{"type":"turn.started"}`, denial("item_1")}, + // Said once the ledger is holding the first, so the second is in the + // pipe and nowhere else: the reader cannot have taken it into its + // own buffer, because it was not there to take. + LateEvents: []string{denial("item_2")}, + LateAfter: late, + Hang: true, + }) + cfg := h.config() + cfg.Refusals = ledger + s, err := h.drv.NewSession(context.Background(), cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + answers := make(chan driver.PromptResult, 1) + go func() { + result, _ := s.Prompt(context.Background(), "Task 1. Event 2.") + answers <- result + }() + + <-ledger.writing + require.NoError(t, os.WriteFile(late, nil, 0o600)) + held := make(chan struct{}) + go func() { + defer close(held) + // Longer than the grace and the drain budget together, and well + // inside the ten seconds a refusal's write is allowed. + time.Sleep(6 * time.Second) + close(ledger.release) + }() + + result := <-answers + <-held + require.NoError(t, s.Close()) + assert.Len(t, ledger.Recorded(), 2, + "the refusal that arrived while the ledger was writing is still read, and the ledger holds %d", len(ledger.Recorded())) + assert.Len(t, result.Refusals, 2, + "the result carries what the ledger carries, and carries %d", len(result.Refusals)) +} + // A session closed under a turn still reports the refusals that turn made. // A worker that has stopped reading its input leaves the prompt's write // blocked, and closing the session fails that write: the turn is then @@ -1276,7 +1348,7 @@ func TestARefusalReadAfterTheUpdatesCloseIsNotAPanic(t *testing.T) { require.NoError(t, s.Close()) require.Eventually(t, func() bool { return strings.Contains(session.StderrTail(), "rejected") }, 10*time.Second, 20*time.Millisecond, "the worker logged its refusal on its way out") - session.stderrRefusals() + session.stderrRefusals(nil) assert.Len(t, recorder.Recorded(), 1, "the refusal is recorded, and the update it carries is dropped rather than sent on a closed channel") @@ -1376,3 +1448,117 @@ func TestACanceledTurnRecordsItsRefusals(t *testing.T) { } assert.Len(t, recorder.Recorded(), 1, "the refusal Codex logged is recorded, not lost with the cancel") } + +// What a refusal keeps of a tool call id is bounded, and so is how many the +// session remembers having refused. The agent writes both, and the redactor +// takes things out of a string without making it shorter — so an agent that +// names a tool call a megabyte long would otherwise have that megabyte kept +// on the turn, in the queue, and in the session's memory of what it refused. +func TestWhatARefusalKeepsOfWhatTheAgentWroteIsBounded(t *testing.T) { + recorder := &drivertest.Refusals{} + h := newHarness(t, scenario{TurnContext: safeTurnContext()}) + cfg := h.config() + cfg.Refusals = recorder + s, err := h.drv.NewSession(context.Background(), cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + session := s.(*session) + tr := &turn{done: make(chan struct{})} + session.mu.Lock() + session.turn = tr + session.mu.Unlock() + + // Separators, so the redactor reads it as an id rather than as a + // credential and replaces the whole thing — which is what a long run of + // one character gets, and would prove nothing about the cut. + long := strings.Repeat("tool.call-", 8*maxToolCallID/10) + session.refused(nil, "item:"+long, long, "mcp__other__write", driver.ToolOther) + require.NoError(t, s.Close()) + + session.mu.Lock() + defer session.mu.Unlock() + require.Len(t, tr.refusals, 1) + assert.LessOrEqual(t, len(tr.refusals[0].ToolCallID), maxToolCallID, + "an id is cut to what this driver keeps of it, and kept %d", len(tr.refusals[0].ToolCallID)) + for key := range session.recorded { + assert.LessOrEqual(t, len(key), maxToolCallID+len("stderr:000:"), + "and what the session remembers is cut too") + } + assert.Len(t, recorder.Recorded(), 1) +} + +// A session that has refused more distinct tool calls than it can remember +// ends, rather than go on unable to tell a repeat from a first — recording +// the same refusal twice would be worse. The acp driver bounds the same +// memory for the same reason. +func TestASessionRemembersSoManyRefusalsAndNoMore(t *testing.T) { + recorder := &drivertest.Refusals{} + h := newHarness(t, scenario{TurnContext: safeTurnContext(), Hang: true}) + cfg := h.config() + cfg.Refusals = recorder + s, err := h.drv.NewSession(context.Background(), cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + session := s.(*session) + + for i := range maxRecorded + 10 { + session.refused(nil, fmt.Sprintf("item:call-%d", i), fmt.Sprintf("call-%d", i), "exec", driver.ToolExecute) + } + session.mu.Lock() + remembered := len(session.recorded) + session.mu.Unlock() + assert.Equal(t, maxRecorded, remembered, "a session remembers so many and no more, and remembered %d", remembered) + waitDone(t, s) +} + +// Two tool call ids that begin alike are two refusals. What a session +// remembers a refusal by is bounded, because the agent writes it — but +// bounding it by cutting would make a prefix the identity, and every later +// refusal sharing that start would be taken for a repeat and recorded +// nowhere. +func TestTwoIdsThatBeginAlikeAreTwoRefusals(t *testing.T) { + recorder := &drivertest.Refusals{} + h := newHarness(t, scenario{TurnContext: safeTurnContext(), Hang: true}) + cfg := h.config() + cfg.Refusals = recorder + s, err := h.drv.NewSession(context.Background(), cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + session := s.(*session) + tr := &turn{done: make(chan struct{})} + session.mu.Lock() + session.turn = tr + session.mu.Unlock() + + same := strings.Repeat("tool.call-", 4*maxToolCallID/10) + session.refused(tr, "item:"+same+"-one", same+"-one", "exec", driver.ToolExecute) + session.refused(tr, "item:"+same+"-two", same+"-two", "exec", driver.ToolExecute) + require.NoError(t, s.Close()) + + session.mu.Lock() + defer session.mu.Unlock() + assert.Len(t, tr.refusals, 2, "two ids alike for longer than a session keeps are still two refusals") + assert.Len(t, recorder.Recorded(), 2, "and the ledger has both") +} + +// The refusal that fills a session's memory is recorded like any other. The +// session ends on it — past it a repeat cannot be told from a first — but it +// ends after taking that refusal, not instead of it. +func TestTheRefusalThatFillsTheMemoryIsStillRecorded(t *testing.T) { + recorder := &drivertest.Refusals{} + h := newHarness(t, scenario{TurnContext: safeTurnContext(), Hang: true}) + cfg := h.config() + cfg.Refusals = recorder + s, err := h.drv.NewSession(context.Background(), cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + session := s.(*session) + + for i := range maxRecorded { + session.refused(nil, fmt.Sprintf("item:call-%d", i), fmt.Sprintf("call-%d", i), "exec", driver.ToolExecute) + } + require.NoError(t, s.Close()) + assert.Len(t, recorder.Recorded(), maxRecorded, + "every refusal up to and including the one that filled the memory, and the ledger has %d", len(recorder.Recorded())) + waitDone(t, s) +} diff --git a/internal/connector/driver/codex/scribe.go b/internal/connector/driver/codex/scribe.go new file mode 100644 index 000000000..d529f9a76 --- /dev/null +++ b/internal/connector/driver/codex/scribe.go @@ -0,0 +1,267 @@ +package codex + +import ( + "sync" + + "github.com/basecamp/basecamp-cli/internal/connector/driver" +) + +// A refusal is read on one goroutine and written on another +// +// Writing a refusal to the ledger is allowed ten seconds and used to run on +// the goroutine reading the worker's output. That made the reading +// unboundable: the read end of the pipe is released when the session stops +// waiting for it, and any bound on that — a clock, a count of bytes — is +// either spent by a ledger write, taking the pipe away with the worker's +// next refusal still in it, or not spent by a descendant writing without +// end, which then holds the session open for as long as it lives. +// +// So the reader only reads, and the scribe only writes. The reader's promise +// to driver.Worker.ReadingDone — that it does nothing slow between reads — +// is what lets the drain bound be a clock and mean what it says. +// +// # What the scribe owes the contract +// +// A refusal is recorded before the update for it is emitted (driver.go). The +// scribe emits it, after its write lands, so that ordering holds exactly +// rather than being narrowed. What is traded is the other half — that a +// refusal is never held only in a session's memory — for the window between +// the reader handing one over and the scribe writing it. A connector crash +// there loses it; a worker exiting, a turn cut short, and a session closing +// do not, because the queue is drained before the session's updates close, +// which is where the dispatcher settles what the ledger would not take. +// +// # When the queue is full +// +// It blocks the reader, and it never drops. Dropping would be the loss this +// exists to end, wearing a new hat. +// +// Blocking is only safe before the reader has been asked to stop, and there +// it is right: the pipe fills, the worker blocks writing, and an agent +// producing refusals faster than the ledger takes them is slowed to the +// ledger's pace. That is what a slow ledger already did to this driver when +// both ran on one goroutine, so it is not a new way to fail. +// +// After the asking it would be fatal — the drain's clock would run while the +// reader sat on a full queue, and the pipe would be abandoned with the +// worker's output still in it — so after the asking the reader never waits. +// It appends, and the queue grows past its mark. That growth is bounded by +// what one drain can produce, which the drain budget bounds in time, and it +// is a burst at the end of a session rather than a leak. +// +// The mark is not derived from the pipe's size or the scanner's buffer. +// Neither is a bound the code has: a worker inherits the pipe descriptor and +// may enlarge it up to the host's pipe-max-size, and the scanner is allowed +// to grow to 64 MiB. It is simply how many refusals are worth holding before +// a live worker is told to slow down. +// +// refusalCap is where the growing stops. A drain bounded only by its clock +// is not bounded in memory: two seconds of reading from a descendant that +// refuses as fast as it can write is as much as the machine will take. At +// the cap the reader stops taking more FROM THE PIPE, and goes on parsing +// what it has already taken — a scanner holds a bufferful that the pipe no +// longer does, and abandoning that would drop lines nobody had seen. +const ( + refusalMark = 256 + refusalCap = 4096 +) + +// pending is a refusal the reader has read and the scribe has not yet +// written, with the update that is emitted once it has. +type pending struct { + refusal driver.Refusal + update driver.Update +} + +// scribe writes refusals to the ledger, one at a time, off the reader. +type scribe struct { + record func(driver.Refusal) + emit func(driver.Update) + + mu sync.Mutex + room *sync.Cond + work *sync.Cond + queue []pending + nowait bool // the reader has been asked to stop: never make it wait + closing bool + finished bool // the scribe has drained and gone + heard bool // the queue reached its cap during a drain + // late counts the writes being made on a caller's own goroutine because + // the scribe had already gone. It lives under the same lock that + // publishes finished, so admitting one and deciding everything is + // written are a single decision rather than two that can cross. + late int + // writes counts the writes the scribe itself has taken off the queue + // and not yet finished, so a drain can tell an empty queue from a + // finished one. + writes int + quiet *sync.Cond + // writing is held across every call to record, wherever it is made + // from. The recorder is promised it is never called concurrently with + // itself (driver.go), and a late write runs on its reader's goroutine + // rather than the scribe's, so the promise needs a lock rather than a + // single goroutine to keep it. + writing sync.Mutex + done chan struct{} +} + +func newScribe(record func(driver.Refusal), emit func(driver.Update)) *scribe { + s := &scribe{record: record, emit: emit, done: make(chan struct{})} + s.room = sync.NewCond(&s.mu) + s.work = sync.NewCond(&s.mu) + s.quiet = sync.NewCond(&s.mu) + go s.run() + return s +} + +// hand gives the scribe a refusal to write. It waits for room while the +// reader can afford to, and never once it cannot. +func (s *scribe) hand(p pending) { + s.mu.Lock() + for len(s.queue) >= refusalMark && !s.nowait && !s.closing && !s.finished { + s.room.Wait() + } + if s.finished { + // The tail: a refusal read from the worker's last word by an ending + // that is not the reader's, after the scribe has drained and gone. + // It is written here rather than lost, and counted under the same + // lock that published finished — so a close either waits for this + // one or has not begun, never decides the ledger is whole while it + // is starting. + s.late++ + s.mu.Unlock() + s.writeNow(p) + s.mu.Lock() + s.late-- + s.quiet.Broadcast() + s.mu.Unlock() + return + } + s.queue = append(s.queue, p) + if s.nowait && len(s.queue) >= refusalCap { + // A drain that has heard as much as it can hold. Nothing read is + // lost; the reader is told to stop reading, which is what its own + // clock was about to do. + s.heard = true + } + s.work.Signal() + s.mu.Unlock() +} + +// enough reports that a drain filled the queue to its cap, so the reader +// should stop reading. Nothing that was read is dropped. +func (s *scribe) enough() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.heard +} + +// noWaiting says the reader has been asked to stop, so it must not be held +// up handing anything over: from here the queue grows rather than blocks. +func (s *scribe) noWaiting() { + s.mu.Lock() + s.nowait = true + s.room.Broadcast() + s.mu.Unlock() +} + +// drain waits for everything handed over so far to be written, without +// saying there is no more. A turn's result is settled after this, so a +// refusal read in that turn is in the ledger before the prompt that made it +// comes back — the ordering callers had when the write was on the reader. +// +// It is called from a turn's ending, which runs away from the reader, so +// waiting here holds up no reading. +func (s *scribe) drain() { + s.mu.Lock() + for (len(s.queue) > 0 || s.writes > 0) && !s.finished { + s.quiet.Wait() + } + s.mu.Unlock() + // A late write is on its caller's own goroutine and counted the same + // way; this waits for those too. + s.mu.Lock() + for s.late > 0 { + s.quiet.Wait() + } + s.mu.Unlock() +} + +// close says there is no more, and waits for what there is to be written. +// Whoever calls it promises the reader is through. It is idempotent. +func (s *scribe) close() { + s.mu.Lock() + if !s.closing { + s.closing = true + s.room.Broadcast() + s.work.Broadcast() + } + s.mu.Unlock() + <-s.done + // # What this barrier promises + // + // When it returns, every refusal handed over so far has been written: + // the ones the scribe took, and the ones a caller wrote itself because + // the scribe had already gone. Neither is dropped and neither races + // this, because admitting a late write and counting it down happen + // under the lock this waits on. + // + // It promises nothing about a refusal handed over AFTER it returns. One + // can be: an ending that is not the reader's may read the worker's last + // word later still. That refusal is written too, by whoever read it, on + // that goroutine — it is simply not waited for here, in the same window + // where an update emitted then is dropped rather than sent. + s.mu.Lock() + for s.late > 0 { + s.quiet.Wait() + } + s.mu.Unlock() +} + +func (s *scribe) run() { + defer close(s.done) + for { + s.mu.Lock() + for len(s.queue) == 0 && !s.closing { + s.work.Wait() + } + if len(s.queue) == 0 { + // Published under the lock a hand takes, so no refusal can be + // queued to a scribe that has decided it is done. + s.finished = true + s.room.Broadcast() + s.quiet.Broadcast() + s.mu.Unlock() + return + } + next := s.queue[0] + s.queue = s.queue[1:] + s.writes++ + s.room.Signal() + s.mu.Unlock() + + // Outside the queue's lock: the write is the slow thing, and the + // reader goes on reading while it happens. + s.writing.Lock() + s.record(next.refusal) + s.writing.Unlock() + // The update comes after the write, which is the contract's own + // order: a refusal is recorded before the update for it is emitted. + s.emit(next.update) + s.mu.Lock() + s.writes-- + s.quiet.Broadcast() + s.mu.Unlock() + } +} + +// writeNow writes a refusal on the caller's own goroutine, for one read +// after the scribe has gone. Two of these can be in flight at once — two +// endings that are not the reader's — so they take the same lock the scribe +// writes under. +func (s *scribe) writeNow(p pending) { + s.writing.Lock() + s.record(p.refusal) + s.writing.Unlock() + s.emit(p.update) +} diff --git a/internal/connector/driver/codex/scribe_test.go b/internal/connector/driver/codex/scribe_test.go new file mode 100644 index 000000000..1882b0629 --- /dev/null +++ b/internal/connector/driver/codex/scribe_test.go @@ -0,0 +1,344 @@ +//go:build unix + +package codex + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/basecamp/basecamp-cli/internal/connector/driver" +) + +// A refusal's update is emitted after its write has landed, never before. +// That is the driver's contract — the recorder is called before the update +// is emitted — and it is the half of it the scribe keeps exactly rather than +// narrows. Emitting at hand-off would publish a refusal the ledger has not +// taken. +func TestTheUpdateComesAfterTheWrite(t *testing.T) { + var mu sync.Mutex + var order []string + writing := make(chan struct{}) + release := make(chan struct{}) + s := newScribe( + func(driver.Refusal) { + mu.Lock() + order = append(order, "write") + mu.Unlock() + close(writing) + <-release + }, + func(driver.Update) { + mu.Lock() + order = append(order, "update") + mu.Unlock() + }, + ) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "one"}}) + + // The write has begun and the update must not have been emitted. + <-writing + mu.Lock() + assert.Equal(t, []string{"write"}, order, "the update waits for the write") + mu.Unlock() + + close(release) + s.close() + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{"write", "update"}, order) +} + +// Handing a refusal over never loses it, whatever the scribe is doing — +// including after it has drained and gone, when whoever read it writes it. +func TestARefusalHandedOverAfterTheScribeHasGoneIsStillWritten(t *testing.T) { + var mu sync.Mutex + var written []string + s := newScribe( + func(r driver.Refusal) { + mu.Lock() + written = append(written, r.ToolCallID) + mu.Unlock() + }, + func(driver.Update) {}, + ) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "before"}}) + s.close() + s.hand(pending{refusal: driver.Refusal{ToolCallID: "after"}}) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{"before", "after"}, written, "nothing is dropped on either side of the close") +} + +// The queue holds a live worker back rather than dropping anything: past the +// mark, handing over waits for room. Dropping would be the loss this exists +// to end. +func TestAFullQueueHoldsTheReaderBack(t *testing.T) { + release := make(chan struct{}) + first := make(chan struct{}) + var once sync.Once + s := newScribe( + func(driver.Refusal) { + once.Do(func() { close(first) }) + <-release + }, + func(driver.Update) {}, + ) + // One is taken and stuck in its write; the rest fill the queue to the + // mark, and the next has to wait. + s.hand(pending{refusal: driver.Refusal{ToolCallID: "stuck"}}) + <-first + for i := range refusalMark { + s.hand(pending{refusal: driver.Refusal{ToolCallID: string(rune('a' + i%26))}}) + } + + waited := make(chan struct{}) + go func() { + defer close(waited) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "over"}}) + }() + select { + case <-waited: + t.Fatal("the queue took more than its mark without holding the reader back") + case <-time.After(200 * time.Millisecond): + } + + close(release) + <-waited + s.close() +} + +// Once the reader has been asked to stop it is never held back, because the +// clock on its drain is running: a reader waiting for room would spend that +// clock and abandon the pipe with the worker's output still in it. +func TestAReaderAskedToStopIsNeverHeldBack(t *testing.T) { + release := make(chan struct{}) + first := make(chan struct{}) + var once sync.Once + s := newScribe( + func(driver.Refusal) { + once.Do(func() { close(first) }) + <-release + }, + func(driver.Update) {}, + ) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "stuck"}}) + <-first + s.noWaiting() + + handed := make(chan struct{}) + go func() { + defer close(handed) + for range refusalMark * 3 { + s.hand(pending{refusal: driver.Refusal{ToolCallID: "more"}}) + } + }() + select { + case <-handed: + case <-time.After(10 * time.Second): + t.Fatal("a reader that has been asked to stop was held back handing refusals over") + } + + close(release) + s.close() +} + +// Everything handed over is written by the time the scribe is closed, which +// is what the session leans on: it closes the scribe before the updates, +// because that is where the dispatcher settles what the ledger would not +// take. +func TestClosingTheScribeWritesEverythingHandedOver(t *testing.T) { + var mu sync.Mutex + count := 0 + s := newScribe( + func(driver.Refusal) { + mu.Lock() + count++ + mu.Unlock() + }, + func(driver.Update) {}, + ) + const handed = 500 + for range handed { + s.hand(pending{refusal: driver.Refusal{ToolCallID: "one"}}) + } + s.close() + + mu.Lock() + defer mu.Unlock() + require.Equal(t, handed, count, "the close is what makes the ledger whole") +} + +// The close waits for a late write that has already begun. A refusal read +// by an ending that is not the reader's can arrive after the scribe has +// drained and gone; whoever read it writes it, and the close either waits +// for that write or has not started — never decides the ledger is whole +// while one is running. +func TestClosingWaitsForALateWriteAlreadyBegun(t *testing.T) { + var mu sync.Mutex + var done []string + writing := make(chan struct{}) + release := make(chan struct{}) + s := newScribe( + func(r driver.Refusal) { + if r.ToolCallID == "late" { + close(writing) + <-release + } + mu.Lock() + done = append(done, r.ToolCallID) + mu.Unlock() + }, + func(driver.Update) {}, + ) + s.close() // the scribe has drained and gone + + handed := make(chan struct{}) + go func() { + defer close(handed) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "late"}}) + }() + <-writing // the late write has begun + + closed := make(chan struct{}) + go func() { + defer close(closed) + s.close() + }() + select { + case <-closed: + t.Fatal("the close decided the ledger was whole while a late write was running") + case <-time.After(200 * time.Millisecond): + } + + close(release) + <-closed + <-handed + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{"late"}, done, "and the write it waited for did happen") +} + +// A drain that hears more refusals than it can hold tells the reader to stop +// reading, rather than growing without limit. Nothing that was read is lost; +// what was not read is abandoned, which is what the drain's clock was about +// to do anyway. +func TestADrainThatFillsTheQueueTellsTheReaderToStop(t *testing.T) { + release := make(chan struct{}) + first := make(chan struct{}) + var once sync.Once + s := newScribe( + func(driver.Refusal) { + once.Do(func() { close(first) }) + <-release + }, + func(driver.Update) {}, + ) + s.hand(pending{refusal: driver.Refusal{ToolCallID: "stuck"}}) + <-first + s.noWaiting() + + require.False(t, s.enough(), "nothing has been heard yet") + for range refusalCap { + s.hand(pending{refusal: driver.Refusal{ToolCallID: "more"}}) + } + assert.True(t, s.enough(), "the reader is told to stop once the queue is full") + + close(release) + s.close() +} + +// The recorder is never called concurrently with itself, which driver.go +// promises an implementation it may rely on. A late write runs on its +// reader's goroutine rather than the scribe's, and two endings that are not +// the reader's can make one each, so the promise needs a lock rather than a +// single goroutine to keep it. +func TestTheRecorderIsNeverCalledTwiceAtOnce(t *testing.T) { + var mu sync.Mutex + inside, peak := 0, 0 + held := make(chan struct{}) + var once sync.Once + s := newScribe( + func(driver.Refusal) { + mu.Lock() + inside++ + peak = max(peak, inside) + mu.Unlock() + once.Do(func() { close(held) }) + time.Sleep(20 * time.Millisecond) + mu.Lock() + inside-- + mu.Unlock() + }, + func(driver.Update) {}, + ) + s.close() // every write from here is a late one, on its caller's goroutine + + var wg sync.WaitGroup + for i := range 8 { + wg.Add(1) + go func() { + defer wg.Done() + s.hand(pending{refusal: driver.Refusal{ToolCallID: string(rune('a' + i))}}) + }() + } + <-held + wg.Wait() + s.close() + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, 1, peak, "the recorder saw one call at a time, and %d at once", peak) +} + +// A turn's refusals are in the ledger before its result comes back. The +// writes are not on the reader any more, so an ending that settled a turn +// without waiting would hand a caller a result whose refusals the ledger did +// not yet hold — and the dispatcher settles an attempt on what it holds. +func TestDrainingWaitsForWhatIsQueuedWithoutClosing(t *testing.T) { + var mu sync.Mutex + written := 0 + release := make(chan struct{}) + first := make(chan struct{}) + var once sync.Once + s := newScribe( + func(driver.Refusal) { + once.Do(func() { close(first) }) + <-release + mu.Lock() + written++ + mu.Unlock() + }, + func(driver.Update) {}, + ) + t.Cleanup(s.close) + for range 4 { + s.hand(pending{refusal: driver.Refusal{ToolCallID: "one"}}) + } + <-first + + drained := make(chan struct{}) + go func() { defer close(drained); s.drain() }() + select { + case <-drained: + t.Fatal("the drain returned with refusals still queued") + case <-time.After(200 * time.Millisecond): + } + + close(release) + <-drained + mu.Lock() + assert.Equal(t, 4, written, "everything handed over is written, and %d was", written) + mu.Unlock() + + // And the scribe is still open afterwards: draining is not closing. + s.hand(pending{refusal: driver.Refusal{ToolCallID: "after"}}) + s.drain() + mu.Lock() + defer mu.Unlock() + assert.Equal(t, 5, written, "draining is not closing") +} diff --git a/internal/connector/driver/driver.go b/internal/connector/driver/driver.go index 437e52126..c75d2b787 100644 --- a/internal/connector/driver/driver.go +++ b/internal/connector/driver/driver.go @@ -65,9 +65,19 @@ // recorded in the ledger, once, at the moment the driver answers the request // — or, for an agent that answers its own requests under a mode the driver // froze (claude -p), at the moment the driver first reads that it was -// refused. It is never held only in a session's memory, because a worker that -// exits before its result, a connector that crashes mid-turn, and a turn cut -// short by a deadline all end the session that memory lives in. +// refused. +// +// A driver may hand the write to a goroutine of its own rather than make it +// where it read the refusal, and the codex driver does: a ledger write is +// allowed ten seconds, and a reader that spends them cannot be given a bound +// on its reading that means anything (driver/codex/scribe.go). What that +// costs is named here rather than left to be discovered. Between the reading +// and the write the refusal is in memory, and a connector that crashes there +// loses it. What it does NOT cost is the rest: the queue is drained before +// the session's updates close, so a worker that exits before its result, a +// turn cut short by a deadline, and a session closed under a turn all still +// find it written — and the update for it is emitted by the writer, after +// the write, so the order below holds exactly. // // 1. The driver calls SessionConfig.Refusals.RecordRefusal before it sends // its answer to the agent, or before it emits the update for a refusal @@ -91,7 +101,10 @@ // recorder could not write, and the ended attempt's count is final. The // session's updates are drained before the attempt is released, and the // recorder is called before an update is emitted, so a worker that exits -// between a refusal and its result has already recorded it. +// between a refusal and its result has already recorded it. That holds +// for a driver that writes on a goroutine of its own too: it is the +// writer that emits, once the write has landed, and its queue is drained +// before the updates close. // // Once-ness is the driver's (a set of tool call ids per session), not a key in // the ledger: it holds for as long as a session lives, which is as long as a @@ -316,8 +329,14 @@ type Refusal struct { } // RefusalRecorder records a refusal at the moment a driver makes or observes -// it (see "Refusals" above). RecordRefusal must not block for long: a driver -// calls it on the goroutine that reads the agent's stream. +// it (see "Refusals" above). +// +// It is called once per refusal, never concurrently with itself, and an +// implementation needs no locking of its own for that. Which goroutine calls +// it is the driver's: most call it on the one reading the agent's stream, so +// it must not block for long — a driver that cannot afford that hands the +// write to a goroutine of its own, as codex does, and then emits the update +// for the refusal only once this has returned. type RefusalRecorder interface { RecordRefusal(ctx context.Context, r Refusal) error } diff --git a/internal/connector/driver/driver_test.go b/internal/connector/driver/driver_test.go index 71740bede..b422a30bc 100644 --- a/internal/connector/driver/driver_test.go +++ b/internal/connector/driver/driver_test.go @@ -5,6 +5,7 @@ package driver import ( "context" "errors" + "io" "os" "os/exec" "path/filepath" @@ -339,3 +340,138 @@ func TestALaterGroupSignalIsNotSentToAGroupTheRecordNoLongerOwns(t *testing.T) { require.NoError(t, signalRecordedGroup(p, syscall.SIGKILL)) assert.Eventually(t, func() bool { return !GroupMembersRemain(p) }, 5*time.Second, 20*time.Millisecond) } + +// A stop does not discard what is in the pipe. Nothing closes the read end +// under a reader, because closing throws away whatever the worker wrote that +// nobody has parsed. Nothing here is concurrent and nothing waits: the lines +// are in the pipe before the stop, so what is asked is only whether the stop +// destroys them. +func TestAStopDoesNotDiscardWhatIsInThePipe(t *testing.T) { + readEnd, writeEnd, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { _ = readEnd.Close(); _ = writeEnd.Close() }) + out := &output{f: readEnd} + + _, err = writeEnd.Write([]byte("one\ntwo\n")) + require.NoError(t, err) + // Asked twice, and by two callers: a close and a cancel can both ask. + out.stop() + out.stop() + + buf := make([]byte, 64) + n, err := out.Read(buf) + require.NoError(t, err) + assert.Equal(t, "one\ntwo\n", string(buf[:n]), "what was in the pipe when the stop came") + + n, err = out.Read(buf) + assert.Zero(t, n) + assert.ErrorIs(t, err, io.EOF, "and then the reader ends, rather than waiting on a pipe nobody will close") +} + +// The drain budget is absolute from the asking, not a window that begins +// again after every read. A descendant that writes a single byte just before +// each window expires makes progress forever without ever letting the pipe +// fall idle, and a per-read window would follow it for as long as it lived — +// hours, for as much output as any byte count would allow. +func TestADescendantDrippingOutputCannotHoldTheReader(t *testing.T) { + readEnd, writeEnd, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { _ = readEnd.Close(); _ = writeEnd.Close() }) + out := &output{f: readEnd} + + dripping := make(chan struct{}) + go func() { + defer close(dripping) + for { + // Just often enough that no window of the reader's own ever + // finds the pipe empty. + time.Sleep(drainWindow / 2) + if _, werr := writeEnd.Write([]byte("x")); werr != nil { + return + } + } + }() + t.Cleanup(func() { _ = readEnd.Close(); <-dripping }) + + out.stop() + started := time.Now() + ended := make(chan error, 1) + go func() { + buf := make([]byte, 4<<10) + for { + if _, rerr := out.Read(buf); rerr != nil { + ended <- rerr + return + } + } + }() + select { + case rerr := <-ended: + assert.ErrorIs(t, rerr, io.EOF, "the reading ends, rather than following a descendant forever") + assert.Less(t, time.Since(started), 4*drainBudget, "and it ends on the budget, not on the descendant") + case <-time.After(30 * time.Second): + t.Fatal("a reader asked to stop never ended while a descendant dripped output into the pipe") + } +} + +// The deadline on the pipe has one owner. Asking a reader to stop sets a +// flag and touches nothing else, so it cannot cut short a window a read has +// opened. That is a property of who writes what, not of a schedule, so it is +// checked as one: a deadline set here survives any number of stops. +func TestAStopDoesNotTouchTheReadDeadline(t *testing.T) { + readEnd, writeEnd, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { _ = readEnd.Close(); _ = writeEnd.Close() }) + out := &output{f: readEnd} + + const mine = 300 * time.Millisecond + require.NoError(t, readEnd.SetReadDeadline(time.Now().Add(mine))) + out.stop() + out.stop() + + started := time.Now() + buf := make([]byte, 64) + n, err := readEnd.Read(buf) // the file itself, not through output + waited := time.Since(started) + + assert.Zero(t, n) + require.ErrorIs(t, err, os.ErrDeadlineExceeded) + assert.GreaterOrEqual(t, waited, mine/2, "the deadline in force is the one set here, not one a stop installed") +} + +// The drain budget runs from the asking, not from when the reader next +// notices it. A reader with nothing arriving sits in an idle window, and a +// stop landing inside one is not seen until that window is out; a budget +// begun at the noticing would outlast what it promises by that much. +// +// This is a question about when the clock is taken, not about how long +// anything runs, so it is asked that way: after the asking, the time to +// measure from exists. Timing it would be a knife-edge — the reader notices +// as soon as anything arrives, so the gap only opens when output begins just +// as a window expires — and a test on a knife-edge proves nothing either +// way. +func TestTheDrainBudgetRunsFromTheAsking(t *testing.T) { + readEnd, writeEnd, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { _ = readEnd.Close(); _ = writeEnd.Close() }) + out := &output{f: readEnd} + + require.True(t, out.askedAt().IsZero(), "nothing has been asked yet") + before := time.Now() + out.stop() + after := time.Now() + + asked := out.askedAt() + require.False(t, asked.IsZero(), "the budget's clock is taken when the stop is asked for, not when a read next looks") + assert.False(t, asked.Before(before), "and it is taken then") + assert.False(t, asked.After(after)) + + // It carries a monotonic reading, so the budget is not at the mercy of + // the wall clock moving under it. + assert.NotEqual(t, asked, asked.Round(0), "the time kept is monotonic, not a wall clock rebuilt from a number") + + // And it is the first asking that counts, so a second does not hand the + // reader a fresh budget. + out.stop() + assert.Equal(t, asked, out.askedAt(), "a later stop does not restart the budget") +} diff --git a/internal/connector/driver/worker.go b/internal/connector/driver/worker.go index 164f92501..43cd0d6ee 100644 --- a/internal/connector/driver/worker.go +++ b/internal/connector/driver/worker.go @@ -11,6 +11,7 @@ import ( "os/exec" "strings" "sync" + "sync/atomic" "syscall" "time" ) @@ -19,6 +20,146 @@ import ( // pipes a stray descendant still holds. const pipeWaitDelay = 2 * time.Second +// drainWindow is how long a reader that has been asked to stop gives the +// pipe to produce what is already on its way, idleWindow is how often a +// reader with nothing looks up to see whether it has been asked, and +// drainBudget is the longest it goes on reading after being asked. The +// windows are opened by the reader for itself: a read that finds nothing in +// one is a pipe with nothing in it. +// +// The two bound different things, and both are needed. +// +// drainWindow bounds WAITING: how long a reader sits with nothing arriving +// before it calls the pipe empty. That is how a stop ordinarily ends, within +// tens of milliseconds, and it needs no clock beyond itself. +// +// drainBudget bounds FOLLOWING: how long a reader goes on taking output that +// keeps arriving. Only one thing produces that after a worker is dead — a +// descendant outside its group writing without stop — and it is the one case +// with no other ending, so a window that begins again after every read is no +// bound at all: a descendant writing a byte before each one expires would +// hold a session open for hours. This one is absolute from the asking. +// +// Absolute is only honest because a reader registered here does nothing slow +// between reads, so the clock cannot run while the reader is working rather +// than waiting. That is a promise, and it is the whole reason this package's +// drivers write to the ledger on a goroutine of their own: a reader that +// wrote to a database here would be inside a write when the clock ran out, +// and would come back to find its own pipe closed with the worker's next +// line still in it. See Worker.ReadingDone. +// +// A count of bytes cannot stand in for either. It is too permissive on time, +// because a descendant can drip them out forever, and too strict on volume, +// because a worker inherits the pipe descriptor and may enlarge it beyond +// any size assumed here. +const ( + drainWindow = 50 * time.Millisecond + idleWindow = 500 * time.Millisecond + drainBudget = 2 * time.Second +) + +// output is the read end of a worker's pipe, and it belongs to whoever reads +// it for as long as they are reading. Nothing closes it underneath them: +// closing discards whatever the worker wrote that nobody has parsed yet, and +// there is no size that could be assumed safe — a worker inherits the +// descriptor and may enlarge the pipe itself, up to the host's +// pipe-max-size. +// +// So a reader is asked to stop rather than cut off, and asking touches +// nothing but a flag. The deadline on the pipe has one owner, the reader, +// which is what makes a window mean what it says: nobody else can shorten +// one, and there is no moment in which a deadline is in force whose reason +// is not yet visible. +// +// A worker is ended before its reader is asked, so everything the worker +// wrote is in the pipe by then. What is given up, past the budget, is only +// what a descendant outside the worker's group goes on writing — which the +// connector had already decided not to wait for, and which is the one case +// that has no other ending. +type output struct { + f *os.File + // stopped is whether the reader has been asked to stop. One field, + // written by one call, so there is no moment in which half the state is + // visible. + stopped atomic.Bool + // stopAt is when the stop was asked for, written by stop before the flag + // that publishes it. The reader takes it from here rather than from when + // it next looks, so the budget runs from the asking as it says it does — + // a reader inside an idle window when the asking comes would otherwise + // start the clock up to that window late. + // + // It is kept as a time.Time under a mutex rather than as nanoseconds in + // an atomic, because a time.Time carries a monotonic reading and a + // number does not. Rebuilt from nanoseconds, the wall clock decides the + // budget: a forward jump ends a drain early with the worker's output + // still buffered, and a backward one holds the session open past it. + stopMu sync.Mutex + stopAt time.Time +} + +func (o *output) Read(p []byte) (int, error) { + for { + stopped := o.stopped.Load() + window := idleWindow + if stopped { + left := drainBudget - time.Since(o.askedAt()) + if left <= 0 { + // Still arriving, and no longer waited for. + return 0, io.EOF + } + // Never past the budget: the last window is whatever is left of + // it, not a whole one begun at the end. + window = min(drainWindow, left) + } + if err := o.f.SetReadDeadline(time.Now().Add(window)); err != nil { + return 0, err + } + if !stopped && o.stopped.Load() { + // Asked between the sample and the deadline. The idle window + // just installed would hide that until it expired, and the + // budget would be most of the way gone by then, so the window + // is opened again as a drain's. + continue + } + n, err := o.f.Read(p) + switch { + case n > 0: + return n, nil + case !errors.Is(err, os.ErrDeadlineExceeded): + return n, err + case stopped: + // A window this read opened, and nothing came. + return 0, io.EOF + } + // Nothing yet, and nobody has asked. Look again. + } +} + +// stop asks the reader to end once it has what is there, and no later than +// the drain budget. It writes one field and nothing else — in particular it +// does not reach for the deadline, so it cannot cut short a window a read +// has opened — and it says the same thing however many times it is called. +func (o *output) stop() { + // The time before the flag, so a reader that sees the flag always finds + // a time to measure from, and the budget runs from the asking. The first + // asking is the one that counts. + o.stopMu.Lock() + if o.stopAt.IsZero() { + o.stopAt = time.Now() + } + o.stopMu.Unlock() + o.stopped.Store(true) +} + +// askedAt is when the stop was asked for, monotonic. +func (o *output) askedAt() time.Time { + o.stopMu.Lock() + defer o.stopMu.Unlock() + return o.stopAt +} + +func (o *output) close() { _ = o.f.Close() } + // # One owner, one release point // // This is the connector's rule for a task's process tree and its ledger @@ -178,13 +319,16 @@ type Worker struct { cmd *exec.Cmd process Process stdin io.WriteCloser - stdout *os.File + stdout *output stderr *tailBuffer done chan struct{} exit Exit killOnce sync.Once releaseOnce sync.Once + + readingMu sync.Mutex + reading <-chan struct{} } // StartWorker launches cmd through cfg's launcher, in cfg's scope, as a new @@ -244,7 +388,7 @@ func StartWorker(ctx context.Context, cfg SessionConfig, cmd Command) (*Worker, return nil, fmt.Errorf("%w: %w", ErrNotStarted, err) } ec.Stdout = writeEnd - w.stdout = readEnd + w.stdout = &output{f: readEnd} if err := ec.Start(); err != nil { // exec.Cmd.Start returns an error only when no process was created: // a missing binary, a bad directory, a failed fork. @@ -308,11 +452,46 @@ func (w *Worker) Stdin() io.WriteCloser { return w.stdin } // Stdout is the worker's standard output. Read it to end of file. func (w *Worker) Stdout() io.Reader { return w.stdout } -// CloseStdout closes the worker's output: a reader blocked on it returns, and -// the descriptor is released. -// For a worker that is gone while a descendant that left its group still -// holds the pipe. -func (w *Worker) CloseStdout() { _ = w.stdout.Close() } +// ReadingDone hands the Worker the signal that its output has been read to +// the end, so the read end of the pipe is released with the reading rather +// than on a clock. +// +// The descriptor is the Worker's to release; whether the reading is over is +// the reader's to say. Without this the Worker has to guess, and a guess is +// destructive: closing the read end discards whatever the worker wrote that +// nobody has parsed yet. +// +// A Worker with no reader registered keeps the clock, because nothing else +// would release it. +// +// Registering is two promises about the goroutine that reads: that it closes +// the channel, and that it does nothing slow between reads. The second is +// what lets the drain budget be a clock — see drainBudget — and it is why a +// driver's ledger writes belong on a goroutine of their own. +func (w *Worker) ReadingDone(done <-chan struct{}) { + w.readingMu.Lock() + defer w.readingMu.Unlock() + w.reading = done +} + +func (w *Worker) readingSignal() <-chan struct{} { + w.readingMu.Lock() + defer w.readingMu.Unlock() + return w.reading +} + +// StopReading asks the worker's reader to end once it has read what is in +// the pipe, and no later than the drain budget. For a worker that is gone +// while a descendant that left its group still holds the output open, so the +// end of file never comes: the reader is asked rather than having the pipe +// taken from under it, which would discard what the worker wrote and nobody +// has parsed. It says the same thing however many times it is called. +func (w *Worker) StopReading() { w.stdout.stop() } + +// CloseStdout closes the worker's output at once, discarding whatever is in +// the pipe: for releasing the descriptor once nobody is reading it. A reader +// that is still reading is asked to stop instead. +func (w *Worker) CloseStdout() { w.stdout.close() } // Done is closed once the process has exited and been reaped. func (w *Worker) Done() <-chan struct{} { return w.done } @@ -360,12 +539,28 @@ func (w *Worker) Terminate(grace time.Duration) { _ = w.cmd.Process.Kill() }) <-w.done - // The output pipe is the Worker's to release as well. Its reader gets the - // same bound Wait gives a stray descendant to finish draining what the - // worker wrote before it went, and then the descriptor is closed whether - // or not the reader closed it. + // The output pipe is the Worker's to release as well, once its reader is + // through it: a reader that says when it is done is waited for, because + // closing the descriptor under it would discard what the worker wrote + // and nobody has parsed. A worker nobody reads this way gets the same + // bound Wait gives a stray descendant, and then the descriptor is closed + // regardless — nothing else would release it. w.releaseOnce.Do(func() { - time.AfterFunc(pipeWaitDelay, w.CloseStdout) + reading := w.readingSignal() + if reading == nil { + time.AfterFunc(pipeWaitDelay, w.CloseStdout) + return + } + // Asked before it is waited for: a descendant outside the group can + // hold the write end open, so the end of file never comes and a + // reader left to itself never ends — which would hold this + // worker's descriptor, and everything waiting on the reading, for + // as long as that descendant lived. + w.stdout.stop() + go func() { + <-reading + w.CloseStdout() + }() }) } diff --git a/internal/connector/driver/worker_other.go b/internal/connector/driver/worker_other.go index df6f39846..6eaf5d34c 100644 --- a/internal/connector/driver/worker_other.go +++ b/internal/connector/driver/worker_other.go @@ -23,6 +23,8 @@ func (*Worker) Process() Process { return Process{} } func (*Worker) Stdin() io.WriteCloser { return nil } func (*Worker) Stdout() io.Reader { return nil } func (*Worker) CloseStdout() {} +func (*Worker) StopReading() {} +func (*Worker) ReadingDone(<-chan struct{}) {} func (*Worker) Done() <-chan struct{} { return nil } func (*Worker) Exit() Exit { return Exit{} } func (*Worker) StderrTail(*Redactor) string { return "" }