From 4c6c89fb56756552b89c83192b23397cfbf960dd Mon Sep 17 00:00:00 2001 From: Rauf Guliyev Date: Mon, 1 Jun 2026 13:32:41 -0700 Subject: [PATCH] feat(clickhouse): implement multi-cluster fan-out for events and stats Adds support for CLICKHOUSE_CONNECTION_STRINGS to allow concurrent writing to multiple clusters. Includes deduplication logic, isolated batchers for additional endpoints, and resource cleanup fixes. --- packages/clickhouse/go.mod | 3 + packages/clickhouse/pkg/clickhouse.go | 18 ++- packages/clickhouse/pkg/events/delivery.go | 10 +- packages/clickhouse/pkg/hoststats/delivery.go | 11 +- .../clickhouse/pkg/hoststats/hoststats.go | 55 +++++++ .../pkg/hoststats/hoststats_test.go | 136 ++++++++++++++++++ packages/clickhouse/pkg/team.go | 8 ++ .../cmd/resume-build/fph_bench.go | 2 + packages/orchestrator/pkg/cfg/model.go | 75 +++++++--- packages/orchestrator/pkg/cfg/model_test.go | 107 ++++++++++++++ packages/orchestrator/pkg/events/events.go | 29 ++-- packages/orchestrator/pkg/factories/run.go | 125 ++++++++++++++-- .../pkg/sandbox/snapshot_metrics.go | 2 + 13 files changed, 530 insertions(+), 51 deletions(-) create mode 100644 packages/clickhouse/pkg/hoststats/hoststats_test.go diff --git a/packages/clickhouse/go.mod b/packages/clickhouse/go.mod index 30f934a0d3..a45e9eefca 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/stretchr/testify v1.11.1 go.opentelemetry.io/otel v1.43.0 go.opentelemetry.io/otel/metric v1.43.0 go.opentelemetry.io/otel/trace v1.43.0 @@ -24,6 +25,7 @@ require ( github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/coder/websocket v1.8.13 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/elastic/go-sysinfo v1.15.4 // indirect @@ -65,6 +67,7 @@ require ( github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/paulmach/orb v0.11.1 // indirect github.com/pierrec/lz4/v4 v4.1.22 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/pressly/goose/v3 v3.26.0 // indirect github.com/prometheus/procfs v0.17.0 // indirect github.com/redis/go-redis/v9 v9.17.3 // indirect diff --git a/packages/clickhouse/pkg/clickhouse.go b/packages/clickhouse/pkg/clickhouse.go index f877df210e..6e7406b3e7 100644 --- a/packages/clickhouse/pkg/clickhouse.go +++ b/packages/clickhouse/pkg/clickhouse.go @@ -2,6 +2,7 @@ package clickhouse import ( "context" + "errors" "fmt" "time" @@ -38,7 +39,6 @@ func NewDriver(connectionString string) (driver.Conn, error) { options.MaxOpenConns = 10 options.MaxIdleConns = 3 - options.TLS = nil conn, err := clickhouse.Open(options) if err != nil { @@ -57,6 +57,22 @@ func New(connectionString string) (*Client, error) { return &Client{conn: conn}, nil } +// EndpointFromDSN returns the credential-stripped host:port for use in logs +// and metric attributes. Never use the raw DSN there — it contains the password. +// On parse failure returns a fixed sentinel so the DSN (which clickhouse-go's +// url.Error embeds verbatim) never reaches a log line. +func EndpointFromDSN(dsn string) (string, error) { + options, err := clickhouse.ParseDSN(dsn) + if err != nil { + return "", errors.New("parse DSN") + } + if len(options.Addr) == 0 { + return "", errors.New("DSN has no addresses") + } + + return options.Addr[0], nil +} + // Close drains the queue and flushes remaining items func (c *Client) Close(context.Context) error { return c.conn.Close() diff --git a/packages/clickhouse/pkg/events/delivery.go b/packages/clickhouse/pkg/events/delivery.go index 94054025a8..a3ae67bfe7 100644 --- a/packages/clickhouse/pkg/events/delivery.go +++ b/packages/clickhouse/pkg/events/delivery.go @@ -53,7 +53,9 @@ type ClickhouseDelivery struct { var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/clickhouse/pkg/events") -func NewDefaultClickhouseSandboxEventsDelivery(ctx context.Context, conn driver.Conn, featureFlags *featureflags.Client) (*ClickhouseDelivery, error) { +const DefaultBatcherName = "sandbox-events" + +func NewDefaultClickhouseSandboxEventsDelivery(ctx context.Context, conn driver.Conn, featureFlags *featureflags.Client, batcherName string) (*ClickhouseDelivery, error) { maxBatchSize := featureFlags.IntFlag(ctx, featureflags.ClickhouseBatcherMaxBatchSize) maxDelay := time.Duration(featureFlags.IntFlag(ctx, featureflags.ClickhouseBatcherMaxDelay)) * time.Millisecond @@ -62,7 +64,7 @@ func NewDefaultClickhouseSandboxEventsDelivery(ctx context.Context, conn driver. return NewClickhouseSandboxEventsDelivery( ctx, conn, batcher.BatcherOptions{ - Name: "sandbox-events", + Name: batcherName, MaxBatchSize: maxBatchSize, MaxDelay: maxDelay, QueueSize: batcherQueueSize, @@ -112,7 +114,8 @@ func (c *ClickhouseDelivery) Publish(_ context.Context, _ string, event events.S }) } -func (c *ClickhouseDelivery) Close(context.Context) error { +// Close waits for queued items to flush. +func (c *ClickhouseDelivery) Close(_ context.Context) error { return c.batcher.Stop() } @@ -128,6 +131,7 @@ func (c *ClickhouseDelivery) batchInserter(ctx context.Context, events []Sandbox return fmt.Errorf("error preparing batch: %w", err) } + defer batch.Close() for _, event := range events { err := batch.Append( diff --git a/packages/clickhouse/pkg/hoststats/delivery.go b/packages/clickhouse/pkg/hoststats/delivery.go index 8582d4bb2f..f5ba863858 100644 --- a/packages/clickhouse/pkg/hoststats/delivery.go +++ b/packages/clickhouse/pkg/hoststats/delivery.go @@ -47,10 +47,13 @@ type ClickhouseDelivery struct { var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/clickhouse/pkg/hoststats") +const DefaultBatcherName = "sandbox-host-stats" + func NewDefaultClickhouseHostStatsDelivery( ctx context.Context, conn driver.Conn, featureFlags *featureflags.Client, + batcherName string, ) (*ClickhouseDelivery, error) { maxBatchSize := featureFlags.IntFlag(ctx, featureflags.ClickhouseBatcherMaxBatchSize) maxDelay := time.Duration(featureFlags.IntFlag(ctx, featureflags.ClickhouseBatcherMaxDelay)) * time.Millisecond @@ -58,7 +61,7 @@ func NewDefaultClickhouseHostStatsDelivery( return NewClickhouseHostStatsDelivery( ctx, conn, batcher.BatcherOptions{ - Name: "sandbox-host-stats", + Name: batcherName, MaxBatchSize: maxBatchSize, MaxDelay: maxDelay, QueueSize: batcherQueueSize, @@ -93,7 +96,10 @@ func (c *ClickhouseDelivery) Push(stat SandboxHostStat) error { return c.batcher.Push(stat) } -func (c *ClickhouseDelivery) Close(context.Context) error { +// 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 *ClickhouseDelivery) Close(_ context.Context) error { return c.batcher.Stop() } @@ -109,6 +115,7 @@ func (c *ClickhouseDelivery) batchInserter(ctx context.Context, stats []SandboxH return fmt.Errorf("error preparing batch: %w", err) } + defer batch.Close() for _, stat := range stats { err := batch.Append( diff --git a/packages/clickhouse/pkg/hoststats/hoststats.go b/packages/clickhouse/pkg/hoststats/hoststats.go index 4fbd6c2b16..2c9e2bfd51 100644 --- a/packages/clickhouse/pkg/hoststats/hoststats.go +++ b/packages/clickhouse/pkg/hoststats/hoststats.go @@ -2,6 +2,8 @@ package hoststats import ( "context" + "errors" + "sync" "time" "github.com/google/uuid" @@ -56,3 +58,56 @@ func NewNoopDelivery() Delivery { func (d *noopDelivery) Push(_ SandboxHostStat) error { return nil } func (d *noopDelivery) Close(_ context.Context) error { return nil } + +// multiDelivery fans out to every target. Push is serial (each target's Push +// is a non-blocking batcher send); Close is parallel so a stalled target +// can't block the others from draining. +type multiDelivery struct { + targets []Delivery +} + +var _ Delivery = (*multiDelivery)(nil) + +// NewMultiDelivery returns noop for 0 targets and the target directly for 1, +// so callers can wrap unconditionally. +func NewMultiDelivery(targets ...Delivery) Delivery { + switch len(targets) { + case 0: + return NewNoopDelivery() + case 1: + return targets[0] + default: + return &multiDelivery{targets: targets} + } +} + +func (m *multiDelivery) Push(stat SandboxHostStat) error { + var err error + for _, t := range m.targets { + if e := t.Push(stat); e != nil { + err = errors.Join(err, e) + } + } + + return err +} + +func (m *multiDelivery) Close(ctx context.Context) error { + var ( + wg sync.WaitGroup + mu sync.Mutex + errs error + ) + for _, t := range m.targets { + wg.Go(func() { + if e := t.Close(ctx); e != nil { + mu.Lock() + errs = errors.Join(errs, e) + mu.Unlock() + } + }) + } + wg.Wait() + + return errs +} diff --git a/packages/clickhouse/pkg/hoststats/hoststats_test.go b/packages/clickhouse/pkg/hoststats/hoststats_test.go new file mode 100644 index 0000000000..075b708581 --- /dev/null +++ b/packages/clickhouse/pkg/hoststats/hoststats_test.go @@ -0,0 +1,136 @@ +package hoststats + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type fakeDelivery struct { + pushed []SandboxHostStat + pushErr error + closed bool + closeErr error +} + +func (f *fakeDelivery) Push(stat SandboxHostStat) error { + f.pushed = append(f.pushed, stat) + + return f.pushErr +} + +func (f *fakeDelivery) Close(_ context.Context) error { + f.closed = true + + return f.closeErr +} + +func TestMultiDelivery_PushFansOutToAllTargets(t *testing.T) { + t.Parallel() + + a, b := &fakeDelivery{}, &fakeDelivery{} + md := NewMultiDelivery(a, b) + + stat := SandboxHostStat{SandboxID: "sbx-1"} + require.NoError(t, md.Push(stat)) + + assert.Equal(t, []SandboxHostStat{stat}, a.pushed) + assert.Equal(t, []SandboxHostStat{stat}, b.pushed) +} + +func TestMultiDelivery_PushContinuesAfterTargetError(t *testing.T) { + t.Parallel() + + bad := &fakeDelivery{pushErr: errors.New("boom")} + good := &fakeDelivery{} + md := NewMultiDelivery(bad, good) + + stat := SandboxHostStat{SandboxID: "sbx-1"} + err := md.Push(stat) + + require.Error(t, err) + require.ErrorContains(t, err, "boom") + // Good target still received the stat — best-effort, no early return. + assert.Equal(t, []SandboxHostStat{stat}, good.pushed) + assert.Equal(t, []SandboxHostStat{stat}, bad.pushed) +} + +func TestMultiDelivery_CloseAllTargets(t *testing.T) { + t.Parallel() + + a := &fakeDelivery{closeErr: errors.New("a closed badly")} + b := &fakeDelivery{} + md := NewMultiDelivery(a, b) + + err := md.Close(context.Background()) + + require.Error(t, err) + require.ErrorContains(t, err, "a closed badly") + assert.True(t, a.closed) + assert.True(t, b.closed) +} + +func TestMultiDelivery_PushJoinsAllErrors(t *testing.T) { + t.Parallel() + + errA := errors.New("push failed A") + errB := errors.New("push failed B") + a := &fakeDelivery{pushErr: errA} + b := &fakeDelivery{pushErr: errB} + md := NewMultiDelivery(a, b) + + stat := SandboxHostStat{SandboxID: "sbx-1"} + err := md.Push(stat) + + require.Error(t, err) + require.ErrorIs(t, err, errA) + require.ErrorIs(t, err, errB) + assert.Equal(t, []SandboxHostStat{stat}, a.pushed) + assert.Equal(t, []SandboxHostStat{stat}, b.pushed) +} + +func TestMultiDelivery_CloseJoinsAllErrors(t *testing.T) { + t.Parallel() + + errA := errors.New("close failed A") + errB := errors.New("close failed B") + a := &fakeDelivery{closeErr: errA} + b := &fakeDelivery{closeErr: errB} + md := NewMultiDelivery(a, b) + + err := md.Close(context.Background()) + + require.Error(t, err) + require.ErrorIs(t, err, errA) + require.ErrorIs(t, err, errB) + assert.True(t, a.closed) + assert.True(t, b.closed) +} + +func TestMultiDelivery_EmptyTargets(t *testing.T) { + t.Parallel() + + md := NewMultiDelivery() + assert.NoError(t, md.Push(SandboxHostStat{})) + assert.NoError(t, md.Close(context.Background())) +} + +func TestMultiDelivery_SingleTargetBypassesWrapper(t *testing.T) { + t.Parallel() + + single := &fakeDelivery{} + md := NewMultiDelivery(single) + + // One-target case returns the underlying delivery directly — no wrapper. + assert.Same(t, single, md) + + stat := SandboxHostStat{SandboxID: "sbx-1"} + require.NoError(t, md.Push(stat)) + assert.Equal(t, []SandboxHostStat{stat}, single.pushed) + + require.NoError(t, md.Close(context.Background())) + assert.True(t, single.closed) +} diff --git a/packages/clickhouse/pkg/team.go b/packages/clickhouse/pkg/team.go index 81d3765f56..21d86d3d95 100644 --- a/packages/clickhouse/pkg/team.go +++ b/packages/clickhouse/pkg/team.go @@ -119,6 +119,10 @@ func (c *Client) QueryMaxStartRateTeamMetrics(ctx context.Context, teamID string // No data -> return 0 if !rows.Next() { + if err := rows.Err(); err != nil { + return MaxTeamMetric{}, fmt.Errorf("error iterating max start rate team metrics: %w", err) + } + return MaxTeamMetric{ Value: 0, Timestamp: time.Now(), @@ -158,6 +162,10 @@ func (c *Client) QueryMaxConcurrentTeamMetrics(ctx context.Context, teamID strin // No data -> return 0 if !rows.Next() { + if err := rows.Err(); err != nil { + return MaxTeamMetric{}, fmt.Errorf("error iterating max concurrent team metrics: %w", err) + } + return MaxTeamMetric{ Value: 0, Timestamp: time.Now(), diff --git a/packages/orchestrator/cmd/resume-build/fph_bench.go b/packages/orchestrator/cmd/resume-build/fph_bench.go index f9bd0b941a..67b8fb2bc2 100644 --- a/packages/orchestrator/cmd/resume-build/fph_bench.go +++ b/packages/orchestrator/cmd/resume-build/fph_bench.go @@ -1,3 +1,5 @@ +//go:build linux + package main import ( diff --git a/packages/orchestrator/pkg/cfg/model.go b/packages/orchestrator/pkg/cfg/model.go index 84325bade6..156d1fcfd8 100644 --- a/packages/orchestrator/pkg/cfg/model.go +++ b/packages/orchestrator/pkg/cfg/model.go @@ -81,29 +81,58 @@ func makePathsAbsolute(c *BuilderConfig) error { type Config struct { BuilderConfig - ClickhouseConnectionString string `env:"CLICKHOUSE_CONNECTION_STRING"` - ForceStop bool `env:"FORCE_STOP"` - GRPCPort uint16 `env:"GRPC_PORT" envDefault:"5008"` - LaunchDarklyAPIKey string `env:"LAUNCH_DARKLY_API_KEY"` - LocalUploadBaseURL string `env:"LOCAL_UPLOAD_BASE_URL"` - NodeIP string `env:"NODE_IP" envDefault:"localhost"` - NodeLabels []string `env:"NODE_LABELS" envSeparator:","` - OrchestratorLockPath string `env:"ORCHESTRATOR_LOCK_PATH" envDefault:"/orchestrator.lock"` - NFSProxyLogging bool `env:"NFS_PROXY_LOGGING" envDefault:"false"` - NFSProxyTracing bool `env:"NFS_PROXY_TRACING" envDefault:"false"` - NFSProxyMetrics bool `env:"NFS_PROXY_METRICS" envDefault:"true"` - NFSProxyRecordHandleCalls bool `env:"NFS_PROXY_RECORD_HANDLE_CALLS" envDefault:"false"` - NFSProxyRecordStatCalls bool `env:"NFS_PROXY_RECORD_STAT_CALLS" envDefault:"false"` - NFSProxyLogLevel nfs.LogLevel `env:"NFS_PROXY_LOG_LEVEL" envDefault:"info"` - ProxyPort uint16 `env:"PROXY_PORT" envDefault:"5007"` - RedisClusterURL string `env:"REDIS_CLUSTER_URL"` - RedisTLSCABase64 string `env:"REDIS_TLS_CA_BASE64"` - RedisURL string `env:"REDIS_URL"` - RedisPoolSize int `env:"REDIS_POOL_SIZE" envDefault:"5"` - RedisMinIdleConns int `env:"REDIS_MIN_IDLE_CONNS" envDefault:"2"` - NBDPoolSize int `env:"NBD_POOL_SIZE" envDefault:"64"` - Services []string `env:"ORCHESTRATOR_SERVICES" envDefault:"orchestrator"` - PersistentVolumeMounts map[string]string `env:"PERSISTENT_VOLUME_MOUNTS"` + ClickhouseConnectionString string `env:"CLICKHOUSE_CONNECTION_STRING"` + ClickhouseConnectionStrings []string `env:"CLICKHOUSE_CONNECTION_STRINGS" envSeparator:";"` + ForceStop bool `env:"FORCE_STOP"` + GRPCPort uint16 `env:"GRPC_PORT" envDefault:"5008"` + LaunchDarklyAPIKey string `env:"LAUNCH_DARKLY_API_KEY"` + LocalUploadBaseURL string `env:"LOCAL_UPLOAD_BASE_URL"` + NodeIP string `env:"NODE_IP" envDefault:"localhost"` + NodeLabels []string `env:"NODE_LABELS" envSeparator:","` + OrchestratorLockPath string `env:"ORCHESTRATOR_LOCK_PATH" envDefault:"/orchestrator.lock"` + NFSProxyLogging bool `env:"NFS_PROXY_LOGGING" envDefault:"false"` + NFSProxyTracing bool `env:"NFS_PROXY_TRACING" envDefault:"false"` + NFSProxyMetrics bool `env:"NFS_PROXY_METRICS" envDefault:"true"` + NFSProxyRecordHandleCalls bool `env:"NFS_PROXY_RECORD_HANDLE_CALLS" envDefault:"false"` + NFSProxyRecordStatCalls bool `env:"NFS_PROXY_RECORD_STAT_CALLS" envDefault:"false"` + NFSProxyLogLevel nfs.LogLevel `env:"NFS_PROXY_LOG_LEVEL" envDefault:"info"` + ProxyPort uint16 `env:"PROXY_PORT" envDefault:"5007"` + RedisClusterURL string `env:"REDIS_CLUSTER_URL"` + RedisTLSCABase64 string `env:"REDIS_TLS_CA_BASE64"` + RedisURL string `env:"REDIS_URL"` + RedisPoolSize int `env:"REDIS_POOL_SIZE" envDefault:"5"` + RedisMinIdleConns int `env:"REDIS_MIN_IDLE_CONNS" envDefault:"2"` + NBDPoolSize int `env:"NBD_POOL_SIZE" envDefault:"64"` + Services []string `env:"ORCHESTRATOR_SERVICES" envDefault:"orchestrator"` + PersistentVolumeMounts map[string]string `env:"PERSISTENT_VOLUME_MOUNTS"` +} + +// AdditionalClickhouseEndpoints returns the non-blank entries from +// CLICKHOUSE_CONNECTION_STRINGS that are *in addition to* the singular +// CLICKHOUSE_CONNECTION_STRING. Order is preserved; first occurrence wins on +// dedup. Returns nil endpoints if nothing remains. +func (c Config) AdditionalClickhouseEndpoints() (endpoints, droppedDuplicates []string) { + singular := strings.TrimSpace(c.ClickhouseConnectionString) + seen := make(map[string]struct{}, len(c.ClickhouseConnectionStrings)) + if singular != "" { + seen[singular] = struct{}{} + } + + for _, raw := range c.ClickhouseConnectionStrings { + s := strings.TrimSpace(raw) + if s == "" { + continue + } + if _, dup := seen[s]; dup { + droppedDuplicates = append(droppedDuplicates, s) + + continue + } + seen[s] = struct{}{} + endpoints = append(endpoints, s) + } + + return endpoints, droppedDuplicates } func (c Config) NodeAddress() *string { diff --git a/packages/orchestrator/pkg/cfg/model_test.go b/packages/orchestrator/pkg/cfg/model_test.go index 5d543778d3..236caec7fb 100644 --- a/packages/orchestrator/pkg/cfg/model_test.go +++ b/packages/orchestrator/pkg/cfg/model_test.go @@ -80,3 +80,110 @@ func TestParse(t *testing.T) { assert.False(t, config.NFSProxyMetrics) }) } + +func TestAdditionalClickhouseEndpoints(t *testing.T) { + t.Run("returns nil when neither var is set", func(t *testing.T) { + c := Config{} + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Nil(t, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("singular only — returns nil", func(t *testing.T) { + c := Config{ClickhouseConnectionString: "A"} + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Nil(t, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("plural only", func(t *testing.T) { + c := Config{ClickhouseConnectionStrings: []string{"B", "C"}} + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"B", "C"}, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("singular filters matching plural entry, reported as dropped", func(t *testing.T) { + c := Config{ + ClickhouseConnectionString: "A", + ClickhouseConnectionStrings: []string{"A", "B"}, + } + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"B"}, endpoints) + assert.Equal(t, []string{"A"}, dropped) + }) + + t.Run("plural internal duplicates deduped, reported as dropped", func(t *testing.T) { + c := Config{ClickhouseConnectionStrings: []string{"A", "B", "A"}} + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"A", "B"}, endpoints) + assert.Equal(t, []string{"A"}, dropped) + }) + + t.Run("blank/whitespace entries dropped silently (not reported)", func(t *testing.T) { + c := Config{ClickhouseConnectionStrings: []string{"A", "", " ", "B"}} + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"A", "B"}, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("whitespace-trimmed duplicate of singular dropped, reported", func(t *testing.T) { + c := Config{ + ClickhouseConnectionString: "A", + ClickhouseConnectionStrings: []string{" A ", "B"}, + } + endpoints, dropped := c.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"B"}, endpoints) + assert.Equal(t, []string{"A"}, dropped) + }) + + t.Run("env-var parse: basic split", func(t *testing.T) { + t.Setenv("CLICKHOUSE_CONNECTION_STRING", "") + t.Setenv("CLICKHOUSE_CONNECTION_STRINGS", "dsn1;dsn2;dsn3") + config, err := Parse() + require.NoError(t, err) + endpoints, dropped := config.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"dsn1", "dsn2", "dsn3"}, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("env-var parse: empty tokens and trailing/leading separators", func(t *testing.T) { + t.Setenv("CLICKHOUSE_CONNECTION_STRING", "") + t.Setenv("CLICKHOUSE_CONNECTION_STRINGS", ";dsn1;;dsn2;") + config, err := Parse() + require.NoError(t, err) + endpoints, dropped := config.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"dsn1", "dsn2"}, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("env-var parse: whitespace around entries trimmed", func(t *testing.T) { + t.Setenv("CLICKHOUSE_CONNECTION_STRING", "") + t.Setenv("CLICKHOUSE_CONNECTION_STRINGS", " dsn1 ; dsn2 ") + config, err := Parse() + require.NoError(t, err) + endpoints, dropped := config.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"dsn1", "dsn2"}, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("env-var parse: all-blank yields nil", func(t *testing.T) { + t.Setenv("CLICKHOUSE_CONNECTION_STRING", "") + t.Setenv("CLICKHOUSE_CONNECTION_STRINGS", ";; ;;") + config, err := Parse() + require.NoError(t, err) + endpoints, dropped := config.AdditionalClickhouseEndpoints() + assert.Nil(t, endpoints) + assert.Nil(t, dropped) + }) + + t.Run("env-var parse: dup of singular reported, dup of earlier plural reported", func(t *testing.T) { + t.Setenv("CLICKHOUSE_CONNECTION_STRING", "dsn1") + t.Setenv("CLICKHOUSE_CONNECTION_STRINGS", "dsn1;dsn2;dsn2") + config, err := Parse() + require.NoError(t, err) + endpoints, dropped := config.AdditionalClickhouseEndpoints() + assert.Equal(t, []string{"dsn2"}, endpoints) + assert.Equal(t, []string{"dsn1", "dsn2"}, dropped) + }) +} diff --git a/packages/orchestrator/pkg/events/events.go b/packages/orchestrator/pkg/events/events.go index b7dd6e1ca2..2e101e4c29 100644 --- a/packages/orchestrator/pkg/events/events.go +++ b/packages/orchestrator/pkg/events/events.go @@ -41,25 +41,32 @@ func (e *EventsService) Publish(ctx context.Context, teamID uuid.UUID, event eve return } - wg := sync.WaitGroup{} for _, target := range e.deliveryTargets { - wg.Go(func() { - if err := target.Publish(ctx, deliveryKey, event); err != nil { - logger.L().Error(ctx, "Failed to publish sandbox event", zap.Error(err), zap.Any("event", event)) - } - }) + if err := target.Publish(ctx, deliveryKey, event); err != nil { + logger.L().Error(ctx, "Failed to publish sandbox event", zap.Error(err), zap.Any("event", event)) + } } - wg.Wait() } +// Close is parallel: a stalled target must not block the others from draining. func (e *EventsService) Close(ctx context.Context) error { - var err error + var ( + wg sync.WaitGroup + mu sync.Mutex + errs error + ) for _, target := range e.deliveryTargets { - closeErr := target.Close(ctx) - err = errors.Join(err, closeErr) + wg.Go(func() { + if closeErr := target.Close(ctx); closeErr != nil { + mu.Lock() + errs = errors.Join(errs, closeErr) + mu.Unlock() + } + }) } + wg.Wait() - return err + return errs } func validateEvent(event events.SandboxEvent) error { diff --git a/packages/orchestrator/pkg/factories/run.go b/packages/orchestrator/pkg/factories/run.go index 513df32b44..f879b8169a 100644 --- a/packages/orchestrator/pkg/factories/run.go +++ b/packages/orchestrator/pkg/factories/run.go @@ -413,10 +413,11 @@ func run(config cfg.Config, opts Options) (success bool) { }}) sbxEventsDeliveryTargets := make([]event.Delivery[event.SandboxEvent], 0) + hostStatsTargets := make([]clickhousehoststats.Delivery, 0, 1+len(config.ClickhouseConnectionStrings)) - hostStatsDelivery := clickhousehoststats.NewNoopDelivery() - - // Clickhouse sandbox events and host stats delivery + // Legacy singular ClickHouse delivery path. Fatal on init error and uses + // the unsuffixed default batcher names to preserve pre-multi-endpoint + // behavior and existing dashboards/alerts. if config.ClickhouseConnectionString != "" { clickhouseConn, err := clickhouse.NewDriver(config.ClickhouseConnectionString) if err != nil { @@ -426,23 +427,121 @@ func run(config cfg.Config, opts Options) (success bool) { return clickhouseConn.Close() }}) - sbxEventsDeliveryClickhouse, err := clickhouseevents.NewDefaultClickhouseSandboxEventsDelivery(ctx, clickhouseConn, featureFlags) + sbxEventsDeliveryClickhouse, err := clickhouseevents.NewDefaultClickhouseSandboxEventsDelivery( + ctx, + clickhouseConn, + featureFlags, + clickhouseevents.DefaultBatcherName, + ) if err != nil { logger.L().Fatal(ctx, "failed to create clickhouse events delivery", zap.Error(err)) } - sbxEventsDeliveryTargets = append(sbxEventsDeliveryTargets, sbxEventsDeliveryClickhouse) - closers = append(closers, closer{"sandbox events delivery for clickhouse", sbxEventsDeliveryClickhouse.Close}) - hostStatsDeliveryClickhouse, err := clickhousehoststats.NewDefaultClickhouseHostStatsDelivery(ctx, clickhouseConn, featureFlags) + hostStatsDeliveryClickhouse, err := clickhousehoststats.NewDefaultClickhouseHostStatsDelivery( + ctx, + clickhouseConn, + featureFlags, + clickhousehoststats.DefaultBatcherName, + ) if err != nil { logger.L().Fatal(ctx, "failed to create clickhouse host stats delivery", zap.Error(err)) } + hostStatsTargets = append(hostStatsTargets, hostStatsDeliveryClickhouse) + } + + // Additional ClickHouse delivery endpoints are best-effort. One driver + + // delivery pair per endpoint keeps connection pools, batcher queues, and + // OTel metrics isolated so a slow/failing endpoint cannot stall the others. + additionalEndpoints, droppedDuplicates := config.AdditionalClickhouseEndpoints() + for _, d := range droppedDuplicates { + endpoint, err := clickhouse.EndpointFromDSN(d) + if err != nil { + logger.L().Info(ctx, "dropped duplicate unparseable ClickHouse endpoint", zap.Error(err)) + + continue + } + logger.L().Info(ctx, "dropped duplicate ClickHouse endpoint", zap.String("endpoint", endpoint)) + } + if len(additionalEndpoints) > 0 { + logger.L().Info(ctx, "resolved additional ClickHouse delivery endpoints", + zap.Int("count", len(additionalEndpoints)), + ) + + logClose := func(what, label string, fn func() error) { + if err := fn(); err != nil { + logger.L().Error(ctx, "failed to close "+what, + zap.String("endpoint", label), zap.Error(err)) + } + } - hostStatsDelivery = hostStatsDeliveryClickhouse - closers = append(closers, closer{"sandbox host stats delivery", hostStatsDeliveryClickhouse.Close}) + for _, dsn := range additionalEndpoints { + label, err := clickhouse.EndpointFromDSN(dsn) + if err != nil { + logger.L().Error(ctx, "failed to parse clickhouse DSN, skipping endpoint", zap.Error(err)) + + continue + } + + clickhouseConn, err := clickhouse.NewDriver(dsn) + if err != nil { + logger.L().Error(ctx, "failed to create clickhouse driver, skipping endpoint", + zap.String("endpoint", label), + zap.Error(err), + ) + + continue + } + + sbxEventsDeliveryClickhouse, err := clickhouseevents.NewDefaultClickhouseSandboxEventsDelivery( + ctx, + clickhouseConn, + featureFlags, + clickhouseevents.DefaultBatcherName+":"+label, + ) + if err != nil { + logger.L().Error(ctx, "failed to create clickhouse events delivery, skipping endpoint", + zap.String("endpoint", label), + zap.Error(err), + ) + logClose("clickhouse connection after events delivery failure", label, clickhouseConn.Close) + + continue + } + + hostStatsDeliveryClickhouse, err := clickhousehoststats.NewDefaultClickhouseHostStatsDelivery( + ctx, + clickhouseConn, + featureFlags, + clickhousehoststats.DefaultBatcherName+":"+label, + ) + if err != nil { + logger.L().Error(ctx, "failed to create clickhouse host stats delivery, skipping endpoint", + zap.String("endpoint", label), + zap.Error(err), + ) + // Events delivery before its underlying driver — it holds a goroutine writing to it. + logClose( + "clickhouse events delivery after host stats delivery failure", + label, + func() error { return sbxEventsDeliveryClickhouse.Close(ctx) }, + ) + logClose("clickhouse connection after host stats delivery failure", label, clickhouseConn.Close) + + continue + } + + sbxEventsDeliveryTargets = append(sbxEventsDeliveryTargets, sbxEventsDeliveryClickhouse) + closers = append(closers, closer{"clickhouse connection " + label, func(context.Context) error { + return clickhouseConn.Close() + }}) + + hostStatsTargets = append(hostStatsTargets, hostStatsDeliveryClickhouse) + } } + hostStatsDelivery := clickhousehoststats.NewMultiDelivery(hostStatsTargets...) + // cgroup manager for resource accounting cgroupManager, err := cgroup.NewManager() if err != nil { @@ -459,9 +558,13 @@ func run(config cfg.Config, opts Options) (success bool) { if redisClient != nil { sbxEventsDeliveryRedis := event.NewRedisStreamsDelivery[event.SandboxEvent](redisClient, event.SandboxEventsStreamName) sbxEventsDeliveryTargets = append(sbxEventsDeliveryTargets, sbxEventsDeliveryRedis) - closers = append(closers, closer{"sandbox events delivery for redis", sbxEventsDeliveryRedis.Close}) } + // Wrapper closers run before per-driver closers (deliveries write through the drivers). + eventsService := events.NewEventsService(sbxEventsDeliveryTargets) + closers = append(closers, closer{"sandbox host stats deliveries (all)", hostStatsDelivery.Close}) + closers = append(closers, closer{"sandbox events deliveries (all)", eventsService.Close}) + // sandbox observer sandboxObserver, err := metrics.NewSandboxObserver(ctx, nodeID, serviceName, commitSHA, version, serviceInstanceID, sandboxes) if err != nil { @@ -568,7 +671,7 @@ func run(config cfg.Config, opts Options) (success bool) { Proxy: sandboxProxy, Persistence: persistence, FeatureFlags: featureFlags, - SbxEventsService: events.NewEventsService(sbxEventsDeliveryTargets), + SbxEventsService: eventsService, PeerRegistry: peerRegistry, Uploads: uploads, }) diff --git a/packages/orchestrator/pkg/sandbox/snapshot_metrics.go b/packages/orchestrator/pkg/sandbox/snapshot_metrics.go index 6b27fb70fe..de6f5c9690 100644 --- a/packages/orchestrator/pkg/sandbox/snapshot_metrics.go +++ b/packages/orchestrator/pkg/sandbox/snapshot_metrics.go @@ -1,3 +1,5 @@ +//go:build linux + package sandbox import (