diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index b026318..aee4d42 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -23,7 +23,9 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-go@v6 with: - go-version: "stable" + # Match the module floor exactly; "stable" makes benchmark history + # change when a new Go release appears without a repository change. + go-version-file: go.mod - name: Run benchmark run: make ensure-build && make bench-run | tee build/output.txt - name: Store benchmark result diff --git a/README.md b/README.md index 3d3b5e6..3ce274d 100644 --- a/README.md +++ b/README.md @@ -39,8 +39,12 @@ func main() { return nil }), ) - // stop the batcher - defer batcher.Close() + // stop the batcher, and report if the drain did not finish in time + defer func() { + if err := batcher.Close(); err != nil { + fmt.Printf("shutdown incomplete: %v\n", err) + } + }() // add operations to the batcher for i := 0; i < 1000; i++ { @@ -171,15 +175,89 @@ if err := batcher.Join(timeout); err != nil { ### Stopping the batcher -To stop the batcher you can use the `StopProcessing` function: +To stop the batcher, use `Close`: ```go -defer batcher.Close() +defer func() { + if err := batcher.Close(); err != nil { + // The drain did not finish within the grace period. Work is still being + // processed in the background, so decide deliberately: wait longer with + // Shutdown, or accept that the process is about to exit with work pending. + log.Printf("batcher shutdown incomplete: %v", err) + } +}() + +// batcher.IsClosed() == true once the drain has completed +``` + +`Close` seals admission and drains work that has already been accepted, waiting up +to 30 seconds by default (configurable with `WithCloseGrace`). It is safe to call +multiple times and from multiple goroutines. + +If the grace period expires, `Close` reports that the drain is incomplete — it does +**not** discard the remaining work, which keeps being processed in the background. +Do not write a bare `defer batcher.Close()`: it discards that report, and if the +process exits immediately afterwards, accepted work is lost without any signal. + +When you need to control the wait, or to keep waiting after a timeout, use +`Shutdown`: + +```go +ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) +defer cancel() + +if err := b.Shutdown(ctx); err != nil { + var incomplete *batcher.ShutdownIncompleteError + if errors.As(err, &incomplete) { + // Still draining. Nothing was lost, and we can wait longer. + log.Printf("%d items still pending", incomplete.Pending) + + err = b.Shutdown(context.Background()) + } +} +``` + +`Shutdown` is resumable: a later call waits on the same drain rather than starting +a new one, so an expired deadline never costs you accepted work. -// batcher.IsClosed() == true after this point +### Back-pressure and rejection + +By default the queue is unbounded, which absorbs bursts well but turns a sustained +overload into unbounded memory growth. Note that shrinking the batch interval does +**not** bound queued work — only bounding the queue does. + +To get back-pressure instead, bound the queue and use `Enqueue`, which reports why +an item was refused: + +```go +b := batcher.New( + batcher.WithProcessor(processor.Process), + batcher.WithMaxQueueSize[Item](10_000), +) + +if err := b.Enqueue(ctx, item); err != nil { + // batcher.ErrClosing -> shutting down + // context.DeadlineExceeded -> queue stayed full until the deadline passed + // context.Canceled -> queue was full and the caller cancelled + return err +} ``` -This function is safe to be called multiple times as it will only stop the processor once. +`Add` remains available as the fast path for best-effort use: it returns no error, +and after shutdown it is a counted no-op rather than a panic. + +### Observing a batcher + +`Stats` returns an allocation-free snapshot suitable for frequent scraping: + +```go +s := b.Stats() + +// s.Queued -> queue depth; the value to alert on +// s.Pending -> accepted work not yet finished, including in-flight batches +// s.Rejected -> refused enqueues +// s.DroppedErrors -> diagnostics lost because Errors() was not drained +``` ### Handling Errors diff --git a/docs/improvements/plan-perf.md b/docs/improvements/plan-perf.md index c64948c..d4243c8 100644 --- a/docs/improvements/plan-perf.md +++ b/docs/improvements/plan-perf.md @@ -373,11 +373,15 @@ conservation check, never an instantaneous API guarantee. is bounded by: ```text -N + BatchSize + (concurrency × BatchSize) + P +N + (1 + concurrency) × BatchSize + P ``` where `P` is the number of publishers currently inside the gate, exposed as -`Stats().PublishersInGate`. At `n = 1` there is no dispatch stage, so the +`Stats().PublishersInGate`. The `(1 + concurrency)` term is deliberate: one batch +is held by the aggregator after leaving the queue, and `concurrency` batches can +be inside processors. Writing it as `BatchSize + concurrency × BatchSize` +understates the peak by one batch, which at `n = 1` is a factor-of-two error on +that term. At `n = 1` there is no dispatch stage, so the `concurrency × BatchSize` term collapses to the single inline batch already counted by `BatchSize`; the formula is therefore conservative rather than wrong at `n = 1`. `P` is bounded by the caller's own concurrency, not @@ -567,9 +571,32 @@ reviewable behaviour-preserving change. - Characterization tests written pre-swap pass unchanged post-swap: batch boundaries, timer-armed-on-first-item behaviour, ordering, error propagation, no-empty-batch guarantee. -- Goroutines per *constructed, auto-started* batcher drop from 6 to 3, verified by - test. (2.2 removes the remaining two by replacing `chann`, reaching 1 per - running batcher and 0 for an unstarted one.) +- Goroutines per *constructed, auto-started* batcher drop from 6 to **5**, + verified by test and by direct measurement (50 batchers: +250 goroutines, 0 + leaked after `Close`). + + The plan first predicted 3, then 4. Both were wrong, and the reasons are worth + recording because they constrain the design: + + 1. **Merging aggregation into the processing loop changes latency.** Measured + with a 50ms processor at 10k items/s, a merged loop inverted the baseline + finding (5ms window p50 28ms versus 100ms window p50 75ms), because the next + batch only began accumulating after the processor returned. Aggregation and + processing must stay separate goroutines here; decoupling them differently + is a Phase 3 decision, not one this milestone may smuggle in. + 2. **Input draining must be isolated from batch bookkeeping.** With a single + goroutine doing both, sequential `Add` regressed 39-50% against the stored + baseline, breaching the ≤+10% gate: the `chann` relay has a small bounded + ingress, and its output was not being read promptly while the aggregator was + busy with timer and slice work. Adding a dedicated forwarding stage restored + parity (no regression; one case 6.7% faster). This is what `rill`'s + `ToChans`/`Batch`/`FromChans` pipeline was buying. + + The five are: `chann` input relay, input forwarder, aggregator, processor loop, + and `chann` error relay. Phase 2.2 removes both `chann` relays and the + forwarder along with them, because an owned queue needs no relay and can be + read directly. The single-goroutine target in the threshold table applies only + once Phase 3 owns the worker model. - Enqueue benchmarks from 1.1 show no regression beyond predeclared thresholds. - `go.mod`/`go.sum` no longer reference `rill`; no other dependency added. - No public API change in this commit. @@ -1032,7 +1059,9 @@ rather than intuition. deadline; **`Len()` now counting accepted-but-unfinished work including in-flight batches, so it can be non-zero with an empty queue**; `Join`'s clarified quiescence-snapshot contract; the removal of the `rill` and `chann` - dependencies; and any `ProvideBatcherInFX` signature change. + dependencies; **the module's minimum Go version moving from 1.22.4 to 1.25.0** + (required by `golang.org/x/sys` v0.46.0, pulled in when dependencies were + modernised in 2.1); and any `ProvideBatcherInFX` signature change. - **`Config()` must stop exposing live mutable state.** It currently returns the internal pointer (`batcher.go:76-78`), letting callers mutate `ProcessorFunc` or `BatchSize` during processing — an exported data race that can invalidate @@ -1135,14 +1164,22 @@ commit leaves the library in a consistent state, so any prefix of the plan is a valid stopping point: - Stopping after **1.x**: measurement only, no behaviour change. -- Stopping after **2.1**: same semantics, `rill` removed, goroutines 6 → 3. +- Stopping after **2.1**: same semantics, `rill` removed, goroutines 6 → 5, + and ~47% less memory per enqueued item. - Stopping after **2.2**: admission is panic-free, accounting is sound, `chann` - is gone (1 goroutine per running batcher, 0 idle), and the internal - seal/gate/`noInput` coordinator is complete. `Close()` keeps its existing - signature, no longer abandons the drain, and a minimal `Stats()` is public; only - the resumable API and typed error are absent. -- Stopping after **2.3**: no accepted work is ever silently discarded. -- Stopping after **2.4**: diagnostics bounded, panics survivable. + is gone (**measured 2 goroutines per running batcher**, down from 6, and enqueue + 54% faster), and the internal seal/gate/`noInput` coordinator is complete. + `Close()` keeps its existing signature, no longer abandons the drain, and a + minimal `Stats()` is public; only the resumable API and typed error are absent. +- Stopping after **2.3**: no accepted work is ever silently discarded. Verified + directly: a partial batch with a 30s interval and a shorter grace now reports 50 + accepted / 50 processed / 0 pending, where it previously reported 50 accepted / + 0 processed / 50 phantom pending. +- Stopping after **2.4**: diagnostics bounded, panics survivable. Phase 2 as a + whole is verified against the original defects: `Add` after `Close` is a counted + rejection instead of a process panic, bounded mode caps the queue and reports + rejections instead of growing the heap without limit, and goroutines per batcher + are 2 with none leaked. - Stopping after **3.x**: small windows are honoured under slow processors. - Stopping after **4.x**: operators can see overload developing. diff --git a/docs/improvements/thresholds.md b/docs/improvements/thresholds.md index dcd5155..5cb8983 100644 --- a/docs/improvements/thresholds.md +++ b/docs/improvements/thresholds.md @@ -12,9 +12,17 @@ reported by the scenario matrix and compared as a trend; it never fails a PR. Thresholds are keyed to the environment. Re-baseline when any of these change. +The Go version is whatever `go.mod` declares: CI resolves it with +`go-version-file: go.mod` rather than a hardcoded string, so this row and the +toolchain cannot drift apart. Phase 2 raised the floor from 1.22.4 to 1.25.0 when +`golang.org/x/sys` was updated. + +Stored baselines record the toolchain they were captured on in their own header, +which is what makes them comparable or not; see Baselines below. + | Field | Value | | ---------- | ---------------------------------------- | -| Go version | 1.22.4 | +| Go version | 1.25.0 (from `go.mod`) | | Runner | `ubuntu-latest` (GitHub-hosted, mutable) | | Arch | `amd64` | @@ -27,12 +35,15 @@ orientation only and are not gates. Allocation counts are stable even on noisy shared runners, which is why they are the blocking signal rather than timings. -| Gate | Threshold | Enforced by | -| --------------------------------------- | --------- | ---------------------------------------- | -| `Add` allocations, unbounded path | exactly 0 | `TestAddAllocatesNothingPerCall` | -| `Stats()` allocations | exactly 0 | added with `Stats()` in Phase 2 | -| Recovery wrapper allocations, non-panic | exactly 0 | added with panic recovery in Phase 2 | +| Gate | Threshold | Enforced by | +| --------------------------------------- | -------------------------------------------- | ------------------------------------------------------------- | +| `Add` allocations, unbounded path | exactly 0 | `TestAddAllocatesNothingPerCall` | +| Recovery wrapper allocations, non-panic | exactly 0 | `TestRecoveredPanicAddsNoSteadyStateAllocations` | | Scenario recorder allocations per item | ≤ 1 total, and must not grow with run length | `TestHarnessRecorderDoesNotAllocatePerItem` (`test/scenario`) | +| Goroutines per running batcher | exactly 2 | `TestGoroutineBudgetPerRunningBatcher` | + +`Stats()` is a fixed set of atomic loads returning a value type, so it has no +allocation gate of its own; the `Add` gate covers the hot path that matters. The recorder threshold is not "exactly 0" because `AllocsPerItem` measures the whole pipeline, including Batcher's own per-batch allocations, not just the @@ -40,12 +51,12 @@ recorder. Measured values are 0.03-0.04 allocations per item and do not grow wit run length, which is the property that matters: a recorder that allocated per item would make every allocation figure it reports a measurement of itself. -`Add` currently allocates 0 per call in the timed region because the item is -constructed by the caller. Phase 2 replaces `chann` with a slice-backed unbounded -queue, at which point the gate becomes "zero allocations per `Add` in steady -state": `append` must allocate when it grows its backing array, so the gate is -measured after a stated warmup with queue capacity retained across drains, and -growth-path allocations are the one named exemption. +`Add` allocates 0 per call in the timed region because the item is constructed by +the caller. Phase 2 replaced `chann` with a slice-backed unbounded queue, so the +gate is now "zero allocations per `Add` in steady state": `append` must allocate +when it grows its backing array, so the gate is measured after a stated warmup +with queue capacity retained across drains, and growth-path allocations are the +one named exemption. ## Throughput gate (advisory until CI baseline exists) @@ -55,18 +66,46 @@ growth-path allocations are the one named exemption. Compare with `benchstat` over `-count=10`. A single run is not evidence. This is advisory until a scheduled run stores an ubuntu-latest/amd64 baseline; the current workflow uploads raw input but intentionally does not fail on this threshold. -## Goroutine gates (blocking, from Phase 2 onward) - -| Gate | Threshold | -| ----------------------------------- | ---------------------------------- | -| Goroutines per idle batcher | exactly 0 | -| Goroutines per running `n=1` | exactly 1 (aggregator) | -| Goroutines per running `n>1` | exactly 1 + n | -| Goroutines after terminal `closed` | equal to pre-construction baseline | - -Current `main` owns 6 goroutines per batcher: 1 aggregator, 2 `rill` pipeline -relays, and 2 `chann` relays plus the pipeline's batch goroutine. Phase 2.1 -removes `rill` (6 → 3) and Phase 2.2 removes `chann` (3 → 1). +## Goroutine gates (blocking) + +These are the gates in force today. They must agree with the allocation table +above, which enforces the same running-batcher count. + +| Gate | Threshold | +| ---------------------------------- | ----------------------------------------- | +| Goroutines per unstarted batcher | exactly 0 | +| Goroutines per running batcher | exactly 2 (aggregator + serial processor) | +| Goroutines after terminal `closed` | equal to pre-construction baseline | + +Enforced by `TestGoroutineBudgetPerRunningBatcher`. + +Phase 3 introduces explicit worker concurrency and will replace the single +running-batcher row with `1 + n` (aggregator plus workers). That row is +deliberately absent here rather than stated as a gate, because a threshold for a +configuration this code cannot express is not enforceable — and two rows claiming +different counts for the same batcher is worse than one row that is merely +incomplete. + +Current `main` owns 6 goroutines per batcher. Phase 2.1 removed `rill` +(**measured 6 → 5**) and Phase 2.2 removed both `chann` relays and the input +forwarder they required (**measured 5 → 2**: aggregator plus processor), enforced +by `TestGoroutineBudgetPerRunningBatcher`. + +Removing the relays also removed a channel hop and a goroutine handoff per item. +Measured against the stored baseline: **-54% geomean sec/op** on the enqueue +microbenchmarks (`Add` 229.8ns → 63.4ns at small batch sizes) and -49% to -90% +bytes/op. This is the one place in the plan where a safety change also made the +hot path materially faster. + +The 6 → 5 figure supersedes earlier 6 → 3 and 6 → 4 estimates. Fewer goroutines +were unreachable without changing observable behaviour: merging aggregation into +the processing loop inverted the documented latency baseline, and merging input +draining into aggregation regressed sequential `Add` by 39-50% because the +`chann` relay's bounded ingress was not drained promptly. + +Both earlier estimates assumed aggregation and processing could share one +goroutine. They cannot without changing observable latency, which is why the +enforced count is 2 rather than 1. ## Conditional gates (Milestone 4.2 only) diff --git a/go.mod b/go.mod index ce6d2b9..d558f26 100644 --- a/go.mod +++ b/go.mod @@ -1,21 +1,18 @@ module github.com/NSXBet/batcher -go 1.22.4 +go 1.25.0 require ( - github.com/destel/rill v0.1.2 - github.com/stretchr/testify v1.9.0 - go.uber.org/atomic v1.11.0 + github.com/stretchr/testify v1.11.1 go.uber.org/fx v1.24.0 - go.uber.org/zap v1.27.0 - golang.design/x/chann v0.1.2 + go.uber.org/zap v1.28.0 ) require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect go.uber.org/dig v1.19.0 // indirect - go.uber.org/multierr v1.10.0 // indirect - golang.org/x/sys v0.0.0-20220412211240-33da011f77ad // indirect + go.uber.org/multierr v1.11.0 // indirect + golang.org/x/sys v0.46.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 580d99e..fbb054a 100644 --- a/go.sum +++ b/go.sum @@ -1,27 +1,23 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/destel/rill v0.1.2 h1:+JK5JoH70GuQf+Z/JOegc4d2//QbwG5q+6dbHXc5UVE= -github.com/destel/rill v0.1.2/go.mod h1:srKuXzvGqINUEGYR5b/iwvW+L9/S35RxVHWGYbXNoO4= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= -github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= -go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= -go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= go.uber.org/dig v1.19.0 h1:BACLhebsYdpQ7IROQ1AGPjrXcP5dF80U3gKoFzbaq/4= go.uber.org/dig v1.19.0/go.mod h1:Us0rSJiThwCv2GteUN0Q7OKvU7n5J4dxZ9JKUXozFdE= go.uber.org/fx v1.24.0 h1:wE8mruvpg2kiiL1Vqd0CC+tr0/24XIB10Iwp2lLWzkg= go.uber.org/fx v1.24.0/go.mod h1:AmDeGyS+ZARGKM4tlH4FY2Jr63VjbEDJHtqXTGP5hbo= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ= -go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= -go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8= -go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= -golang.design/x/chann v0.1.2 h1:eHF9wjuQnpp+j4ryWhyxC/pFuYzbvMAkudA/I5ALovY= -golang.design/x/chann v0.1.2/go.mod h1:Rh5KhCAp+0qu9+FfKPymHpu8onmjl89sFwMeiw3SK14= -golang.org/x/sys v0.0.0-20220412211240-33da011f77ad h1:ntjMns5wyP/fN65tdBD4g8J5w8n015+iIIs9rtjXkY0= -golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= +go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= +go.uber.org/zap v1.28.0 h1:IZzaP1Fv73/T/pBMLk4VutPl36uNC+OSUh3JLG3FIjo= +go.uber.org/zap v1.28.0/go.mod h1:rDLpOi171uODNm/mxFcuYWxDsqWSAVkFdX4XojSKg/Q= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/pkg/batcher/admission_test.go b/pkg/batcher/admission_test.go new file mode 100644 index 0000000..6228009 --- /dev/null +++ b/pkg/batcher/admission_test.go @@ -0,0 +1,554 @@ +package batcher_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/NSXBet/batcher/pkg/batcher" + "github.com/stretchr/testify/require" +) + +// Admission and drain protocol tests. +// +// These cover the interleavings that make shutdown safe. Each one exists because +// a specific plausible implementation is wrong in a way that ordinary tests do not +// catch: a lost wakeup hangs forever, a missed reservation drops accepted work, and +// a closed input channel panics in the caller's goroutine. + +// TestSealWhileProducersRaceNeverPanics is the blunt instrument: hammer admission +// while sealing, many times, and require that nothing panics and nothing is lost. +// +// The input queue is never closed precisely so this cannot panic. A design that +// closed it would fail here with "send on closed channel" in a caller's goroutine, +// which is the worst place for a library to fail. +func TestSealWhileProducersRaceNeverPanics(t *testing.T) { + t.Parallel() + + const ( + trials = 400 + producers = 8 + perActor = 25 + ) + + for trial := range trials { + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](16), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + var ( + wg sync.WaitGroup + start = make(chan struct{}) + ) + + for range producers { + wg.Add(1) + + go func() { + defer wg.Done() + + <-start + + for range perActor { + b.Add(1) + } + }() + } + + closed := make(chan error, 1) + + go func() { + <-start + closed <- b.Close() + }() + + close(start) + wg.Wait() + + require.NoError(t, <-closed, "trial %d: shutdown must complete", trial) + + stats := b.Stats() + + require.GreaterOrEqual(t, stats.Pending, int64(0), + "trial %d: pending must never go negative", trial) + require.Zero(t, stats.Pending, + "trial %d: pending must converge to zero after drain", trial) + require.Zero(t, stats.IntakePending, + "trial %d: intake must be exhausted after drain", trial) + require.Zero(t, stats.PublishersInGate, + "trial %d: no publisher may remain in the gate", trial) + + // Every accepted item reached a terminal outcome, and nothing was invented. + require.Equal(t, stats.Accepted, stats.Completed+stats.Failed+stats.Panicked, + "trial %d: accepted must equal terminal outcomes at quiescence", trial) + require.Equal(t, int64(stats.Accepted), processed.Load(), + "trial %d: every accepted item must be processed", trial) + } +} + +// TestLastPublisherLeavesBeforeSealCompletesShutdown covers the lost-wakeup +// interleaving where the final publisher exits the gate *before* sealing begins. +// +// leave() only signals quiescence when it observes the seal, so if nothing else +// checked, no signal would ever arrive and Close would hang forever. The +// coordinator's own post-seal gate check is what prevents that, and this test is +// what proves the check is load-bearing: remove it and this hangs. +func TestLastPublisherLeavesBeforeSealCompletesShutdown(t *testing.T) { + t.Parallel() + + const trials = 300 + + for trial := range trials { + b := batcher.New( + batcher.WithBatchSize[int](4), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + // Publish and return, so the gate is empty well before sealing starts. + b.Add(1) + + done := make(chan error, 1) + + go func() { done <- b.Close() }() + + select { + case err := <-done: + require.NoError(t, err, "trial %d", trial) + case <-time.After(10 * time.Second): + t.Fatalf("trial %d: Close hung; the coordinator's post-seal gate check is missing", trial) + } + } +} + +// TestShutdownWithQueuedWorkProcessesEverything is the final-intake completeness +// case: seal immediately after publishing, repeatedly, and require the item is +// always processed. +// +// An implementation that decided "intake is finished" by observing an empty queue +// would fail here, because an item can be published between that observation and +// the decision. Draining by intake accounting is what makes it correct. +func TestShutdownWithQueuedWorkProcessesEverything(t *testing.T) { + t.Parallel() + + const trials = 300 + + for trial := range trials { + var processed atomic.Int64 + + b := batcher.New( + // Large batch and long interval: only shutdown can flush this. + batcher.WithBatchSize[int](10_000), + batcher.WithBatchInterval[int](time.Hour), + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + b.Add(1) + + require.NoError(t, b.Close(), "trial %d", trial) + require.Equal(t, int64(1), processed.Load(), + "trial %d: the queued item must be processed, not discarded", trial) + } +} + +// TestAddAfterSealIsRejectedNotPanicking pins that a late Add is a counted +// rejection. +// +// Previously this panicked with "send on closed channel", turning a graceful +// shutdown into a crash in the caller's goroutine. Silently succeeding would be +// just as bad, because the caller would believe the item was queued. +func TestAddAfterSealIsRejectedNotPanicking(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[int](4), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + require.NoError(t, b.Close()) + + before := b.Stats() + + require.NotPanics(t, func() { b.Add(1) }, + "Add after shutdown must not panic") + + after := b.Stats() + + require.Equal(t, before.Accepted, after.Accepted, + "a rejected Add must not count as accepted") + require.Equal(t, before.Rejected+1, after.Rejected, + "a rejected Add must be counted") + require.Zero(t, after.Pending, + "a rejected Add must leave no pending obligation") +} + +// TestEnqueueAfterSealReportsClosing pins that Enqueue tells the caller why the +// item was refused, rather than failing silently the way Add must. +func TestEnqueueAfterSealReportsClosing(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[int](4), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + require.NoError(t, b.Close()) + + err := b.Enqueue(context.Background(), 1) + + require.ErrorIs(t, err, batcher.ErrClosing) + require.Zero(t, b.Stats().Pending, "a refused Enqueue must leave no obligation") +} + +// TestEnqueueRespectsContextOnFullQueue pins back-pressure with a deadline: a +// bounded queue that is full must let the caller give up, and giving up must not +// leave a phantom obligation that would stall shutdown forever. +func TestEnqueueRespectsContextOnFullQueue(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + defer close(release) + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Hour), + batcher.WithMaxQueueSize[int](2), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + saturate(t, b) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + before := b.Stats() + err := b.Enqueue(ctx, 99) + + require.ErrorIs(t, err, context.DeadlineExceeded, + "a full bounded queue must honour the caller's deadline") + + after := b.Stats() + + require.Equal(t, before.Accepted, after.Accepted, + "a cancelled Enqueue must not count as accepted") + require.Equal(t, before.Pending, after.Pending, + "a cancelled Enqueue must roll its reservation back") + require.Equal(t, before.Rejected+1, after.Rejected, + "a cancelled Enqueue must be counted as rejected") +} + +// TestBlockedEnqueueIsReleasedByShutdown pins that sealing releases a publisher +// parked on a full bounded queue. +// +// This is not a nicety. A parked publisher cannot leave the admission gate, and +// shutdown waits for the gate to empty, so without this release the two would +// deadlock permanently. +func TestBlockedEnqueueIsReleasedByShutdown(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + defer close(release) + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Hour), + batcher.WithMaxQueueSize[int](2), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + saturate(t, b) + + parked := make(chan error, 1) + + go func() { + parked <- b.Enqueue(context.Background(), 99) + }() + + // Give the publisher time to actually park inside the queue. + require.Eventually(t, func() bool { + return b.Stats().PublishersInGate > 0 + }, 5*time.Second, time.Millisecond) + + go func() { + _ = b.Close() + }() + + select { + case err := <-parked: + require.ErrorIs(t, err, batcher.ErrClosing, + "a parked publisher must be released with ErrClosing when sealing begins") + case <-time.After(10 * time.Second): + t.Fatal("a publisher parked on a full queue was never released by shutdown") + } +} + +// TestStartIsIdempotent pins that repeated and concurrent Start calls produce +// exactly one processing lifecycle. +// +// Previously each call spawned another loop whose deferred cleanup closed the same +// channels, so a second Start turned into a double-close panic. +func TestStartIsIdempotent(t *testing.T) { + t.Parallel() + + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithSkipAutoStart[int](), + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + var wg sync.WaitGroup + + for range 8 { + wg.Add(1) + + go func() { + defer wg.Done() + + b.Start() + }() + } + + wg.Wait() + + b.Add(1) + + require.NoError(t, b.Close()) + require.Equal(t, int64(1), processed.Load(), + "exactly one lifecycle must process the item exactly once") +} + +// TestShutdownIsIdempotent pins that concurrent and repeated shutdowns agree. +func TestShutdownIsIdempotent(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[int](4), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + b.Add(1) + + var ( + wg sync.WaitGroup + errs = make([]error, 4) + ) + + for i := range errs { + wg.Add(1) + + go func(idx int) { + defer wg.Done() + + errs[idx] = b.Close() + }(i) + } + + wg.Wait() + + for i, err := range errs { + require.NoError(t, err, "caller %d must observe a completed drain", i) + } + + require.True(t, b.IsClosed(), "the batcher must be terminal after shutdown") + require.True(t, b.IsClosing(), "IsClosing must remain true once closed") +} + +// TestErrorsChannelClosesExactlyOnceAfterShutdown pins that diagnostics terminate +// cleanly, so a consumer ranging over Errors() does not block forever and the +// channel is never closed twice. +func TestErrorsChannelClosesExactlyOnceAfterShutdown(t *testing.T) { + t.Parallel() + + failure := errors.New("processor failed") + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func([]int) error { + return failure + }), + ) + + drained := make(chan int) + + go func() { + count := 0 + + for range b.Errors() { + count++ + } + + drained <- count + }() + + b.Add(1) + + require.NoError(t, b.Close()) + + select { + case count := <-drained: + require.Positive(t, count, "the processor error must be observable") + case <-time.After(10 * time.Second): + t.Fatal("Errors() was not closed after shutdown; a consumer would block forever") + } + + stats := b.Stats() + + require.Equal(t, uint64(1), stats.Failed, "a processor error counts as Failed") + require.Zero(t, stats.Completed, "a failed batch must not also count as Completed") + require.Zero(t, stats.Pending, "a failed batch must still release its obligation") +} + +// TestIdleBatcherDoesNotSpin pins that steady-state idling costs nothing. +// +// The aggregator waits in select. If it polled instead, an idle batcher would burn +// CPU continuously, which matters because callers create many of them. +func TestIdleBatcherDoesNotSpin(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[int](1_000), + batcher.WithBatchInterval[int](10*time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + t.Cleanup(func() { _ = b.Close() }) + + // An idle batcher must not invoke the processor at all, no matter how many + // intervals elapse. + time.Sleep(200 * time.Millisecond) + + stats := b.Stats() + + require.Zero(t, stats.Accepted) + require.Zero(t, stats.Completed, "an idle batcher must never flush an empty batch") + require.Zero(t, stats.Queued) +} + +// saturate fills a bounded batcher whose processor is blocked, so that a further +// enqueue is guaranteed to park. +// +// It enqueues with a short deadline until admission refuses, which is the only +// reliable signal that the queue is genuinely full: the aggregator drains greedily, +// so simply watching Queued can race with it. +func saturate(t *testing.T, b *batcher.Batcher[int]) { + t.Helper() + + for attempt := range 1_000 { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + err := b.Enqueue(ctx, attempt) + + cancel() + + if err != nil { + require.ErrorIs(t, err, context.DeadlineExceeded, + "saturation should end with a deadline, not %v", err) + + return + } + } + + t.Fatal("bounded queue never filled; the capacity bound is not being enforced") +} + +// TestAcceptedWorkBoundIncludesHeldAndInFlightBatches pins the documented capacity +// bound, which is the number callers use to size memory. +// +// The bound is N + 2*BatchSize + publishers-in-gate at Concurrency 1, not +// N + BatchSize: the aggregator holds a partial batch that has already left the +// queue, and a second batch can be inside the processor simultaneously. Documenting +// one batch understated the peak by a whole batch, so this asserts both that the +// real bound is respected and that it genuinely exceeds the smaller figure. +func TestAcceptedWorkBoundIncludesHeldAndInFlightBatches(t *testing.T) { + t.Parallel() + + const ( + maxQueueSize = 4 + batchSize = 3 + ) + + release := make(chan struct{}) + + var releaseOnce sync.Once + + releaseAll := func() { releaseOnce.Do(func() { close(release) }) } + + b := batcher.New( + batcher.WithBatchSize[int](batchSize), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithMaxQueueSize[int](maxQueueSize), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + t.Cleanup(func() { + // Release first: cleanup runs after any failed assertion, and Close would + // otherwise wait its full grace for the blocked processor. + releaseAll() + _ = b.Close() + }) + + // Fill using deadline-bounded enqueues so no publisher stays parked in the gate, + // isolating the queue + held + in-flight terms. Stop at the first refusal: that + // is the saturation point, and continuing would only add timeout latency. + for range 64 { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond) + err := b.Enqueue(ctx, 1) + + cancel() + + if err != nil { + break + } + } + + stats := b.Stats() + + documented := int64(maxQueueSize + 2*batchSize) + + require.LessOrEqual(t, stats.Pending, documented+stats.PublishersInGate, + "accepted work must stay within N + 2*BatchSize + gate "+ + "(pending=%d queued=%d gate=%d)", + stats.Pending, stats.Queued, stats.PublishersInGate) + + require.Greater(t, stats.Pending, int64(maxQueueSize+batchSize), + "the peak must exceed N + BatchSize; if it did not, the documented bound "+ + "would be describing a batcher that cannot hold a batch and process one "+ + "at the same time") +} diff --git a/pkg/batcher/batcher.go b/pkg/batcher/batcher.go index 9345e7b..25fd33c 100644 --- a/pkg/batcher/batcher.go +++ b/pkg/batcher/batcher.go @@ -1,16 +1,20 @@ package batcher import ( - "fmt" - "math" + "context" + "errors" + "runtime/debug" "sync" "time" - - "github.com/destel/rill" - "golang.design/x/chann" ) -var ErrTimeout = fmt.Errorf("timeout waiting for batches to complete") +// ErrTimeout is returned when a wait for pending work to finish expires. +var ErrTimeout = errors.New("timeout waiting for batches to complete") + +// ErrClosing is returned by Enqueue when admission has been sealed, and by a +// blocking enqueue that is released by shutdown. It means the item was not +// accepted and no work was queued for it. +var ErrClosing = errors.New("batcher is closing") type Processor[T any] func([]T) error @@ -19,48 +23,71 @@ func NoOpProcessor[T any]([]T) error { } type Config[T any] struct { - SkipAutoStart bool - BatchSize int - BatchInterval time.Duration - Concurrency int - ProcessorFunc Processor[T] + SkipAutoStart bool + BatchSize int + BatchInterval time.Duration + Concurrency int + MaxQueueSize int + CloseGrace time.Duration + ErrorBufferSize int + ProcessorFunc Processor[T] } type Batcher[T any] struct { config *Config[T] - isClosed bool - closeOnce sync.Once - itemCount *AtomicCounter - doneChan chan struct{} - errorsChan *chann.Chann[error] - batchInputChan *chann.Chann[rill.Try[T]] - batchesChan <-chan rill.Try[[]T] + gate *admissionGate + counters counters + + // input is never closed, for the life of the batcher. A publisher can + // therefore never send on a closed channel, which removes an entire class of + // shutdown panic. Intake completion is signalled by noInput instead. + input *queue[T] + + // errorsChan is a plain buffered channel closed exactly once, by the drain + // coordinator, after every possible sender has exited. + errorsChan chan error + + // noInput is closed by the coordinator once no publisher can ever send again. + // It is what lets the aggregator conclude that intake is finished, since the + // queue itself is never closed. + noInput chan struct{} + stopped chan struct{} + startOnce sync.Once + drainOnce sync.Once + + // shutdownDone closes only after the independent coordinator has finished + // draining accepted work and closed diagnostics. A Close caller may time out, + // but the coordinator never does: reporting an incomplete drain is not the + // same as abandoning it. + shutdownOnce sync.Once + shutdownDone chan struct{} } // New creates a new Batcher with the given options. func New[T any](options ...Option[T]) *Batcher[T] { b := &Batcher[T]{ config: &Config[T]{ - BatchSize: DefaultBatchSize, - BatchInterval: DefaultBatchInterval, - Concurrency: DefaultConcurrency, - ProcessorFunc: NoOpProcessor[T], + BatchSize: DefaultBatchSize, + BatchInterval: DefaultBatchInterval, + Concurrency: DefaultConcurrency, + CloseGrace: DefaultCloseGrace, + ErrorBufferSize: DefaultErrorBufferSize, + ProcessorFunc: NoOpProcessor[T], }, - itemCount: NewAtomicCounter(), - doneChan: make(chan struct{}), - batchInputChan: chann.New[rill.Try[T]](chann.Cap(-1)), + gate: newAdmissionGate(), + noInput: make(chan struct{}), + stopped: make(chan struct{}), + shutdownDone: make(chan struct{}), } for _, option := range options { option(b) } - b.errorsChan = chann.New[error](chann.Cap(-1)) - - batchOutput := rill.Batch(b.batchInputChan.Out(), b.config.BatchSize, b.config.BatchInterval) - b.batchesChan = batchOutput + b.input = newQueue[T](b.config.MaxQueueSize) + b.errorsChan = make(chan error, b.config.ErrorBufferSize) if !b.config.SkipAutoStart { b.Start() @@ -69,23 +96,121 @@ func New[T any](options ...Option[T]) *Batcher[T] { return b } +// Start begins processing. It is idempotent: repeated or concurrent calls start +// exactly one processing lifecycle. +// +// Start always launches the consumer, even if shutdown has already sealed +// admission. That is deliberate: a sealed batcher may still hold queued work, and +// the drain needs a consumer to make progress. Declining to start here would leave +// the drain waiting on a consumer that never exists, which deadlocks shutdown when +// Start and Shutdown race. func (b *Batcher[T]) Start() { - go b.startProcessing() + b.startOnce.Do(func() { + b.gate.state.CompareAndSwap(stateNew, stateRunning) + + go b.run() + }) } +// Config returns the batcher's configuration. func (b *Batcher[T]) Config() *Config[T] { return b.config } +// Add enqueues an item. It is the compatibility fast path and returns no error. +// +// In unbounded mode it never blocks. In bounded mode it blocks while the queue is +// full, exactly as it did before bounding existed, and shutdown releases it. After +// admission is sealed it is a no-op that counts a rejection: it neither panics nor +// silently pretends to have accepted the item. +// +// Callers that need to observe rejection or bound their wait should use Enqueue. func (b *Batcher[T]) Add(item T) { - b.batchInputChan.In() <- rill.Try[T]{Value: item} - b.itemCount.Add(1) + _ = b.publish(context.Background(), item, true) +} + +// Enqueue enqueues an item, reporting why it could not be accepted. +// +// It returns ErrClosing when admission is sealed, ctx.Err() when the caller's +// context expires while waiting for space in a bounded queue, or nil on success. +// A non-nil error means nothing was queued. +func (b *Batcher[T]) Enqueue(ctx context.Context, item T) error { + return b.publish(ctx, item, false) } +// publish is the single admission path shared by Add and Enqueue. +// +// The ordering here is the protocol, and each step exists for a reason: +// +// 1. enter the gate, so shutdown cannot conclude that intake is finished while +// this publisher is still deciding; +// 2. reserve the drain obligation BEFORE publishing, so the drain cannot +// terminate in the window between reserving and publishing; +// 3. publish; +// 4. on failure, roll the reservation back BEFORE leaving the gate, so no +// phantom obligation can outlive the gate; +// 5. leave the gate. +func (b *Batcher[T]) publish(ctx context.Context, item T, nonBlockingWhenFull bool) error { + if !b.gate.enter() { + b.counters.rejected.Add(1) + + return ErrClosing + } + + defer b.gate.leave() + + b.counters.reserve() + + var err error + + if nonBlockingWhenFull && b.config.MaxQueueSize <= 0 { + // Unbounded fast path: publication cannot fail. + _ = b.input.tryPush(item) + } else { + err = b.input.push(ctx, item, b.gate.sealCh) + } + + if err != nil { + b.counters.rollback() + + return err + } + + b.counters.accept() + + return nil +} + +// Len reports accepted-or-reserved work that has not reached a terminal outcome. +// +// Note this includes work currently inside a processor call, so it can be non-zero +// while the queue is empty. Use Stats().Queued for queue depth. func (b *Batcher[T]) Len() int { - return int(b.itemCount.Read()) + return int(b.counters.pending.Load()) } +// Stats returns an O(1), allocation-free snapshot of the batcher's counters. See +// Stats for its consistency guarantees. +func (b *Batcher[T]) Stats() Stats { + return Stats{ + Pending: b.counters.pending.Load(), + IntakePending: b.counters.intakePending.Load(), + PublishersInGate: b.gate.inGate(), + Queued: int64(b.input.length()), + Accepted: b.counters.accepted.Load(), + Completed: b.counters.completed.Load(), + Failed: b.counters.failed.Load(), + Panicked: b.counters.panicked.Load(), + Rejected: b.counters.rejected.Load(), + DroppedErrors: b.counters.droppedErrors.Load(), + } +} + +// Join waits until no work is pending, or the timeout expires. +// +// This is a quiescence snapshot, not a barrier: a concurrent Add can make work +// pending again the instant it returns. It is only authoritative after admission +// has been sealed. func (b *Batcher[T]) Join(timeout time.Duration) error { deadline := time.Now().Add(timeout) @@ -102,84 +227,346 @@ func (b *Batcher[T]) Join(timeout time.Duration) error { } } -func (b *Batcher[T]) startProcessing() { - // Close channels in the correct order to prevent goroutine leaks: - // 1. First close the input channel to stop the rill pipeline from producing more batches - // 2. Then drain the output channel to allow rill goroutines to exit cleanly - // 3. Finally close the errors channel - defer b.errorsChan.Close() - defer func() { - // Drain any remaining batches from the rill pipeline to ensure - // all rill goroutines can exit cleanly. This prevents goroutine leaks. - for range b.batchesChan { - // Just drain, don't process during shutdown +// Errors returns the diagnostics channel. It is closed once shutdown completes, +// so a consumer ranging over it terminates. +func (b *Batcher[T]) Errors() <-chan error { + return b.errorsChan +} + +// run aggregates items into batches and hands each one to the processing loop. +// +// Batching rules, unchanged from the previous implementation and pinned by the +// characterization tests: +// +// - The interval timer is armed when the first item enters an empty batch, not +// on a periodic tick. A sparse producer waits one interval per item, and an +// idle batcher does no work at all. +// - A batch reaching BatchSize flushes immediately. +// - Empty batches are never emitted. +// - Shutdown flushes the pending partial batch rather than discarding it. +// +// Aggregation runs here and processing runs in its own goroutine, connected by an +// unbuffered channel. That separation is load-bearing: it lets the next batch +// accumulate while the current one is being processed. Merging the two measurably +// changes latency — with a 50ms processor at 10k items/s a merged loop made a 5ms +// window faster than a 100ms one, inverting the documented baseline — and that is +// a Phase 3 decision, not this milestone's. +func (b *Batcher[T]) run() { + defer close(b.stopped) + + batches := make(chan []T) + + processing := make(chan struct{}) + + go func() { + defer close(processing) + + for items := range batches { + b.process(items) } }() - defer b.batchInputChan.Close() - for { - select { - case <-b.doneChan: + defer func() { + close(batches) + <-processing + }() + + var ( + batch []T + timer *time.Timer + timerC <-chan time.Time + ) + + stopTimer := func() { + timerC = nil + + if timer == nil { return - case batch, ok := <-b.batchesChan: + } + + if !timer.Stop() { + // Drain a tick that fired while we were flushing, so the next batch + // cannot inherit a stale expiry and flush early. + select { + case <-timer.C: + default: + } + } + } + + flush := func() { + if len(batch) == 0 { + return + } + + items := batch + batch = nil + + stopTimer() + + batches <- items + } + + take := func(item T) { + if len(batch) == 0 { + batch = make([]T, 0, b.config.BatchSize) + + if timer == nil { + timer = time.NewTimer(b.config.BatchInterval) + } else { + timer.Reset(b.config.BatchInterval) + } + + timerC = timer.C + } + + batch = append(batch, item) + + if len(batch) >= b.config.BatchSize { + flush() + } + } + + // drainReady empties whatever is currently queued. The aggregator drains + // greedily rather than one item per wakeup, so a burst costs one signal. + drainReady := func() { + for { + item, ok := b.input.pop() if !ok { - // Channel closed, exit gracefully return } - if batch.Error != nil { - b.errorsChan.In() <- batch.Error + b.counters.received(1) + take(item) + } + } - continue + for { + select { + case <-b.input.ready(): + drainReady() + case <-timerC: + flush() + case <-b.noInput: + // Intake is finished: no publisher can ever send again. Drain by + // accounting rather than by observing an empty queue, so an item that was + // reserved-then-published during shutdown cannot be missed. + for b.counters.intakePending.Load() > 0 { + drainReady() } - if err := b.config.ProcessorFunc(batch.Value); err != nil { - b.errorsChan.In() <- err - } + flush() - b.itemCount.Add(int64(-len(batch.Value))) + return } } } -func (b *Batcher[T]) Errors() <-chan error { - return b.errorsChan.Out() +// process invokes the processor for one batch and records exactly one terminal +// outcome for it. +// +// A panic in the processor is recovered and scoped to the batch. Batcher owns +// goroutines the caller cannot see, so it must not let a caller's bug become a +// process-wide crash that also destroys every other queued item: the caller has no +// way to wrap our goroutine in its own recover. The panic value and stack are +// reported as a diagnostic so the bug stays debuggable. +// +// Recovery is scoped to the batch rather than the loop, so one poison batch cannot +// silently take the consumer down and stall every later batch. +// +// The ordering is fixed and load-bearing: whatever happens, exactly one terminal +// category is counted and the drain obligation is released exactly once. Leaving it +// unreleased would inflate Pending permanently and hang shutdown forever. +func (b *Batcher[T]) process(items []T) { + outcome := outcomeCompleted + + defer func() { + b.counters.terminal(outcome, len(items)) + }() + + defer func() { + recovered := recover() + if recovered == nil { + return + } + + outcome = outcomePanicked + + b.publishError(&ProcessorPanicError{ + Value: recovered, + Stack: debug.Stack(), + }) + }() + + if err := b.config.ProcessorFunc(items); err != nil { + outcome = outcomeFailed + + b.publishError(err) + } +} + +// publishError reports a diagnostic without ever blocking the pipeline. +// +// Dropping is deliberate, and the alternative is worse. A bounded channel with a +// blocking send would deadlock the pipeline whenever a processor fails and nobody +// is reading Errors() — a self-inflicted outage caused by a diagnostic. An +// unbounded channel would instead grow without limit during an outage, which is +// exactly when memory is scarcest. +// +// Dropping the newest error rather than the oldest keeps the earliest errors of a +// storm, which usually carry the root cause, and avoids racing the consumer. +// DroppedErrors makes the loss visible: a non-zero value means "you are not +// draining Errors() fast enough", which is itself worth alerting on. +func (b *Batcher[T]) publishError(err error) { + select { + case b.errorsChan <- err: + default: + b.counters.droppedErrors.Add(1) + } } +// Close seals admission and drains accepted work, waiting up to the configured +// grace period (WithCloseGrace, default 30s). +// +// It reports an incomplete drain if the grace expires, but it does NOT abandon the +// drain: accepted work keeps being processed in the background. Reporting an +// incomplete drain and causing one are very different things, and the previous +// implementation silently discarded accepted items. +// +// Callers who want to control the wait, or to keep waiting after a timeout, should +// use Shutdown. func (b *Batcher[T]) Close() error { - var clErr error + ctx, cancel := context.WithTimeout(context.Background(), b.config.CloseGrace) + defer cancel() - b.closeOnce.Do(func() { - // Calculate timeout based on pending items, with a maximum of 10 seconds - timeout := time.Duration(2*2*math.Ceil(float64(b.Len())/float64(b.config.BatchSize))) * - b.config.BatchInterval + return b.Shutdown(ctx) +} - // Cap the timeout at 10 seconds to prevent indefinite blocking - maxTimeout := 10 * time.Second - if timeout > maxTimeout { - timeout = maxTimeout - } +// Shutdown seals admission and waits for accepted work to drain. +// +// It is resumable. If ctx expires first it returns *ShutdownIncompleteError and the +// drain continues; calling Shutdown again with a fresh context waits on that same +// drain rather than starting a new one. That is why there is no continuation +// token: the Batcher itself is the continuation, and it is sealed against new work +// from the first call onward. +// +// Concurrent callers are independent. A caller with a short deadline receives its +// own timeout error; it is never stored as the batcher's terminal result, so it +// cannot poison a later caller who waits longer. +// +// Processor errors and recovered panics are NOT returned here. They remain on +// Errors(), because a shutdown result describes the drain, not the work. +// +// A processor that never returns holds the batcher in draining indefinitely. This +// library cannot cancel user code, so Shutdown reports the condition rather than +// pretending to resolve it. +func (b *Batcher[T]) Shutdown(ctx context.Context) error { + b.beginSealing() + + // An unstarted batcher may still hold queued work. Start the consumer now, or + // the drain could never make progress. This is the one case where shutdown + // starts a lifecycle, and it must happen before any wait: for a bounded queue + // with a parked publisher, nothing else can unblock the gate. + b.startDrainConsumer() + + // The coordinator runs independently of any caller's patience. + b.shutdownOnce.Do(func() { + go b.coordinateDrain() + }) + + // A completed drain wins over an expired context. A select chooses randomly + // when both cases are ready; reporting ShutdownIncompleteError with Pending=0 + // for a batcher that is already closed would be false. + select { + case <-b.shutdownDone: + return nil + default: + } - // Ensure a minimum timeout of 100ms - if timeout < 100*time.Millisecond { - timeout = 100 * time.Millisecond + select { + case <-b.shutdownDone: + return nil + case <-ctx.Done(): + return &ShutdownIncompleteError{ + Pending: b.Len(), + PublishersInGate: int(b.gate.inGate()), + Cause: ctx.Err(), } + } +} - clErr = b.Join(timeout) +// IsClosing reports whether admission has been sealed. The drain may still be in +// progress. +func (b *Batcher[T]) IsClosing() bool { + return b.gate.state.Load() >= stateSealing +} - // Signal shutdown to the processing goroutine - close(b.doneChan) +// IsClosed reports whether shutdown has fully completed and every owned goroutine +// has exited. +func (b *Batcher[T]) IsClosed() bool { + return b.gate.state.Load() == stateClosed +} + +// coordinateDrain performs the drain exactly once, however many callers are +// waiting and however soon they give up. +func (b *Batcher[T]) coordinateDrain() { + defer close(b.shutdownDone) - // Give the goroutine a moment to exit cleanly - // This allows the deferred cleanup to run - time.Sleep(50 * time.Millisecond) + // Wait for publishers to leave the publication window. The aggregator keeps + // draining throughout, which is what lets a publisher parked on a full bounded + // queue make progress and exit the gate. + <-b.gate.empty() - b.isClosed = true + b.gate.state.CompareAndSwap(stateSealing, stateDraining) + + // Only now can the aggregator conclude that intake is finished: no publisher + // can ever send again. + b.drainOnce.Do(func() { + close(b.noInput) }) - return clErr + <-b.stopped + + b.finish() } -func (b *Batcher[T]) IsClosed() bool { - return b.isClosed +// beginSealing closes admission and performs the coordinator's own quiescence +// check. +func (b *Batcher[T]) beginSealing() { + for { + current := b.gate.state.Load() + if current >= stateSealing { + break + } + + if b.gate.state.CompareAndSwap(current, stateSealing) { + break + } + } + + b.gate.seal() + b.gate.checkEmpty() +} + +// startDrainConsumer ensures a consumer exists to drain queued work during +// shutdown, even if the batcher was never started. +// +// It shares startOnce with Start, so exactly one consumer ever runs regardless of +// which path gets there first. +func (b *Batcher[T]) startDrainConsumer() { + b.Start() +} + +// finish performs terminal cleanup exactly once: close diagnostics after every +// sender has exited, then mark the batcher closed. +func (b *Batcher[T]) finish() { + if b.gate.state.CompareAndSwap(stateDraining, stateClosed) { + close(b.errorsChan) + + return + } + + // Another caller may already be finishing; make sure the terminal state is + // visible before returning. + b.gate.state.CompareAndSwap(stateSealing, stateClosed) } diff --git a/pkg/batcher/batcher_test.go b/pkg/batcher/batcher_test.go index 3f60a63..9a07bac 100644 --- a/pkg/batcher/batcher_test.go +++ b/pkg/batcher/batcher_test.go @@ -9,7 +9,7 @@ import ( "github.com/NSXBet/batcher/internal/test" "github.com/NSXBet/batcher/pkg/batcher" "github.com/stretchr/testify/require" - "go.uber.org/atomic" + "sync/atomic" ) func TestCanCreateBatcherWithDefaultConfig(t *testing.T) { @@ -234,8 +234,8 @@ func TestFlushOnClose(t *testing.T) { func TestProcessesEntireBatchesIfTimerHasNotExpired(t *testing.T) { // ARRANGE - batches := atomic.NewInt32(0) - processed := atomic.NewInt32(0) + var batches atomic.Int32 + var processed atomic.Int32 b := batcher.New( batcher.WithBatchSize[test.BatchItem](5), diff --git a/pkg/batcher/characterization_test.go b/pkg/batcher/characterization_test.go new file mode 100644 index 0000000..4cea0c7 --- /dev/null +++ b/pkg/batcher/characterization_test.go @@ -0,0 +1,410 @@ +package batcher_test + +import ( + "sync" + "testing" + "time" + + "github.com/NSXBet/batcher/internal/test" + "github.com/NSXBet/batcher/pkg/batcher" + "github.com/stretchr/testify/require" +) + +// Characterization tests for the batching engine. +// +// These were written against the rill-backed implementation and must keep +// passing, unchanged, after the aggregation loop is brought in-repo. They exist +// to pin observable behaviour that no other test asserts, so that replacing the +// engine cannot quietly change semantics. +// +// Each test states the behaviour it locks down and why it matters. Where the +// behaviour is a consequence of rill's design rather than a deliberate contract, +// that is called out, because those are exactly the places a reimplementation +// could drift. + +// recordingProcessor captures the batches it receives, in order. +type recordingProcessor struct { + mu sync.Mutex + batches [][]test.BatchItem + flushed chan int +} + +func newRecordingProcessor(capacity int) *recordingProcessor { + return &recordingProcessor{flushed: make(chan int, capacity)} +} + +func (r *recordingProcessor) process(items []test.BatchItem) error { + // Copy: the engine owns the slice it passes in, and a future implementation + // may reuse its backing array. + batch := make([]test.BatchItem, len(items)) + copy(batch, items) + + r.mu.Lock() + r.batches = append(r.batches, batch) + r.mu.Unlock() + + r.flushed <- len(batch) + + return nil +} + +func (r *recordingProcessor) snapshot() [][]test.BatchItem { + r.mu.Lock() + defer r.mu.Unlock() + + out := make([][]test.BatchItem, len(r.batches)) + copy(out, r.batches) + + return out +} + +// waitForBatches blocks until n batches have been processed, or fails. +func (r *recordingProcessor) waitForBatches(t *testing.T, n int, within time.Duration) { + t.Helper() + + deadline := time.After(within) + + for range n { + select { + case <-r.flushed: + case <-deadline: + t.Fatalf("timed out waiting for %d batches; got %d", n, len(r.snapshot())) + } + } +} + +// TestCharacterizeTimerStartsOnFirstItemOfBatch pins the single most surprising +// property of the engine: the batch interval is not a periodic flush tick. The +// timer starts when the first item enters an empty batch, so a sparse producer +// waits the full interval for every item. +// +// This is what makes the configured window a per-hop latency cost, and it is the +// behaviour the whole performance plan is built around. A reimplementation that +// used a periodic ticker instead would change latency for every user. +func TestCharacterizeTimerStartsOnFirstItemOfBatch(t *testing.T) { + t.Parallel() + + const interval = 100 * time.Millisecond + + proc := newRecordingProcessor(8) + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](1_000), // never size-triggered + batcher.WithBatchInterval[test.BatchItem](interval), + batcher.WithProcessor(proc.process), + ) + + t.Cleanup(func() { _ = b.Close() }) + + // Let the batcher idle well past one interval with nothing queued. + time.Sleep(3 * interval) + + require.Empty(t, proc.snapshot(), + "an idle batcher must not emit empty batches; the timer is not a periodic tick") + + // The interval is measured from this item, not from any earlier tick. + start := time.Now() + b.Add(test.BatchItem{Key: "first"}) + + proc.waitForBatches(t, 1, 5*interval) + elapsed := time.Since(start) + + require.GreaterOrEqual(t, elapsed, interval/2, + "flush must wait for the interval that started with the first item, not fire immediately") + + batches := proc.snapshot() + require.Len(t, batches, 1) + require.Len(t, batches[0], 1, "a timer flush emits exactly what accumulated") +} + +// TestCharacterizeNoEmptyBatchesEverEmitted pins that the processor is never +// invoked with an empty slice, even across many idle intervals. User processors +// are entitled to assume len(items) > 0. +func TestCharacterizeNoEmptyBatchesEverEmitted(t *testing.T) { + t.Parallel() + + const interval = 20 * time.Millisecond + + var invocations, emptyInvocations int + + var mu sync.Mutex + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](10), + batcher.WithBatchInterval[test.BatchItem](interval), + batcher.WithProcessor(func(items []test.BatchItem) error { + mu.Lock() + defer mu.Unlock() + + invocations++ + + if len(items) == 0 { + emptyInvocations++ + } + + return nil + }), + ) + + t.Cleanup(func() { _ = b.Close() }) + + time.Sleep(10 * interval) + + mu.Lock() + defer mu.Unlock() + + require.Zero(t, invocations, "an idle batcher must never invoke the processor") + require.Zero(t, emptyInvocations, "the processor must never receive an empty batch") +} + +// TestCharacterizeSizeTriggeredFlushIgnoresInterval pins that a full batch is +// emitted immediately rather than waiting for the timer, which is what makes +// batching efficient under load. +func TestCharacterizeSizeTriggeredFlushIgnoresInterval(t *testing.T) { + t.Parallel() + + const ( + batchSize = 5 + interval = 10 * time.Second // must not be reached + ) + + proc := newRecordingProcessor(8) + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](batchSize), + batcher.WithBatchInterval[test.BatchItem](interval), + batcher.WithProcessor(proc.process), + ) + + t.Cleanup(func() { _ = b.Close() }) + + for i := range batchSize { + b.Add(test.BatchItem{Key: keyFor(i)}) + } + + proc.waitForBatches(t, 1, 2*time.Second) + + batches := proc.snapshot() + require.Len(t, batches, 1) + require.Len(t, batches[0], batchSize, + "a full batch must flush on size without waiting for the interval") +} + +// TestCharacterizeBatchBoundariesAndOrdering pins exact batch boundaries and +// ordering for a sequential producer: items are grouped in arrival order, in +// batches of exactly BatchSize until the remainder. +// +// Ordering here is a property of a single producer's program order. Concurrent +// producers have no defined relative order, and this test deliberately does not +// claim otherwise. +func TestCharacterizeBatchBoundariesAndOrdering(t *testing.T) { + t.Parallel() + + const ( + batchSize = 4 + total = 10 // two full batches plus a remainder of two + ) + + proc := newRecordingProcessor(8) + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](batchSize), + batcher.WithBatchInterval[test.BatchItem](25*time.Millisecond), + batcher.WithProcessor(proc.process), + ) + + for i := range total { + b.Add(test.BatchItem{Key: keyFor(i)}) + } + + require.NoError(t, b.Join(5*time.Second)) + require.NoError(t, b.Close()) + + batches := proc.snapshot() + + var ( + sizes []int + keys []string + ) + + for _, batch := range batches { + sizes = append(sizes, len(batch)) + + for _, item := range batch { + keys = append(keys, item.Key) + } + } + + require.Equal(t, []int{batchSize, batchSize, total - 2*batchSize}, sizes, + "batches must be exactly BatchSize until the remainder") + + expected := make([]string, total) + for i := range expected { + expected[i] = keyFor(i) + } + + require.Equal(t, expected, keys, + "a single producer's items must be processed in arrival order") +} + +// TestCharacterizePartialBatchFlushesOnClose pins that shutdown flushes an +// under-full batch rather than discarding it. This is the guarantee Phase 2.3 +// later strengthens; pinning it now means the engine swap cannot regress it. +func TestCharacterizePartialBatchFlushesOnClose(t *testing.T) { + t.Parallel() + + proc := newRecordingProcessor(4) + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](1_000), // never size-triggered + batcher.WithBatchInterval[test.BatchItem](25*time.Millisecond), + batcher.WithProcessor(proc.process), + ) + + b.Add(test.BatchItem{Key: "only"}) + + require.NoError(t, b.Close()) + + batches := proc.snapshot() + require.Len(t, batches, 1, "a partial batch must be flushed during shutdown") + require.Len(t, batches[0], 1) + require.Equal(t, "only", batches[0][0].Key) +} + +// TestCharacterizeProcessorErrorsSurfaceOnErrorsChannel pins error propagation: +// a processor error is published without stopping the pipeline, and the item is +// still accounted for so Join can complete. +func TestCharacterizeProcessorErrorsSurfaceOnErrorsChannel(t *testing.T) { + t.Parallel() + + failure := errCharacterization + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](1), + batcher.WithBatchInterval[test.BatchItem](10*time.Millisecond), + batcher.WithProcessor(func([]test.BatchItem) error { + return failure + }), + ) + + t.Cleanup(func() { _ = b.Close() }) + + b.Add(test.BatchItem{Key: "fails"}) + + select { + case err := <-b.Errors(): + require.ErrorIs(t, err, failure) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for a processor error") + } + + // A failed batch still counts as completed work: Join must not hang. + require.NoError(t, b.Join(5*time.Second), + "a failed batch must still decrement pending work") + + // The pipeline must remain usable after an error. + b.Add(test.BatchItem{Key: "also-fails"}) + + select { + case err := <-b.Errors(): + require.ErrorIs(t, err, failure) + case <-time.After(5 * time.Second): + t.Fatal("pipeline stopped processing after an error") + } +} + +// TestCharacterizeErrorsChannelClosesAfterShutdown pins that Errors() is closed +// on shutdown, so a consumer ranging over it terminates instead of blocking. +func TestCharacterizeErrorsChannelClosesAfterShutdown(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](10), + batcher.WithBatchInterval[test.BatchItem](10*time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[test.BatchItem]), + ) + + drained := make(chan struct{}) + + go func() { + defer close(drained) + + for range b.Errors() { + } + }() + + require.NoError(t, b.Close()) + + select { + case <-drained: + case <-time.After(5 * time.Second): + t.Fatal("Errors() was not closed on shutdown; a consumer would block forever") + } +} + +// TestCharacterizeLenReflectsAcceptedButUnfinishedWork pins Len's meaning: it +// counts accepted work that has not finished processing, and returns to zero +// once the pipeline drains. +func TestCharacterizeLenReflectsAcceptedButUnfinishedWork(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + entered := make(chan struct{}, 1) + + b := batcher.New( + batcher.WithBatchSize[test.BatchItem](2), + batcher.WithBatchInterval[test.BatchItem](10*time.Millisecond), + batcher.WithProcessor(func([]test.BatchItem) error { + select { + case entered <- struct{}{}: + default: + } + + <-release + + return nil + }), + ) + + t.Cleanup(func() { + close(release) + _ = b.Close() + }) + + b.Add(test.BatchItem{Key: "a"}) + b.Add(test.BatchItem{Key: "b"}) + + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("processor was never invoked") + } + + require.Positive(t, b.Len(), + "work handed to a blocked processor must still count as pending") +} + +func keyFor(i int) string { + return "item-" + itoa(i) +} + +func itoa(i int) string { + if i == 0 { + return "0" + } + + var digits []byte + + for i > 0 { + digits = append([]byte{byte('0' + i%10)}, digits...) + i /= 10 + } + + return string(digits) +} + +var errCharacterization = errCharacterizationType("characterization failure") + +type errCharacterizationType string + +func (e errCharacterizationType) Error() string { return string(e) } diff --git a/pkg/batcher/constants.go b/pkg/batcher/constants.go index 443e15f..80110ea 100644 --- a/pkg/batcher/constants.go +++ b/pkg/batcher/constants.go @@ -11,4 +11,16 @@ const ( // DefaultConcurrency is the default concurrency. DefaultConcurrency = 3 + + // DefaultCloseGrace is how long Close waits for the drain to finish before + // reporting that it is incomplete. The drain continues regardless. + // + // A flat value replaces the previous formula derived from queue size and batch + // interval, which was impossible to reason about and was the direct cause of + // accepted work being discarded when the interval exceeded its 10s cap. + DefaultCloseGrace = 30 * time.Second + + // DefaultErrorBufferSize bounds retained diagnostics. The channel must never + // block the pipeline, so a full buffer drops rather than waits. + DefaultErrorBufferSize = 1024 ) diff --git a/pkg/batcher/diagnostics_test.go b/pkg/batcher/diagnostics_test.go new file mode 100644 index 0000000..6dad2a7 --- /dev/null +++ b/pkg/batcher/diagnostics_test.go @@ -0,0 +1,187 @@ +package batcher_test + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/NSXBet/batcher/pkg/batcher" + "github.com/stretchr/testify/require" +) + +// Diagnostics and panic-recovery tests. +// +// Errors are diagnostic, not payload. The tests here make the deliberate trade +// explicit: Batcher must keep processing when diagnostics are not consumed, and +// must report how many diagnostics it had to drop rather than turning an outage +// into a deadlock or an OOM. + +// TestErrorStormWithoutConsumerIsBoundedAndNonBlocking proves the error channel +// cannot grow without limit or stall the processor when nobody reads it. +func TestErrorStormWithoutConsumerIsBoundedAndNonBlocking(t *testing.T) { + t.Parallel() + + const ( + bufferSize = 3 + items = 100 + ) + + failure := errors.New("downstream unavailable") + + var calls atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithErrorBufferSize[int](bufferSize), + batcher.WithProcessor(func([]int) error { + calls.Add(1) + + return failure + }), + ) + + for i := range items { + b.Add(i) + } + + // If publishError blocked on the full channel, this wait would hang after + // bufferSize failures. It must instead drain all work. + require.NoError(t, b.Join(10*time.Second)) + + stats := b.Stats() + + require.Equal(t, int64(items), calls.Load(), + "a full diagnostics buffer must not stop later batches") + require.Equal(t, uint64(items), stats.Failed) + require.Equal(t, uint64(items-bufferSize), stats.DroppedErrors, + "drop-newest must retain the first bufferSize diagnostics and count the rest") + require.Zero(t, stats.Pending, "errors must still release their drain obligations") + + // Closing must not hang even though nobody has drained Errors(). + require.NoError(t, b.Close()) +} + +// TestProcessorPanicIsReportedAndPipelineSurvives proves a poison batch cannot +// crash the process, take the consumer down, or strand every item behind it. +func TestProcessorPanicIsReportedAndPipelineSurvives(t *testing.T) { + t.Parallel() + + var calls atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func([]int) error { + if calls.Add(1) == 1 { + panic("poison batch") + } + + return nil + }), + ) + + b.Add(1) // panics + b.Add(2) // must still be processed + + var diagnostic error + + select { + case diagnostic = <-b.Errors(): + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for recovered panic diagnostic") + } + + var panicErr *batcher.ProcessorPanicError + + require.ErrorAs(t, diagnostic, &panicErr) + require.Equal(t, "poison batch", panicErr.Value) + require.NotEmpty(t, panicErr.Stack, "the recovered panic must retain a stack trace") + require.True(t, strings.Contains(string(panicErr.Stack), "TestProcessorPanic"), + "the stack trace must identify the caller that panicked") + + require.NoError(t, b.Join(5*time.Second)) + + stats := b.Stats() + + require.Equal(t, int64(2), calls.Load(), "the consumer must survive the first panic") + require.Equal(t, uint64(1), stats.Panicked) + require.Equal(t, uint64(1), stats.Completed) + require.Zero(t, stats.Failed, "a panic is not a processor error") + require.Zero(t, stats.Pending, "a panic must not strand its drain obligation") + + require.NoError(t, b.Shutdown(context.Background()), + "a recovered panic must not make the shutdown result fail") +} + +// TestPanicInFinalShutdownBatchDrains proves recovery works for the batch most +// likely to regress: the partial batch flushed only by shutdown. +func TestPanicInFinalShutdownBatchDrains(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithBatchSize[int](100), + batcher.WithBatchInterval[int](time.Hour), + batcher.WithProcessor(func([]int) error { + panic("shutdown batch") + }), + ) + + b.Add(1) + + // Drain the diagnostic concurrently, so this test also covers the final + // diagnostic publication racing the coordinator's error-channel close. + diagnostic := make(chan error, 1) + + go func() { diagnostic <- <-b.Errors() }() + + require.NoError(t, b.Shutdown(context.Background())) + + var err error + + select { + case err = <-diagnostic: + case <-time.After(5 * time.Second): + t.Fatal("the final panic diagnostic was never published") + } + + var panicErr *batcher.ProcessorPanicError + require.ErrorAs(t, err, &panicErr) + + stats := b.Stats() + + require.Equal(t, uint64(1), stats.Panicked) + require.Zero(t, stats.Pending, "the final panic must not hang shutdown") + require.True(t, b.IsClosed()) +} + +// TestRecoveredPanicAddsNoSteadyStateAllocations makes the recovery guard's cost +// explicit. A defer/recover guard must not allocate when nothing panics; otherwise +// a safety feature would tax every batch. +func TestRecoveredPanicAddsNoSteadyStateAllocations(t *testing.T) { + // Deliberately not parallel: AllocsPerRun requires exclusive GOMAXPROCS. + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + t.Cleanup(func() { _ = b.Close() }) + + // Warm up the queue and processing path. + for range 1_000 { + b.Add(1) + } + + require.NoError(t, b.Join(10*time.Second)) + + // This measures Add's producer-side cost, which must remain allocation-free + // despite batches being guarded by recover in another goroutine. + allocs := testing.AllocsPerRun(2_000, func() { b.Add(1) }) + + require.Zero(t, allocs, + "the normal non-panic path must not allocate because recovery is enabled") +} diff --git a/pkg/batcher/fx.go b/pkg/batcher/fx.go index e74d7ca..a9fab4a 100644 --- a/pkg/batcher/fx.go +++ b/pkg/batcher/fx.go @@ -38,12 +38,17 @@ func ProvideBatcherInFX[T any]( return nil })) - lifecycle.Append(fx.StopHook(func(context.Context) error { - if err := batcher.Close(); err != nil { - return err - } - - return nil + // Forward the stop-hook context instead of discarding it, so the + // application's shutdown deadline governs the drain. Previously this + // used Close's own timeout, which meant an app with a longer grace + // period could not use it, and one with a shorter deadline could not + // bound it. + // + // If the hook's context expires the drain continues in the background + // and the error describes what remained, rather than silently + // discarding accepted work. + lifecycle.Append(fx.StopHook(func(ctx context.Context) error { + return batcher.Shutdown(ctx) })) }, ), diff --git a/pkg/batcher/fx_test.go b/pkg/batcher/fx_test.go index e63248b..3502463 100644 --- a/pkg/batcher/fx_test.go +++ b/pkg/batcher/fx_test.go @@ -1,11 +1,11 @@ package batcher_test import ( + "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" - "go.uber.org/atomic" "go.uber.org/fx" "go.uber.org/fx/fxtest" @@ -18,11 +18,11 @@ type BatchItem struct { } type Processor struct { - counter *atomic.Uint32 + counter atomic.Uint32 } func NewProcessor() *Processor { - return &Processor{counter: atomic.NewUint32(0)} + return &Processor{} } func (p *Processor) Process(items []*BatchItem) error { diff --git a/pkg/batcher/gate.go b/pkg/batcher/gate.go new file mode 100644 index 0000000..c8e8905 --- /dev/null +++ b/pkg/batcher/gate.go @@ -0,0 +1,127 @@ +package batcher + +import ( + "sync" + "sync/atomic" +) + +// lifecycle state values. Exactly one holds at any time. +const ( + stateNew int32 = iota // constructed, processing not started + stateRunning // processing started, admitting + stateSealing // admission closed, publishers still leaving the gate + stateDraining // gate empty, accepted work still processing + stateClosed // terminal: all owned goroutines have exited +) + +// admissionGate serialises "is admission still open?" against shutdown without +// putting a mutex on the enqueue path. +// +// The gate counts publishers currently inside the publication window. Its purpose +// is to let the shutdown coordinator prove that no publisher can still be +// mid-publish, which is what makes it safe to declare intake finished. +// +// The protocol is subtle in two places, and both have dedicated tests: +// +// 1. enter() increments, then RE-CHECKS the seal. Without the re-check, a +// publisher that read "not sealed" could increment after the coordinator had +// already observed an empty gate, and publish into a queue nobody will drain. +// 2. every decrement routes through leave(), including enter()'s own rejection +// path. A bare decrement there loses the wakeup: the coordinator can observe a +// non-zero gate and commit to waiting, and the rejecting publisher can then +// drop the count to zero without signalling, hanging shutdown forever. +type admissionGate struct { + state atomic.Int32 + sealed atomic.Bool + + // afterIncrement is used only by package-internal tests to force the race + // between incrementing the gate and re-checking the seal. It is nil in + // production and has no hot-path effect beyond one nil check. + afterIncrement func() + + // publishers counts goroutines between a successful enter() and its leave(). + publishers atomic.Int64 + + // sealCh is closed once when sealing begins, releasing publishers parked on a + // full bounded queue. Closing is used rather than sending so every parked + // publisher observes it. + sealCh chan struct{} + sealOnce sync.Once + emptyCh chan struct{} + emptyOnce sync.Once +} + +func newAdmissionGate() *admissionGate { + return &admissionGate{ + sealCh: make(chan struct{}), + emptyCh: make(chan struct{}), + } +} + +// enter attempts to join the publication window. +// +// A successful enter must be paired with exactly one leave. +func (g *admissionGate) enter() bool { + if g.sealed.Load() { + return false + } + + g.publishers.Add(1) + + if g.afterIncrement != nil { + g.afterIncrement() + } + + if g.sealed.Load() { + // Sealing raced our increment. Route the decrement through leave() so the + // coordinator cannot miss the transition to zero. + g.leave() + + return false + } + + return true +} + +// leave exits the publication window, signalling quiescence when it is the last +// publisher out after sealing. +func (g *admissionGate) leave() { + if g.publishers.Add(-1) == 0 && g.sealed.Load() { + g.signalEmpty() + } +} + +// seal closes admission. It is idempotent. +// +// The caller must follow this with checkEmpty, which covers the interleaving where +// the last publisher left before sealed was stored and therefore never signalled. +func (g *admissionGate) seal() { + g.sealed.Store(true) + g.sealOnce.Do(func() { close(g.sealCh) }) +} + +// checkEmpty is the coordinator's own quiescence check. +// +// This is mandatory, not an optimisation: leave() only signals when it observes +// sealed, so if the final publisher left before sealing, nobody would ever signal. +func (g *admissionGate) checkEmpty() { + if g.publishers.Load() == 0 && g.sealed.Load() { + g.signalEmpty() + } +} + +// signalEmpty arms the quiescence latch exactly once. It closes a channel rather +// than sending, so it never blocks and every waiter observes it. +func (g *admissionGate) signalEmpty() { + g.emptyOnce.Do(func() { close(g.emptyCh) }) +} + +// empty returns the latch that closes once no publisher remains in the window. +func (g *admissionGate) empty() <-chan struct{} { + return g.emptyCh +} + +// inGate reports how many publishers are inside the publication window. +func (g *admissionGate) inGate() int64 { + return g.publishers.Load() +} diff --git a/pkg/batcher/gate_internal_test.go b/pkg/batcher/gate_internal_test.go new file mode 100644 index 0000000..42aedfa --- /dev/null +++ b/pkg/batcher/gate_internal_test.go @@ -0,0 +1,64 @@ +package batcher + +import ( + "testing" + "time" +) + +// TestGateRejectionPathSignalsQuiescence forces the lost-wakeup interleaving that +// only the rejection path can produce: +// +// 1. an entrant increments the gate; +// 2. the coordinator seals, observes a NON-ZERO gate, and commits to waiting; +// 3. the entrant's post-increment re-check sees the seal and decrements to zero. +// +// If that decrement bypasses leave(), nobody signals quiescence and the waiter +// blocks forever. This is distinct from the last-leaver-before-seal case, because +// there the coordinator's own check finds the gate empty and signals; here it has +// already found the gate occupied, so only leave() can signal. +// +// The interleaving is forced with a hook rather than hoped for: a test that only +// sometimes exercises the hazard is not a regression test. Sabotaging enter()'s +// leave() call must make this hang. +func TestGateRejectionPathSignalsQuiescence(t *testing.T) { + t.Parallel() + + const trials = 500 + + for trial := range trials { + g := newAdmissionGate() + + incremented := make(chan struct{}) + sealed := make(chan struct{}) + + // Pause the entrant between its increment and its re-check. + g.afterIncrement = func() { + close(incremented) + <-sealed + } + + entered := make(chan bool, 1) + + go func() { entered <- g.enter() }() + + <-incremented + + // The coordinator seals while the gate is occupied, so its own check cannot + // signal quiescence. Only the entrant's exit can. + g.seal() + g.checkEmpty() + + close(sealed) + + if <-entered { + t.Fatalf("trial %d: entry must be rejected once sealing has begun", trial) + } + + select { + case <-g.empty(): + case <-time.After(2 * time.Second): + t.Fatalf("trial %d: quiescence was never signalled; the rejection path's "+ + "gate decrement bypassed leave()", trial) + } + } +} diff --git a/pkg/batcher/goroutine_budget_test.go b/pkg/batcher/goroutine_budget_test.go new file mode 100644 index 0000000..9a285a9 --- /dev/null +++ b/pkg/batcher/goroutine_budget_test.go @@ -0,0 +1,103 @@ +package batcher_test + +import ( + "runtime" + "testing" + "time" + + "github.com/NSXBet/batcher/internal/test" + "github.com/NSXBet/batcher/pkg/batcher" + "github.com/stretchr/testify/require" +) + +// goroutinesPerRunningBatcher is the owned-goroutine budget after Phase 2.2: +// +// 1. the aggregator +// 2. the processing loop +// +// The input queue and diagnostics queue are owned data structures, not relay +// goroutines. Aggregation and processing remain separate because merging them +// changes observable latency; Phase 3 may change this topology when it introduces +// explicit worker concurrency. +const goroutinesPerRunningBatcher = 2 + +// TestGoroutineBudgetPerRunningBatcher pins the owned-goroutine count. +// +// The count started at 6 (rill plus chann relays), dropped to 5 when rill was +// replaced in 2.1, and reaches 2 in 2.2 once both relay-backed queues become owned +// data structures. Per-batcher goroutine cost matters because callers create one +// batcher per tenant or per key, so this is asserted rather than left to a +// measurement someone runs by hand. +func TestGoroutineBudgetPerRunningBatcher(t *testing.T) { + // Not parallel: goroutine counting requires no other test starting or + // stopping batchers concurrently. + const batchers = 20 + + baseline := stableGoroutineCount(t) + + live := make([]*batcher.Batcher[test.BatchItem], 0, batchers) + + for range batchers { + live = append(live, batcher.New( + batcher.WithBatchSize[test.BatchItem](100), + batcher.WithBatchInterval[test.BatchItem](time.Hour), // stay idle + batcher.WithProcessor(batcher.NoOpProcessor[test.BatchItem]), + )) + } + + // Give every owned goroutine time to be scheduled and counted. + running := 0 + + for range 100 { + running = runtime.NumGoroutine() + + if running >= baseline+batchers*goroutinesPerRunningBatcher { + break + } + + time.Sleep(20 * time.Millisecond) + } + + perBatcher := float64(running-baseline) / float64(batchers) + + require.InDelta(t, float64(goroutinesPerRunningBatcher), perBatcher, 0.25, + "expected %d goroutines per running batcher, measured %.2f "+ + "(baseline %d, running %d). Changing this budget is a design decision: "+ + "update the constant and the plan's goroutine table together.", + goroutinesPerRunningBatcher, perBatcher, baseline, running) + + for _, b := range live { + require.NoError(t, b.Close()) + } + + // Every owned goroutine must exit, or a long-lived process leaks one set per + // batcher it ever created. + settled := stableGoroutineCount(t) + + require.LessOrEqual(t, settled, baseline, + "goroutines must return to the pre-construction baseline after Close "+ + "(baseline %d, settled %d)", baseline, settled) +} + +// stableGoroutineCount waits for the goroutine count to stop changing, so a +// scheduler that has not yet reaped finished goroutines cannot be mistaken for a +// leak. +func stableGoroutineCount(t *testing.T) int { + t.Helper() + + previous := -1 + + for range 100 { + runtime.GC() + time.Sleep(20 * time.Millisecond) + + current := runtime.NumGoroutine() + if current == previous { + return current + } + + previous = current + } + + return previous +} diff --git a/pkg/batcher/goroutine_leak_test.go b/pkg/batcher/goroutine_leak_test.go index 633be92..87610ab 100644 --- a/pkg/batcher/goroutine_leak_test.go +++ b/pkg/batcher/goroutine_leak_test.go @@ -10,13 +10,13 @@ import ( "github.com/NSXBet/batcher/pkg/batcher" "github.com/stretchr/testify/require" - "go.uber.org/atomic" + "sync/atomic" ) // TestBatcher_GoroutineCleanupOnStop is the main regression test for goroutine leaks. // It reproduces the issue where goroutines remain blocked after Close() is called. func TestBatcher_GoroutineCleanupOnStop(t *testing.T) { - processCount := atomic.NewInt32(0) + var processCount atomic.Int32 processor := func(items []string) error { count := processCount.Add(1) @@ -85,7 +85,7 @@ func TestBatcher_GoroutineCleanupOnStop(t *testing.T) { // TestBatcher_StopWithPendingMessages tests that pending messages don't block goroutine cleanup. func TestBatcher_StopWithPendingMessages(t *testing.T) { - processedBatches := atomic.NewInt32(0) + var processedBatches atomic.Int32 processor := func(items []string) error { processedBatches.Add(1) @@ -139,7 +139,7 @@ func TestBatcher_StopWithPendingMessages(t *testing.T) { // TestBatcher_StopWhileProcessing tests cleanup when the processor is actively working. func TestBatcher_StopWhileProcessing(t *testing.T) { - processing := atomic.NewBool(false) + var processing atomic.Bool processingDone := make(chan struct{}) processor := func(items []string) error { @@ -206,7 +206,7 @@ func TestBatcher_StopWhileProcessing(t *testing.T) { // TestBatcher_StopWithFailingProcessor tests cleanup when the processor consistently fails. func TestBatcher_StopWithFailingProcessor(t *testing.T) { - processCount := atomic.NewInt32(0) + var processCount atomic.Int32 processor := func(items []string) error { processCount.Add(1) @@ -319,9 +319,24 @@ func TestBatcher_RapidStartStopCycles(t *testing.T) { // TestBatcher_CloseTimeoutBehavior tests that Close() doesn't block forever. func TestBatcher_CloseTimeoutBehavior(t *testing.T) { - // Processor that never returns + // A processor that never returns. The library cannot cancel user code, so the + // contract is that Close reports an incomplete drain rather than hanging + // forever or pretending the work finished. + release := make(chan struct{}) + defer close(release) + + // Signalled by the processor itself. Waiting on Stats().Pending would only prove + // work is pending, not that the processor has been entered, so Close could race + // ahead of processing and the in-flight case would go uncovered. + entered := make(chan struct{}) + + var enterOnce sync.Once + processor := func(items []string) error { - time.Sleep(10 * time.Minute) // Intentionally long + enterOnce.Do(func() { close(entered) }) + + <-release + return nil } @@ -329,17 +344,22 @@ func TestBatcher_CloseTimeoutBehavior(t *testing.T) { batcher.WithBatchSize[string](5), batcher.WithBatchInterval[string](10*time.Millisecond), batcher.WithProcessor(processor), + // Keep the test fast: the grace period is what bounds Close, and its exact + // value is configuration rather than behaviour. + batcher.WithCloseGrace[string](500*time.Millisecond), ) - // Add items to trigger processing for i := 0; i < 5; i++ { b.Add(fmt.Sprintf("msg-%d", i)) } - // Wait for processing to start - time.Sleep(100 * time.Millisecond) + // Wait for the processor to be entered, so the batch really is in flight. + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("the processor was never entered") + } - // Close() should return within a reasonable time, even though processor is blocked closeStart := time.Now() closeDone := make(chan error, 1) @@ -352,10 +372,15 @@ func TestBatcher_CloseTimeoutBehavior(t *testing.T) { closeDuration := time.Since(closeStart) t.Logf("Close() completed in %v with error: %v", closeDuration, err) - // Close should complete within a reasonable timeout (not 10 minutes) - if closeDuration > 10*time.Second { - t.Errorf("Close() took too long: %v", closeDuration) - } + require.ErrorIs(t, err, batcher.ErrTimeout, + "a blocked processor must be reported as an incomplete drain") + require.Less(t, closeDuration, 5*time.Second, + "Close must be bounded by its grace period, not by the processor") + + // The drain was reported incomplete, not abandoned: the batcher is still + // draining, so it must not claim to be closed. + require.True(t, b.IsClosing(), "admission must be sealed") + require.False(t, b.IsClosed(), "a batcher with work still in flight is not closed") case <-time.After(15 * time.Second): t.Fatal("Close() blocked for more than 15 seconds - this is unacceptable") diff --git a/pkg/batcher/options.go b/pkg/batcher/options.go index be88ee1..4771f9e 100644 --- a/pkg/batcher/options.go +++ b/pkg/batcher/options.go @@ -41,3 +41,61 @@ func WithSkipAutoStart[T any]() Option[T] { b.config.SkipAutoStart = true } } + +// WithMaxQueueSize bounds the number of queued items. +// +// The default is unbounded, which absorbs bursts but converts a sustained +// overload into unbounded memory growth. Bounding it gives back-pressure instead: +// Add blocks while the queue is full, and Enqueue can report the condition. +// +// Note this bounds queued items only. Total accepted-but-unfinished work is +// bounded by: +// +// N + 2*BatchSize + publishers currently inside the publication window +// +// Two batches, not one: the aggregator holds a partial batch that has left the +// queue but not yet been dispatched, and a second batch can be inside the +// processor at the same time. Sizing memory from N alone therefore understates +// the peak by up to two batches. +func WithMaxQueueSize[T any](maxQueueSize int) Option[T] { + return func(b *Batcher[T]) { + if maxQueueSize < 0 { + maxQueueSize = 0 + } + + b.config.MaxQueueSize = maxQueueSize + } +} + +// WithCloseGrace sets how long Close waits for the drain to complete before +// reporting ErrTimeout. The drain is never abandoned when the grace expires. +func WithCloseGrace[T any](grace time.Duration) Option[T] { + return func(b *Batcher[T]) { + if grace <= 0 { + grace = DefaultCloseGrace + } + + b.config.CloseGrace = grace + } +} + +// WithErrorBufferSize sets how many diagnostics the Errors() channel buffers. +// +// When the buffer is full, already-retained errors are kept and newer ones are +// dropped. Keeping the oldest is deliberate: the first errors of a storm usually +// carry the root cause, and dropping the oldest instead would mean racing the +// consumer for the right to discard. +// +// The buffer is bounded on purpose. A blocking error channel would deadlock the +// pipeline whenever a processor fails and nobody reads Errors(), and an unbounded +// one would grow without limit during an outage. Stats().DroppedErrors reports the +// loss. +func WithErrorBufferSize[T any](size int) Option[T] { + return func(b *Batcher[T]) { + if size <= 0 { + size = DefaultErrorBufferSize + } + + b.config.ErrorBufferSize = size + } +} diff --git a/pkg/batcher/queue.go b/pkg/batcher/queue.go new file mode 100644 index 0000000..209ff9c --- /dev/null +++ b/pkg/batcher/queue.go @@ -0,0 +1,186 @@ +package batcher + +import ( + "context" + "sync" +) + +// queue is Batcher's owned intake queue. +// +// It replaces golang.design/x/chann for two reasons, one correctness and one +// cost: +// +// - A chann unbounded queue owns a relay goroutine whose only exit path closes +// its own ingress channel — the very channel publishers send to. That makes +// "never close the input" and "leak no goroutines" mutually exclusive. +// - Owning the queue removes a channel hop and a goroutine handoff per item. +// +// One implementation serves both modes so there is a single code path to reason +// about: capacity 0 means unbounded, capacity N bounds queued items at N. +// +// The queue is NEVER closed. Shutdown is signalled out of band, by sealing +// admission and by intake accounting, so a publisher can never send on a closed +// channel. Readiness is published on capacity-1 signal channels rather than by +// handing out the storage itself, which is what lets the same type back both a +// bounded and an unbounded mode without a relay. +type queue[T any] struct { + mu sync.Mutex + items []T + // head is the read cursor. Popping advances head instead of reslicing from + // the front, so the backing array is reused and steady-state enqueue does not + // allocate. The array is only reset once fully drained. + head int + capacity int + + // notEmpty and notFull are capacity-1 latches, not queues: a pending signal + // means "there may be work" or "there may be space". Consumers must re-check + // the real state after waking, which they do by calling pop or retrying push. + notEmpty chan struct{} + notFull chan struct{} +} + +// newQueue creates a queue. capacity <= 0 means unbounded. +func newQueue[T any](capacity int) *queue[T] { + if capacity < 0 { + capacity = 0 + } + + return &queue[T]{ + capacity: capacity, + notEmpty: make(chan struct{}, 1), + notFull: make(chan struct{}, 1), + } +} + +// ready returns the readiness latch the aggregator selects on. +func (q *queue[T]) ready() <-chan struct{} { + return q.notEmpty +} + +// push appends an item, blocking only when a bounded queue is full. +// +// While blocked it selects on ctx and sealCh, so a parked publisher is always +// releasable: ctx for callers that supplied a deadline, sealCh for shutdown. +// Without the sealCh arm, a publisher parked on a full queue could never leave +// the admission gate and shutdown would deadlock. +// +// An unbounded queue never blocks, so ctx and sealCh are not consulted. +func (q *queue[T]) push(ctx context.Context, item T, sealCh <-chan struct{}) error { + for { + q.mu.Lock() + + if q.capacity == 0 || len(q.items)-q.head < q.capacity { + q.items = append(q.items, item) + q.mu.Unlock() + + signal(q.notEmpty) + + return nil + } + + q.mu.Unlock() + + select { + case <-q.notFull: + // Space may be available; retry. + case <-ctx.Done(): + return ctx.Err() + case <-sealCh: + return ErrClosing + } + } +} + +// tryPush appends without ever blocking, reporting whether it fit. +func (q *queue[T]) tryPush(item T) bool { + q.mu.Lock() + + if q.capacity != 0 && len(q.items)-q.head >= q.capacity { + q.mu.Unlock() + + return false + } + + q.items = append(q.items, item) + q.mu.Unlock() + + signal(q.notEmpty) + + return true +} + +// pop removes the oldest item, reporting false when the queue is empty. +func (q *queue[T]) pop() (T, bool) { + var zero T + + q.mu.Lock() + + if q.head >= len(q.items) { + // Fully drained: reset the cursor and keep the backing array, so the next + // batch of pushes reuses capacity instead of allocating. + if q.head > 0 { + q.items = q.items[:0] + q.head = 0 + } + + q.mu.Unlock() + + return zero, false + } + + item := q.items[q.head] + // Clear the slot so a popped item is not kept alive by the backing array. + q.items[q.head] = zero + q.head++ + + remaining := len(q.items) - q.head + + // Reclaim the consumed prefix once it dominates the slice. + // + // Without this, the backing array grows with total throughput rather than with + // queue depth: push always appends, and the full-drain reset above only fires + // when the queue is observed completely empty. A steady producer that keeps even + // one item resident never triggers it, so the array grows without bound. + // Measured before this compaction: a queue holding a single item reached 219,136 + // slots after 200,000 pushes. + // + // Compacting when head >= remaining keeps the copy cost amortised O(1) per item, + // because each compaction halves the live region's offset and moves at most as + // many items as have been consumed since the last one. + if q.head >= remaining { + copy(q.items, q.items[q.head:]) + clear(q.items[remaining:]) + + q.items = q.items[:remaining] + q.head = 0 + } + + q.mu.Unlock() + + if remaining > 0 { + // Keep the latch armed: the aggregator drains greedily, but re-arming means + // a missed wakeup cannot strand items. + signal(q.notEmpty) + } + + signal(q.notFull) + + return item, true +} + +// length reports queued items. +func (q *queue[T]) length() int { + q.mu.Lock() + defer q.mu.Unlock() + + return len(q.items) - q.head +} + +// signal arms a capacity-1 latch without blocking. A latch that is already armed +// needs no second signal, because a waiter re-checks real state after waking. +func signal(ch chan struct{}) { + select { + case ch <- struct{}{}: + default: + } +} diff --git a/pkg/batcher/queue_test.go b/pkg/batcher/queue_test.go new file mode 100644 index 0000000..8b1817d --- /dev/null +++ b/pkg/batcher/queue_test.go @@ -0,0 +1,76 @@ +package batcher + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestQueueCompactsConsumedPrefix pins a memory bound that is easy to miss: +// queue backing storage must be proportional to queue depth, not the total number +// of items ever pushed. +// +// Before compaction, this test's one-item queue reached a capacity of 219,136 +// after 200,000 push/pop cycles. The queue was never observed empty, so its +// full-drain reset never fired, push kept appending, and the backing array grew +// without bound even though the live depth stayed at one. +func TestQueueCompactsConsumedPrefix(t *testing.T) { + t.Parallel() + + q := newQueue[int](0) + sealCh := make(chan struct{}) + + require.NoError(t, q.push(context.Background(), 0, sealCh)) + + const cycles = 200_000 + + for i := range cycles { + require.NoError(t, q.push(context.Background(), i, sealCh)) + + _, ok := q.pop() + require.True(t, ok, "cycle %d: queue unexpectedly empty", i) + } + + q.mu.Lock() + defer q.mu.Unlock() + + require.Equal(t, 1, len(q.items)-q.head, "one item must remain live") + require.LessOrEqual(t, cap(q.items), 8, + "backing capacity must remain proportional to depth; got cap=%d after %d cycles", + cap(q.items), cycles) +} + +// TestQueueCompactionPreservesOrder verifies that moving the unconsumed suffix to +// the front does not reorder items. +func TestQueueCompactionPreservesOrder(t *testing.T) { + t.Parallel() + + q := newQueue[int](0) + sealCh := make(chan struct{}) + + for i := range 100 { + require.NoError(t, q.push(context.Background(), i, sealCh)) + } + + // Consume enough to trigger at least one prefix compaction while leaving a + // substantial suffix behind. + for i := range 75 { + item, ok := q.pop() + require.True(t, ok) + require.Equal(t, i, item) + } + + for i := 100; i < 150; i++ { + require.NoError(t, q.push(context.Background(), i, sealCh)) + } + + for want := 75; want < 150; want++ { + item, ok := q.pop() + require.True(t, ok) + require.Equal(t, want, item, "queue order must survive compaction") + } + + _, ok := q.pop() + require.False(t, ok) +} diff --git a/pkg/batcher/shutdown.go b/pkg/batcher/shutdown.go new file mode 100644 index 0000000..804a306 --- /dev/null +++ b/pkg/batcher/shutdown.go @@ -0,0 +1,75 @@ +package batcher + +import ( + "errors" + "fmt" +) + +// ShutdownIncompleteError reports that a shutdown wait expired before the drain +// finished. +// +// It does NOT mean work was discarded. The drain continues in the background, and +// a caller who wants to keep waiting can call Shutdown again with a fresh context; +// that second call waits on the same drain rather than restarting it. +type ShutdownIncompleteError struct { + // Pending is a conservative drain obligation, not a queue depth. + // + // Three things inflate it relative to "items still queued", and callers who + // treat it as a precise figure will be misled: + // + // - it counts items currently inside a processor call; + // - it counts publishers that have reserved but not yet published, which may + // yet abort; + // - it is sampled from atomics without a lock, so it is a snapshot. + // + // It is an exact count of accepted-but-unfinished work only once + // PublishersInGate is zero. For queue depth, use Stats().Queued. + Pending int + + // PublishersInGate is how many publishers were still inside the publication + // window when the wait expired. Zero means Pending is exact. + PublishersInGate int + + // Cause is the caller's context error: context.DeadlineExceeded or + // context.Canceled. + Cause error +} + +func (e *ShutdownIncompleteError) Error() string { + return fmt.Sprintf( + "batcher shutdown incomplete: %d pending, %d publishers still admitting: %v", + e.Pending, e.PublishersInGate, e.Cause, + ) +} + +// Unwrap exposes the context error, so errors.Is(err, context.DeadlineExceeded) +// works as callers expect. +func (e *ShutdownIncompleteError) Unwrap() error { + return e.Cause +} + +// Is reports ErrTimeout equivalence, so callers written against Close's original +// error keep working after moving to Shutdown. +func (e *ShutdownIncompleteError) Is(target error) bool { + return errors.Is(target, ErrTimeout) +} + +// ProcessorPanicError reports that a processor panicked and the panic was +// recovered. +// +// Batcher recovers processor panics because it owns goroutines the caller cannot +// reach: an unrecovered panic there would crash the whole process and discard every +// other queued item, and the caller has no way to install its own recover. The +// panic value and stack are preserved so the bug remains debuggable rather than +// being silently swallowed. +type ProcessorPanicError struct { + // Value is whatever was passed to panic(). + Value any + + // Stack is the stack trace captured at recovery. + Stack []byte +} + +func (e *ProcessorPanicError) Error() string { + return fmt.Sprintf("batcher processor panicked: %v", e.Value) +} diff --git a/pkg/batcher/shutdown_test.go b/pkg/batcher/shutdown_test.go new file mode 100644 index 0000000..8b744ef --- /dev/null +++ b/pkg/batcher/shutdown_test.go @@ -0,0 +1,402 @@ +package batcher_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/NSXBet/batcher/pkg/batcher" + "github.com/stretchr/testify/require" +) + +// Shutdown semantics tests. +// +// The contract these pin is "drain or report, never discard". Before this +// milestone, a shutdown whose timeout expired closed the pipeline and threw away +// accepted work, which is the worst possible outcome: the caller was told +// shutdown finished and the items were gone. + +// TestShutdownIsResumable is the core of the contract. A caller that gives up +// early must not affect the drain, and a later caller must be able to keep waiting +// on the same drain rather than starting a new one. +func TestShutdownIsResumable(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func(items []int) error { + <-release + processed.Add(int64(len(items))) + + return nil + }), + ) + + b.Add(1) + + require.Eventually(t, func() bool { + return b.Stats().Pending > 0 + }, 5*time.Second, time.Millisecond) + + // First attempt: give up almost immediately. + shortCtx, cancelShort := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancelShort() + + err := b.Shutdown(shortCtx) + + var incomplete *batcher.ShutdownIncompleteError + + require.ErrorAs(t, err, &incomplete, + "an expired wait must report an incomplete drain") + require.Positive(t, incomplete.Pending, + "the report must say how much work remained") + require.ErrorIs(t, err, context.DeadlineExceeded, + "the cause must be visible to errors.Is") + require.ErrorIs(t, err, batcher.ErrTimeout, + "callers written against Close's original error must keep working") + + // The drain was not abandoned: the batcher is sealed but not closed. + require.True(t, b.IsClosing(), "admission must be sealed after the first attempt") + require.False(t, b.IsClosed(), "the drain must still be in progress") + + // Let the processor finish, then resume waiting with a fresh context. + close(release) + + require.NoError(t, b.Shutdown(context.Background()), + "a later Shutdown must wait on the same drain and observe completion") + + require.True(t, b.IsClosed(), "the batcher must be terminal once the drain completes") + require.Equal(t, int64(1), processed.Load(), + "the accepted item must be processed, not discarded by the first timeout") +} + +// TestShutdownDeadlineDoesNotPoisonLaterCallers pins that a per-caller timeout is +// per-caller. A short deadline must not be stored as the batcher's terminal +// result, or one impatient caller would corrupt every other caller's answer. +func TestShutdownDeadlineDoesNotPoisonLaterCallers(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + b.Add(1) + + require.Eventually(t, func() bool { + return b.Stats().Pending > 0 + }, 5*time.Second, time.Millisecond) + + shortCtx, cancelShort := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancelShort() + + longDone := make(chan error, 1) + + go func() { + longDone <- b.Shutdown(context.Background()) + }() + + shortErr := b.Shutdown(shortCtx) + require.Error(t, shortErr, "the impatient caller must get its own timeout") + + close(release) + + select { + case err := <-longDone: + require.NoError(t, err, + "the patient caller must observe a completed drain, not the other caller's timeout") + case <-time.After(10 * time.Second): + t.Fatal("the patient caller never observed drain completion") + } + + // A shutdown after terminal completion returns the stored terminal result. + require.NoError(t, b.Shutdown(context.Background()), + "shutdown after completion must return the terminal result") +} + +// TestShutdownBeforeStartDrainsQueuedWork pins that work accepted by a batcher +// that was never started is still drained. +// +// Without this, a caller using WithSkipAutoStart who enqueued and then shut down +// would silently lose everything, because nothing was ever consuming. +func TestShutdownBeforeStartDrainsQueuedWork(t *testing.T) { + t.Parallel() + + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](10_000), // only shutdown can flush this + batcher.WithBatchInterval[int](time.Hour), + batcher.WithSkipAutoStart[int](), + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + for i := range 5 { + b.Add(i) + } + + require.NoError(t, b.Shutdown(context.Background())) + require.Equal(t, int64(5), processed.Load(), + "work queued before Start must still be drained") +} + +// TestShutdownOnEmptyUnstartedBatcherCompletes pins the trivial case: an unused +// batcher must shut down promptly rather than waiting for a consumer that never +// had work. +func TestShutdownOnEmptyUnstartedBatcherCompletes(t *testing.T) { + t.Parallel() + + b := batcher.New( + batcher.WithSkipAutoStart[int](), + batcher.WithProcessor(batcher.NoOpProcessor[int]), + ) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + require.NoError(t, b.Shutdown(ctx)) + require.True(t, b.IsClosed()) +} + +// TestShutdownWithParkedPublisherOnFullBoundedQueue pins the specific ordering +// hazard the protocol exists to prevent. +// +// A publisher parked on a full bounded queue cannot leave the admission gate, and +// shutdown waits for the gate to empty. If shutdown waited before ensuring a +// consumer was draining, or failed to release parked publishers, the two would +// deadlock permanently. This is run repeatedly because it is an ordering race. +func TestShutdownWithParkedPublisherOnFullBoundedQueue(t *testing.T) { + t.Parallel() + + const trials = 300 + + for trial := range trials { + release := make(chan struct{}) + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Hour), + batcher.WithMaxQueueSize[int](2), + batcher.WithSkipAutoStart[int](), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + // Fill the bounded queue while nothing is consuming. + for range 2 { + require.NoError(t, b.Enqueue(context.Background(), 1)) + } + + // This publisher must park: the queue is full and unstarted. + parked := make(chan error, 1) + + go func() { parked <- b.Enqueue(context.Background(), 99) }() + + shutdownDone := make(chan error, 1) + + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + shutdownDone <- b.Shutdown(ctx) + }() + + // The publisher must be released rather than held forever. It may return + // ErrClosing, or it may win the final race with the newly available slot: + // it entered the admission gate before sealing, and once the shutdown + // consumer drains a slot, both sealCh and notFull are ready. The protocol + // permits that pre-seal publisher to succeed, provided the drain accounts + // for it. Enqueue calls that START after sealing are covered separately by + // TestEnqueueAfterSealReportsClosing and must return ErrClosing. + select { + case err := <-parked: + require.True(t, err == nil || errors.Is(err, batcher.ErrClosing), + "trial %d: a pre-seal parked publisher may succeed or be released "+ + "with ErrClosing, got %v", trial, err) + case <-time.After(10 * time.Second): + t.Fatalf("trial %d: parked publisher was never released", trial) + } + + close(release) + + select { + case err := <-shutdownDone: + require.NoError(t, err, + "trial %d: the drain must complete, not merely return", trial) + case <-time.After(10 * time.Second): + t.Fatalf("trial %d: shutdown deadlocked with a parked publisher", trial) + } + } +} + +// TestPartialBatchWithLongIntervalIsNotDiscarded pins the data-loss bug directly. +// +// With an interval far longer than the shutdown grace, the old implementation +// reported a timeout and threw the batch away: 50 items accepted, 0 processed, and +// Len still claiming 50. Now the batch must be processed, or the shortfall must be +// reported; never silently dropped. +func TestPartialBatchWithLongIntervalIsNotDiscarded(t *testing.T) { + t.Parallel() + + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](100), + batcher.WithBatchInterval[int](30*time.Second), // far beyond the grace period + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + const accepted = 50 + + for i := range accepted { + b.Add(i) + } + + require.NoError(t, b.Close(), + "shutdown must flush the partial batch rather than waiting for the interval") + require.Equal(t, int64(accepted), processed.Load(), + "every accepted item must be processed") + require.Zero(t, b.Len(), "no phantom pending count may remain") +} + +// TestStartRacingShutdownYieldsOneLifecycle pins that the two cannot combine into +// two consumers or a double close. +func TestStartRacingShutdownYieldsOneLifecycle(t *testing.T) { + t.Parallel() + + const trials = 200 + + for trial := range trials { + var processed atomic.Int64 + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithSkipAutoStart[int](), + batcher.WithProcessor(func(items []int) error { + processed.Add(int64(len(items))) + + return nil + }), + ) + + b.Add(1) + + var wg sync.WaitGroup + + wg.Add(2) + + go func() { + defer wg.Done() + + b.Start() + }() + + go func() { + defer wg.Done() + + _ = b.Shutdown(context.Background()) + }() + + wg.Wait() + + require.NoError(t, b.Shutdown(context.Background()), "trial %d", trial) + require.Equal(t, int64(1), processed.Load(), + "trial %d: the item must be processed exactly once", trial) + } +} + +// TestShutdownResultIgnoresProcessorErrors pins the separation between "the drain +// finished" and "the work succeeded". Conflating them would make a failing +// downstream look like a broken shutdown. +func TestShutdownResultIgnoresProcessorErrors(t *testing.T) { + t.Parallel() + + failure := errors.New("downstream unavailable") + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func([]int) error { + return failure + }), + ) + + // Drain diagnostics so the buffer cannot fill. + go func() { + for range b.Errors() { + } + }() + + b.Add(1) + + require.NoError(t, b.Shutdown(context.Background()), + "a processor error must not be reported as a shutdown failure") + + stats := b.Stats() + + require.Equal(t, uint64(1), stats.Failed) + require.Zero(t, stats.Pending) +} + +// TestNonReturningProcessorKeepsBatcherDraining documents the one case the library +// cannot resolve: user code that never returns. The honest behaviour is to report +// it, not to pretend the drain finished. +func TestNonReturningProcessorKeepsBatcherDraining(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + defer close(release) + + b := batcher.New( + batcher.WithBatchSize[int](1), + batcher.WithBatchInterval[int](time.Millisecond), + batcher.WithProcessor(func([]int) error { + <-release + + return nil + }), + ) + + b.Add(1) + + require.Eventually(t, func() bool { + return b.Stats().Pending > 0 + }, 5*time.Second, time.Millisecond) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + err := b.Shutdown(ctx) + + require.Error(t, err, "a blocked processor must be reported, not waited on forever") + require.True(t, b.IsClosing(), "admission must be sealed") + require.False(t, b.IsClosed(), + "a batcher whose processor never returns is draining, not closed") +} diff --git a/pkg/batcher/stats.go b/pkg/batcher/stats.go new file mode 100644 index 0000000..d455e7b --- /dev/null +++ b/pkg/batcher/stats.go @@ -0,0 +1,136 @@ +package batcher + +import "sync/atomic" + +// Stats is a point-in-time snapshot of a Batcher's counters. +// +// Reading it is O(1) and allocation-free: it loads a fixed set of atomics and +// never reads a channel or walks the queue. +// +// One field is not lock-free. Queued is read from the queue itself, which takes +// the queue's mutex for the duration of a single length check, so a scrape costs +// about the same contention as one push or pop. That is a deliberate trade: +// maintaining Queued as an atomic would add a fourth producer-path RMW to every +// enqueue, and measurement showed contended producer atomics dominate enqueue +// cost. Scraping frequently is fine; scraping in a tight loop is not. +// +// The snapshot is eventually consistent, NOT transactional. Fields are individual +// atomic loads, so a state transition in flight can be observed partially. Two +// consequences worth stating plainly: +// +// - Accepted == Completed + Failed + Panicked + Pending holds exactly only at +// publisher quiescence (after shutdown observes an empty gate), not while load +// is in flight. +// - Pending is a conservative drain obligation. It counts publishers that have +// reserved but not yet published, so it can briefly exceed accepted work, and +// it is not a queue depth. Use Queued for depth. +type Stats struct { + // Pending is accepted-or-reserved work that has not reached a terminal + // outcome. The drain waits for this to reach zero. + Pending int64 + + // IntakePending is work the aggregator has not yet received. It is the + // counter that decides when intake is exhausted during shutdown. + IntakePending int64 + + // PublishersInGate is how many goroutines are currently inside the + // publication window. Zero means no publisher can still be mid-publish. + PublishersInGate int64 + + // Queued is items sitting in the intake queue: published, not yet received by + // the aggregator. This is the queue depth to alert on. + Queued int64 + + // Accepted counts successful publications. A rejected or cancelled enqueue + // never increments it. + Accepted uint64 + + // Completed, Failed and Panicked are mutually exclusive terminal outcomes, + // counted in items. Failed means the processor returned a non-nil error; + // Panicked means it panicked and the panic was recovered. + Completed uint64 + Failed uint64 + Panicked uint64 + + // Rejected counts enqueue attempts refused because admission was sealed, a + // bounded queue was full for a non-blocking Add, or the caller's context + // expired. + Rejected uint64 + + // DroppedErrors counts diagnostics discarded because the Errors() buffer was + // full. A non-zero value means diagnostics are being lost, not that batches + // failed: it is a signal that Errors() is not being drained fast enough. + DroppedErrors uint64 +} + +// counters holds the live atomics behind Stats. +// +// Placement is a performance decision. Only the reservation and acceptance +// counters are touched per item on the enqueue path; every terminal counter is +// updated once per batch on the processing path, where the cost amortises. Four +// contended atomics per Add measured up to 4x the cost of one, so nothing is added +// to the hot path without a reason. +type counters struct { + pending atomic.Int64 + intakePending atomic.Int64 + accepted atomic.Uint64 + completed atomic.Uint64 + failed atomic.Uint64 + panicked atomic.Uint64 + rejected atomic.Uint64 + droppedErrors atomic.Uint64 +} + +// reserve claims a drain obligation before publication. +// +// Reserving first is what stops the drain from concluding while a publisher is +// still between "decided to publish" and "published". +func (c *counters) reserve() { + c.pending.Add(1) + c.intakePending.Add(1) +} + +// rollback releases a reservation that never became an accepted item. +// +// This MUST run before the publisher leaves the admission gate. If it ran after, +// the coordinator could see an empty gate while a phantom intake obligation +// remained, and the final drain would wait for an item that will never arrive. +func (c *counters) rollback() { + c.pending.Add(-1) + c.intakePending.Add(-1) + c.rejected.Add(1) +} + +// accept records a successful publication. +func (c *counters) accept() { + c.accepted.Add(1) +} + +// received records the aggregator taking an item out of the queue. +func (c *counters) received(n int) { + c.intakePending.Add(int64(-n)) +} + +// terminal records exactly one outcome for a finished batch and releases its +// drain obligation. +func (c *counters) terminal(outcome outcomeKind, n int) { + switch outcome { + case outcomeCompleted: + c.completed.Add(uint64(n)) + case outcomeFailed: + c.failed.Add(uint64(n)) + case outcomePanicked: + c.panicked.Add(uint64(n)) + } + + c.pending.Add(int64(-n)) +} + +// outcomeKind enumerates the mutually exclusive terminal outcomes for a batch. +type outcomeKind int + +const ( + outcomeCompleted outcomeKind = iota + outcomeFailed + outcomePanicked +)