feat(batch): configurable context-cancellation mode (WithCancelMode)

[MasterOfBinary]

Summary

Adds a configurable context-cancellation policy so callers choose how a Batch reacts to cancellation — settling the recurring drain-vs-stop review debate with an explicit knob instead of one hardcoded behavior.

Stacked on #65 → #64. Merge order: #64 → #65 → this. GitHub retargets automatically as each lands.

API (backwards-compatible — zero value == today's behavior)

type CancelMode int

const (
    CancelDrain CancelMode = iota // default
    CancelStop
)

func (b *Batch[T]) WithCancelMode(m CancelMode) *Batch[T]

• CancelDrain (default): on cancel, keep processing items already read from the Source, relying on the Source to stop and close its channels. Nothing already read into the pipeline is dropped. This is the existing behavior, so current users see no change.
• CancelStop: on cancel, stop reading promptly instead of waiting for the Source. Items already buffered are still processed; items not yet read from the Source may be dropped.
• The deadlock-safe ctx-aware error sends (from fix(batch): panic recovery, deadlock-safe error sends, bounded prealloc, Done() race #65) stay active in both modes.

Implementation

• New cancelMode field + WithCancelMode builder, mirroring WithBufferConfig (settable only before Go(); panics if called after).
• CancelStop enables a case <-ctx.Done() arm in doReader's select via the nil-channel trick — the case is inert (nil channel) in CancelDrain, so the default path is byte-for-byte unchanged.

Testing (TDD)

• CancelStop stops on an uncooperative (ctx-ignoring) source; CancelDrain does not (relies on the Source) — both RED tests hung for 2 s before wiring, then passed.
• CancelStop still processes already-buffered items (drops nothing already read).
• WithCancelMode panics after Go(); default mode is CancelDrain.
• go test -race ./... green; existing hardening/wedge tests unaffected.

Backwards compatibility

Purely additive. CancelDrain is the zero value, so a Batch with no WithCancelMode call behaves identically to before.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a configurable CancelMode for the batch processing engine, allowing users to choose between draining already-read items (CancelDrain, the default) or stopping promptly upon context cancellation (CancelStop). The feedback focuses on adding input validation for the CancelMode parameter to prevent invalid integer values, and refactoring the tests to use a single, robust releasableSource to eliminate goroutine leaks caused by the uncooperative test source blocking indefinitely.

[gemini-code-assist]

To prevent bugs and ensure defensive programming, it is highly recommended to validate the CancelMode input in WithCancelMode. Since CancelMode is an integer type, an invalid value (e.g., 42) could be passed, which would silently fall back to CancelDrain behavior and make debugging difficult.

func (b *Batch[T]) WithCancelMode(m CancelMode) *Batch[T] {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.running {
		panic("batch: WithCancelMode cannot be called after Go() has started")
	}

	if m != CancelDrain && m != CancelStop {
		panic(fmt.Sprintf("batch: invalid cancel mode %d", m))
	}

	b.cancelMode = m
	return b
}

[gemini-code-assist]

The uncooperativeSource intentionally leaks a goroutine because it blocks forever on <-block or out <- i after the reader stops reading. Leaking goroutines in tests can lead to resource exhaustion, flaky tests, or test failures (especially if leak detectors are used).

We can completely eliminate uncooperativeSource and these goroutine leaks by using releasableSource everywhere, provided we update releasableSource to support interruption during channel sends (by selecting on s.release).

Let's define a single, robust releasableSource here:

// releasableSource emits a fixed number of items, then blocks until Release() is called.
// It does not honor ctx.Done(), simulating an uncooperative source, but can be
// explicitly released to prevent goroutine leaks in tests.
type releasableSource struct {
	emit      int
	delivered *uint32
	started   chan struct{}
	release   chan struct{}
	startOnce sync.Once
}

func newReleasableSource(emit int, delivered *uint32) *releasableSource {
	return &releasableSource{
		emit:      emit,
		delivered: delivered,
		started:   make(chan struct{}),
		release:   make(chan struct{}),
	}
}

func (s *releasableSource) Release() {
	s.startOnce.Do(func() { close(s.started) })
	close(s.release)
}

func (s *releasableSource) Read(ctx context.Context) (<-chan any, <-chan error) {
	out := make(chan any)
	errs := make(chan error)

	go func() {
		defer close(out)
		defer close(errs)
		s.startOnce.Do(func() { close(s.started) })
		for i := 0; i < s.emit; i++ {
			select {
			case out <- i:
			case <-s.release:
				return
			}
			if s.delivered != nil {
				atomic.AddUint32(s.delivered, 1)
			}
		}
		<-s.release
	}()

	return out, errs
}

[gemini-code-assist]

Update TestCancelStop_StopsOnUncooperativeSource to use releasableSource and call Release() to prevent goroutine leaks.

func TestCancelStop_StopsOnUncooperativeSource(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	src := newReleasableSource(2, nil)
	defer src.Release()

	b := New[any](nil).WithCancelMode(CancelStop)
	IgnoreErrors(b.Go(ctx, src, &noopProc{}))

	// Make sure the source is actually running before we cancel, so we are
	// exercising the cancel-while-reading path rather than a pre-cancel race.
	<-src.started
	cancel()

	select {
	case <-b.Done():
		// Pipeline stopped promptly on cancel despite the Source never closing.
	case <-time.After(2 * time.Second):
		t.Fatal("CancelStop: Done() did not close within 2s on an uncooperative source")
	}
}

[gemini-code-assist]

Update TestCancelDrain_DefaultRelaysOnSource to use releasableSource and call Release() to prevent goroutine leaks and clean up the unused _ = ctx variable.

func TestCancelDrain_DefaultRelaysOnSource(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	src := newReleasableSource(2, nil)
	defer src.Release()

	// No WithCancelMode call => default CancelDrain.
	b := New[any](nil)
	IgnoreErrors(b.Go(ctx, src, &noopProc{}))

	<-src.started
	cancel()

	// In drain mode the engine must NOT complete on its own while the Source
	// keeps its channels open.
	select {
	case <-b.Done():
		t.Fatal("CancelDrain: Done() closed on its own; default must rely on the Source to stop")
	case <-time.After(200 * time.Millisecond):
		// Expected: still waiting on the Source.
	}
}

[gemini-code-assist]

Since we redefined releasableSource at the top of the file to replace uncooperativeSource, we can delete this duplicate definition.

[gemini-code-assist]

Update TestCancelStop_ProcessesBufferedItems to use releasableSource and call Release() to prevent goroutine leaks.

func TestCancelStop_ProcessesBufferedItems(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const emit = 5
	var delivered uint32
	var processed uint32

	src := newReleasableSource(emit, &delivered)
	defer src.Release()

	proc := &countingProc{counter: &processed}

	// Large item buffer so all delivered items sit buffered ahead of the
	// processor, and MinItems=1 so each item forms its own batch and is
	// processed promptly.
	b := New[any](NewConstantConfig(&ConfigValues{MinItems: 1})).
		WithCancelMode(CancelStop).
		WithBufferConfig(BufferConfig{ItemBufferSize: 100, ErrorBufferSize: 100})
	IgnoreErrors(b.Go(ctx, src, proc))

	// Wait until the Source has delivered all `emit` items into the pipeline.
	deadline := time.After(2 * time.Second)
	for atomic.LoadUint32(&delivered) < emit {
		select {
		case <-deadline:
			t.Fatalf("source only delivered %d/%d items before timeout",
				atomic.LoadUint32(&delivered), emit)
		default:
			time.Sleep(time.Millisecond)
		}
	}

	cancel()

	select {
	case <-b.Done():
	case <-time.After(2 * time.Second):
		t.Fatal("CancelStop: Done() did not close within 2s")
	}

	if got := atomic.LoadUint32(&processed); got != emit {
		t.Fatalf("CancelStop dropped buffered work: processed %d, want %d", got, emit)
	}
}