diff --git a/packages/clickhouse/go.mod b/packages/clickhouse/go.mod index e6e0069dfd..e37f1db498 100644 --- a/packages/clickhouse/go.mod +++ b/packages/clickhouse/go.mod @@ -10,6 +10,7 @@ require ( github.com/ClickHouse/clickhouse-go/v2 v2.40.1 github.com/e2b-dev/infra/packages/shared v0.0.0 github.com/google/uuid v1.6.0 + github.com/launchdarkly/go-server-sdk/v7 v7.13.0 github.com/stretchr/testify v1.11.1 go.opentelemetry.io/otel v1.43.0 go.opentelemetry.io/otel/metric v1.43.0 @@ -57,7 +58,6 @@ require ( github.com/launchdarkly/go-sdk-events/v3 v3.5.0 // indirect github.com/launchdarkly/go-semver v1.0.3 // indirect github.com/launchdarkly/go-server-sdk-evaluation/v3 v3.0.1 // indirect - github.com/launchdarkly/go-server-sdk/v7 v7.13.0 // indirect github.com/mailru/easyjson v0.9.1 // indirect github.com/mattn/go-isatty v0.0.22 // indirect github.com/mfridman/interpolate v0.0.2 // indirect diff --git a/packages/clickhouse/pkg/events/delivery.go b/packages/clickhouse/pkg/events/delivery.go index a3ae67bfe7..dd395ca799 100644 --- a/packages/clickhouse/pkg/events/delivery.go +++ b/packages/clickhouse/pkg/events/delivery.go @@ -51,6 +51,12 @@ type ClickhouseDelivery struct { conn driver.Conn } +type GatedClickhouseDelivery struct { + *ClickhouseDelivery + + ff *featureflags.Client +} + var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/clickhouse/pkg/events") const DefaultBatcherName = "sandbox-events" @@ -75,6 +81,10 @@ func NewDefaultClickhouseSandboxEventsDelivery(ctx context.Context, conn driver. ) } +func NewGatedDelivery(inner *ClickhouseDelivery, featureFlags *featureflags.Client) *GatedClickhouseDelivery { + return &GatedClickhouseDelivery{ClickhouseDelivery: inner, ff: featureFlags} +} + func NewClickhouseSandboxEventsDelivery(ctx context.Context, conn driver.Conn, opts batcher.BatcherOptions) (*ClickhouseDelivery, error) { var err error @@ -114,7 +124,15 @@ func (c *ClickhouseDelivery) Publish(_ context.Context, _ string, event events.S }) } -// Close waits for queued items to flush. +func (c *GatedClickhouseDelivery) Publish(ctx context.Context, key string, event events.SandboxEvent) error { + if c.ff != nil && c.ff.BoolFlag(ctx, featureflags.ClickhouseWriteFanoutFlag) { + return c.ClickhouseDelivery.Publish(ctx, key, 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() } diff --git a/packages/clickhouse/pkg/events/delivery_test.go b/packages/clickhouse/pkg/events/delivery_test.go new file mode 100644 index 0000000000..b94fe919f6 --- /dev/null +++ b/packages/clickhouse/pkg/events/delivery_test.go @@ -0,0 +1,62 @@ +package events + +import ( + "context" + "testing" + + "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" + "github.com/stretchr/testify/require" + + sharedevents "github.com/e2b-dev/infra/packages/shared/pkg/events" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" +) + +func newTestFeatureFlags(t *testing.T) (*featureflags.Client, *ldtestdata.TestDataSource) { + t.Helper() + + source := ldtestdata.DataSource() + ff, err := featureflags.NewClientWithDatasource(source) + require.NoError(t, err) + t.Cleanup(func() { _ = ff.Close(context.Background()) }) + + return ff, source +} + +func setWriteFanoutFlag(t *testing.T, source *ldtestdata.TestDataSource, value bool) { + t.Helper() + + source.Update(source.Flag(featureflags.ClickhouseWriteFanoutFlag.Key()).VariationForAll(value)) +} + +func TestGatedClickhouseDelivery_PublishFlagOffDrops(t *testing.T) { + t.Parallel() + + ff, source := newTestFeatureFlags(t) + setWriteFanoutFlag(t, source, false) + d := &GatedClickhouseDelivery{ClickhouseDelivery: &ClickhouseDelivery{}, ff: ff} + + err := d.Publish(context.Background(), "key", sharedevents.SandboxEvent{SandboxID: "sbx-1"}) + require.NoError(t, err) +} + +func TestClickhouseDelivery_PublishSkipsFlagCheck(t *testing.T) { + t.Parallel() + + _, source := newTestFeatureFlags(t) + // Even with the flag off, ungated deliveries write unconditionally. + setWriteFanoutFlag(t, source, false) + d := &ClickhouseDelivery{} + + require.Panics(t, func() { + _ = d.Publish(context.Background(), "key", sharedevents.SandboxEvent{}) + }, "ungated delivery should bypass the flag and reach batcher.Push (panics on nil batcher)") +} + +func TestGatedClickhouseDelivery_PublishNilFeatureFlagsDrops(t *testing.T) { + t.Parallel() + + d := &GatedClickhouseDelivery{ClickhouseDelivery: &ClickhouseDelivery{}, ff: nil} + + err := d.Publish(context.Background(), "key", sharedevents.SandboxEvent{}) + require.NoError(t, err) +} diff --git a/packages/clickhouse/pkg/hoststats/delivery.go b/packages/clickhouse/pkg/hoststats/delivery.go index f5ba863858..deab83c0f4 100644 --- a/packages/clickhouse/pkg/hoststats/delivery.go +++ b/packages/clickhouse/pkg/hoststats/delivery.go @@ -45,6 +45,12 @@ type ClickhouseDelivery struct { conn driver.Conn } +type GatedClickhouseDelivery struct { + *ClickhouseDelivery + + ff *featureflags.Client +} + var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/clickhouse/pkg/hoststats") const DefaultBatcherName = "sandbox-host-stats" @@ -72,6 +78,10 @@ func NewDefaultClickhouseHostStatsDelivery( ) } +func NewGatedDelivery(inner *ClickhouseDelivery, featureFlags *featureflags.Client) *GatedClickhouseDelivery { + return &GatedClickhouseDelivery{ClickhouseDelivery: inner, ff: featureFlags} +} + func NewClickhouseHostStatsDelivery( ctx context.Context, conn driver.Conn, @@ -96,9 +106,15 @@ func (c *ClickhouseDelivery) Push(stat SandboxHostStat) error { return c.batcher.Push(stat) } -// Close waits for queued items to flush. ctx is ignored: honoring it would -// leak the flush goroutine onto a connection the caller is about to tear -// down. Real deadlines require plumbing ctx through batcher.Stop. +func (c *GatedClickhouseDelivery) Push(stat SandboxHostStat) error { + if c.ff != nil && c.ff.BoolFlag(context.Background(), featureflags.ClickhouseWriteFanoutFlag) { + return c.ClickhouseDelivery.Push(stat) + } + + 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() } diff --git a/packages/clickhouse/pkg/hoststats/delivery_test.go b/packages/clickhouse/pkg/hoststats/delivery_test.go new file mode 100644 index 0000000000..79f256a82d --- /dev/null +++ b/packages/clickhouse/pkg/hoststats/delivery_test.go @@ -0,0 +1,61 @@ +package hoststats + +import ( + "context" + "testing" + + "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" +) + +func newTestFeatureFlags(t *testing.T) (*featureflags.Client, *ldtestdata.TestDataSource) { + t.Helper() + + source := ldtestdata.DataSource() + ff, err := featureflags.NewClientWithDatasource(source) + require.NoError(t, err) + t.Cleanup(func() { _ = ff.Close(context.Background()) }) + + return ff, source +} + +func setWriteFanoutFlag(t *testing.T, source *ldtestdata.TestDataSource, value bool) { + t.Helper() + + source.Update(source.Flag(featureflags.ClickhouseWriteFanoutFlag.Key()).VariationForAll(value)) +} + +func TestGatedClickhouseDelivery_PushFlagOffDrops(t *testing.T) { + t.Parallel() + + ff, source := newTestFeatureFlags(t) + setWriteFanoutFlag(t, source, false) + d := &GatedClickhouseDelivery{ClickhouseDelivery: &ClickhouseDelivery{}, ff: ff} + + err := d.Push(SandboxHostStat{SandboxID: "sbx-1"}) + require.NoError(t, err) +} + +func TestClickhouseDelivery_PushSkipsFlagCheck(t *testing.T) { + t.Parallel() + + _, source := newTestFeatureFlags(t) + // Even with the flag off, ungated deliveries write unconditionally. + setWriteFanoutFlag(t, source, false) + d := &ClickhouseDelivery{} + + require.Panics(t, func() { + _ = d.Push(SandboxHostStat{}) + }, "ungated delivery should bypass the flag and reach batcher.Push (panics on nil batcher)") +} + +func TestGatedClickhouseDelivery_PushNilFeatureFlagsDrops(t *testing.T) { + t.Parallel() + + d := &GatedClickhouseDelivery{ClickhouseDelivery: &ClickhouseDelivery{}, ff: nil} + + err := d.Push(SandboxHostStat{}) + require.NoError(t, err) +} diff --git a/packages/orchestrator/pkg/factories/run.go b/packages/orchestrator/pkg/factories/run.go index 945ee4ec36..ec32c3c464 100644 --- a/packages/orchestrator/pkg/factories/run.go +++ b/packages/orchestrator/pkg/factories/run.go @@ -563,6 +563,7 @@ func run(config cfg.Config, opts Options) (success bool) { continue } + sbxGatedEventsDeliveryClickhouse := clickhouseevents.NewGatedDelivery(sbxEventsDeliveryClickhouse, featureFlags) hostStatsDeliveryClickhouse, err := clickhousehoststats.NewDefaultClickhouseHostStatsDelivery( ctx, @@ -585,13 +586,14 @@ func run(config cfg.Config, opts Options) (success bool) { continue } + hostStatsGatedDeliveryClickhouse := clickhousehoststats.NewGatedDelivery(hostStatsDeliveryClickhouse, featureFlags) - sbxEventsDeliveryTargets = append(sbxEventsDeliveryTargets, sbxEventsDeliveryClickhouse) + sbxEventsDeliveryTargets = append(sbxEventsDeliveryTargets, sbxGatedEventsDeliveryClickhouse) closers = append(closers, closer{"clickhouse connection " + label, func(context.Context) error { return clickhouseConn.Close() }}) - hostStatsTargets = append(hostStatsTargets, hostStatsDeliveryClickhouse) + hostStatsTargets = append(hostStatsTargets, hostStatsGatedDeliveryClickhouse) } } diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 6a079db557..bd79d4f1f5 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -488,6 +488,11 @@ var ( // "" (empty) → singular CLICKHOUSE_CONNECTION_STRING (self-managed default). // "0", "1", ... → index into CLICKHOUSE_CONNECTION_STRINGS ClickhouseReadEndpointFlag = NewStringFlag("clickhouse-read-endpoint", "") + + // ClickhouseWriteFanoutFlag: when false, drop writes to alternate + // ClickHouse endpoints (CLICKHOUSE_CONNECTION_STRINGS). Default DSN + // is unaffected. + ClickhouseWriteFanoutFlag = NewBoolFlag("clickhouse-write-fanout", false) ) // ResolveFirecrackerVersion resolves the firecracker version using the FirecrackerVersions feature flag.