fix(clickhouse): stop silently dropping terminal lifecycle events under burst

packages/clickhouse/pkg/batcher/batcher.go

@@ -68,6 +68,10 @@ type Batcher[T any] struct {
 	ch      chan T
 	doneCh  chan struct{}
 	started bool
+	// flushCancel aborts the in-flight (and any subsequent) flush. It is set in
+	// Start and invoked by StopWithContext once its drain budget expires, so a
+	// stalled flush at shutdown cannot outlast the scheduler's kill window.
+	flushCancel context.CancelFunc
 
 	flags *featureflags.Client
 	name  string
@@ -158,16 +162,34 @@ func (b *Batcher[T]) Start(ctx context.Context) error {
 	b.doneCh = make(chan struct{})
 	b.started = true
 
+	// flushCtx lets a bounded StopWithContext abort a stalled flush during the
+	// final drain. In steady state it is never cancelled, so flushes are
+	// unbounded as before.
+	flushCtx, flushCancel := context.WithCancel(ctx)
+	b.flushCancel = flushCancel
+
 	go func() {
-		b.processBatches(ctx)
+		b.processBatches(flushCtx)
 		close(b.doneCh)
 	}()
 
 	return nil
 }
 
-// Stop stops batch processing and flushes remaining items.
+// Stop stops batch processing and flushes remaining items, waiting as long as it
+// takes. Use StopWithContext to bound the final drain.
 func (b *Batcher[T]) Stop() error {
+	return b.StopWithContext(context.Background())
+}
+
+// StopWithContext stops batch processing and flushes remaining items, bounding
+// the drain (including a flush already in progress and its retries) by ctx. When
+// ctx is done before the drain finishes, the in-flight flush is cancelled and
+// the remaining batch is dropped rather than retried, so a graceful stop stays
+// within the scheduler's kill window instead of being SIGKILLed mid-drain (which
+// would lose the whole queue). It always waits on doneCh, so the flush goroutine
+// is never leaked.
+func (b *Batcher[T]) StopWithContext(ctx context.Context) error {
 	b.mu.Lock()
 	if !b.started {
 		b.mu.Unlock()
@@ -176,9 +198,24 @@ func (b *Batcher[T]) Stop() error {
 	}
 
 	b.started = false
+	flushCancel := b.flushCancel
 	close(b.ch)
 	b.mu.Unlock()
 
+	// Cancel the drain's flush when the budget expires. Stop early once the
+	// processor is done so we don't leak this goroutine on an unbounded ctx.
+	stopped := make(chan struct{})
+	defer close(stopped)
+	if flushCancel != nil {
+		go func() {
+			select {
+			case <-ctx.Done():
+				flushCancel()
+			case <-stopped:
+			}
+		}()
+	}
+
 	<-b.doneCh
 
 	return nil
@@ -206,6 +243,37 @@ func (b *Batcher[T]) Push(item T) error {
 	}
 }
 
+// PushWithContext enqueues an item, blocking until the queue has room or ctx is
+// done. Unlike Push (which drops immediately when the queue is full), it gives
+// bounded back-pressure so a caller can avoid silently losing an item it cares
+// about. On ctx cancellation/timeout the item is dropped and ErrBatcherQueueFull
+// is returned, matching Push's drop accounting.
+//
+// The RLock is held for the duration of the blocking send. Stop takes the write
+// lock before closing the channel, so it cannot close ch while a send is in
+// flight here; it instead waits (bounded by ctx) for this to return. That both
+// prevents a send-on-closed panic and lets in-flight terminal events settle
+// during a graceful stop.
+func (b *Batcher[T]) PushWithContext(ctx context.Context, item T) error {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+
+	if !b.started {
+		return ErrBatcherNotStarted
+	}
+
+	select {
+	case b.ch <- item:
+		mQueueLen.Record(ctx, int64(len(b.ch)), b.attrs)
+
+		return nil
+	case <-ctx.Done():
+		mItemsDropped.Add(ctx, 1, b.attrs)
+
+		return ErrBatcherQueueFull
+	}
+}
+
 // QueueLen returns the number of items pending in the queue.
 func (b *Batcher[T]) QueueLen() int {
 	return len(b.ch)
@@ -219,6 +287,7 @@ func (b *Batcher[T]) processBatches(ctx context.Context) {
 
 	// Runtime limits belong to this goroutine; public fields remain startup options.
 	maxBatchSize, maxDelay := b.limits(ctx, b.MaxBatchSize, b.MaxDelay)
+	maxRetries := b.flushRetries(ctx)
 	recordConfig := func() {
 		mMaxBatchSize.Record(ctx, int64(maxBatchSize), b.attrs)
 		mMaxDelay.Record(ctx, maxDelay.Milliseconds(), b.attrs)
@@ -242,7 +311,7 @@ func (b *Batcher[T]) processBatches(ctx context.Context) {
 		mFlushWaitDuration.Record(ctx, time.Since(batchStartTime).Milliseconds(), b.attrs)
 		start := time.Now()
 		result := "success"
-		if err := b.Func(ctx, batch); err != nil {
+		if err := b.flushWithRetry(ctx, batch, maxRetries); err != nil {
 			result = "error"
 			b.ErrorHandler(err)
 		}
@@ -278,6 +347,7 @@ func (b *Batcher[T]) processBatches(ctx context.Context) {
 		case <-refreshC:
 			previousDelay := maxDelay
 			maxBatchSize, maxDelay = b.limits(ctx, maxBatchSize, maxDelay)
+			maxRetries = b.flushRetries(ctx)
 			if maxDelay != previousDelay {
 				remaining := maxDelay
 				if len(batch) > 0 {
@@ -312,3 +382,53 @@ func (b *Batcher[T]) limits(ctx context.Context, size int, delay time.Duration)
 
 	return size, delay
 }
+
+// flushRetries is the number of extra flush attempts after the first failure.
+// Static batchers (no feature-flag client) do not retry, preserving legacy
+// drop-on-first-error behavior. A negative flag value is treated as zero.
+func (b *Batcher[T]) flushRetries(ctx context.Context) int {
+	if b.flags == nil {
+		return 0
+	}
+	if n := b.flags.IntFlag(ctx, featureflags.ClickhouseBatcherFlushMaxRetries, featureflags.BatcherContext(b.name)); n > 0 {
+		return n
+	}
+
+	return 0
+}
+
+// flushWithRetry calls Func, retrying up to maxRetries extra times with a short
+// capped exponential backoff when it fails, so a transient ClickHouse blip does
+// not discard an entire batch. It returns the last error if every attempt
+// fails (the caller then drops the batch, as before). Each retry increments
+// batcher.flushes{result="retry"} so repeated failures are visible before loss.
+// Retrying here blocks the batch loop, which is the intended back-pressure; the
+// total added delay is bounded (maxRetries * capped backoff).
+func (b *Batcher[T]) flushWithRetry(ctx context.Context, batch []T, maxRetries int) error {
+	const (
+		baseBackoff = 50 * time.Millisecond
+		maxBackoff  = 2 * time.Second
+	)
+
+	var err error
+	for attempt := 0; ; attempt++ {
+		if err = b.Func(ctx, batch); err == nil {
+			return nil
+		}
+		if attempt >= maxRetries {
+			return err
+		}
+
+		mFlushes.Add(ctx, 1, b.attrs, metric.WithAttributes(attribute.String("result", "retry")))
+
+		backoff := baseBackoff << attempt
+		if backoff <= 0 || backoff > maxBackoff {
+			backoff = maxBackoff
+		}
+		select {
+		case <-time.After(backoff):
+		case <-ctx.Done():
+			return err
+		}
+	}
+}

packages/clickhouse/pkg/batcher/batcher_terminal_test.go

@@ -0,0 +1,196 @@
+package batcher
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/launchdarkly/go-sdk-common/v3/ldvalue"
+	"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+)
+
+// blockFirstFlush returns a BatcherFunc that, on its first call only, signals
+// entered and then blocks until release is closed; every later flush returns
+// immediately. This lets a test pin the processor on a full queue without the
+// Func re-blocking on subsequent batches.
+func blockFirstFlush[T any](entered chan<- struct{}, release <-chan struct{}) func(context.Context, []T) error {
+	var once sync.Once
+
+	return func(context.Context, []T) error {
+		once.Do(func() {
+			close(entered)
+			<-release
+		})
+
+		return nil
+	}
+}
+
+// PushWithContext blocks while the queue is full and succeeds once the
+// processor drains a slot, instead of dropping like the non-blocking Push.
+func TestBatcherPushWithContextBlocksThenSucceeds(t *testing.T) {
+	t.Parallel()
+	release := make(chan struct{})
+	entered := make(chan struct{})
+	b, err := NewBatcher(blockFirstFlush[int](entered, release),
+		BatcherOptions{MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 1})
+	require.NoError(t, err)
+	require.NoError(t, b.Start(t.Context()))
+	defer b.Stop()
+
+	// First item is picked up by the processor and blocks it in Func.
+	require.NoError(t, b.Push(1))
+	<-entered
+	// Queue (cap 1) now fills.
+	require.NoError(t, b.Push(2))
+	require.ErrorIs(t, b.Push(3), ErrBatcherQueueFull, "non-blocking Push drops on a full queue")
+
+	// A bounded-blocking push on the full queue succeeds once Func returns and
+	// the processor pulls the queued item, freeing a slot.
+	done := make(chan error, 1)
+	go func() { done <- b.PushWithContext(t.Context(), 4) }()
+	time.Sleep(20 * time.Millisecond)
+	select {
+	case err := <-done:
+		t.Fatalf("PushWithContext returned early (%v); it should have blocked on the full queue", err)
+	default:
+	}
+	close(release)
+	require.NoError(t, <-done, "PushWithContext should succeed after a slot frees up")
+}
+
+// PushWithContext drops (and does not block forever) when its context expires
+// while the queue stays full.
+func TestBatcherPushWithContextTimesOut(t *testing.T) {
+	t.Parallel()
+	release := make(chan struct{})
+	entered := make(chan struct{})
+	b, err := NewBatcher(blockFirstFlush[int](entered, release),
+		BatcherOptions{MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 1})
+	require.NoError(t, err)
+	require.NoError(t, b.Start(t.Context()))
+	defer func() { close(release); b.Stop() }()
+
+	require.NoError(t, b.Push(1))
+	<-entered
+	require.NoError(t, b.Push(2)) // queue full
+
+	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
+	defer cancel()
+	require.ErrorIs(t, b.PushWithContext(ctx, 3), ErrBatcherQueueFull,
+		"a terminal push whose deadline elapses on a stuck queue is dropped, not blocked forever")
+}
+
+func TestBatcherPushWithContextNotStarted(t *testing.T) {
+	t.Parallel()
+	b, err := NewBatcher[int](func(context.Context, []int) error { return nil }, BatcherOptions{})
+	require.NoError(t, err)
+	require.ErrorIs(t, b.PushWithContext(t.Context(), 1), ErrBatcherNotStarted)
+}
+
+// StopWithContext bounds the final drain: when the drain ctx is already done, a
+// persistently failing flush is not retried forever — Stop returns promptly and
+// the goroutine is not leaked.
+func TestBatcherStopWithContextBoundsDrainRetry(t *testing.T) {
+	t.Parallel()
+	source := ldtestdata.DataSource()
+	// Many retries with a long backoff would hang for minutes if unbounded.
+	source.Update(source.Flag(featureflags.ClickhouseBatcherFlushMaxRetries.Key()).ValueForAll(ldvalue.Int(1_000_000)))
+	source.Update(source.Flag(featureflags.ClickhouseBatcherMaxBatchSize.Key()).ValueForAll(ldvalue.Int(1)))
+	ff, err := featureflags.NewClientWithDatasource(source)
+	require.NoError(t, err)
+	defer ff.Close(context.WithoutCancel(t.Context()))
+
+	var calls atomic.Int32
+	b, err := NewBatcher(func(_ context.Context, _ []int) error {
+		calls.Add(1)
+
+		return errors.New("clickhouse down")
+	}, BatcherOptions{Name: "drain-test", FeatureFlags: ff, MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 4})
+	require.NoError(t, err)
+	require.NoError(t, b.Start(t.Context()))
+
+	require.NoError(t, b.Push(1))
+	time.Sleep(20 * time.Millisecond) // let the processor pick it up and start failing
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel() // drain budget already exhausted
+
+	done := make(chan error, 1)
+	go func() { done <- b.StopWithContext(ctx) }()
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("StopWithContext did not return; drain retried past its cancelled ctx")
+	}
+}
+
+// A transient flush failure is retried (per the flag) and the batch is
+// delivered rather than dropped on the first error.
+func TestBatcherFlushRetriesThenSucceeds(t *testing.T) {
+	t.Parallel()
+	source := ldtestdata.DataSource()
+	source.Update(source.Flag(featureflags.ClickhouseBatcherFlushMaxRetries.Key()).ValueForAll(ldvalue.Int(3)))
+	source.Update(source.Flag(featureflags.ClickhouseBatcherMaxBatchSize.Key()).ValueForAll(ldvalue.Int(1)))
+	ff, err := featureflags.NewClientWithDatasource(source)
+	require.NoError(t, err)
+	defer ff.Close(context.WithoutCancel(t.Context()))
+
+	var calls atomic.Int32
+	delivered := make(chan int, 1)
+	b, err := NewBatcher(func(_ context.Context, batch []int) error {
+		if calls.Add(1) < 3 {
+			return errors.New("transient clickhouse error")
+		}
+		delivered <- len(batch)
+
+		return nil
+	}, BatcherOptions{Name: "retry-test", FeatureFlags: ff, MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 4})
+	require.NoError(t, err)
+	require.NoError(t, b.Start(t.Context()))
+	defer b.Stop()
+
+	require.NoError(t, b.Push(1))
+	select {
+	case n := <-delivered:
+		require.Equal(t, 1, n)
+	case <-time.After(5 * time.Second):
+		t.Fatal("batch was not delivered within timeout; retry did not recover the transient failure")
+	}
+	require.GreaterOrEqual(t, calls.Load(), int32(3), "expected the flush to be retried until success")
+}
+
+// With retries disabled (no flag client), a flush failure drops the batch after
+// a single attempt — the legacy behavior is preserved.
+func TestBatcherFlushNoRetryByDefault(t *testing.T) {
+	t.Parallel()
+	var calls atomic.Int32
+	handled := make(chan error, 1)
+	b, err := NewBatcher(func(_ context.Context, _ []int) error {
+		calls.Add(1)
+
+		return errors.New("boom")
+	}, BatcherOptions{
+		MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 4,
+		ErrorHandler: func(err error) { handled <- err },
+	})
+	require.NoError(t, err)
+	require.NoError(t, b.Start(t.Context()))
+	defer b.Stop()
+
+	require.NoError(t, b.Push(1))
+	select {
+	case <-handled:
+	case <-time.After(5 * time.Second):
+		t.Fatal("ErrorHandler was not called")
+	}
+	time.Sleep(50 * time.Millisecond)
+	require.Equal(t, int32(1), calls.Load(), "static batcher must not retry (legacy drop-on-first-error)")
+}