diff --git a/packages/clickhouse/pkg/batcher/batcher.go b/packages/clickhouse/pkg/batcher/batcher.go index 2ba3553119..d77c20a38e 100644 --- a/packages/clickhouse/pkg/batcher/batcher.go +++ b/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 + } + } +} diff --git a/packages/clickhouse/pkg/batcher/batcher_terminal_test.go b/packages/clickhouse/pkg/batcher/batcher_terminal_test.go new file mode 100644 index 0000000000..d50c85e552 --- /dev/null +++ b/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)") +} diff --git a/packages/clickhouse/pkg/events/delivery.go b/packages/clickhouse/pkg/events/delivery.go index 1aec57dbfd..08d129bdbb 100644 --- a/packages/clickhouse/pkg/events/delivery.go +++ b/packages/clickhouse/pkg/events/delivery.go @@ -5,12 +5,14 @@ import ( "database/sql" "encoding/json" "fmt" + "time" "github.com/ClickHouse/clickhouse-go/v2" "github.com/ClickHouse/clickhouse-go/v2/lib/driver" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" "go.uber.org/zap" @@ -19,6 +21,7 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/events" "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/e2b-dev/infra/packages/shared/pkg/logger" + "github.com/e2b-dev/infra/packages/shared/pkg/utils" ) const InsertSandboxEventQuery = `INSERT INTO sandbox_events @@ -64,6 +67,18 @@ type GatedClickhouseDelivery struct { var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/clickhouse/pkg/events") +var meter = otel.Meter("github.com/e2b-dev/infra/packages/clickhouse/pkg/events") + +// mEventDropped counts sandbox lifecycle events dropped before reaching +// ClickHouse, broken out by event type and whether the type is terminal, so a +// loss of terminal (billing/audit) events is directly alertable rather than +// hidden inside the generic batcher.items.dropped counter. +var mEventDropped = utils.Must(meter.Int64Counter( + "sandbox_events.delivery.dropped", + metric.WithDescription("Sandbox lifecycle events dropped before reaching ClickHouse, by event type and terminality"), + metric.WithUnit("{event}"), +)) + const DefaultBatcherName = "sandbox-events" func NewDefaultClickhouseSandboxEventsDelivery(ctx context.Context, conn driver.Conn, featureFlags *featureflags.Client, batcherName string) (*ClickhouseDelivery, error) { @@ -101,7 +116,7 @@ func NewClickhouseSandboxEventsDelivery(ctx context.Context, conn driver.Conn, o return delivery, nil } -func (c *ClickhouseDelivery) Publish(_ context.Context, _ string, event events.SandboxEvent) error { +func (c *ClickhouseDelivery) Publish(ctx context.Context, _ string, event events.SandboxEvent) error { eventDataJson, err := json.Marshal(event.EventData) if err != nil { return err @@ -117,7 +132,7 @@ func (c *ClickhouseDelivery) Publish(_ context.Context, _ string, event events.S ttlDays = events.MaxEventsTTLDays } - return c.batcher.Push(SandboxEvent{ + row := SandboxEvent{ Version: event.Version, ID: event.ID, Type: event.Type, @@ -130,7 +145,66 @@ func (c *ClickhouseDelivery) Publish(_ context.Context, _ string, event events.S SandboxTeamID: event.SandboxTeamID, SandboxExecutionID: event.SandboxExecutionID, EventsTTLDays: ttlDays, - }) + } + + if err := c.enqueue(ctx, event.Type, row); err != nil { + mEventDropped.Add(ctx, 1, metric.WithAttributes( + attribute.String("event_type", event.Type), + attribute.Bool("terminal", isTerminalEvent(event.Type)), + )) + + return err + } + + return nil +} + +// enqueue routes an event into the batcher. Terminal lifecycle events (sandbox +// killed/paused/checkpointed) are the terminal, billing/audit-relevant surfaces, +// so they use a bounded-blocking push that applies back-pressure instead of +// dropping the event the instant the queue is full. All other (high-volume, +// non-terminal) events keep the non-blocking push, so the hot path is unchanged. +// The publish runs in a TrackWork-guarded goroutine, so bounded blocking here +// does not block the Delete RPC and is awaited on graceful drain. +func (c *ClickhouseDelivery) enqueue(ctx context.Context, eventType string, row SandboxEvent) error { + if !isTerminalEvent(eventType) { + return c.batcher.Push(row) + } + + timeout := terminalPushTimeout(ctx, c.ff) + if timeout <= 0 { + return c.batcher.Push(row) + } + + pushCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) + defer cancel() + + return c.batcher.PushWithContext(pushCtx, row) +} + +// isTerminalEvent reports whether an event type is the terminal/irrecoverable +// lifecycle surface whose loss corrupts billing or audit. Scoped to killed — +// the event the sandbox cannot be running after and the one the incident in +// #3770 lost. paused/checkpointed also stop billing and may warrant the same +// treatment, but they are recoverable (a paused sandbox can resume) and can be +// higher volume, so they are intentionally left on the non-blocking path for +// now; widen this set (ideally behind a per-type flag) if they prove lossy too. +func isTerminalEvent(eventType string) bool { + return eventType == events.SandboxKilledEvent +} + +// terminalPushTimeout reads the bounded-blocking push timeout for terminal +// events. Non-positive disables blocking (legacy non-blocking drop). +func terminalPushTimeout(ctx context.Context, ff *featureflags.Client) time.Duration { + if ff == nil { + return 0 + } + ms := ff.IntFlag(ctx, featureflags.ClickhouseBatcherTerminalPushTimeoutMs) + if ms <= 0 { + return 0 + } + + return time.Duration(ms) * time.Millisecond } func (c *GatedClickhouseDelivery) Publish(ctx context.Context, key string, event events.SandboxEvent) error { @@ -141,9 +215,39 @@ func (c *GatedClickhouseDelivery) Publish(ctx context.Context, key string, event return nil } -// Close drains the batcher. ctx is ignored to avoid leaking the flush goroutine. -func (c *ClickhouseDelivery) Close(_ context.Context) error { - return c.batcher.Stop() +// Close drains the batcher, bounding the final flush by the shutdown ctx and the +// drain-timeout flag so a ClickHouse stall during a graceful stop cannot outlast +// the scheduler's kill window. A non-positive timeout keeps the legacy unbounded +// drain. The flush goroutine is always awaited (never leaked); ctx only caps how +// long the last batch is retried before it is dropped. +func (c *ClickhouseDelivery) Close(ctx context.Context) error { + timeout := drainTimeout(ctx, c.ff) + if timeout <= 0 { + // Unbounded drain, but still honor ctx: a forced shutdown (closeCtx + // already cancelled) ends the drain immediately instead of hanging. + return c.batcher.StopWithContext(ctx) + } + + // On a forced shutdown closeCtx is already cancelled, so the drain ends + // immediately instead of waiting the full budget; on a graceful stop the + // timeout is the bound. + drainCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + return c.batcher.StopWithContext(drainCtx) +} + +// drainTimeout reads the shutdown drain budget. Non-positive disables bounding. +func drainTimeout(ctx context.Context, ff *featureflags.Client) time.Duration { + if ff == nil { + return 0 + } + ms := ff.IntFlag(ctx, featureflags.ClickhouseBatcherDrainTimeoutMs) + if ms <= 0 { + return 0 + } + + return time.Duration(ms) * time.Millisecond } func (c *ClickhouseDelivery) batchInserter(ctx context.Context, events []SandboxEvent) error { diff --git a/packages/clickhouse/pkg/events/delivery_terminal_test.go b/packages/clickhouse/pkg/events/delivery_terminal_test.go new file mode 100644 index 0000000000..4e9cf37e20 --- /dev/null +++ b/packages/clickhouse/pkg/events/delivery_terminal_test.go @@ -0,0 +1,94 @@ +package events + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/launchdarkly/go-sdk-common/v3/ldvalue" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/clickhouse/pkg/batcher" + sharedevents "github.com/e2b-dev/infra/packages/shared/pkg/events" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" +) + +func TestIsTerminalEvent(t *testing.T) { + t.Parallel() + + require.True(t, isTerminalEvent(sharedevents.SandboxKilledEvent), "killed should be terminal") + + // Scoped to killed for now: paused/checkpointed are recoverable and kept on + // the non-blocking path (see isTerminalEvent doc). + nonTerminal := []string{ + sharedevents.SandboxPausedEvent, + sharedevents.SandboxCheckpointedEvent, + sharedevents.SandboxCreatedEvent, + sharedevents.SandboxResumedEvent, + sharedevents.SandboxUpdatedEvent, + "", + "something.else", + } + for _, et := range nonTerminal { + require.Falsef(t, isTerminalEvent(et), "%s should not be terminal", et) + } +} + +func TestTerminalPushTimeout(t *testing.T) { + t.Parallel() + + require.Zero(t, terminalPushTimeout(t.Context(), nil), "nil flags disables blocking") + + ff, source := newTestFeatureFlags(t) + source.Update(source.Flag(featureflags.ClickhouseBatcherTerminalPushTimeoutMs.Key()).ValueForAll(ldvalue.Int(0))) + require.Zero(t, terminalPushTimeout(t.Context(), ff), "non-positive flag disables blocking") + + source.Update(source.Flag(featureflags.ClickhouseBatcherTerminalPushTimeoutMs.Key()).ValueForAll(ldvalue.Int(250))) + require.Equal(t, 250*time.Millisecond, terminalPushTimeout(t.Context(), ff)) +} + +// On a full queue, a terminal event applies bounded back-pressure (blocks up to +// the configured timeout) while a non-terminal event drops immediately. +func TestClickhouseDelivery_TerminalBlocksNonTerminalDrops(t *testing.T) { + t.Parallel() + + const timeout = 150 * time.Millisecond + ff, source := newTestFeatureFlags(t) + source.Update(source.Flag(featureflags.ClickhouseBatcherTerminalPushTimeoutMs.Key()).ValueForAll(ldvalue.Int(int(timeout.Milliseconds())))) + + release := make(chan struct{}) + entered := make(chan struct{}) + var once sync.Once + b, err := batcher.NewBatcher(func(_ context.Context, _ []SandboxEvent) error { + once.Do(func() { + close(entered) + <-release + }) + + return nil + }, batcher.BatcherOptions{Name: "term-test", MaxDelay: time.Hour, MaxBatchSize: 1, QueueSize: 1}) + require.NoError(t, err) + require.NoError(t, b.Start(t.Context())) + defer func() { close(release); _ = b.Stop() }() + + d := &ClickhouseDelivery{batcher: b, ff: ff, batcherName: "term-test"} + + // Fill the pipeline: first event is pulled into Func (blocked), second fills + // the cap-1 queue. + require.NoError(t, d.Publish(t.Context(), "k", sharedevents.SandboxEvent{Type: sharedevents.SandboxCreatedEvent, SandboxID: "s1"})) + <-entered + require.NoError(t, d.Publish(t.Context(), "k", sharedevents.SandboxEvent{Type: sharedevents.SandboxCreatedEvent, SandboxID: "s2"})) + + // Non-terminal on a full queue: dropped immediately. + start := time.Now() + err = d.Publish(t.Context(), "k", sharedevents.SandboxEvent{Type: sharedevents.SandboxCreatedEvent, SandboxID: "s3"}) + require.ErrorIs(t, err, batcher.ErrBatcherQueueFull) + require.Less(t, time.Since(start), timeout/2, "non-terminal event must not block") + + // Terminal on a still-full queue: blocks ~timeout, then drops. + start = time.Now() + err = d.Publish(t.Context(), "k", sharedevents.SandboxEvent{Type: sharedevents.SandboxKilledEvent, SandboxID: "s4"}) + require.ErrorIs(t, err, batcher.ErrBatcherQueueFull) + require.GreaterOrEqual(t, time.Since(start), timeout*3/4, "terminal event must apply bounded back-pressure before dropping") +} diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 592bd2e195..810a109ce6 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -622,6 +622,27 @@ var ( ClickhouseBatcherMaxBatchSize = NewIntFlag("clickhouse-batcher-max-batch-size", 1000) ClickhouseBatcherMaxDelay = NewIntFlag("clickhouse-batcher-max-delay", 1000) // 1s in milliseconds ClickhouseBatcherQueueSize = NewIntFlag("clickhouse-batcher-queue-size", 1000) + + // ClickhouseBatcherTerminalPushTimeoutMs bounds how long a terminal lifecycle + // event (sandbox killed) blocks waiting for room in the batcher queue before + // it is dropped. 0 or negative keeps the legacy non-blocking behavior (drop + // immediately when the queue is full), which is the default so the behavior + // is dark-launched and enabled per-environment via LaunchDarkly. Only + // terminal events block; high-volume non-terminal events are unaffected. + ClickhouseBatcherTerminalPushTimeoutMs = NewIntFlag("clickhouse-batcher-terminal-push-timeout-ms", 0) + // ClickhouseBatcherFlushMaxRetries is the number of extra attempts to flush a + // batch to ClickHouse after the first attempt fails, before the batch is + // dropped. 0 keeps the legacy behavior (one attempt, drop the batch on error) + // and is the default so the behavior is dark-launched via LaunchDarkly. + ClickhouseBatcherFlushMaxRetries = NewIntFlag("clickhouse-batcher-flush-max-retries", 0) + // ClickhouseBatcherDrainTimeoutMs bounds the final drain flush at shutdown so + // a ClickHouse stall during a graceful stop cannot exceed the scheduler's + // kill window (after which the process is SIGKILLed and the whole in-memory + // queue is lost anyway). It caps the retry/flush wait once the batcher is + // stopping; it does not affect steady-state flushes. 0 or negative keeps the + // legacy unbounded drain. + ClickhouseBatcherDrainTimeoutMs = NewIntFlag("clickhouse-batcher-drain-timeout-ms", 10000) + BestOfKSampleSize = NewIntFlag("best-of-k-sample-size", 3) // Default K=3 BestOfKMaxOvercommit = NewIntFlag("best-of-k-max-overcommit", 400) // Default R=4 (stored as percentage, max over-commit ratio) BestOfKAlpha = NewIntFlag("best-of-k-alpha", 50) // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)