Skip to content

Commit 891890e

Browse files
authored
Merge pull request #6 from StructuPath/fix/trace-buffer-bound-on-failed-batch
fix(worker): keep the trace buffer within its cap across a failed send
2 parents 3ea75d3 + b463e7a commit 891890e

2 files changed

Lines changed: 123 additions & 6 deletions

File tree

‎internal/worker/events.go‎

Lines changed: 59 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,38 @@ func (s *TraceStream) buffer(event protocol.Event) {
394394
s.pending = append(s.pending, event)
395395
}
396396

397+
// enforceBound drops the oldest undelivered events until the buffer is back
398+
// within its cap. It exists for the one path that can exceed the cap in a
399+
// single step — a failed batch restored at the head — where dropping once is
400+
// not necessarily enough.
401+
//
402+
// The loop stops if a drop frees nothing rather than spinning: dropOldest
403+
// declines to act on a buffer too short to leave a gap marker behind, which
404+
// is reachable only with a cap below two. Caller holds the mutex.
405+
func (s *TraceStream) enforceBound() {
406+
for len(s.pending) > s.maxBuffered {
407+
before := len(s.pending)
408+
s.dropOldest()
409+
if len(s.pending) >= before {
410+
return
411+
}
412+
}
413+
}
414+
415+
// gapCoveredTo reads the upper bound of the range a gap marker declares. The
416+
// marker's own Seq is its lower bound, so an unreadable payload degrades to
417+
// "covers itself only" — narrower than the truth, never wider, which keeps a
418+
// merge from silently shrinking a declared range.
419+
func gapCoveredTo(event protocol.Event) int64 {
420+
var declared struct {
421+
ToSeq int64 `json:"to_seq"`
422+
}
423+
if err := json.Unmarshal(event.Payload, &declared); err != nil {
424+
return event.Seq
425+
}
426+
return declared.ToSeq
427+
}
428+
397429
// isGapMarker reports whether an event is one this stream synthesized for a
398430
// dropped run, so consecutive overflows merge into one marker instead of
399431
// producing a flood of them.
@@ -425,11 +457,25 @@ func (s *TraceStream) dropOldest() {
425457
return
426458
}
427459
discarded := s.pending[:drop]
428-
first, last := discarded[0], discarded[drop-1]
460+
first := discarded[0]
429461
lost := drop
430-
if merging {
431-
// The head was our own marker, not a lost event.
432-
lost--
462+
// The range the replacement marker must declare is the widest any
463+
// discarded element covered — not simply the last one's seq. A marker
464+
// swallowed by this drop already stood for events that are gone, and
465+
// letting the new marker end short of that range would leave them
466+
// undeclared: silent loss, which is the one outcome this whole mechanism
467+
// exists to prevent. Markers can sit anywhere in the run, not only at the
468+
// head, because a failed batch is restored in front of whatever was
469+
// buffered while it was in flight.
470+
covered := discarded[drop-1].Seq
471+
for _, event := range discarded {
472+
if !isGapMarker(event) {
473+
continue
474+
}
475+
lost-- // our own marker, not a lost event
476+
if to := gapCoveredTo(event); to > covered {
477+
covered = to
478+
}
433479
}
434480
s.dropped += int64(lost)
435481

@@ -444,7 +490,7 @@ func (s *TraceStream) dropOldest() {
444490
"reason": "the worker's event buffer overflowed while ingestion was unavailable",
445491
"dropped": s.dropped,
446492
"from_seq": first.Seq,
447-
"to_seq": last.Seq,
493+
"to_seq": covered,
448494
"buffer_cap": s.maxBuffered,
449495
"note": "the complete trace is in the attempt-local JSONL record",
450496
})
@@ -461,7 +507,7 @@ func (s *TraceStream) dropOldest() {
461507
s.logger.Warn("trace_gap_record_failed", "attempt_id", s.attemptID, "error", err)
462508
}
463509
s.logger.Warn("trace_buffer_overflow", "attempt_id", s.attemptID,
464-
"dropped", lost, "from_seq", first.Seq, "to_seq", last.Seq)
510+
"dropped", lost, "from_seq", first.Seq, "to_seq", covered)
465511
}
466512

