feat(clickhouse): flag async_insert for sandbox host stats flushes

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
GitOrigin-RevId: cb1a9ba40f69575ba48a4c0155002777c3b449cd

Author: 2 people

packages/clickhouse/pkg/hoststats/delivery.go

@@ -5,6 +5,7 @@ import (
 	"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"
@@ -43,6 +44,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
 type ClickhouseDelivery struct {
 	batcher *batcher.Batcher[SandboxHostStat]
 	conn    driver.Conn
+	ff      *featureflags.Client
 }
 
 type GatedClickhouseDelivery struct {
@@ -75,6 +77,7 @@ func NewDefaultClickhouseHostStatsDelivery(
 				logger.L().Error(ctx, "error batching sandbox host stats", zap.Error(err))
 			},
 		},
+		featureFlags,
 	)
 }
 
@@ -86,8 +89,9 @@ func NewClickhouseHostStatsDelivery(
 	ctx context.Context,
 	conn driver.Conn,
 	opts batcher.BatcherOptions,
+	featureFlags *featureflags.Client,
 ) (*ClickhouseDelivery, error) {
-	delivery := &ClickhouseDelivery{conn: conn}
+	delivery := &ClickhouseDelivery{conn: conn, ff: featureFlags}
 
 	var err error
 	delivery.batcher, err = batcher.NewBatcher(delivery.batchInserter, opts)
@@ -119,11 +123,25 @@ func (c *ClickhouseDelivery) Close(_ context.Context) error {
 	return c.batcher.Stop()
 }
 
+// insertSettings returns the per-query ClickHouse settings for one flush, or
+// nil when the server defaults apply.
+func (c *ClickhouseDelivery) insertSettings(ctx context.Context) clickhouse.Settings {
+	if c.ff == nil || !c.ff.BoolFlag(ctx, featureflags.ClickhouseHostStatsAsyncInsertFlag) {
+		return nil
+	}
+
+	return clickhouse.Settings{"async_insert": 1}
+}
+
 func (c *ClickhouseDelivery) batchInserter(ctx context.Context, stats []SandboxHostStat) error {
 	attrs := trace.WithAttributes(attribute.Int("batch.size", len(stats)))
 	ctx, span := tracer.Start(ctx, "Flush host stats batch to Clickhouse", attrs)
 	defer span.End()
 
+	if settings := c.insertSettings(ctx); settings != nil {
+		ctx = clickhouse.Context(ctx, clickhouse.WithSettings(settings))
+	}
+
 	batch, err := c.conn.PrepareBatch(ctx, InsertSandboxHostStatQuery, driver.WithReleaseConnection())
 	if err != nil {
 		span.RecordError(err)

packages/clickhouse/pkg/hoststats/delivery_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
 	"github.com/stretchr/testify/require"
 
@@ -59,3 +60,46 @@ func TestGatedClickhouseDelivery_PushNilFeatureFlagsDrops(t *testing.T) {
 	err := d.Push(SandboxHostStat{})
 	require.NoError(t, err)
 }
+
+func setAsyncInsertFlag(t *testing.T, source *ldtestdata.TestDataSource, value bool) {
+	t.Helper()
+
+	source.Update(source.Flag(featureflags.ClickhouseHostStatsAsyncInsertFlag.Key()).VariationForAll(value))
+}
+
+func TestClickhouseDelivery_InsertSettingsFlagOnEnablesAsyncInsert(t *testing.T) {
+	t.Parallel()
+
+	ff, source := newTestFeatureFlags(t)
+	setAsyncInsertFlag(t, source, true)
+	d := &ClickhouseDelivery{ff: ff}
+
+	require.Equal(t, clickhouse.Settings{"async_insert": 1}, d.insertSettings(t.Context()))
+}
+
+func TestClickhouseDelivery_InsertSettingsFlagOffKeepsServerDefaults(t *testing.T) {
+	t.Parallel()
+
+	ff, source := newTestFeatureFlags(t)
+	setAsyncInsertFlag(t, source, false)
+	d := &ClickhouseDelivery{ff: ff}
+
+	require.Nil(t, d.insertSettings(t.Context()))
+}
+
+func TestClickhouseDelivery_InsertSettingsUnsetFlagKeepsServerDefaults(t *testing.T) {
+	t.Parallel()
+
+	ff, _ := newTestFeatureFlags(t)
+	d := &ClickhouseDelivery{ff: ff}
+
+	require.Nil(t, d.insertSettings(t.Context()), "flag must fall back to off so a deploy changes nothing until flipped")
+}
+
+func TestClickhouseDelivery_InsertSettingsNilFeatureFlagsKeepsServerDefaults(t *testing.T) {
+	t.Parallel()
+
+	d := &ClickhouseDelivery{ff: nil}
+
+	require.Nil(t, d.insertSettings(t.Context()))
+}

packages/shared/pkg/featureflags/flags.go

@@ -862,6 +862,14 @@ var (
 	// ClickHouse endpoints (CLICKHOUSE_CONNECTION_STRINGS). Default DSN
 	// is unaffected.
 	ClickhouseWriteFanoutFlag = NewBoolFlag("clickhouse-write-fanout", false)
+
+	// ClickhouseHostStatsAsyncInsertFlag sets async_insert=1 on every
+	// sandbox_host_stats flush. Each orchestrator flushes its own small batch,
+	// so without it the server writes one tiny part per node per flush; with
+	// it the server buffers those inserts and writes one part per buffer flush.
+	// wait_for_async_insert stays at the server default, so a rejected flush
+	// still fails the batch and reaches the error handler.
+	ClickhouseHostStatsAsyncInsertFlag = NewBoolFlag("clickhouse-host-stats-async-insert", false)
 )
 
 // LogsWriteConfigFlag controls where sandbox/external logs are written, so