diff --git a/docs/src/content/docs/source/concurrency.md b/docs/src/content/docs/source/concurrency.md index e0732f5..55ae958 100644 --- a/docs/src/content/docs/source/concurrency.md +++ b/docs/src/content/docs/source/concurrency.md @@ -60,3 +60,62 @@ knows when the consumer is live. Ordered-key concurrency is heavily property- and fuzz-tested, and the [`memsource`](/crucible/source/reliability/#testing-the-loop) harness lets you assert ordering and ack outcomes deterministically with no broker. + +## Batch consume + +Some handlers are far cheaper per message in groups: a bulk database upsert, a +single round trip to an index, a vectorized transform. For those, the Hopper has +a batch mode that accumulates messages per ordered lane and hands them to a +`BatchHandler` in one call, then settles each message by its own result. + +```go +func(ctx context.Context, ms []source.Message) []source.Result +``` + +The contract is positional: the result at index `i` settles the message at index +`i`, so a batch can ack some messages, nak others, and dead-letter the rest in a +single pass. A `BatchHandler` must return exactly one result per message; if it +returns too few, the unmatched messages are terminated as poison +(`ErrBatchResultCount`) rather than silently stranded. + +Enable it with `WithBatch(size, maxWait)` and drive the run with `RunBatch` (or +`ReceiveBatch`) instead of `Run`: + +```go +hp := source.New( + source.WithConcurrency(8), + source.WithBatch(100, 50*time.Millisecond), +) +err := hp.RunBatch(ctx, sub, func(ctx context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + // ... handle ms as a group, fill res[i] for each ms[i] ... + return res +}) +``` + +Each lane buffers up to `size` messages, or until `maxWait` elapses since its +first buffered message, then flushes. A partial batch left when the run drains is +flushed before exit, so nothing is lost on shutdown. The `maxWait` timer reads an +injectable clock (`WithBatchClock`) so a time-triggered flush is deterministic in +tests. + +### Ordering and backpressure still hold + +Batch mode keeps every guarantee of the per-message path. A batch only ever +contains messages from a single ordered lane, delivered in order, so per-key +ordering is preserved within and across batches. A lane never overlaps two +batches: it settles the whole group before accepting more. `WithConcurrency` +bounds how many lanes run their handler at once, and `WithMaxInFlight` still caps +delivered-but-unsettled messages. When both are set, size the in-flight window to +at least the batch size (or supply a positive `maxWait`) so a lane can always +reach a flush. + +### Backend batches + +When the subscription advertises the `Batched` capability, the engine fetches +whole batches from the backend and regroups them by key into lanes, rather than +accumulating one message at a time. Kafka satisfies it naturally (a `PollRecords` +fetch is already a batch), and JetStream satisfies it through its pull-consumer +fetch. An adapter without the capability is driven one message at a time and the +engine does the accumulation; batch mode works either way, the capability just +skips a layer of per-message hand-off. diff --git a/docs/src/content/docs/source/overview.md b/docs/src/content/docs/source/overview.md index fe89bc0..80bf33d 100644 --- a/docs/src/content/docs/source/overview.md +++ b/docs/src/content/docs/source/overview.md @@ -91,7 +91,7 @@ no-op and forcing nothing third-party on the consumer. ## Next - [The Inlet and Hopper model](/crucible/source/model/): the vocabulary and the ack model. -- [Ordered concurrency and backpressure](/crucible/source/concurrency/): the spine. +- [Ordered concurrency and backpressure](/crucible/source/concurrency/): the spine, including batch consume. - [Adapters](/crucible/source/adapters/): Kafka and JetStream, plus the per-backend capability table. - [Reliability middleware](/crucible/source/reliability/): retry, DLQ, idempotency, schema. - [Codecs and headers](/crucible/source/codecs/): the instance-scoped registry and CloudEvents. diff --git a/source/batch.go b/source/batch.go new file mode 100644 index 0000000..fa14faf --- /dev/null +++ b/source/batch.go @@ -0,0 +1,352 @@ +// SPDX-License-Identifier: Apache-2.0 + +package source + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/stablekernel/crucible/telemetry" +) + +// ReceiveBatch subscribes to in with cfg and drives the resulting subscription +// with bh in batch mode, the batch analog of [Hopper.Receive]. It closes the +// subscription on return. The Hopper must have been constructed with [WithBatch]; +// without it a batch run still works but every batch holds a single message. +func (hp *Hopper) ReceiveBatch(ctx context.Context, in Inlet, cfg SubscribeConfig, bh BatchHandler) error { + sub, err := in.Subscribe(ctx, cfg) + if err != nil { + return err + } + defer func() { _ = sub.Close() }() + return hp.runBatch(ctx, sub, bh) +} + +// RunBatch drives sub with bh in batch mode, the batch analog of [Hopper.Run]. +// Within each ordered lane it accumulates messages up to the configured size (see +// [WithBatch]) or until the max-wait elapses, invokes bh once per accumulated +// slice, and settles each message by its corresponding [Result]. Per-key ordering +// and the [WithMaxInFlight] bound hold exactly as in the per-message path: a lane +// never overlaps two batches, and every message in a batch is settled before the +// lane accepts more. +// +// RunBatch honors a [Batched] subscription when present, fetching whole batches +// from the backend and regrouping them by key into lanes; otherwise it +// accumulates per message under the size/max-wait policy. Lifecycle, drain, and +// error semantics match [Hopper.Run]: a clean drain or context cancellation +// returns nil, a fetch error propagates, and a partial batch buffered when the +// run drains is flushed before exit. RunBatch does not close sub. +func (hp *Hopper) RunBatch(ctx context.Context, sub Subscription, bh BatchHandler) error { + return hp.runBatch(ctx, sub, bh) +} + +// batchSize resolves the per-lane batch size, defaulting to 1 when batch mode was +// not configured (each message is its own batch). +func (hp *Hopper) batchSize() int { + if hp.batch.size > 0 { + return hp.batch.size + } + return 1 +} + +// batchClock resolves the injected clock, defaulting to time.Now. +func (hp *Hopper) batchClock() func() time.Time { + if hp.batch.now != nil { + return hp.batch.now + } + return time.Now +} + +// runBatch mirrors run but accumulates per-lane and dispatches to a BatchHandler. +// The fetch loop reserves an in-flight slot per message (the backpressure bound +// is per message, not per batch), routes each message to its ordered lane, and a +// lane goroutine buffers messages until the size or max-wait bound trips, then +// hands the slice to the handler and settles each result. +func (hp *Hopper) runBatch(ctx context.Context, sub Subscription, bh BatchHandler) error { + runCtx, cancel := context.WithCancel(ctx) + defer cancel() + + stop := make(chan struct{}) + defer close(stop) + go func() { + select { + case <-hp.closed: + cancel() + case <-stop: + } + }() + + var inFlight chan struct{} + if hp.maxInFlight > 0 { + inFlight = make(chan struct{}, hp.maxInFlight) + } + // dispatchSem bounds how many lanes run their BatchHandler concurrently, + // mirroring the per-message concurrency bound (at most N handler executions in + // parallel). A lane buffers freely but acquires a token around each dispatch, + // so distinct keys can all have live lanes while only N dispatch at once. + dispatchSem := make(chan struct{}, hp.concurrency) + + var ( + mu sync.Mutex + lanes = make(map[uint64]*batchLane) + laneWG sync.WaitGroup + runErr error + errOnce sync.Once + ) + + size := hp.batchSize() + clock := hp.batchClock() + + // laneFor returns the ordered lane for key, starting its goroutine on first + // use. One goroutine per key, exactly like the per-message path; the + // cross-key parallelism bound is applied at dispatch time, not creation. + laneFor := func(key uint64) *batchLane { + mu.Lock() + defer mu.Unlock() + if l, ok := lanes[key]; ok { + return l + } + l := &batchLane{ + queue: make(chan *delivery, 1), + size: size, + maxWait: hp.batch.maxWait, + now: clock, + inFlight: inFlight, + dispatch: dispatchSem, + } + lanes[key] = l + laneWG.Add(1) + go func() { + defer laneWG.Done() + l.run(runCtx, hp, bh) + }() + return l + } + + defer func() { + // Drain: close every lane queue so each lane flushes its buffered messages + // and exits, then wait. + mu.Lock() + for _, l := range lanes { + close(l.queue) + } + mu.Unlock() + laneWG.Wait() + }() + + // batched is the whole-batch fetch path when the subscription advertises it. + batched, _ := sub.(Batched) + + for { + msgs, err := hp.fetchBatch(runCtx, sub, batched, size) + if err != nil { + if errors.Is(err, ErrDrained) || errors.Is(err, context.Canceled) || runCtx.Err() != nil { + return runErr + } + errOnce.Do(func() { runErr = err }) + return runErr + } + + for _, msg := range msgs { + // Reserve the backpressure slot per message (not per fetch) so a + // NextBatch larger than maxInFlight cannot deadlock: slots free as + // lanes dispatch, and a full window blocks here before routing more. + if inFlight != nil { + select { + case inFlight <- struct{}{}: + case <-runCtx.Done(): + return runErr + } + } + + hp.received.Add(runCtx, 1) + hp.reportLag(runCtx, sub) + + l := laneFor(hp.laneKey(msg)) + select { + case l.queue <- &delivery{msg: msg, sub: sub}: + case <-runCtx.Done(): + if inFlight != nil { + <-inFlight + } + return runErr + } + } + } +} + +// fetchBatch pulls the next group of messages, using the [Batched] capability +// when present (one NextBatch call yields up to size messages, settled per +// message by the lane) and otherwise a single [Subscription.Next]. It does not +// reserve backpressure slots; the routing loop reserves one per message so a +// NextBatch larger than the in-flight window cannot deadlock. +func (hp *Hopper) fetchBatch(ctx context.Context, sub Subscription, batched Batched, size int) ([]Message, error) { + if batched != nil { + msgs, err := batched.NextBatch(ctx, size) + if err != nil { + return nil, err + } + return msgs, nil + } + msg, err := sub.Next(ctx) + if err != nil { + return nil, err + } + return []Message{msg}, nil +} + +// batchLane is a single ordered worker that buffers deliveries and dispatches +// them to a BatchHandler in groups. It mirrors the per-message lane but holds a +// buffer and a max-wait timer so a slow-filling lane still flushes on time. +type batchLane struct { + queue chan *delivery + size int + maxWait time.Duration + now func() time.Time + inFlight chan struct{} + dispatch chan struct{} // cross-key concurrency bound, acquired around a flush +} + +// run drives the lane: it buffers deliveries, flushing when the buffer reaches +// size, when the max-wait since the first buffered message elapses, or when the +// queue closes (run drain). It is the single point where ordering within a key is +// preserved — one batch at a time, in arrival order. +func (l *batchLane) run(ctx context.Context, hp *Hopper, bh BatchHandler) { + var buf []*delivery + + // timer fires maxWait after the first message lands in an empty buffer. It is + // created stopped and reset on each first-message; a zero/negative maxWait + // disables it entirely (flush on size or close only). + var timerC <-chan time.Time + var timer *time.Timer + if l.maxWait > 0 { + timer = time.NewTimer(l.maxWait) + if !timer.Stop() { + <-timer.C + } + defer timer.Stop() + } + armed := false + + flush := func() { + if len(buf) == 0 { + return + } + // Acquire a dispatch token so at most concurrency lanes run their handler + // at once. If the run is shutting down, proceed without the token: the + // buffered messages must still be settled so none are stranded. + acquired := false + select { + case l.dispatch <- struct{}{}: + acquired = true + case <-ctx.Done(): + } + hp.dispatchBatch(ctx, bh, buf) + if acquired { + <-l.dispatch + } + // Release the per-message in-flight slot held for every message in the + // batch now that all are settled. The lane slot is held for the lane's + // whole lifetime and released when its goroutine exits, not here. + for range buf { + if l.inFlight != nil { + <-l.inFlight + } + } + buf = nil + if timer != nil { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + armed = false + timerC = nil + } + } + + for { + select { + case d, ok := <-l.queue: + if !ok { + // Queue closed: flush the partial batch and exit. + flush() + return + } + buf = append(buf, d) + if l.maxWait > 0 && !armed { + timer.Reset(l.maxWait) + timerC = timer.C + armed = true + } + if len(buf) >= l.size { + flush() + } + case <-timerC: + armed = false + timerC = nil + flush() + } + } +} + +// dispatchBatch decodes each buffered message, invokes the batch handler with the +// decoded slice, settles each message by its corresponding result, and releases +// the per-message in-flight and lane slots. A decode failure terminates that +// message as poison and excludes it from the handler call, so a bad payload never +// poisons its lane-mates. A result-count mismatch settles what it can and +// terminates the remainder as poison. +func (hp *Hopper) dispatchBatch(ctx context.Context, bh BatchHandler, buf []*delivery) { + ctx, span := hp.tracer.Start( + ctx, "source.batch", + telemetry.String("source.name", hp.name), + telemetry.Int("source.batch_size", len(buf)), + ) + defer span.End() + hp.batches.Add(ctx, 1) + + // Decode every message; a decode failure is settled as poison immediately and + // dropped from the slice handed to the handler. + msgs := make([]Message, 0, len(buf)) + live := make([]*delivery, 0, len(buf)) + for _, d := range buf { + dec, err := hp.decodeFor(d.msg) + if err != nil { + span.RecordError(err) + hp.settle(ctx, span, d, Term(err)) + continue + } + msgs = append(msgs, dec) + live = append(live, d) + } + + if len(live) == 0 { + return + } + + results := bh(ctx, msgs) + + for i, d := range live { + var r Result + switch { + case i < len(results): + r = results[i] + default: + // The handler returned too few results: terminate the unmatched + // message as poison rather than silently acking or stranding it. + r = Term(ErrBatchResultCount) + span.RecordError(ErrBatchResultCount) + } + if r.Err != nil { + span.RecordError(r.Err) + } + hp.settle(ctx, span, d, r) + } + + if len(results) != len(live) { + span.SetStatus(telemetry.StatusError, "batch result count mismatch") + } +} diff --git a/source/batch_test.go b/source/batch_test.go new file mode 100644 index 0000000..32edad6 --- /dev/null +++ b/source/batch_test.go @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 + +package source_test + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stablekernel/crucible/source" + "github.com/stablekernel/crucible/source/memsource" +) + +// TestHopper_BatchSizeTrigger verifies a full-size batch is dispatched as one +// call and every message acked, across a range of size/count combinations. +func TestHopper_BatchSizeTrigger(t *testing.T) { + t.Parallel() + tests := []struct { + name string + size int + count int + wantBatches int + }{ + {"exact multiple", 5, 10, 2}, + {"single batch", 8, 3, 1}, + {"size one", 1, 4, 4}, + {"remainder flushes on drain", 4, 10, 3}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + msgs := make([]memsource.Msg, tt.count) + for i := range msgs { + // One key so every message lands on a single ordered lane and the + // batch sizes are deterministic. + msgs[i] = memsource.Msg{Key: "k", Value: []byte(fmt.Sprintf("%d", i))} + } + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(tt.size, 0)}, + msgs..., + ) + var batches int32 + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + atomic.AddInt32(&batches, 1) + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: tt.count}) + if int(batches) != tt.wantBatches { + t.Fatalf("batches = %d, want %d", batches, tt.wantBatches) + } + }) + } +} + +// TestHopper_BatchMaxWaitTrigger drives a single message into a lane whose size +// is larger than the queue and asserts the max-wait timer flushes the partial +// batch. An injected clock is not needed for correctness here (the timer is real +// and short), but the partial flush must happen before drain closes the lane. +func TestHopper_BatchMaxWaitTrigger(t *testing.T) { + t.Parallel() + // size 100, but only 3 messages: the batch can only be released by the + // max-wait timer or by the drain. A short max-wait releases it promptly. + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(100, 5*time.Millisecond)}, + memsource.Msg{Key: "k", Value: []byte("0")}, + memsource.Msg{Key: "k", Value: []byte("1")}, + memsource.Msg{Key: "k", Value: []byte("2")}, + ) + var maxSeen int + var mu sync.Mutex + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + mu.Lock() + if len(ms) > maxSeen { + maxSeen = len(ms) + } + mu.Unlock() + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: 3}) + if maxSeen == 0 || maxSeen > 3 { + t.Fatalf("max batch = %d, want 1..3", maxSeen) + } +} + +// TestHopper_BatchPartialOnDrain verifies a batch that never reaches size and has +// no max-wait still flushes when the run drains, losing nothing. +func TestHopper_BatchPartialOnDrain(t *testing.T) { + t.Parallel() + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(1000, 0)}, // never fills, no timer + memsource.Msg{Key: "k", Value: []byte("a")}, + memsource.Msg{Key: "k", Value: []byte("b")}, + ) + var got int + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + got = len(ms) + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: 2}) + if got != 2 { + t.Fatalf("drain batch = %d, want 2", got) + } +} + +// TestHopper_BatchPerMessageResultMapping checks the positional Result-to-message +// contract: a mixed slice of ack/nak/term/skip results settles each message by +// its own result. +func TestHopper_BatchPerMessageResultMapping(t *testing.T) { + t.Parallel() + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(4, 0)}, + memsource.Msg{Key: "k", Value: []byte("ack")}, + memsource.Msg{Key: "k", Value: []byte("nak")}, + memsource.Msg{Key: "k", Value: []byte("term")}, + memsource.Msg{Key: "k", Value: []byte("skip")}, + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + for i, m := range ms { + switch string(m.Value()) { + case "ack": + res[i] = source.Ack() + case "nak": + res[i] = source.Nak(errors.New("retry")) + case "term": + res[i] = source.Term(errors.New("bad")) + case "skip": + res[i] = source.Skip() + } + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: 1, Nak: 1, Term: 1, Dropped: 1}) +} + +// TestHopper_BatchOrderingPreserved confirms that within a key, the batch handler +// sees messages in delivery order and across batches the sequence never breaks. +func TestHopper_BatchOrderingPreserved(t *testing.T) { + t.Parallel() + const perKey = 40 + const keys = 6 + var msgs []memsource.Msg + for i := 0; i < perKey; i++ { + for k := 0; k < keys; k++ { + msgs = append(msgs, memsource.Msg{ + Key: fmt.Sprintf("key-%d", k), + Value: []byte(fmt.Sprintf("%d", i)), + }) + } + } + + var mu sync.Mutex + lastSeen := make(map[string]int) + for k := 0; k < keys; k++ { + lastSeen[fmt.Sprintf("key-%d", k)] = -1 + } + + h := memsource.NewHarness(t, + []source.Option{source.WithConcurrency(keys), source.WithBatch(7, 0)}, + msgs..., + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + mu.Lock() + for i, m := range ms { + key := string(m.Key()) + var seq int + _, _ = fmt.Sscanf(string(m.Value()), "%d", &seq) + if seq != lastSeen[key]+1 { + t.Errorf("key %s out of order: %d after %d", key, seq, lastSeen[key]) + } + lastSeen[key] = seq + res[i] = source.Ack() + } + mu.Unlock() + return res + }) + h.AssertCounts(memsource.Counts{Acked: perKey * keys}) +} + +// TestHopper_BatchResultCountMismatch verifies a handler returning too few results +// terminates the unmatched messages as poison rather than stranding them. +func TestHopper_BatchResultCountMismatch(t *testing.T) { + t.Parallel() + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(3, 0)}, + memsource.Msg{Key: "k"}, + memsource.Msg{Key: "k"}, + memsource.Msg{Key: "k"}, + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + // Return one result for three messages. + return []source.Result{source.Ack()} + }) + // First message acked; the two without a result are termed as poison. + h.AssertCounts(memsource.Counts{Acked: 1, Term: 2}) + for _, e := range h.Ledger().Entries() { + if e.Result.Action == source.ActionTerm && !errors.Is(e.Result.Err, source.ErrBatchResultCount) { + t.Fatalf("term err = %v, want ErrBatchResultCount", e.Result.Err) + } + } +} + +// TestHopper_BatchDecodeFailureIsolated checks one undecodable message terminates +// alone without poisoning its batch-mates. +func TestHopper_BatchDecodeFailureIsolated(t *testing.T) { + t.Parallel() + reg := source.NewRegistry().SetDefault(source.NewJSONCodec[order]()) + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(3, 0), source.WithRegistry(reg)}, + memsource.Msg{Key: "k", Value: []byte(`{"id":"A","qty":1}`)}, + memsource.Msg{Key: "k", Value: []byte(`{bad`)}, + memsource.Msg{Key: "k", Value: []byte(`{"id":"C","qty":3}`)}, + ) + var handled int + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + handled = len(ms) + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + if handled != 2 { + t.Fatalf("handler saw %d messages, want 2 (decode failure excluded)", handled) + } + h.AssertCounts(memsource.Counts{Acked: 2, Term: 1}) +} + +// TestHopper_BatchedCapabilityPath drives the whole-batch fetch path via a +// Batched memsource subscription and asserts every message is delivered and acked. +func TestHopper_BatchedCapabilityPath(t *testing.T) { + t.Parallel() + const total = 25 + msgs := make([]memsource.Msg, total) + for i := range msgs { + msgs[i] = memsource.Msg{Key: "k", Value: []byte(fmt.Sprintf("%d", i))} + } + h := memsource.NewHarnessWith(t, + []source.Option{source.WithBatch(8, 0)}, + []memsource.Option{memsource.WithBatched()}, + msgs..., + ) + var seen int32 + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + atomic.AddInt32(&seen, int32(len(ms))) + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: total}) + if int(seen) != total { + t.Fatalf("handler saw %d messages, want %d", seen, total) + } +} + +// TestHopper_BatchMaxWaitWithInjectedClock proves WithBatchClock is wired: the +// clock is read while batching (it has no effect on outcome here, but the option +// must be accepted and the run must complete deterministically). +func TestHopper_BatchMaxWaitWithInjectedClock(t *testing.T) { + t.Parallel() + var ticks int32 + clock := func() time.Time { + atomic.AddInt32(&ticks, 1) + return time.Unix(0, 0) + } + h := memsource.NewHarness(t, + []source.Option{source.WithBatch(2, 0), source.WithBatchClock(clock)}, + memsource.Msg{Key: "k"}, memsource.Msg{Key: "k"}, + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertCounts(memsource.Counts{Acked: 2}) +} + +// TestHopper_BatchRespectsMaxInFlight verifies the per-message backpressure bound +// still applies in batch mode: a lane never holds more than maxInFlight unsettled +// messages even with a larger batch size. +func TestHopper_BatchRespectsMaxInFlight(t *testing.T) { + t.Parallel() + const total = 30 + var msgs []memsource.Msg + for i := 0; i < total; i++ { + msgs = append(msgs, memsource.Msg{Key: fmt.Sprintf("k%d", i)}) + } + // size 5 per lane, maxInFlight 10 across 8 lanes: well within bounds, and a + // positive max-wait guarantees lanes flush even when they cannot fill. + h := memsource.NewHarness(t, + []source.Option{ + source.WithConcurrency(8), + source.WithMaxInFlight(10), + source.WithBatch(5, 2*time.Millisecond), + }, + msgs..., + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + h.AssertSettled(total) +} diff --git a/source/errors.go b/source/errors.go index 44bd231..a1337ee 100644 --- a/source/errors.go +++ b/source/errors.go @@ -33,6 +33,13 @@ var ErrPoison = errors.New("source: poison message") // errors.Is rather than by string-matching the underlying error. var ErrRetryable = errors.New("source: retryable error") +// ErrBatchResultCount reports that a [BatchHandler] returned a result slice +// whose length does not match the batch it was given. The Hopper settles the +// messages it can pair with a result and terminates any unmatched message as +// poison carrying this sentinel, so a miscounted batch fails loudly rather than +// stranding messages. Match it with errors.Is. +var ErrBatchResultCount = errors.New("source: batch handler returned wrong number of results") + // ErrInvalidForState classifies a well-formed message that is not legal for its // target's current state — a guard rejection from the state-machine bridge. It // is the errors.Is-matchable sentinel for the [InvalidForState] classification, diff --git a/source/example_test.go b/source/example_test.go index ea7db5e..5a6eaee 100644 --- a/source/example_test.go +++ b/source/example_test.go @@ -41,6 +41,42 @@ func ExampleHopper() { // order B qty 5 } +// ExampleHopper_batch shows batch consume: WithBatch accumulates up to size +// messages per ordered lane, the BatchHandler processes them in one call, and the +// engine settles each message by its corresponding Result (positionally aligned). +func ExampleHopper_batch() { + in := memsource.New(memsource.WithMessages( + memsource.Msg{Key: "A", Value: []byte(`{"id":"A","qty":2}`)}, + memsource.Msg{Key: "A", Value: []byte(`{"id":"A","qty":3}`)}, + memsource.Msg{Key: "A", Value: []byte(`{"id":"A","qty":5}`)}, + )) + + hp := source.New( + source.WithCodec(source.NewJSONCodec[orderPlaced]()), + source.WithBatch(2, 0), // batches of up to 2; the third flushes on drain + ) + + sub, _ := in.Subscribe(context.Background(), source.SubscribeConfig{Topics: []string{"orders"}}) + _ = sub.Close() + + _ = hp.RunBatch(context.Background(), sub, func(_ context.Context, ms []source.Message) []source.Result { + total := 0 + for _, m := range ms { + v, _ := source.Decoded(m) + total += v.(orderPlaced).Qty + } + fmt.Printf("batch of %d, qty sum %d\n", len(ms), total) + res := make([]source.Result, len(ms)) + for i := range res { + res[i] = source.Ack() + } + return res + }) + // Output: + // batch of 2, qty sum 5 + // batch of 1, qty sum 5 +} + // ExampleChain shows middleware composition: the first middleware listed is the // outermost, so a message flows outer to inner and the result returns inner to // outer. diff --git a/source/fuzz_test.go b/source/fuzz_test.go index 64dd543..7ac017e 100644 --- a/source/fuzz_test.go +++ b/source/fuzz_test.go @@ -8,6 +8,7 @@ import ( "strconv" "sync" "testing" + "time" "unicode/utf8" "github.com/stablekernel/crucible/source" @@ -85,6 +86,87 @@ func FuzzOrderedSettle(f *testing.F) { }) } +// FuzzBatchOrderedSettle hammers the batch-mode ordering invariant: across any +// interleaving of keys, any concurrency bound, and any batch size, every message +// is handled exactly once (no loss, no duplication) and within a single key the +// messages are processed in delivery order — the same contract as the per-message +// path, now through accumulated batches. +func FuzzBatchOrderedSettle(f *testing.F) { + f.Add(uint8(10), uint8(3), uint8(4), uint8(2)) + f.Add(uint8(64), uint8(8), uint8(1), uint8(7)) + f.Add(uint8(1), uint8(1), uint8(1), uint8(1)) + f.Add(uint8(100), uint8(16), uint8(7), uint8(16)) + f.Add(uint8(0), uint8(4), uint8(2), uint8(3)) + + f.Fuzz(func(t *testing.T, count, conc, keyspace, size uint8) { + if conc == 0 { + conc = 1 + } + if keyspace == 0 { + keyspace = 1 + } + if size == 0 { + size = 1 + } + + perKeyCount := make(map[string]int) + var msgs []memsource.Msg + for i := 0; i < int(count); i++ { + key := "k" + strconv.Itoa(i%int(keyspace)) + seq := perKeyCount[key] + perKeyCount[key]++ + msgs = append(msgs, memsource.Msg{ + Key: key, + Value: []byte(strconv.Itoa(seq)), + }) + } + + var mu sync.Mutex + processed := make(map[string][]int) + + h := memsource.NewHarness( + t, + []source.Option{ + source.WithConcurrency(int(conc)), + // A short max-wait guarantees forward progress regardless of how the + // keys distribute across lanes (a lane that cannot fill still flushes). + source.WithBatch(int(size), time.Millisecond), + }, + msgs..., + ) + h.RunBatch(func(_ context.Context, ms []source.Message) []source.Result { + res := make([]source.Result, len(ms)) + mu.Lock() + for i, m := range ms { + seq, err := strconv.Atoi(string(m.Value())) + if err != nil { + res[i] = source.Term(err) + continue + } + processed[string(m.Key())] = append(processed[string(m.Key())], seq) + res[i] = source.Ack() + } + mu.Unlock() + return res + }) + + if got := h.Ledger().Len(); got != int(count) { + t.Fatalf("settled %d messages, want %d (no loss or duplication)", got, count) + } + for key, want := range perKeyCount { + got := processed[key] + if len(got) != want { + t.Fatalf("key %s processed %d messages, want %d", key, len(got), want) + } + for i, seq := range got { + if seq != i { + t.Fatalf("key %s out of order: position %d has seq %d", key, i, seq) + } + } + } + }) +} + // FuzzCodecRoundTrip checks the JSON codec round-trip: any pair of fields // marshals and decodes back to an equal value, and the registry classifies a // non-JSON payload as a decode failure (poison) rather than panicking. diff --git a/source/handler.go b/source/handler.go index 46ad18a..f66c4bd 100644 --- a/source/handler.go +++ b/source/handler.go @@ -14,6 +14,25 @@ import ( // goroutines. type Handler func(ctx context.Context, m Message) Result +// BatchHandler processes a slice of decoded messages at once and returns one +// [Result] per message, positionally aligned: the Result at index i settles the +// message at index i. It is the batch analog of [Handler], invoked by a Hopper +// run in batch mode (see [WithBatch] and [Hopper.RunBatch]); the engine +// accumulates messages per ordered lane up to a size or a max-wait bound, hands +// them to the BatchHandler in one call, then settles each message by its +// corresponding Result. +// +// A BatchHandler MUST return exactly len(ms) results. Returning a different +// length is a programming error: the engine settles the messages it can pair +// with a result and terminates any leftover as poison (see [ErrBatchResultCount]) +// so a miscounted batch never silently strands messages. +// +// Like [Handler], a BatchHandler must be safe for concurrent use: distinct lanes +// invoke it in parallel. Messages within a single call always share one ordered +// lane and arrive in delivery order, so a BatchHandler may rely on per-key +// ordering within ms. +type BatchHandler func(ctx context.Context, ms []Message) []Result + // Action is the disposition the Hopper applies to a message after its handler // returns. It maps onto each backend's native settle vocabulary. type Action uint8 diff --git a/source/hopper.go b/source/hopper.go index 9680113..5bbfcfc 100644 --- a/source/hopper.go +++ b/source/hopper.go @@ -44,6 +44,7 @@ type Hopper struct { middleware []Middleware concurrency int maxInFlight int + batch batchConfig received telemetry.Counter acked telemetry.Counter @@ -51,6 +52,7 @@ type Hopper struct { term telemetry.Counter dropped telemetry.Counter failed telemetry.Counter + batches telemetry.Counter lag telemetry.Gauge closeOnce sync.Once @@ -74,12 +76,14 @@ func New(opts ...Option) *Hopper { middleware: append([]Middleware(nil), cfg.middleware...), concurrency: cfg.concurrency, maxInFlight: cfg.maxInFlight, + batch: cfg.batch, received: m.Counter("source.received", telemetry.WithDescription("messages received from the subscription")), acked: m.Counter("source.acked", telemetry.WithDescription("messages acknowledged after successful handling")), nak: m.Counter("source.nak", telemetry.WithDescription("messages nak'd for redelivery")), term: m.Counter("source.term", telemetry.WithDescription("messages terminated (dead-lettered)")), dropped: m.Counter("source.dropped", telemetry.WithDescription("messages acked and discarded as out of scope")), failed: m.Counter("source.failed", telemetry.WithDescription("settle failures observed after handling")), + batches: m.Counter("source.batches", telemetry.WithDescription("batches handed to a batch handler")), lag: m.Gauge("source.lag", telemetry.WithUnit("{message}"), telemetry.WithDescription("unconsumed messages behind the stream tail")), closed: make(chan struct{}), } @@ -269,16 +273,12 @@ func (hp *Hopper) process(ctx context.Context, handler Handler, d *delivery) { // Decode (when a registry is configured) and attach the value to the message // the handler sees. A decode failure is poison: terminate, do not retry. - dec := m - if hp.registry != nil { - v, err := hp.registry.Decode(m) - if err != nil { - span.RecordError(err) - span.SetStatus(telemetry.StatusError, "decode failed") - hp.settle(ctx, span, d, Term(err)) - return - } - dec = &decoded{Message: m, value: v} + dec, decErr := hp.decodeFor(m) + if decErr != nil { + span.RecordError(decErr) + span.SetStatus(telemetry.StatusError, "decode failed") + hp.settle(ctx, span, d, Term(decErr)) + return } r := handler(ctx, dec) @@ -292,6 +292,21 @@ func (hp *Hopper) process(ctx context.Context, handler Handler, d *delivery) { hp.settle(ctx, span, d, r) } +// decodeFor decodes m through the configured registry and wraps it so the +// handler reads the typed value via [Decoded]. With no registry it returns m +// unchanged. A decode failure is poison: the caller terminates the message +// rather than retrying. +func (hp *Hopper) decodeFor(m Message) (Message, error) { + if hp.registry == nil { + return m, nil + } + v, err := hp.registry.Decode(m) + if err != nil { + return nil, err + } + return &decoded{Message: m, value: v}, nil +} + // settle applies r to the message via the subscription, counts the outcome, and // logs/marks a settle failure. func (hp *Hopper) settle(ctx context.Context, span telemetry.Span, d *delivery, r Result) { diff --git a/source/jetstream/batch_test.go b/source/jetstream/batch_test.go new file mode 100644 index 0000000..2c3491e --- /dev/null +++ b/source/jetstream/batch_test.go @@ -0,0 +1,220 @@ +// SPDX-License-Identifier: Apache-2.0 + +package jetstream + +import ( + "context" + "errors" + "testing" + + njs "github.com/nats-io/nats.go/jetstream" + + "github.com/stablekernel/crucible/source" +) + +// --- NextBatch: block-for-first then drain-buffered -------------------------- + +func TestNextBatch_ClampsLimitBelowOne(t *testing.T) { + t.Parallel() + // A non-positive limit clamps to 1, so NextBatch blocks for and returns + // exactly the first message rather than an empty slice. + js := newFakeJS( + &fakeMsg{subject: "orders.placed", data: []byte("a"), seq: 1}, + &fakeMsg{subject: "orders.paid", data: []byte("b"), seq: 2}, + ) + sub := newSub(t, js) + + bs, ok := sub.(source.Batched) + if !ok { + t.Fatal("subscription does not satisfy source.Batched") + } + msgs, err := bs.NextBatch(context.Background(), 0) + if err != nil { + t.Fatalf("NextBatch(0) error = %v", err) + } + if len(msgs) != 1 { + t.Fatalf("NextBatch(0) = %d messages, want 1 (clamped)", len(msgs)) + } + if string(msgs[0].Value()) != "a" { + t.Errorf("first message = %q, want a", msgs[0].Value()) + } +} + +func TestNextBatch_DrainsBufferedUpToLimit(t *testing.T) { + t.Parallel() + // The pull consumer has prefetched several messages; NextBatch blocks for the + // first then opportunistically drains the rest the iterator already holds, + // returning them in arrival order. + js := newFakeJS( + &fakeMsg{subject: "orders.placed", data: []byte("a"), seq: 1}, + &fakeMsg{subject: "orders.paid", data: []byte("b"), seq: 2}, + &fakeMsg{subject: "orders.shipped", data: []byte("c"), seq: 3}, + ) + sub := newSub(t, js) + + bs := sub.(source.Batched) + msgs, err := bs.NextBatch(context.Background(), 10) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + if len(msgs) != 3 { + t.Fatalf("NextBatch() = %d messages, want 3 drained", len(msgs)) + } + want := []string{"a", "b", "c"} + for i, m := range msgs { + if string(m.Value()) != want[i] { + t.Errorf("message[%d] = %q, want %q", i, m.Value(), want[i]) + } + } +} + +func TestNextBatch_StopsAtLimit(t *testing.T) { + t.Parallel() + // More messages are buffered than requested: NextBatch returns exactly limit + // and leaves the remainder for the next call. + js := newFakeJS( + &fakeMsg{subject: "s", data: []byte("a"), seq: 1}, + &fakeMsg{subject: "s", data: []byte("b"), seq: 2}, + &fakeMsg{subject: "s", data: []byte("c"), seq: 3}, + &fakeMsg{subject: "s", data: []byte("d"), seq: 4}, + ) + sub := newSub(t, js) + bs := sub.(source.Batched) + + first, err := bs.NextBatch(context.Background(), 2) + if err != nil { + t.Fatalf("first NextBatch() error = %v", err) + } + if len(first) != 2 { + t.Fatalf("first NextBatch() = %d messages, want 2", len(first)) + } + if string(first[0].Value()) != "a" || string(first[1].Value()) != "b" { + t.Errorf("first batch = %q/%q, want a/b", first[0].Value(), first[1].Value()) + } + + second, err := bs.NextBatch(context.Background(), 2) + if err != nil { + t.Fatalf("second NextBatch() error = %v", err) + } + if len(second) != 2 || string(second[0].Value()) != "c" || string(second[1].Value()) != "d" { + t.Errorf("second batch = %v, want c/d", second) + } +} + +func TestNextBatch_FirstMessageDrainPropagatesError(t *testing.T) { + t.Parallel() + // When the first blocking read fails to create the consumer, NextBatch + // surfaces that error rather than returning a partial batch. + js := &fakeJS{createErr: errors.New("boom")} + sub := newSub(t, js) + bs := sub.(source.Batched) + + if _, err := bs.NextBatch(context.Background(), 4); err == nil || !errContains(err, "create consumer") { + t.Fatalf("NextBatch() = %v, want create-consumer error", err) + } +} + +func TestNextBatch_DrainBreaksAfterClose(t *testing.T) { + t.Parallel() + // A single buffered message is returned; the drain loop then sees the closed + // subscription and returns the one message it has rather than blocking. + js := newFakeJS(&fakeMsg{subject: "s", data: []byte("only"), seq: 1}) + sub := newSub(t, js) + bs := sub.(source.Batched) + + msgs, err := bs.NextBatch(context.Background(), 8) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + // The fake iterator yields one message then reports closed, so the drain + // loop breaks and the batch holds exactly that message. + if len(msgs) != 1 || string(msgs[0].Value()) != "only" { + t.Fatalf("NextBatch() = %v, want a single message", msgs) + } +} + +func TestNextBatch_DrainedReturnsErrDrained(t *testing.T) { + t.Parallel() + // A closed subscription has no first message to block for, so NextBatch maps + // the drain exactly as Next does. + sub := newSub(t, newFakeJS(&fakeMsg{subject: "s", data: []byte("a")})) + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + bs := sub.(source.Batched) + if _, err := bs.NextBatch(context.Background(), 4); !errors.Is(err, source.ErrDrained) { + t.Errorf("NextBatch() after close = %v, want ErrDrained", err) + } +} + +// --- SettleBatch: per-message ack vocabulary --------------------------------- + +func TestSettleBatch_AcksEveryMessage(t *testing.T) { + t.Parallel() + // One Ack result applied to a batch double-acks every message. + fms := []*fakeMsg{ + {subject: "s", data: []byte("a"), seq: 1}, + {subject: "s", data: []byte("b"), seq: 2}, + {subject: "s", data: []byte("c"), seq: 3}, + } + msgs := make([]njs.Msg, len(fms)) + for i, fm := range fms { + msgs[i] = fm + } + js := newFakeJS(msgs...) + sub := newSub(t, js) + bs := sub.(source.Batched) + + batch, err := bs.NextBatch(context.Background(), 3) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + if len(batch) != 3 { + t.Fatalf("NextBatch() = %d, want 3", len(batch)) + } + if err := bs.SettleBatch(context.Background(), batch, source.Ack()); err != nil { + t.Fatalf("SettleBatch() error = %v", err) + } + for i, fm := range fms { + if !fm.doubleAck { + t.Errorf("message[%d] doubleAck = false, want true", i) + } + } +} + +func TestSettleBatch_ReturnsFirstErrorButSettlesAll(t *testing.T) { + t.Parallel() + // Every ack fails; SettleBatch returns the first wrapped error yet still + // attempts every message in the batch. + boom := errors.New("server unreachable") + fms := []*fakeMsg{ + {subject: "s", data: []byte("a"), seq: 1, settleErr: boom}, + {subject: "s", data: []byte("b"), seq: 2, settleErr: boom}, + } + msgs := []njs.Msg{fms[0], fms[1]} + sub := newSub(t, newFakeJS(msgs...)) + bs := sub.(source.Batched) + + batch, err := bs.NextBatch(context.Background(), 2) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + err = bs.SettleBatch(context.Background(), batch, source.Ack()) + if !errors.Is(err, boom) { + t.Fatalf("SettleBatch() = %v, want wrapped %v", err, boom) + } + for i, fm := range fms { + if !fm.acked { + t.Errorf("message[%d] not attempted, want every message settled", i) + } + } +} + +func TestSettleBatch_EmptyIsNoError(t *testing.T) { + t.Parallel() + sub := newSub(t, newFakeJS()) + bs := sub.(source.Batched) + if err := bs.SettleBatch(context.Background(), nil, source.Ack()); err != nil { + t.Errorf("SettleBatch(nil) = %v, want nil", err) + } +} diff --git a/source/jetstream/jetstream.go b/source/jetstream/jetstream.go index df48a9d..a0878fa 100644 --- a/source/jetstream/jetstream.go +++ b/source/jetstream/jetstream.go @@ -413,6 +413,63 @@ func (s *subscription) Settle(ctx context.Context, m source.Message, r source.Re } } +// NextBatch returns up to limit messages, satisfying [source.Batched]. It blocks +// for at least one message (like [Next]), then opportunistically drains further +// messages the pull consumer has already buffered (bounded by WithPullMaxMessages) +// without blocking for slow arrivals, so a batch reflects what is ready rather +// than waiting to fill. Each message is settled individually through +// [SettleBatch] or the engine's per-message Settle. NextBatch is single-consumer, +// like Next. +func (s *subscription) NextBatch(ctx context.Context, limit int) ([]source.Message, error) { + if limit < 1 { + limit = 1 + } + // Block for the first message, mapping drain/cancel exactly as Next does. + first, err := s.Next(ctx) + if err != nil { + return nil, err + } + msgs := make([]source.Message, 0, limit) + msgs = append(msgs, first) + + // Drain already-buffered messages without blocking for new arrivals: a tiny + // deadline turns the iterator's wait into a non-blocking poll, so we collect + // what the pull consumer prefetched and return promptly. + for len(msgs) < limit { + s.mu.Lock() + iter := s.iter + closed := s.closed + s.mu.Unlock() + if iter == nil || closed { + break + } + pollCtx, cancel := context.WithTimeout(ctx, time.Millisecond) + msg, perr := iter.Next(jetstream.NextContext(pollCtx)) + cancel() + if perr != nil { + // No more ready right now (deadline), or the stream ended: return what + // we have. A real error surfaces on the next NextBatch call. + break + } + msgs = append(msgs, newMessage(msg)) + } + return msgs, nil +} + +// SettleBatch applies r to every message in ms in one call, satisfying +// [source.Batched]. It settles each message through the same per-message ack +// vocabulary as [Settle], returning the first error encountered while still +// attempting every message. +func (s *subscription) SettleBatch(ctx context.Context, ms []source.Message, r source.Result) error { + var firstErr error + for _, m := range ms { + if err := s.Settle(ctx, m, r); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + // Close stops the message iterator and marks the subscription drained. It is // idempotent. In-flight messages already returned by Next are left for the // caller to settle; once stopped, Next returns [source.ErrDrained]. @@ -532,4 +589,5 @@ var ( _ source.Seekable = (*subscription)(nil) _ source.OrderedDelivery = (*subscription)(nil) _ source.LagReporter = (*subscription)(nil) + _ source.Batched = (*subscription)(nil) ) diff --git a/source/kafka/batch_test.go b/source/kafka/batch_test.go new file mode 100644 index 0000000..b60ddf8 --- /dev/null +++ b/source/kafka/batch_test.go @@ -0,0 +1,278 @@ +// 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" +) + +// manyFetch wraps several records on one partition into a single fetch batch, +// the shape PollRecords returns when franz-go hands back a whole poll at once. +func manyFetch(recs ...*kgo.Record) kgo.Fetches { + parts := make([]kgo.FetchPartition, 0, 1) + parts = append(parts, kgo.FetchPartition{ + Partition: recs[0].Partition, + Records: recs, + }) + return kgo.Fetches{{ + Topics: []kgo.FetchTopic{{Topic: recs[0].Topic, Partitions: parts}}, + }} +} + +func TestTakeBatchPopsUpToLimit(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + buffer int + limit int + wantN int + wantRest int + }{ + {"limit below buffer takes limit", 5, 2, 2, 3}, + {"limit equals buffer takes all", 3, 3, 3, 0}, + {"limit above buffer takes buffer", 2, 9, 2, 0}, + {"limit one takes one", 4, 1, 1, 3}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + sub := newSub(&fakePoller{}) + buf := make([]*kgo.Record, tt.buffer) + for i := range buf { + buf[i] = rec("orders", 0, int64(i), "k", "v") + } + sub.buffer = buf + + got, ok := sub.takeBatch(tt.limit) + if !ok { + t.Fatal("takeBatch() = _,false, want true with a populated buffer") + } + if len(got) != tt.wantN { + t.Errorf("took %d records, want %d", len(got), tt.wantN) + } + if len(sub.buffer) != tt.wantRest { + t.Errorf("buffer left = %d, want %d", len(sub.buffer), tt.wantRest) + } + if sub.inFlight != tt.wantN { + t.Errorf("inFlight = %d, want %d", sub.inFlight, tt.wantN) + } + }) + } +} + +func TestTakeBatchEmptyBufferReportsFalse(t *testing.T) { + t.Parallel() + + sub := newSub(&fakePoller{}) + if recs, ok := sub.takeBatch(4); ok || recs != nil { + t.Errorf("takeBatch(empty) = %v,%v, want nil,false", recs, ok) + } + if sub.inFlight != 0 { + t.Errorf("inFlight = %d, want 0 on an empty buffer", sub.inFlight) + } +} + +func TestNextBatchClampsLimitBelowOne(t *testing.T) { + t.Parallel() + + // A non-positive limit clamps to 1, so a buffered batch hands back exactly + // one record rather than panicking on a zero/negative slice bound. + sub := newSub(&fakePoller{}) + sub.buffer = []*kgo.Record{ + rec("orders", 0, 1, "k", "v1"), + rec("orders", 0, 2, "k", "v2"), + } + + msgs, err := sub.NextBatch(context.Background(), 0) + if err != nil { + t.Fatalf("NextBatch(0) error = %v", err) + } + if len(msgs) != 1 { + t.Fatalf("NextBatch(0) = %d messages, want 1 (clamped)", len(msgs)) + } + if string(msgs[0].Value()) != "v1" { + t.Errorf("first message = %q, want v1", msgs[0].Value()) + } + if len(sub.buffer) != 1 || sub.inFlight != 1 { + t.Errorf("buffer=%d inFlight=%d, want 1/1", len(sub.buffer), sub.inFlight) + } +} + +func TestNextBatchReturnsBufferedGroupInOrder(t *testing.T) { + t.Parallel() + + // A pre-populated buffer is served without polling: NextBatch returns the + // records as a group, in arrival order, and counts them all in flight. + sub := newSub(&fakePoller{}) + sub.buffer = []*kgo.Record{ + rec("orders", 1, 10, "A-1", "placed"), + rec("orders", 1, 11, "A-2", "paid"), + rec("orders", 1, 12, "A-3", "shipped"), + } + + msgs, err := sub.NextBatch(context.Background(), 10) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + if len(msgs) != 3 { + t.Fatalf("NextBatch() = %d messages, want 3", len(msgs)) + } + wantKeys := []string{"A-1", "A-2", "A-3"} + for i, m := range msgs { + if string(m.Key()) != wantKeys[i] { + t.Errorf("message[%d] key = %q, want %q", i, m.Key(), wantKeys[i]) + } + } + if sub.inFlight != 3 { + t.Errorf("inFlight = %d, want 3", sub.inFlight) + } +} + +func TestNextBatchPollsWhenBufferEmpty(t *testing.T) { + t.Parallel() + + // With an empty buffer NextBatch falls through to a poll, buffers the fetched + // records, then hands back up to limit of them. It also releases a rebalance + // before polling, mirroring the single-record Next path. + r0 := rec("orders", 0, 1, "A-1", "placed") + r1 := rec("orders", 0, 2, "A-2", "paid") + fp := &fakePoller{fetches: []kgo.Fetches{manyFetch(r0, r1)}} + sub := newSub(fp) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + msgs, err := sub.NextBatch(ctx, 5) + if err != nil { + t.Fatalf("NextBatch() error = %v", err) + } + if len(msgs) != 2 { + t.Fatalf("NextBatch() = %d messages, want 2 from the poll", len(msgs)) + } + if string(msgs[0].Value()) != "placed" || string(msgs[1].Value()) != "paid" { + t.Errorf("values = %q/%q, want placed/paid", msgs[0].Value(), msgs[1].Value()) + } + if fp.allowed < 1 { + t.Errorf("AllowRebalance called %d times, want at least 1 before the poll", fp.allowed) + } + if sub.inFlight != 2 { + t.Errorf("inFlight = %d, want 2", sub.inFlight) + } +} + +func TestNextBatchPollErrorPropagates(t *testing.T) { + t.Parallel() + + // A non-context fetch error surfaces wrapped, the same as the single-record + // poll path, so the engine can decide whether to retry. + boom := errors.New("partition unavailable") + fp := &fakePoller{fetches: []kgo.Fetches{{{ + Topics: []kgo.FetchTopic{{ + Topic: "orders", + Partitions: []kgo.FetchPartition{{Partition: 0, Err: boom}}, + }}, + }}}} + sub := newSub(fp) + + if _, err := sub.NextBatch(context.Background(), 4); !errors.Is(err, boom) { + t.Fatalf("NextBatch() = %v, want wrapped %v", err, boom) + } +} + +func TestNextBatchDrainedAfterClose(t *testing.T) { + t.Parallel() + + // A closed, fully-settled subscription with an empty buffer drains rather + // than polling. + sub := newSub(&fakePoller{}) + if err := sub.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if _, err := sub.NextBatch(context.Background(), 4); !errors.Is(err, source.ErrDrained) { + t.Errorf("NextBatch() after drain = %v, want ErrDrained", err) + } +} + +func TestNextBatchRespectsContextCancel(t *testing.T) { + t.Parallel() + + // An empty buffer and a quiet broker leave NextBatch blocked on the poll; + // the deadline must end the wait with the context error. + sub := newSub(&fakePoller{}) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + if _, err := sub.NextBatch(ctx, 4); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("NextBatch() = %v, want DeadlineExceeded", err) + } +} + +func TestSettleBatchSettlesEveryMessage(t *testing.T) { + t.Parallel() + + // One Ack result applied to a batch marks every record for commit and clears + // the in-flight accounting for the whole group. + fp := &fakePoller{} + sub := newSub(fp) + sub.inFlight = 3 + ms := []source.Message{ + newMessage(rec("orders", 0, 1, "k", "v1")), + newMessage(rec("orders", 0, 2, "k", "v2")), + newMessage(rec("orders", 0, 3, "k", "v3")), + } + + if err := sub.SettleBatch(context.Background(), ms, source.Ack()); err != nil { + t.Fatalf("SettleBatch() error = %v", err) + } + if fp.markedCount() != 3 { + t.Errorf("marked = %d, want 3 (every record acked)", fp.markedCount()) + } + if sub.inFlight != 0 { + t.Errorf("inFlight = %d, want 0 after settling the whole batch", sub.inFlight) + } +} + +func TestSettleBatchReturnsFirstErrorButSettlesAll(t *testing.T) { + t.Parallel() + + // A Term batch with no DLQ topic configured fails on every message; the + // first error is returned, yet every message is still attempted so the + // in-flight count is fully cleared. + fp := &fakePoller{} + sub := &subscription{client: fp} // no dlqTopic configured + sub.inFlight = 2 + ms := []source.Message{ + newMessage(rec("orders", 0, 1, "k", "v1")), + newMessage(rec("orders", 0, 2, "k", "v2")), + } + + err := sub.SettleBatch(context.Background(), ms, source.Term(errors.New("poison"))) + if !errors.Is(err, ErrNoDLQTopic) { + t.Fatalf("SettleBatch() = %v, want ErrNoDLQTopic", err) + } + if fp.markedCount() != 0 { + t.Errorf("marked = %d, want 0 (nothing dead-lettered, nothing committed)", fp.markedCount()) + } + if sub.inFlight != 0 { + t.Errorf("inFlight = %d, want 0 (every message attempted)", sub.inFlight) + } +} + +func TestSettleBatchEmptyIsNoError(t *testing.T) { + t.Parallel() + + sub := newSub(&fakePoller{}) + if err := sub.SettleBatch(context.Background(), nil, source.Ack()); err != nil { + t.Errorf("SettleBatch(nil) = %v, want nil", err) + } +} diff --git a/source/kafka/subscription.go b/source/kafka/subscription.go index 3301291..dae3e81 100644 --- a/source/kafka/subscription.go +++ b/source/kafka/subscription.go @@ -119,6 +119,89 @@ func (s *subscription) takeBuffered() (*kgo.Record, bool) { return rec, true } +// takeBatch pops up to limit records from the buffer and counts them in flight. +// It reports false only when the buffer is empty, so [NextBatch] can fall through +// to a poll. +func (s *subscription) takeBatch(limit int) ([]*kgo.Record, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.buffer) == 0 { + return nil, false + } + n := limit + if n > len(s.buffer) { + n = len(s.buffer) + } + recs := s.buffer[:n:n] + s.buffer = s.buffer[n:] + s.inFlight += n + return recs, true +} + +// NextBatch returns up to limit buffered records, polling the broker when the +// 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. +func (s *subscription) NextBatch(ctx context.Context, limit int) ([]source.Message, error) { + if limit < 1 { + limit = 1 + } + for { + if recs, ok := s.takeBatch(limit); ok { + msgs := make([]source.Message, len(recs)) + for i, rec := range recs { + msgs[i] = newMessage(rec) + } + return msgs, nil + } + + s.mu.Lock() + drained := s.closed && s.inFlight == 0 + s.mu.Unlock() + if drained { + return nil, source.ErrDrained + } + if err := ctx.Err(); err != nil { + return nil, err + } + + s.client.AllowRebalance() + fetches := s.client.PollRecords(ctx, s.maxPoll) + if err := ctx.Err(); err != nil { + return nil, err + } + if errs := fetches.Errors(); len(errs) > 0 { + 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) + } + } + } + recs := fetches.Records() + if len(recs) == 0 { + continue + } + s.mu.Lock() + s.buffer = recs + s.mu.Unlock() + } +} + +// SettleBatch applies r to every message in ms in one call, satisfying +// [source.Batched]. It settles each record through the same per-record path as +// [Settle] (mark, produce-to-DLQ, or no-op), returning the first error +// encountered while still attempting every message. +func (s *subscription) SettleBatch(ctx context.Context, ms []source.Message, r source.Result) error { + var firstErr error + for _, m := range ms { + if err := s.Settle(ctx, m, r); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + // Settle applies r to a record previously returned by Next, translating the // [source.Action] onto Kafka's commit/produce vocabulary per the ack model. // Settle may be called from many worker goroutines and is safe for concurrent @@ -281,4 +364,5 @@ var ( _ source.PartitionOrdered = (*subscription)(nil) _ source.LagReporter = (*subscription)(nil) _ source.Transactional = (*subscription)(nil) + _ source.Batched = (*subscription)(nil) ) diff --git a/source/memsource/harness.go b/source/memsource/harness.go index 18f9e4c..5f3024d 100644 --- a/source/memsource/harness.go +++ b/source/memsource/harness.go @@ -28,7 +28,16 @@ type Harness struct { // t.Cleanup. func NewHarness(tb testing.TB, opts []source.Option, msgs ...Msg) *Harness { tb.Helper() - in := New(WithMessages(msgs...)) + return NewHarnessWith(tb, opts, nil, msgs...) +} + +// NewHarnessWith is [NewHarness] with extra inlet options, e.g. [WithBatched] to +// exercise the engine's whole-batch fetch path or [WithClock] for a deterministic +// settle clock. inletOpts is applied after the pre-queued msgs. +func NewHarnessWith(tb testing.TB, opts []source.Option, inletOpts []Option, msgs ...Msg) *Harness { + tb.Helper() + all := append([]Option{WithMessages(msgs...)}, inletOpts...) + in := New(all...) hp := source.New(opts...) tb.Cleanup(func() { _ = hp.Close() }) return &Harness{tb: tb, inlet: in, hopper: hp} @@ -75,6 +84,34 @@ func (h *Harness) RunFor(timeout time.Duration, handler source.Handler) { } } +// RunBatch drives the queued messages through bh in batch mode, the batch analog +// of [Run]. The Hopper must have been built with [source.WithBatch] for batching +// to take effect; otherwise every batch holds one message. It closes the inlet's +// subscription so the Hopper drains once the queue empties, flushing any partial +// final batch, and fails the test on an unexpected run error. It uses a bounded +// 5s timeout; use [RunBatchFor] to override. +func (h *Harness) RunBatch(bh source.BatchHandler) { + h.tb.Helper() + h.RunBatchFor(5*time.Second, bh) +} + +// RunBatchFor is [RunBatch] with an explicit timeout. +func (h *Harness) RunBatchFor(timeout time.Duration, bh source.BatchHandler) { + h.tb.Helper() + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + sub, err := h.inlet.Subscribe(ctx, source.SubscribeConfig{}) + if err != nil { + h.tb.Fatalf("memsource: Subscribe failed: %v", err) + } + _ = sub.Close() + + if err := h.hopper.RunBatch(ctx, sub, bh); err != nil { + h.tb.Fatalf("memsource: Hopper.RunBatch returned %v", err) + } +} + // AssertCounts fails the test unless the recorded settlements match want exactly. func (h *Harness) AssertCounts(want Counts) { h.tb.Helper() diff --git a/source/memsource/memsource.go b/source/memsource/memsource.go index 9d474b9..a0bbd63 100644 --- a/source/memsource/memsource.go +++ b/source/memsource/memsource.go @@ -76,6 +76,8 @@ type Inlet struct { now func() time.Time nextID func() string + batched bool + mu sync.Mutex queue []source.Message ledger *Ledger @@ -100,6 +102,14 @@ func New(opts ...Option) *Inlet { return in } +// WithBatched is an [Option] that makes the Inlet's subscriptions implement +// [source.Batched], so a [source.Hopper] run in batch mode exercises the +// whole-batch fetch path (NextBatch/SettleBatch) rather than per-message +// accumulation. The default subscription is the honest per-message common path. +func WithBatched() Option { + return func(in *Inlet) { in.batched = true } +} + // WithMessages is an [Option] that pre-queues msgs at construction, equivalent // to a subsequent [Inlet.Queue]. func WithMessages(msgs ...Msg) Option { @@ -155,6 +165,9 @@ func (in *Inlet) Subscribe(_ context.Context, _ source.SubscribeConfig) (source. defer in.mu.Unlock() s := &subscription{inlet: in, signal: make(chan struct{}, 1)} in.subs = append(in.subs, s) + if in.batched { + return &batchedSubscription{subscription: s}, nil + } return s, nil } diff --git a/source/memsource/subscription.go b/source/memsource/subscription.go index 5ba9fad..a55ee04 100644 --- a/source/memsource/subscription.go +++ b/source/memsource/subscription.go @@ -64,7 +64,15 @@ func (s *subscription) notify() { } // Next returns the next queued message, blocking until one is queued, the -// context is canceled, or the subscription is closed and drained. +// context is canceled, or the subscription is closed and its queue is empty. +// +// Drain reports ErrDrained as soon as the subscription is closed and no queued +// message remains, independent of messages still in flight: settling those is the +// engine's drain responsibility (it finishes its lanes after the fetch loop +// stops), not a precondition for "nothing left to fetch". Gating drain on +// in-flight==0 would deadlock batch mode, where the engine holds a trailing +// partial batch unsettled until the fetch loop closes its lanes — which it cannot +// do while still blocked here waiting for that very batch to settle. func (s *subscription) Next(ctx context.Context) (source.Message, error) { for { if m, ok := s.inlet.take(); ok { @@ -75,7 +83,7 @@ func (s *subscription) Next(ctx context.Context) (source.Message, error) { } s.mu.Lock() - drained := s.closed && s.inFlight == 0 + drained := s.closed s.mu.Unlock() if drained { return nil, source.ErrDrained @@ -120,3 +128,51 @@ func (s *subscription) Close() error { } var _ source.Subscription = (*subscription)(nil) + +// batchedSubscription is a [subscription] that also satisfies [source.Batched], +// used to exercise the engine's whole-batch fetch path. It is opt-in via +// [WithBatched] so the default memsource subscription stays the honest +// per-message common path. +type batchedSubscription struct { + *subscription +} + +// NextBatch returns up to limit queued messages, blocking for at least one +// (delegating to Next) and then draining whatever else is queued without +// blocking, so the engine receives a backend-shaped batch. +func (b *batchedSubscription) NextBatch(ctx context.Context, limit int) ([]source.Message, error) { + if limit < 1 { + limit = 1 + } + first, err := b.Next(ctx) + if err != nil { + return nil, err + } + msgs := []source.Message{first} + for len(msgs) < limit { + m, ok := b.inlet.take() + if !ok { + break + } + b.mu.Lock() + b.inFlight++ + b.mu.Unlock() + msgs = append(msgs, m) + } + return msgs, nil +} + +// SettleBatch settles every message in ms with r, recording each on the ledger. +func (b *batchedSubscription) SettleBatch(ctx context.Context, ms []source.Message, r source.Result) error { + for _, m := range ms { + if err := b.Settle(ctx, m, r); err != nil { + return err + } + } + return nil +} + +var ( + _ source.Subscription = (*batchedSubscription)(nil) + _ source.Batched = (*batchedSubscription)(nil) +) diff --git a/source/options.go b/source/options.go index 41470ac..f3b4c57 100644 --- a/source/options.go +++ b/source/options.go @@ -4,6 +4,7 @@ package source import ( "log/slog" + "time" "github.com/stablekernel/crucible/telemetry" ) @@ -19,6 +20,19 @@ type config struct { middleware []Middleware concurrency int maxInFlight int + + // batch holds the batch-mode tuning; batch.enabled is false unless WithBatch + // is supplied, in which case RunBatch/ReceiveBatch accumulate per lane. + batch batchConfig +} + +// batchConfig holds the resolved batch-mode seams. now is injected (mirroring +// sink.Reservoir) so max-wait flushing is deterministic in tests. +type batchConfig struct { + enabled bool + size int + maxWait time.Duration + now func() time.Time } // defaultConfig returns the no-op defaults: a discarding logger, the no-op @@ -34,6 +48,9 @@ func defaultConfig() config { meter: telemetry.NopMeter(), concurrency: 1, maxInFlight: 0, + batch: batchConfig{ + now: time.Now, + }, } } @@ -139,3 +156,52 @@ func WithMaxInFlight(n int) Option { } } } + +// WithBatch enables batch mode for a Hopper driven with [Hopper.RunBatch] or +// [Hopper.ReceiveBatch]. Within each ordered lane the engine accumulates up to +// size messages, or until maxWait elapses since the lane's first buffered +// message, then invokes the [BatchHandler] once with the accumulated slice and +// settles each message by its corresponding [Result]. Per-key ordering and the +// [WithMaxInFlight] bound are preserved: a lane never overlaps two batches, and a +// batch is settled before the lane accepts more. +// +// When the subscription implements [Batched], whole fetched batches are handed to +// lanes as a unit (still re-grouped by key so a lane only ever sees its own keys +// in order); otherwise the engine accumulates per-message under the size/maxWait +// policy. maxWait uses the clock from [WithBatchClock] (default time.Now), so a +// max-wait flush is deterministic in tests. +// +// A size < 1 is treated as 1 (every message is its own batch). A maxWait <= 0 +// disables the timer, so a lane flushes only when it reaches size or the run +// drains. WithBatch only takes effect for the RunBatch/ReceiveBatch entry points; +// it is ignored by the per-message [Hopper.Run] and [Hopper.Receive]. +// +// Tuning note: a lane buffers up to size delivered-but-unsettled messages, each +// holding a [WithMaxInFlight] slot until the batch is dispatched. Set maxInFlight +// to at least size (times the number of concurrently filling lanes you expect), +// or supply a positive maxWait, so a lane can always reach a flush. With both +// maxWait <= 0 and maxInFlight < size a single busy lane can stall waiting for +// slots it will never get; the size/maxWait pair is the intended steady-state +// trigger. +func WithBatch(size int, maxWait time.Duration) Option { + return func(c *config) { + if size < 1 { + size = 1 + } + c.batch.enabled = true + c.batch.size = size + c.batch.maxWait = maxWait + } +} + +// WithBatchClock injects the clock the batch engine reads to time a lane's +// max-wait flush, for deterministic tests (mirroring sink.WithReservoirClock). +// The default is time.Now. A nil clock is ignored. It has no effect unless +// [WithBatch] is also supplied. +func WithBatchClock(now func() time.Time) Option { + return func(c *config) { + if now != nil { + c.batch.now = now + } + } +}