From 2637ae91a7af4791869c103b1a75b639a93341b7 Mon Sep 17 00:00:00 2001 From: Joshua Temple Date: Sat, 22 Aug 2026 21:25:42 -0400 Subject: [PATCH 1/2] fix(source/kafka): honor drain contract after Close; make Nak redeliver in-session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0-1: Next/NextBatch re-checked closed&&inFlight==0 before each poll, so a closed subscription with records still in flight kept polling and yielded NEW records — violating the source.Subscription drain contract when driven outside the Hopper. Once closed they now never poll: already-buffered records are yielded, then the call blocks on a settle-wakeup channel until the last in-flight record settles, then returns ErrDrained (ctx cancellation honored). P0-2 (review option a): plain Nak previously only declined to mark, so the record was skipped for the live group until restart/rebalance while the core contract claimed universal at-least-once. Plain Nak now re-seeks through the same machinery as NakAfter with delay zero (pause partition, seek to the record's offset, resume), making redelivery deterministic within a live subscription. Cross-restart redelivery rides committed offsets and stays documented best-effort. Contract text aligned in one voice across source/doc.go, inlet.go, handler.go ActionNak, kafka package doc, and kafka README divergence table; requeueWithDelay cost model documented (head-of-line pause, buffered-ahead records, concurrent-mark override); poll-error record discard documented. Tests: close-then-Next/NextBatch regressions against the real subscription (fake poller), plain-Nak immediate re-seek, NakAfter delay, duplicate-settle safety, Hopper redelivery storm (bounded goroutines, >=501 deliveries). --- source/doc.go | 25 ++- source/handler.go | 9 +- source/hopper_test.go | 92 ++++++++++ source/inlet.go | 5 +- source/kafka/README.md | 26 ++- source/kafka/drain_test.go | 332 +++++++++++++++++++++++++++++++++++ source/kafka/kafka.go | 10 +- source/kafka/subscription.go | 121 ++++++++++--- 8 files changed, 569 insertions(+), 51 deletions(-) create mode 100644 source/kafka/drain_test.go diff --git a/source/doc.go b/source/doc.go index e16857c..83bb96b 100644 --- a/source/doc.go +++ b/source/doc.go @@ -23,15 +23,22 @@ // // # Contract // -// Delivery is at-least-once by default: a message is acked only after its -// handler reports success, never before processing. A handler returns a -// [Result] — [Ack], [Nak], [Term], [InProgress], or [Manual] — and the Hopper -// applies it to the backend. Backends differ (Kafka commits offsets per -// partition; JetStream acks per message), so capabilities a backend may or may -// not have — replay, consumer groups, transactions — are optional interfaces -// discovered by type assertion ([Seekable], [ConsumerGroups], [Transactional], -// …) rather than a lowest-common-denominator API that lies about what a backend -// can do. +// Delivery is at-least-once within a live subscription: a message is acked only +// after its handler reports success, never before processing, and a [Nak] +// redelivers — JetStream naks natively (with optional delay), Kafka pauses and +// re-seeks the record's partition so it is fetched again. Across process +// restarts and rebalances, backends whose redelivery rides persisted offsets +// (Kafka) resume from the last committed position, which a concurrently +// committed higher offset can advance past a nacked record; those backends +// therefore redeliver exactly within a session and best-effort across restarts, +// and each adapter documents the precise semantics in its own module. A handler +// returns a [Result] — [Ack], [Nak], [Term], [InProgress], or [Manual] — and +// the Hopper applies it to the backend. Backends differ (Kafka commits offsets +// per partition; JetStream acks per message), so capabilities a backend may or +// may not have — replay, consumer groups, transactions — are optional +// interfaces discovered by type assertion ([Seekable], [ConsumerGroups], +// [Transactional], …) rather than a lowest-common-denominator API that lies +// about what a backend can do. // // # No forced dependencies // diff --git a/source/handler.go b/source/handler.go index f66c4bd..a788765 100644 --- a/source/handler.go +++ b/source/handler.go @@ -43,8 +43,13 @@ const ( // means "ack". ActionAck Action = iota // ActionNak asks for redelivery: the message failed transiently and should - // be tried again (JetStream Nak/NakWithDelay; Kafka declines to commit so the - // record is re-read). Result.Requeue is an optional backoff delay. + // be tried again (JetStream naks natively with optional delay; Kafka pauses + // and re-seeks the record's partition so it is fetched again). Redelivery + // is guaranteed within a live subscription; across process restarts or + // rebalances, backends whose redelivery rides persisted offsets (Kafka) + // honor the last committed position, so a concurrently committed higher + // offset can pass a nacked record — each adapter documents its exact + // semantics. Result.Requeue is an optional backoff delay. ActionNak // ActionTerm rejects the message permanently: it must not be redelivered // (JetStream Term; Kafka routes it to a dead-letter topic, then commits). diff --git a/source/hopper_test.go b/source/hopper_test.go index 0a7fe8c..fcee6fe 100644 --- a/source/hopper_test.go +++ b/source/hopper_test.go @@ -6,6 +6,7 @@ import ( "context" "errors" "fmt" + "runtime" "sync" "sync/atomic" "testing" @@ -499,3 +500,94 @@ func TestHopper_MaxInFlightBackpressure(t *testing.T) { } h.AssertSettled(total) } + +// stormSub is a minimal Subscription whose Settle requeues nacked messages, so +// a handler that keeps Naking drives Hopper.run through a real redelivery loop +// without a broker. +type stormSub struct { + queue chan source.Message +} + +func newStormSub(m source.Message) *stormSub { + q := make(chan source.Message, 8) + q <- m + return &stormSub{queue: q} +} + +func (s *stormSub) Next(ctx context.Context) (source.Message, error) { + select { + case m := <-s.queue: + return m, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (s *stormSub) Settle(_ context.Context, m source.Message, r source.Result) error { + if r.Action == source.ActionNak { + s.queue <- m // redeliver + } + return nil +} + +func (s *stormSub) Close() error { return nil } + +// TestHopper_RedeliveryStormBoundedResources drives 500 consecutive Naks of a +// single-key message through a real Hopper.run redelivery loop and pins the +// resource contract: Run terminates, the delivery loop really cycled (>= 501 +// deliveries for 500 naks + 1 ack), and every lane/fetch goroutine unwinds so +// the goroutine count returns to its pre-run baseline. +func TestHopper_RedeliveryStormBoundedResources(t *testing.T) { + t.Parallel() + + sub := newStormSub(testMsg{key: []byte("k"), value: []byte("v")}) + hp := source.New(source.WithConcurrency(4)) + t.Cleanup(func() { _ = hp.Close() }) + + var deliveries atomic.Int64 + acked := make(chan struct{}) + var ackOnce sync.Once + handler := func(context.Context, source.Message) source.Result { + if deliveries.Add(1) <= 500 { + return source.Nak(errors.New("transient")) + } + ackOnce.Do(func() { close(acked) }) + return source.Ack() + } + + baseline := runtime.NumGoroutine() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- hp.Run(ctx, sub, handler) }() + + select { + case <-acked: + case <-time.After(5 * time.Second): + t.Fatal("handler never reached its first ack after the redelivery storm") + } + cancel() + select { + case err := <-done: + if err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("Run = %v, want nil or context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run did not return after cancel") + } + + if got := deliveries.Load(); got < 501 { + t.Fatalf("deliveries = %d, want >= 501 (500 naks + 1 ack through the redelivery loop)", got) + } + + // Bounded resources: the fetch loop and lane goroutines must all unwind + // once Run returns, returning the process to its pre-run goroutine count. + deadline := time.Now().Add(2 * time.Second) + for runtime.NumGoroutine() > baseline+2 { + if time.Now().After(deadline) { + t.Fatalf("goroutines after run = %d, want back within +2 of baseline %d", runtime.NumGoroutine(), baseline) + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/source/inlet.go b/source/inlet.go index 300dd8c..a66acdd 100644 --- a/source/inlet.go +++ b/source/inlet.go @@ -49,7 +49,10 @@ type Inlet interface { type Subscription interface { // Next returns the next message. It blocks until one is available, returns // ctx.Err() if ctx is canceled, or returns ErrDrained once the subscription - // has been closed and all delivered messages settled. + // has been closed and all delivered messages settled. After Close, a + // backend must not deliver new messages: only messages already returned to + // the buffer before Close may still be yielded, and once none remain, + // Next blocks until in-flight settles finish and then reports ErrDrained. Next(ctx context.Context) (Message, error) // Settle applies a handler [Result] to a message previously returned by Next: // ack/commit, schedule redelivery, route to dead-letter, or extend the diff --git a/source/kafka/README.md b/source/kafka/README.md index d5951f5..36ff876 100644 --- a/source/kafka/README.md +++ b/source/kafka/README.md @@ -33,19 +33,27 @@ and the marked offsets are committed on graceful drain and on rebalance | Result | Kafka behavior | | ------------- | ----------------------------------------------------------- | | `Ack` | mark the record for commit (commit-after-process) | -| `Nak` | do **not** mark; the record is re-read on restart/rebalance | -| `NakAfter(d)` | best-effort: pause partition, wait `d`, re-seek, resume | +| `Nak` | never mark; pause, re-seek to the record's offset, resume — the record is fetched again in this session | +| `NakAfter(d)` | same as `Nak`, waiting out `d` between pause and re-seek (best-effort) | | `Term` | produce the record to the dead-letter topic, then mark | | `InProgress` | no-op (Kafka has no per-message ack deadline) | | `Manual` | no-op (the handler settled via `Message.As` + the client) | -### Divergence: Nak delay - -Kafka has no native per-message redelivery delay. `NakAfter(d)` is emulated by -pausing the record's partition, waiting out `d` (or the context), re-seeking to -the record's own offset, and resuming. A plain `Nak` simply declines to mark, -so the record is re-delivered on the next restart or rebalance. This is the -documented divergence from JetStream's native delayed nak. +### Divergence: Nak delay and cross-restart redelivery + +Kafka has no native per-message redelivery. A `Nak` (plain or delayed) is +emulated by pausing the record's partition, waiting out the requested delay +(zero for a plain `Nak`), re-seeking the partition to the record's own offset, +and resuming — so the record is redelivered deterministically **within the live +subscription**. Costs to know about: the delay head-of-line-blocks the paused +partition; records already fetched but not yet yielded are still delivered +before the redelivered record; and concurrent settles on the same partition can +commit past the nacked offset before the re-seek lands. + +Across process restarts and rebalances, redelivery rides committed offsets: a +concurrently committed higher offset can advance past a nacked record, so +cross-restart redelivery is **best-effort**, not guaranteed. This is the one +divergence from JetStream's native nak semantics. ## Capabilities diff --git a/source/kafka/drain_test.go b/source/kafka/drain_test.go new file mode 100644 index 0000000..56766f2 --- /dev/null +++ b/source/kafka/drain_test.go @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 + +package kafka + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/twmb/franz-go/pkg/kgo" + + "github.com/stablekernel/crucible/source" +) + +// recsFetch wraps several records into a one-partition fetch batch (the +// multi-record analog of oneFetch). +func recsFetch(recs ...*kgo.Record) kgo.Fetches { + return kgo.Fetches{{ + Topics: []kgo.FetchTopic{{ + Topic: recs[0].Topic, + Partitions: []kgo.FetchPartition{{ + Partition: recs[0].Partition, + Records: recs, + }}, + }}, + }} +} + +// fetchIx reports how many scripted fetch batches the poller has served, so +// drain tests can prove Close stops polling. +func (f *fakePoller) fetchIxLocked() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.fetchIx +} + +// nextResult carries a Next outcome across a goroutine boundary. +type nextResult struct { + msg source.Message + err error +} + +// TestNextAfterCloseYieldsOnlyBufferedThenDrained pins the P0-1 drain contract +// against the real subscription: after Close, Next yields only already-fetched +// records, never polls again, and blocks while records are in flight — returning +// ErrDrained only once the last delivered record settles. +func TestNextAfterCloseYieldsOnlyBufferedThenDrained(t *testing.T) { + t.Parallel() + + r0 := rec("orders", 2, 10, "k", "v0") + r1 := rec("orders", 2, 11, "k", "v1") + r2 := rec("orders", 2, 12, "k", "v2") + fp := &fakePoller{fetches: []kgo.Fetches{recsFetch(r0, r1, r2)}} + sub := newSub(fp) + ctx := context.Background() + + m0, err := sub.Next(ctx) + if err != nil { + t.Fatalf("Next() error = %v", err) + } + if string(m0.Value()) != "v0" { + t.Fatalf("first Next value = %q, want v0", m0.Value()) + } + + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + + m1, err := sub.Next(ctx) + if err != nil { + t.Fatalf("Next() after close error = %v", err) + } + m2, err := sub.Next(ctx) + if err != nil { + t.Fatalf("Next() after close error = %v", err) + } + if string(m1.Value()) != "v1" || string(m2.Value()) != "v2" { + t.Errorf("buffered values = %q/%q, want v1/v2", m1.Value(), m2.Value()) + } + if got := fp.fetchIxLocked(); got != 1 { + t.Fatalf("polls served = %d, want 1 (Next must never poll after Close)", got) + } + + // Buffer empty, one record in flight: Next must block, not poll and not + // drain early. + resCh := make(chan nextResult, 1) + go func() { + m, err := sub.Next(ctx) + resCh <- nextResult{msg: m, err: err} + }() + select { + case res := <-resCh: + t.Fatalf("Next returned before the last settle: msg=%v err=%v", res.msg != nil, res.err) + case <-time.After(50 * time.Millisecond): + } + + // Settle every delivered record; only the last settle may release the + // blocked Next. + for _, m := range []source.Message{m0, m1, m2} { + if err := sub.Settle(ctx, m, source.Ack()); err != nil { + t.Fatalf("Settle(ack) error = %v", err) + } + } + select { + case res := <-resCh: + if !errors.Is(res.err, source.ErrDrained) { + t.Fatalf("blocked Next = %v, want ErrDrained", res.err) + } + if res.msg != nil { + t.Fatalf("drained Next yielded %+v, want no message", res.msg) + } + case <-time.After(time.Second): + t.Fatal("blocked Next did not return after the last settle") + } + if got := fp.markedCount(); got != 3 { + t.Errorf("marked = %d, want 3", got) + } +} + +// TestNextBatchAfterCloseYieldsOnlyBufferedThenDrained pins the P0-1 drain +// contract for the batch path: after Close, NextBatch serves only the buffer +// without polling and drains only after every settled delivery. +func TestNextBatchAfterCloseYieldsOnlyBufferedThenDrained(t *testing.T) { + t.Parallel() + + r0 := rec("orders", 2, 10, "k", "v0") + r1 := rec("orders", 2, 11, "k", "v1") + r2 := rec("orders", 2, 12, "k", "v2") + fp := &fakePoller{fetches: []kgo.Fetches{recsFetch(r0, r1, r2)}} + sub := newSub(fp) + ctx := context.Background() + + first, err := sub.NextBatch(ctx, 2) + if err != nil { + t.Fatalf("NextBatch(2) error = %v", err) + } + if len(first) != 2 || string(first[0].Value()) != "v0" || string(first[1].Value()) != "v1" { + t.Fatalf("first batch values = %q/%q, want v0/v1", first[0].Value(), first[1].Value()) + } + + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + + second, err := sub.NextBatch(ctx, 2) + if err != nil { + t.Fatalf("NextBatch(2) after close error = %v", err) + } + if len(second) != 1 || string(second[0].Value()) != "v2" { + t.Fatalf("second batch = %d records (%q), want exactly [v2]", len(second), second[0].Value()) + } + if got := fp.fetchIxLocked(); got != 1 { + t.Fatalf("polls served = %d, want 1 (NextBatch must never poll after Close)", got) + } + + // Settle the first batch, then hold a third call while v2 is in flight. + for _, m := range first { + if err := sub.Settle(ctx, m, source.Ack()); err != nil { + t.Fatalf("Settle(batch, ack) error = %v", err) + } + } + resCh := make(chan nextResult, 1) + go func() { + ms, err := sub.NextBatch(ctx, 2) + if err != nil { + resCh <- nextResult{err: err} + return + } + _ = ms + resCh <- nextResult{} + }() + select { + case res := <-resCh: + t.Fatalf("NextBatch returned before the last settle: err=%v", res.err) + case <-time.After(50 * time.Millisecond): + } + + if err := sub.Settle(ctx, second[0], source.Ack()); err != nil { + t.Fatalf("Settle(v2, ack) error = %v", err) + } + select { + case res := <-resCh: + if !errors.Is(res.err, source.ErrDrained) { + t.Fatalf("blocked NextBatch = %v, want ErrDrained", res.err) + } + case <-time.After(time.Second): + t.Fatal("blocked NextBatch did not return after the last settle") + } +} + +// TestNextAfterCloseDoesNotPollWithEmptyBufferAndNoInFlight pins the simplest +// drain shape: Close on a quiet subscription drains immediately without a +// single poll. +func TestNextAfterCloseDoesNotPollWithEmptyBufferAndNoInFlight(t *testing.T) { + t.Parallel() + + fp := &fakePoller{} // no scripted fetches + sub := newSub(fp) + + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + + resCh := make(chan nextResult, 1) + go func() { + m, err := sub.Next(context.Background()) + resCh <- nextResult{msg: m, err: err} + }() + select { + case res := <-resCh: + if !errors.Is(res.err, source.ErrDrained) { + t.Fatalf("Next() after close = %v, want ErrDrained", res.err) + } + case <-time.After(time.Second): + t.Fatal("Next() did not drain after Close on an empty subscription") + } + if got := fp.fetchIxLocked(); got != 0 { + t.Errorf("polls served = %d, want 0 (a closed subscription never polls)", got) + } +} + +// TestPlainNakReseeksPartitionImmediately pins the P0-2 decision: a plain Nak +// (zero Requeue) redelivers in-session via the same pause/re-seek/resume +// machinery as a delayed nak, and never marks the record. +func TestPlainNakReseeksPartitionImmediately(t *testing.T) { + t.Parallel() + + fp := &fakePoller{} + sub := newSub(fp) + r := rec("orders", 5, 88, "k", "v") + + start := time.Now() + if err := sub.Settle(context.Background(), newMessage(r), source.Nak(errors.New("boom"))); err != nil { + t.Fatalf("Settle(nak) error = %v", err) + } + if elapsed := time.Since(start); elapsed > 50*time.Millisecond { + t.Errorf("plain Nak took %v, want effectively immediate", elapsed) + } + if got := fp.markedCount(); got != 0 { + t.Errorf("marked = %d, want 0 (a nak never commits)", got) + } + if len(fp.paused) != 1 || len(fp.resumed) != 1 { + t.Fatalf("paused=%d resumed=%d, want 1/1", len(fp.paused), len(fp.resumed)) + } + if len(fp.setOffsets) != 1 { + t.Fatalf("setOffsets calls = %d, want 1 (re-seek to the record offset)", len(fp.setOffsets)) + } + eo, ok := fp.setOffsets[0]["orders"][5] + if !ok || eo.Offset != 88 || eo.Epoch != -1 { + t.Errorf("re-seek = %+v, want offset 88 epoch -1 on orders/5", eo) + } +} + +// TestNakAfterDelaysThenReseeks pins the delayed-nak path end to end: the +// requested delay elapses before the re-seek, and the same +// pause/re-seek/resume choreography runs. +func TestNakAfterDelaysThenReseeks(t *testing.T) { + t.Parallel() + + fp := &fakePoller{} + sub := newSub(fp) + r := rec("orders", 5, 88, "k", "v") + + start := time.Now() + if err := sub.Settle(context.Background(), newMessage(r), source.NakAfter(60*time.Millisecond, errors.New("boom"))); err != nil { + t.Fatalf("Settle(nak-after) error = %v", err) + } + if elapsed := time.Since(start); elapsed < 55*time.Millisecond { + t.Errorf("NakAfter(60ms) returned after %v, want >= ~55ms", elapsed) + } + if got := fp.markedCount(); got != 0 { + t.Errorf("marked = %d, want 0 (a nak never commits)", got) + } + if len(fp.paused) != 1 || len(fp.resumed) != 1 { + t.Fatalf("paused=%d resumed=%d, want 1/1", len(fp.paused), len(fp.resumed)) + } + eo, ok := fp.setOffsets[0]["orders"][5] + if !ok || eo.Offset != 88 || eo.Epoch != -1 { + t.Errorf("re-seek = %+v, want offset 88 epoch -1 on orders/5", eo) + } +} + +// TestDuplicateSettleIsSafe defines duplicate-settle behavior: settling the +// same message twice is accepted (the backend mark is idempotent at the broker), +// the drain accounting stays honest, and the subscription still drains cleanly. +func TestDuplicateSettleIsSafe(t *testing.T) { + t.Parallel() + + r0 := rec("orders", 2, 10, "k", "v0") + r1 := rec("orders", 2, 11, "k", "v1") + fp := &fakePoller{fetches: []kgo.Fetches{recsFetch(r0, r1)}} + sub := newSub(fp) + ctx := context.Background() + + m0, err := sub.Next(ctx) + if err != nil { + t.Fatalf("Next() error = %v", err) + } + m1, err := sub.Next(ctx) + if err != nil { + t.Fatalf("Next() error = %v", err) + } + + for i, m := range []source.Message{m0, m1} { + for j := 0; j < 2; j++ { + if err := sub.Settle(ctx, m, source.Ack()); err != nil { + t.Fatalf("duplicate Settle(%d, pass %d) error = %v", i, j, err) + } + } + } + if got := fp.markedCount(); got != 4 { + t.Errorf("marked = %d, want 4 (each Settle marks; duplicate marks are broker-idempotent)", got) + } + + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + resCh := make(chan nextResult, 1) + go func() { + m, err := sub.Next(ctx) + resCh <- nextResult{msg: m, err: err} + }() + select { + case res := <-resCh: + if !errors.Is(res.err, source.ErrDrained) { + t.Fatalf("Next() after duplicate settles = %v, want ErrDrained", res.err) + } + case <-time.After(time.Second): + t.Fatal("subscription did not drain after duplicate settles (in-flight accounting corrupted)") + } +} diff --git a/source/kafka/kafka.go b/source/kafka/kafka.go index c0cac48..f272e1d 100644 --- a/source/kafka/kafka.go +++ b/source/kafka/kafka.go @@ -15,10 +15,12 @@ // Each handler [source.Result] maps onto Kafka as follows: // // - Ack marks the record for commit (commit-after-process). -// - Nak does NOT mark the record, so it is re-read on restart or rebalance. -// A requeue delay is best-effort, applied by pausing and re-seeking the -// record's partition; this is a documented divergence from JetStream's -// native delayed nak (Kafka has no per-message redelivery delay). +// - Nak never marks the record and redelivers it in-session: the partition is +// paused, re-seeked to the record's offset so it is fetched again, then +// resumed; a requeue delay waits out the pause. Redelivery across restarts +// or rebalances rides committed offsets, so a concurrently committed higher +// offset can pass the nacked record — a documented best-effort divergence, +// not an in-session one. // - Term produces the record to the configured dead-letter topic, then marks // it for commit so it is not re-read. // - InProgress is a no-op: Kafka has no per-message ack deadline to extend. diff --git a/source/kafka/subscription.go b/source/kafka/subscription.go index 2625284..030c84c 100644 --- a/source/kafka/subscription.go +++ b/source/kafka/subscription.go @@ -52,6 +52,11 @@ type subscription struct { buffer []*kgo.Record inFlight int closed bool + // notify wakes a Next/NextBatch that is waiting (after Close) for the last + // in-flight records to settle. Capacity one: sends are best-effort wakeups, + // and every waiter re-checks the drain state under mu after waking. Guarded + // by mu; created lazily by signal() so a zero-value subscription stays safe. + notify chan struct{} // onAssignedFn/onRevokedFn are the engine-registered hooks the franz-go // rebalance trampolines forward to. Guarded by hookMu. @@ -63,8 +68,10 @@ type subscription struct { // Next returns the next buffered record, polling the broker when the buffer is // empty. It blocks until a record is available, returns ctx.Err() on // cancellation, or [source.ErrDrained] once the subscription is closed and all -// delivered records are settled. Next is single-consumer (the engine's fetch -// loop). +// delivered records are settled. After Close it never polls for new records: +// only records already fetched into the buffer are yielded, and once the buffer +// is empty Next blocks until the remaining in-flight records settle, then +// returns ErrDrained. Next is single-consumer (the engine's fetch loop). func (s *subscription) Next(ctx context.Context) (source.Message, error) { for { if rec, ok := s.takeBuffered(); ok { @@ -72,11 +79,21 @@ func (s *subscription) Next(ctx context.Context) (source.Message, error) { } s.mu.Lock() - drained := s.closed && s.inFlight == 0 + closed, drained := s.closed, s.inFlight == 0 s.mu.Unlock() - if drained { + if closed && drained { return nil, source.ErrDrained } + if closed { + // Draining with records still in flight: hold here — do not poll — + // until the last settle wakes us, or the caller gives up. + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-s.signal(): + } + continue + } if err := ctx.Err(); err != nil { return nil, err } @@ -92,7 +109,10 @@ func (s *subscription) Next(ctx context.Context) (source.Message, error) { } if errs := fetches.Errors(); len(errs) > 0 { // Surface the first non-context fetch error; the engine decides - // whether to retry the loop. + // whether to retry the loop. Records fetched alongside an error are + // deliberately discarded rather than buffered: none of them have + // been delivered, so nothing can be settled or marked from them, + // and the next successful poll refetches their offsets. for _, fe := range errs { if fe.Err != nil && fe.Err != context.Canceled && fe.Err != context.DeadlineExceeded { return nil, fmt.Errorf("source/kafka: poll %s[%d]: %w", fe.Topic, fe.Partition, fe.Err) @@ -110,6 +130,29 @@ func (s *subscription) Next(ctx context.Context) (source.Message, error) { } } +// signal returns the wakeup channel for drain waiters, creating it on first +// use so a zero-value subscription stays safe. +func (s *subscription) signal() <-chan struct{} { + s.mu.Lock() + defer s.mu.Unlock() + if s.notify == nil { + s.notify = make(chan struct{}, 1) + } + return s.notify +} + +// wake nudges any Next/NextBatch waiting for in-flight settles after Close. +// Best-effort: a coalesced wakeup is fine because every waiter re-checks the +// drain state after receiving. Call with s.mu held. +func (s *subscription) wake() { + if s.notify != nil { + select { + case s.notify <- struct{}{}: + default: + } + } +} + // takeBuffered pops one record from the buffer and counts it in flight. func (s *subscription) takeBuffered() (*kgo.Record, bool) { s.mu.Lock() @@ -146,7 +189,10 @@ func (s *subscription) takeBatch(limit int) ([]*kgo.Record, bool) { // buffer is empty, satisfying [source.Batched]. franz-go already fetches in // batches (PollRecords), so this exposes a poll's records as a group rather than // draining them one at a time; the engine settles each through SettleBatch. -// NextBatch is single-consumer, like Next. +// After Close it never polls for new records: only already-buffered records are +// yielded, and once the buffer is empty it blocks until the remaining in-flight +// records settle, then returns [source.ErrDrained] — the same drain contract as +// [Next]. NextBatch is single-consumer, like Next. func (s *subscription) NextBatch(ctx context.Context, limit int) ([]source.Message, error) { if limit < 1 { limit = 1 @@ -161,11 +207,21 @@ func (s *subscription) NextBatch(ctx context.Context, limit int) ([]source.Messa } s.mu.Lock() - drained := s.closed && s.inFlight == 0 + closed, drained := s.closed, s.inFlight == 0 s.mu.Unlock() - if drained { + if closed && drained { return nil, source.ErrDrained } + if closed { + // Draining with records still in flight: hold here — do not poll — + // until the last settle wakes us, or the caller gives up. + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-s.signal(): + } + continue + } if err := ctx.Err(); err != nil { return nil, err } @@ -226,13 +282,14 @@ func (s *subscription) Settle(ctx context.Context, m source.Message, r source.Re return nil case source.ActionNak: - // Do NOT mark: the record stays uncommitted and is re-read on the next - // restart or rebalance. A requeue delay is best-effort: pause the - // partition, sleep, re-seek to this record's offset, resume. - if r.Requeue > 0 { - return s.requeueWithDelay(ctx, rec, r.Requeue) - } - return nil + // Redeliver: never mark the record. The partition is paused, re-seeked + // to this record's offset so it (and anything after it) is fetched + // again in this session, then resumed; a Requeue delay waits out the + // pause, and a plain Nak re-seeks immediately (delay zero). Redelivery + // across restarts/rebalances rides committed offsets, so a concurrent + // ack that commits past this record can skip it after such a restart — + // a documented best-effort divergence, not an in-session one. + return s.requeueWithDelay(ctx, rec, r.Requeue) case source.ActionTerm: // Produce to the dead-letter topic, then mark so it is not re-read. @@ -255,21 +312,31 @@ func (s *subscription) Settle(ctx context.Context, m source.Message, r source.Re } } -// settled decrements the in-flight count and is deferred from Settle so it runs -// on every path, including errors, keeping the drain accounting honest. +// settled decrements the in-flight count and wakes any drain waiter; deferred +// from Settle so it runs on every path, including errors, keeping the drain +// accounting honest. func (s *subscription) settled() { s.mu.Lock() if s.inFlight > 0 { s.inFlight-- } + s.wake() s.mu.Unlock() } -// requeueWithDelay implements the best-effort Nak delay: pause the record's -// partition so no further records are fetched from it, wait out the delay (or -// the context), re-seek delivery to the record's own offset so it is re-read, -// then resume the partition. This is a documented divergence — Kafka has no -// native per-message redelivery delay. +// requeueWithDelay implements Nak redelivery: pause the record's partition so +// no further records are fetched from it, wait out the delay (zero for a plain +// Nak), re-seek delivery to the record's own offset so it is re-read, then +// resume the partition. +// +// Cost model, documented as best-effort by design: the delay blocks only the +// calling settle goroutine, but the pause head-of-line-blocks the whole +// partition for the delay's duration (Kafka has no native per-message +// redelivery delay), and records fetched before the seek but not yet yielded +// are still delivered ahead of the redelivered record. A concurrent ack on the +// same partition can also commit past the nacked offset before the re-seek +// lands; the commit advances the persisted position while the live fetch keeps +// reading from the seeked offset until the next rebalance or restart. func (s *subscription) requeueWithDelay(ctx context.Context, rec *kgo.Record, d time.Duration) error { tp := map[string][]int32{rec.Topic: {rec.Partition}} s.client.PauseFetchPartitions(tp) @@ -337,14 +404,16 @@ const ( dlqHeaderError = "crucible-error" ) -// Close begins a graceful drain: Next stops fetching new records once the -// current buffer is exhausted, commits whatever has been marked, and once -// in-flight records settle, Next returns [source.ErrDrained]. Close is -// idempotent. +// Close begins a graceful drain: from the moment it returns, Next/NextBatch +// stop fetching new records and yield only what was already buffered; once the +// in-flight records settle, Next returns [source.ErrDrained]. Marked offsets +// are committed best-effort so a clean shutdown does not re-read +// already-processed records. Close is idempotent. func (s *subscription) Close() error { s.mu.Lock() already := s.closed s.closed = true + s.wake() s.mu.Unlock() if already { return nil From ce6996667c95f4217f5cbd939483c377791f524e Mon Sep 17 00:00:00 2001 From: Joshua Temple Date: Sat, 22 Aug 2026 22:29:49 -0400 Subject: [PATCH 2/2] fix(source/kafka): scope Close error declarations correctly in drain tests --- source/kafka/drain_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/kafka/drain_test.go b/source/kafka/drain_test.go index 56766f2..12bb12a 100644 --- a/source/kafka/drain_test.go +++ b/source/kafka/drain_test.go @@ -63,7 +63,7 @@ func TestNextAfterCloseYieldsOnlyBufferedThenDrained(t *testing.T) { t.Fatalf("first Next value = %q, want v0", m0.Value()) } - if err := sub.Close(); err != nil { + if err = sub.Close(); err != nil { t.Fatalf("Close() error = %v", err) } @@ -139,7 +139,7 @@ func TestNextBatchAfterCloseYieldsOnlyBufferedThenDrained(t *testing.T) { t.Fatalf("first batch values = %q/%q, want v0/v1", first[0].Value(), first[1].Value()) } - if err := sub.Close(); err != nil { + if err = sub.Close(); err != nil { t.Fatalf("Close() error = %v", err) }