467513
// nudge wakes the sender without blocking. Caller holds the mutex.
@@ -557,7 +603,14 @@ func (s *TraceStream) sendOnce(ctx context.Context) (int, error) {
557603
// Back at the head, where they were: the batch is the oldest run, so
558604
// restoring it in front of whatever arrived meanwhile keeps the buffer
559605
// in seq order.
606+
//
607+
// Then re-apply the bound. Emit kept buffering while this batch was
608+
// out — the bound it enforced counted only what was still pending —
609+
// so restoring the batch can push the buffer past its cap. Without
610+
// this the cap is really maxBuffered plus one batch, and "explicitly
611+
// bounded" would be a promise the code does not keep.
560612
s.pending = append(batch, s.pending...)
613+
s.enforceBound()
561614
return 0, err
562615
}
563616
return count, nil

‎internal/worker/events_test.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,3 +427,67 @@ func TestRedactorMatchesJSONEscapedValues(t *testing.T) {
427427
t.Fatalf("redaction broke the payload JSON: %v (%s)", err, got.Payload)
428428
}
429429
}
430+
431+
// gatedIngester holds its first send in flight until released, which is what
432+
// makes the interleaving below deterministic instead of a 1-in-15 flake.
433+
type gatedIngester struct {
434+
entered chan struct{}
435+
release chan struct{}
436+
once sync.Once
437+
}
438+
439+
func (g *gatedIngester) IngestEvents(
440+
_ context.Context, _ string, _ protocol.EventBatch,
441+
) (protocol.EventBatchResult, error) {
442+
g.once.Do(func() { close(g.entered) })
443+
<-g.release
444+
return protocol.EventBatchResult{}, errors.New("control plane is unreachable")
445+
}
446+
447+
// The bound holds across a failed send, not only across Emit. A batch in
448+
// flight has already left the buffer, so Emit's check counts only what
449+
// remains and is satisfied — then restoring that batch at the head carries
450+
// the buffer past its cap. The package comment calls the buffer explicitly
451+
// bounded, and this is the path that made that untrue.
452+
func TestARestoredFailedBatchStaysWithinTheBufferBound(t *testing.T) {
453+
gate := &gatedIngester{entered: make(chan struct{}), release: make(chan struct{})}
454+
stream, _ := newTraceFixture(t, gate, 6)
455+
// Drive the sender by hand: the background loop would race the
456+
// interleaving this test exists to pin down.
457+
stream.stop()
458+
<-stream.done
459+
460+
for seq := int64(1); seq <= 5; seq++ {
461+
if err := stream.Emit(event(seq, protocol.EventToolCall, "Write", nil)); err != nil {
462+
t.Fatalf("emit: %v", err)
463+
}
464+
}
465+
sent := make(chan struct{})
466+
go func() {
467+
defer close(sent)
468+
_, _ = stream.sendOnce(context.Background())
469+
}()
470+
<-gate.entered
471+
472+
// Five events are out. Emit now fills the buffer to its cap, satisfied by
473+
// a count that cannot see them.
474+
for seq := int64(6); seq <= 30; seq++ {
475+
if err := stream.Emit(event(seq, protocol.EventToolCall, "Write", nil)); err != nil {
476+
t.Fatalf("emit: %v", err)
477+
}
478+
}
479+
if pending, _ := stream.Undelivered(); pending > 6 {
480+
t.Fatalf("buffer was past its cap before the restore: %d", pending)
481+
}
482+
483+
close(gate.release)
484+
<-sent
485+
486+
if pending, _ := stream.Undelivered(); pending > 6 {
487+
t.Fatalf("the restored batch carried the buffer to %d, past its cap of 6", pending)
488+
}
489+
// The drop is accounted for, not silent: the JSONL keeps the full record.
490+
if _, dropped := stream.Undelivered(); dropped == 0 {
491+
t.Fatal("events left the buffer without being counted as dropped")
492+
}
493+
}

0 commit comments

Comments
 (0)