batcher.go

package jetstream

import "sync"

// batcher groups events into batches by count, with a max-latency flush so a
// low-volume live tail does not hold events indefinitely (the design's
// tens-of-milliseconds delivery goal) and a final flush for the partial tail
// when the stream ends.
//
// It is safe for concurrent use: add is called from the backfill goroutine and
// (after cutover) the live goroutine, while a periodic flusher may fire from a
// timer goroutine. The mutex also serializes the downstream emit so the
// iterator's yield is never called concurrently.
type batcher struct {
	mu      sync.Mutex
	size    int
	buf     []Event
	emit    func(batch []Event) bool
	emitErr func(error) bool // serialized recoverable-error emission
	stop    bool             // set once emit returned false; further adds are no-ops
	onStop  func()           // fired exactly once when the batcher first stops
	onced   bool             // guards onStop
}

func newBatcher(size int, emit func(batch []Event) bool, emitErr func(error) bool) *batcher {
	if size < 1 {
		size = 1
	}
	// The preallocation is an optimization, not the logical batch bound: buf
	// grows by append past it. Clamp it so an absurd WithBatchSize (e.g.
	// MaxInt) cannot panic make with an out-of-range cap.
	return &batcher{size: size, emit: emit, emitErr: emitErr, buf: make([]Event, 0, min(size, 4096))}
}

// setOnStop registers a callback fired exactly once when the consumer first
// asks to stop (emit returns false). The engine uses it to cancel the live
// tail so a quiet steady-state stream still unwinds promptly: the periodic
// flusher's yield returns false even when no live event is arriving.
//
// Register it BEFORE any goroutine (e.g. the flusher) can drive a stopping
// emit: fireStopLocked latches onced=true the first time the batcher stops, so
// a stop that races ahead of setOnStop would leave onStop permanently unfired.
func (b *batcher) setOnStop(fn func()) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.onStop = fn
}

// fireStopLocked invokes onStop once. Caller holds b.mu.
func (b *batcher) fireStopLocked() {
	if b.onced {
		return
	}
	b.onced = true
	if b.onStop != nil {
		go b.onStop()
	}
}

// add appends ev, emitting a full batch when size is reached. Returns false if
// the consumer asked to stop (emit returned false), after which the batcher is
// inert.
func (b *batcher) add(ev Event) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.stop {
		return false
	}
	b.buf = append(b.buf, ev)
	if len(b.buf) >= b.size {
		return b.flushLocked()
	}
	return true
}

// emitError emits a recoverable error downstream, serialized against batch
// emission by b.mu so the consumer's yield is never called concurrently (the
// live goroutine reports errors while the flusher goroutine emits batches). Any
// buffered events are flushed first so an error never jumps ahead of events
// already accepted. Returns false if the consumer asked to stop, after which
// the batcher is inert (mirroring a batch-emit stop, including onStop).
func (b *batcher) emitError(err error) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.stop {
		return false
	}
	if !b.flushLocked() {
		return false
	}
	if b.emitErr == nil {
		return true
	}
	if !b.emitErr(err) {
		b.stop = true
		b.fireStopLocked()
		return false
	}
	return true
}

// flush emits any buffered events as one batch. A no-op when empty. Returns
// false if the consumer asked to stop.
func (b *batcher) flush() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.flushLocked()
}

func (b *batcher) flushLocked() bool {
	if b.stop {
		return false
	}
	if len(b.buf) == 0 {
		return true
	}
	batch := b.buf
	// Same clamped preallocation as newBatcher: b.size may be absurdly large
	// (any positive WithBatchSize is accepted) and buf grows by append anyway.
	b.buf = make([]Event, 0, min(b.size, 4096))
	if !b.emit(batch) {
		b.stop = true
		b.fireStopLocked()
		return false
	}
	return true
}

// stopped reports whether the consumer asked to stop.
func (b *batcher) stopped() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.stop
}