From 7f4fcd36815d8824d972a499e5d3248aa57b2503 Mon Sep 17 00:00:00 2001 From: Rauf Guliyev Date: Mon, 22 Jun 2026 12:51:07 -0700 Subject: [PATCH 1/2] feat(api): LD-gated ClickHouse read switcher --- packages/api/internal/cfg/model.go | 3 +- packages/api/internal/handlers/store.go | 29 +- packages/clickhouse/pkg/switcher.go | 93 ++++ packages/shared/pkg/featureflags/flags.go | 5 + .../shared/pkg/utils/switching/switching.go | 409 ++++++++++++++++++ .../pkg/utils/switching/switching_test.go | 391 +++++++++++++++++ 6 files changed, 919 insertions(+), 11 deletions(-) create mode 100644 packages/clickhouse/pkg/switcher.go create mode 100644 packages/shared/pkg/utils/switching/switching.go create mode 100644 packages/shared/pkg/utils/switching/switching_test.go diff --git a/packages/api/internal/cfg/model.go b/packages/api/internal/cfg/model.go index f710e8ad9b..f841ebddc0 100644 --- a/packages/api/internal/cfg/model.go +++ b/packages/api/internal/cfg/model.go @@ -32,7 +32,8 @@ type Config struct { AnalyticsCollectorAPIToken string `env:"ANALYTICS_COLLECTOR_API_TOKEN"` AnalyticsCollectorHost string `env:"ANALYTICS_COLLECTOR_HOST"` - ClickhouseConnectionString string `env:"CLICKHOUSE_CONNECTION_STRING"` + ClickhouseConnectionString string `env:"CLICKHOUSE_CONNECTION_STRING"` + ClickhouseConnectionStrings []string `env:"CLICKHOUSE_CONNECTION_STRINGS" envSeparator:";"` LokiPassword string `env:"LOKI_PASSWORD"` LokiURL string `env:"LOKI_URL,required"` diff --git a/packages/api/internal/handlers/store.go b/packages/api/internal/handlers/store.go index a0f1aac938..06a6293f5b 100644 --- a/packages/api/internal/handlers/store.go +++ b/packages/api/internal/handlers/store.go @@ -109,16 +109,19 @@ func NewAPIStore(ctx context.Context, tel *telemetry.Client, redisClient redis.U logger.L().Info(ctx, "Created database client") - var clickhouseStore clickhouse.Clickhouse - - clickhouseConnectionString := config.ClickhouseConnectionString - if clickhouseConnectionString == "" { - clickhouseStore = clickhouse.NewNoopClient() - } else { - clickhouseStore, err = clickhouse.New(clickhouseConnectionString) - if err != nil { - logger.L().Fatal(ctx, "initializing ClickHouse store", zap.Error(err)) - } + // LD-gated switcher: empty flag → singular DSN (self-managed); "0", "1", … + // → alternates from CLICKHOUSE_CONNECTION_STRINGS. Lets reads shift between + // clusters per-query without restarts. Empty singular DSN falls back to a + // noop client. + clickhouseStore, err := clickhouse.NewSwitchingClient( + ctx, + featureFlags, + config.ClickhouseConnectionString, + config.ClickhouseConnectionStrings, + clickhouse.WithAllowNoopDefault(true), + ) + if err != nil { + logger.L().Fatal(ctx, "initializing ClickHouse switching client", zap.Error(err)) } posthogClient, posthogErr := analyticscollector.NewPosthogClient(ctx, config.PosthogAPIKey) @@ -318,6 +321,12 @@ func (a *APIStore) Close(ctx context.Context) error { errs = append(errs, fmt.Errorf("closing snapshot cache: %w", err)) } + if a.clickhouseStore != nil { + if err := a.clickhouseStore.Close(ctx); err != nil { + errs = append(errs, fmt.Errorf("closing ClickHouse store: %w", err)) + } + } + return errors.Join(errs...) } diff --git a/packages/clickhouse/pkg/switcher.go b/packages/clickhouse/pkg/switcher.go new file mode 100644 index 0000000000..970198694d --- /dev/null +++ b/packages/clickhouse/pkg/switcher.go @@ -0,0 +1,93 @@ +package clickhouse + +import ( + "context" + "time" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + "github.com/e2b-dev/infra/packages/shared/pkg/utils/switching" +) + +// SwitchingClient is a Clickhouse client that routes each read to one of +// several DSNs based on the clickhouse-read-endpoint LaunchDarkly flag. +// Each call delegates to switcher.Resolve(ctx), so the active endpoint can +// change between calls without restarting. See +// packages/shared/pkg/utils/switching for the underlying mechanism. +type SwitchingClient struct { + switcher *switching.Switcher[Clickhouse] +} + +var _ Clickhouse = (*SwitchingClient)(nil) + +// NewSwitchingClient builds N+1 clients (one for defaultDSN, one per alternate) +// and selects between them per-call using ClickhouseReadEndpointFlag. An empty +// flag value (the LD default) selects the default client; "0", "1", … select +// alternateDSNs[i]. Invalid values fall back to default + rate-limited warning. +func NewSwitchingClient( + ctx context.Context, + ff *featureflags.Client, + defaultDSN string, + alternateDSNs []string, + opts ...Option, +) (*SwitchingClient, error) { + var sOpts []switching.Option[Clickhouse] + for _, opt := range opts { + if opt != nil { + sOpts = append(sOpts, switching.Option[Clickhouse](opt)) + } + } + + s, err := switching.New[Clickhouse]( + ctx, + ff, + featureflags.ClickhouseReadEndpointFlag, + defaultDSN, + alternateDSNs, + func(dsn string) (Clickhouse, error) { return New(dsn) }, + append(sOpts, switching.WithNoopFactory(func() (Clickhouse, error) { + return NewNoopClient(), nil + }))..., + ) + if err != nil { + return nil, err + } + + return &SwitchingClient{switcher: s}, nil +} + +// Option mirrors switching.Option for caller convenience. +type Option switching.Option[Clickhouse] + +// WithAllowNoopDefault enables falling back to a noop client when the +// default DSN is empty. +func WithAllowNoopDefault(allow bool) Option { + return Option(switching.WithAllowNoopDefault[Clickhouse](allow)) +} + +func (s *SwitchingClient) Close(ctx context.Context) error { + return s.switcher.Close(ctx) +} + +func (s *SwitchingClient) QuerySandboxTimeRange(ctx context.Context, sandboxID, teamID string) (time.Time, time.Time, error) { + return s.switcher.Resolve(ctx).QuerySandboxTimeRange(ctx, sandboxID, teamID) +} + +func (s *SwitchingClient) QuerySandboxMetrics(ctx context.Context, sandboxID, teamID string, start, end time.Time, step time.Duration) ([]Metrics, error) { + return s.switcher.Resolve(ctx).QuerySandboxMetrics(ctx, sandboxID, teamID, start, end, step) +} + +func (s *SwitchingClient) QueryLatestMetrics(ctx context.Context, sandboxIDs []string, teamID string) ([]Metrics, error) { + return s.switcher.Resolve(ctx).QueryLatestMetrics(ctx, sandboxIDs, teamID) +} + +func (s *SwitchingClient) QueryTeamMetrics(ctx context.Context, teamID string, start, end time.Time, step time.Duration) ([]TeamMetrics, error) { + return s.switcher.Resolve(ctx).QueryTeamMetrics(ctx, teamID, start, end, step) +} + +func (s *SwitchingClient) QueryMaxStartRateTeamMetrics(ctx context.Context, teamID string, start, end time.Time, step time.Duration) (MaxTeamMetric, error) { + return s.switcher.Resolve(ctx).QueryMaxStartRateTeamMetrics(ctx, teamID, start, end, step) +} + +func (s *SwitchingClient) QueryMaxConcurrentTeamMetrics(ctx context.Context, teamID string, start, end time.Time) (MaxTeamMetric, error) { + return s.switcher.Resolve(ctx).QueryMaxConcurrentTeamMetrics(ctx, teamID, start, end) +} diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 35f0889d40..11a3f9c8ff 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -433,6 +433,11 @@ var ( DefaultPersistentVolumeType = NewStringFlag("default-persistent-volume-type", "") BuildNodeInfo = NewJSONFlag("preferred-build-node", ldvalue.Null()) FirecrackerVersions = NewJSONFlag("firecracker-versions", ldvalue.FromJSONMarshal(FirecrackerVersionMap)) + + // ClickhouseReadEndpointFlag selects which ClickHouse DSN to use for reads. + // "" (empty) → singular CLICKHOUSE_CONNECTION_STRING (self-managed default). + // "0", "1", ... → index into CLICKHOUSE_CONNECTION_STRINGS + ClickhouseReadEndpointFlag = NewStringFlag("clickhouse-read-endpoint", "") ) // ResolveFirecrackerVersion resolves the firecracker version using the FirecrackerVersions feature flag. diff --git a/packages/shared/pkg/utils/switching/switching.go b/packages/shared/pkg/utils/switching/switching.go new file mode 100644 index 0000000000..cd56c843bc --- /dev/null +++ b/packages/shared/pkg/utils/switching/switching.go @@ -0,0 +1,409 @@ +// Package switching provides a generic, LaunchDarkly-gated resource switcher. +// +// A Switcher holds one default Resource and N alternates, all built up-front +// from DSNs via a caller-supplied factory. Each call to Resolve reads a string +// flag and routes to the default (empty value) or one of the alternates +// (numeric index "0", "1", ...). Invalid values or out-of-range / unavailable +// indexes fall back to the default and log a rate-limited warning, so a +// misconfigured flag never takes traffic down. +// +// Intended for shifting read traffic between endpoints per-query without +// restarting the service. +package switching + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "net/url" + "reflect" + "regexp" + "slices" + "strconv" + "strings" + "sync" + "sync/atomic" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "go.uber.org/zap" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + "github.com/e2b-dev/infra/packages/shared/pkg/logger" +) + +// Resource represents a component that can be gracefully closed. +type Resource interface { + Close(ctx context.Context) error +} + +// Factory builds a Resource from a DSN. +type Factory[T Resource] func(dsn string) (T, error) + +type Option[T Resource] func(*options[T]) + +type options[T Resource] struct { + allowNoopDefault bool + defaultClient T + defaultClientSet bool + factory Factory[T] + noopFactory func() (T, error) + warnCap int64 + meter metric.Meter +} + +// WithAllowNoopDefault enables falling back to a noop client if the default DSN +// is empty. Requires WithNoopFactory to be set. +func WithAllowNoopDefault[T Resource](allow bool) Option[T] { + return func(o *options[T]) { + o.allowNoopDefault = allow + } +} + +// WithNoopFactory provides a constructor for a noop client used when the +// default DSN is empty and noop is allowed. +func WithNoopFactory[T Resource](factory func() (T, error)) Option[T] { + return func(o *options[T]) { + o.noopFactory = factory + } +} + +// WithDefaultClient provides an explicitly constructed default client. +// Note: Passing a client here transfers ownership to the Switcher; the +// switcher will close this client when Switcher.Close() is called. +func WithDefaultClient[T Resource](defaultClient T) Option[T] { + return func(o *options[T]) { + o.defaultClient = defaultClient + o.defaultClientSet = true + } +} + +// WithWarnCap sets the maximum number of unique invalid flag value warnings +// logged before suppression. Defaults to 16. +func WithWarnCap[T Resource](warnCap int64) Option[T] { + return func(o *options[T]) { + o.warnCap = warnCap + } +} + +// WithMeter enables observability metrics for the switcher. +func WithMeter[T Resource](meter metric.Meter) Option[T] { + return func(o *options[T]) { + o.meter = meter + } +} + +type Switcher[T Resource] struct { + defaultClient T + alternates []T + // alternateNonNil caches the results of isNil(alternates[i]) to avoid + // reflection on the hot path (Resolve). + alternateNonNil []bool + closed atomic.Bool + ff *featureflags.Client + flag featureflags.StringFlag + warnSeen sync.Map + warnCount atomic.Int64 + warnCap int64 + + resolveCounter metric.Int64Counter + // flagKeyAttr is pre-computed to avoid per-call allocations in Resolve. + flagKeyAttr attribute.KeyValue +} + +// New creates a new Switcher that routes queries between a default client +// and a list of alternates based on a LaunchDarkly string flag. +// +// The flag value "" (empty) selects the default client. +// Numeric values "0", "1", ... select the corresponding index in alternateDSNs. +// Any other value or an out-of-range index falls back to the default client and +// logs a rate-limited warning. +func New[T Resource]( + ctx context.Context, + ff *featureflags.Client, + flag featureflags.StringFlag, + defaultDSN string, + alternateDSNs []string, + factory Factory[T], + opts ...Option[T], +) (*Switcher[T], error) { + options := options[T]{ + factory: factory, + warnCap: 16, + } + for _, opt := range opts { + if opt != nil { + opt(&options) + } + } + + if ff == nil { + return nil, errors.New("switcher requires a feature flags client") + } + if options.factory == nil { + return nil, errors.New("switcher requires a client factory") + } + if options.warnCap < 0 { + return nil, errors.New("switcher warn cap must not be negative") + } + + var defaultClient T + if options.defaultClientSet { + if isNil(options.defaultClient) { + return nil, errors.New("provided default client must not be nil") + } + defaultClient = options.defaultClient + } else { + switch { + case strings.TrimSpace(defaultDSN) != "": + client, err := options.factory(strings.TrimSpace(defaultDSN)) + if err != nil { + return nil, fmt.Errorf("failed to create default client: %w", err) + } + if isNil(client) { + return nil, errors.New("factory returned nil default client") + } + defaultClient = client + case options.allowNoopDefault: + if options.noopFactory == nil { + return nil, errors.New("allowNoopDefault is true but no noopFactory was provided") + } + client, err := options.noopFactory() + if err != nil { + return nil, fmt.Errorf("failed to create noop client: %w", err) + } + if isNil(client) { + return nil, errors.New("noopFactory returned nil default client") + } + defaultClient = client + default: + return nil, errors.New("default DSN is required (or enable allowNoopDefault)") + } + } + + clonedDSNs := slices.Clone(alternateDSNs) + alternates := make([]T, len(clonedDSNs)) + alternateNonNil := make([]bool, len(clonedDSNs)) + alternateHosts := make([]string, 0, len(clonedDSNs)) + var nonBlankAlternates, initializedAlternates int + var failedIndexes []int + + for i := range clonedDSNs { + dsn := strings.TrimSpace(clonedDSNs[i]) + if dsn == "" { + continue + } + + nonBlankAlternates++ + host := sanitizeHost(dsn) + alternateHosts = append(alternateHosts, fmt.Sprintf("%d=%s", i, host)) + + client, err := options.factory(dsn) + if err != nil { + logger.L().Error(ctx, "failed to create alternate client, skipping entry", + zap.Int("index", i), + zap.String("host", host), + zap.Error(sanitizeDriverErr(err)), + ) + failedIndexes = append(failedIndexes, i) + + continue + } + if isNil(client) { + logger.L().Error(ctx, "factory returned nil alternate client, skipping entry", + zap.Int("index", i), + zap.String("host", host), + ) + failedIndexes = append(failedIndexes, i) + + continue + } + + alternates[i] = client + alternateNonNil[i] = true + initializedAlternates++ + } + + s := &Switcher[T]{ + defaultClient: defaultClient, + alternates: alternates, + alternateNonNil: alternateNonNil, + ff: ff, + flag: flag, + warnCap: options.warnCap, + flagKeyAttr: attribute.String("flag_key", flag.Key()), + } + + if options.meter != nil { + counter, err := options.meter.Int64Counter( + "switcher.resolve_count", + metric.WithDescription("Number of times the switcher resolved a client"), + ) + if err != nil { + logger.L().Error(ctx, "failed to create switcher metrics", zap.Error(err)) + } else { + s.resolveCounter = counter + } + } + + logger.L().Info(ctx, "initialized switching client", + zap.String("flag_key", flag.Key()), + zap.Int("alternate_count_configured", nonBlankAlternates), + zap.Int("alternate_count_initialized", initializedAlternates), + zap.Ints("alternate_failed_indexes", failedIndexes), + zap.Strings("alternate_hosts", alternateHosts), + ) + + return s, nil +} + +func (s *Switcher[T]) Resolve(ctx context.Context) T { + v := strings.TrimSpace(s.ff.StringFlag(ctx, s.flag)) + if v == "" { + s.recordMetric(ctx, "default") + + return s.defaultClient + } + + idx, err := strconv.Atoi(v) + if err != nil || idx < 0 || idx >= len(s.alternates) || !s.alternateNonNil[idx] { + s.warnInvalid(ctx, v) + s.recordMetric(ctx, "fallback_default") + + return s.defaultClient + } + + // Cardinality bound for "target" label: len(alternateDSNs) + 2 + // (default + fallback_default + indices). + s.recordMetric(ctx, v) + + return s.alternates[idx] +} + +func (s *Switcher[T]) recordMetric(ctx context.Context, target string) { + if s.resolveCounter != nil { + s.resolveCounter.Add(ctx, 1, metric.WithAttributes( + s.flagKeyAttr, + attribute.String("target", target), + )) + } +} + +func (s *Switcher[T]) warnInvalid(ctx context.Context, v string) { + h := hashFlagValue(v) + if _, ok := s.warnSeen.Load(h); ok { + return + } + + for { + count := s.warnCount.Load() + if count >= s.warnCap { + return + } + if _, ok := s.warnSeen.Load(h); ok { + return + } + if !s.warnCount.CompareAndSwap(count, count+1) { + continue + } + if _, loaded := s.warnSeen.LoadOrStore(h, struct{}{}); loaded { + s.warnCount.Add(-1) + + return + } + + logger.L().Warn(ctx, "invalid read-endpoint flag value, falling back to default", + zap.String("flag_key", s.flag.Key()), + zap.String("value_hash", h), + zap.Int("value_len", len(v)), + ) + + return + } +} + +func (s *Switcher[T]) Close(ctx context.Context) error { + if !s.closed.CompareAndSwap(false, true) { + return nil + } + + var errs []error + if !isNil(s.defaultClient) { + errs = append(errs, s.defaultClient.Close(ctx)) + } + for _, client := range s.alternates { + if !isNil(client) { + errs = append(errs, client.Close(ctx)) + } + } + + return errors.Join(errs...) +} + +func isNil(i any) bool { + if i == nil { + return true + } + v := reflect.ValueOf(i) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Map, reflect.Pointer, reflect.UnsafePointer, reflect.Interface, reflect.Slice: + return v.IsNil() + } + + return false +} + +func hashFlagValue(v string) string { + sum := sha256.Sum256([]byte(v)) + + return hex.EncodeToString(sum[:])[:8] +} + +func sanitizeHost(dsn string) string { + u, err := url.Parse(dsn) + if err != nil || u.Host == "" { + return "unparseable-dsn-" + hashFlagValue(dsn) + } + + host := u.Hostname() + if host == "" { + return "unparseable-dsn-" + hashFlagValue(dsn) + } + // url.Hostname strips brackets from IPv6 literals; re-add them so the + // host:port form stays unambiguous (e.g. [::1]:9000, not ::1:9000). + if strings.Contains(host, ":") { + host = "[" + host + "]" + } + if port := u.Port(); port != "" { + return host + ":" + port + } + + return host +} + +var ( + dsnRE = regexp.MustCompile(`(?i)\b(clickhouses?|https?|tcp)://[^\s]+`) + passwordRE = regexp.MustCompile(`(?i)(password|pass|pwd)=([^&\s]+)`) + userinfoRE = regexp.MustCompile(`(?i)://([^/@\s:]+):([^/@\s]+)@`) +) + +func sanitizeDriverErr(err error) error { + if err == nil { + return nil + } + + msg := err.Error() + msg = userinfoRE.ReplaceAllString(msg, "://@") + msg = passwordRE.ReplaceAllString(msg, "$1=") + msg = dsnRE.ReplaceAllStringFunc(msg, func(raw string) string { + if idx := strings.Index(raw, "://"); idx >= 0 { + return raw[:idx+3] + "" + } + + return "" + }) + + return errors.New(msg) +} diff --git a/packages/shared/pkg/utils/switching/switching_test.go b/packages/shared/pkg/utils/switching/switching_test.go new file mode 100644 index 0000000000..e68650ac24 --- /dev/null +++ b/packages/shared/pkg/utils/switching/switching_test.go @@ -0,0 +1,391 @@ +package switching + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + "github.com/launchdarkly/go-sdk-common/v3/ldvalue" + "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/metric/noop" +) + +type fakeResource struct { + id string + closeCalls atomic.Int32 +} + +func (f *fakeResource) Close(context.Context) error { + f.closeCalls.Add(1) + + return nil +} + +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 setFlag(t *testing.T, source *ldtestdata.TestDataSource, key string, value string) { + t.Helper() + source.Update(source.Flag(key).ValueForAll(ldvalue.String(value))) +} + +func TestNew_Validation(t *testing.T) { + t.Parallel() + ff, _ := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + factory := func(dsn string) (*fakeResource, error) { return &fakeResource{id: dsn}, nil } + + tests := []struct { + name string + ff *featureflags.Client + factory Factory[*fakeResource] + dsn string + opts []Option[*fakeResource] + wantErr string + }{ + {"missing ff", nil, factory, "default", nil, "requires a feature flags client"}, + {"missing factory", ff, nil, "default", nil, "requires a client factory"}, + {"missing dsn", ff, factory, "", nil, "default DSN is required"}, + {"blank dsn", ff, factory, " ", nil, "default DSN is required"}, + {"noop allowed but no factory", ff, factory, "", []Option[*fakeResource]{WithAllowNoopDefault[*fakeResource](true)}, "no noopFactory was provided"}, + {"negative warn cap", ff, factory, "default", []Option[*fakeResource]{WithWarnCap[*fakeResource](-1)}, "must not be negative"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := New(t.Context(), tt.ff, flag, tt.dsn, nil, tt.factory, tt.opts...) + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestNew_AlternateFactoryFailures(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + // Test case: one alternate returns error, one returns typed nil. + // Both should be skipped and Resolve should fall back to default. + s, err := New[*fakeResource](t.Context(), ff, flag, "default", []string{"good", "bad-err", "typed-nil"}, func(dsn string) (*fakeResource, error) { + switch dsn { + case "bad-err": + return nil, errors.New("factory error") + case "typed-nil": + var nilResource *fakeResource + + return nilResource, nil + default: + return &fakeResource{id: dsn}, nil + } + }) + require.NoError(t, err) + require.Len(t, s.alternates, 3) + require.NotNil(t, s.alternates[0]) + require.Nil(t, s.alternates[1]) + require.Nil(t, s.alternates[2]) + require.True(t, s.alternateNonNil[0]) + require.False(t, s.alternateNonNil[1]) + require.False(t, s.alternateNonNil[2]) + + ctx := context.Background() + + // Should fall back to default for skipped alternates + setFlag(t, source, flag.Key(), "1") // bad-err + require.Equal(t, "default", s.Resolve(ctx).id) + + setFlag(t, source, flag.Key(), "2") // typed-nil + require.Equal(t, "default", s.Resolve(ctx).id) +} + +func TestSwitcher_DSNInvariants(t *testing.T) { + t.Parallel() + ff, _ := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + // Test case: Trimming and Cloning + dsns := []string{" alt-0 ", "alt-1"} + s, err := New[*fakeResource](t.Context(), ff, flag, " default ", dsns, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }) + require.NoError(t, err) + + // Mutate caller slice - should not affect switcher (cloning check) + dsns[1] = "mutated" + + require.Equal(t, "default", s.defaultClient.id) + require.Equal(t, "alt-0", s.alternates[0].id) + require.Equal(t, "alt-1", s.alternates[1].id) +} + +func TestSwitcher_Resolve_EdgeCases(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + flagVal string + expected string + warns int64 + }{ + {"empty", "", "default", 0}, + {"valid with space", " 0 ", "alt-0", 0}, + {"out of range", "99", "default", 1}, + {"negative", "-1", "default", 1}, + {"non-numeric", "bad", "default", 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + s, err := New[*fakeResource](t.Context(), ff, flag, "default", []string{"alt-0"}, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }) + require.NoError(t, err) + + ctx := context.Background() + setFlag(t, source, flag.Key(), tt.flagVal) + require.Equal(t, tt.expected, s.Resolve(ctx).id) + require.Equal(t, tt.warns, s.warnCount.Load()) + }) + } +} + +func TestSwitcher_WarnCap_Deduplication(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + s, err := New[*fakeResource](t.Context(), ff, flag, "default", nil, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }, WithWarnCap[*fakeResource](2)) + require.NoError(t, err) + + ctx := context.Background() + + // Warn 1 + setFlag(t, source, flag.Key(), "bad-1") + s.Resolve(ctx) + require.Equal(t, int64(1), s.warnCount.Load()) + + // Same value again - should not increment count + s.Resolve(ctx) + require.Equal(t, int64(1), s.warnCount.Load()) + + // Warn 2 + setFlag(t, source, flag.Key(), "bad-2") + s.Resolve(ctx) + require.Equal(t, int64(2), s.warnCount.Load()) + + // Warn 3 - capped + setFlag(t, source, flag.Key(), "bad-3") + s.Resolve(ctx) + require.Equal(t, int64(2), s.warnCount.Load()) +} + +func TestSwitcher_Sanitization(t *testing.T) { + t.Parallel() + tests := []struct { + name string + err string + contains []string + notContains []string + }{ + { + name: "nil", + err: "", + contains: nil, + }, + { + name: "simple", + err: "failed to connect", + contains: []string{"failed to connect"}, + }, + { + name: "password in params", + err: "host=localhost password=secret port=5432", + contains: []string{"password="}, + notContains: []string{"secret"}, + }, + { + name: "userinfo in url", + err: "clickhouse://user:pass@localhost:9000/db", + contains: []string{"clickhouse://"}, + notContains: []string{"user", "pass"}, + }, + { + name: "full dsn", + err: "error with clickhouse://localhost:9000?user=default", + contains: []string{"clickhouse://"}, + notContains: []string{"default"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if tt.name == "nil" { + require.NoError(t, sanitizeDriverErr(nil)) + + return + } + got := sanitizeDriverErr(errors.New(tt.err)).Error() + for _, c := range tt.contains { + require.Contains(t, got, c) + } + for _, nc := range tt.notContains { + require.NotContains(t, got, nc) + } + }) + } +} + +func TestSwitcher_Close_Idempotence(t *testing.T) { + t.Parallel() + ff, _ := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + def := &fakeResource{id: "def"} + + s, err := New[*fakeResource](t.Context(), ff, flag, "", nil, func(string) (*fakeResource, error) { return nil, nil }, WithDefaultClient[*fakeResource](def)) + require.NoError(t, err) + + ctx := context.Background() + require.NoError(t, s.Close(ctx)) + require.Equal(t, int32(1), def.closeCalls.Load()) + + // Second call should be no-op + require.NoError(t, s.Close(ctx)) + require.Equal(t, int32(1), def.closeCalls.Load()) +} + +func TestSwitcher_WithDefaultClient(t *testing.T) { + t.Parallel() + ff, _ := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test-flag", "") + def := &fakeResource{id: "explicit-default"} + + s, err := New[*fakeResource]( + t.Context(), ff, flag, "", nil, + func(string) (*fakeResource, error) { return nil, errors.New("factory should not be called") }, + WithDefaultClient[*fakeResource](def), + ) + require.NoError(t, err) + require.Equal(t, def, s.Resolve(context.Background())) + + // P1: Test WithDefaultClient(nil) vulnerability + _, err = New[*fakeResource]( + t.Context(), ff, flag, "", nil, + func(string) (*fakeResource, error) { return nil, nil }, + WithDefaultClient[*fakeResource](nil), + ) + require.Error(t, err) + require.Contains(t, err.Error(), "provided default client must not be nil") +} + +func TestSwitcher_WithMeter(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + meter := noop.NewMeterProvider().Meter("test") + s, err := New[*fakeResource](t.Context(), ff, flag, "default", nil, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }, WithMeter[*fakeResource](meter)) + require.NoError(t, err) + require.NotNil(t, s.resolveCounter) + + // Exercise Resolve with metrics enabled + setFlag(t, source, flag.Key(), "") + s.Resolve(context.Background()) +} + +func TestSwitcher_RecoveryAfterInvalid(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test-flag", "") + + s, err := New[*fakeResource](t.Context(), ff, flag, "default", []string{"alt-0"}, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }) + require.NoError(t, err) + + ctx := context.Background() + + // 1. Start with invalid + setFlag(t, source, flag.Key(), "bad") + require.Equal(t, "default", s.Resolve(ctx).id) + + // 2. Switch to valid - should recover immediately + setFlag(t, source, flag.Key(), "0") + require.Equal(t, "alt-0", s.Resolve(ctx).id) + + // 3. Switch back to invalid + setFlag(t, source, flag.Key(), "99") + require.Equal(t, "default", s.Resolve(ctx).id) + + // 4. Switch to valid again + setFlag(t, source, flag.Key(), "0") + require.Equal(t, "alt-0", s.Resolve(ctx).id) +} + +func TestSwitcher_Concurrency_LiveUpdates(t *testing.T) { + t.Parallel() + ff, source := newTestFeatureFlags(t) + flag := featureflags.NewStringFlag("test", "") + + alts := make([]string, 10) + for i := range 10 { + alts[i] = fmt.Sprintf("alt-%d", i) + } + + s, err := New[*fakeResource](t.Context(), ff, flag, "default", alts, func(dsn string) (*fakeResource, error) { + return &fakeResource{id: dsn}, nil + }) + require.NoError(t, err) + + var wg sync.WaitGroup + ctx := context.Background() + + // Hammer resolve + for i := range 20 { + wg.Add(1) + go func(id int) { + defer wg.Done() + for range 500 { + res := s.Resolve(ctx) + assert.NotNil(t, res) + if id%2 == 0 { + // Some readers should see updates quickly + _ = res.id + } + } + }(i) + } + + // Mutate flag simultaneously + wg.Go(func() { + for i := range 100 { + setFlag(t, source, flag.Key(), fmt.Sprintf("%d", i%10)) + } + // Set to invalid + setFlag(t, source, flag.Key(), "bad") + }) + + wg.Wait() +} From d417f1cb59d4d3bb2b015a73a4e0172a6b22a910 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 22 Jun 2026 20:21:35 +0000 Subject: [PATCH 2/2] chore: auto-commit generated changes --- packages/dashboard-api/go.mod | 11 ++++++++ packages/dashboard-api/go.sum | 26 +++++++++++++++++++ .../pkg/utils/switching/switching_test.go | 5 ++-- 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/packages/dashboard-api/go.mod b/packages/dashboard-api/go.mod index 3afedadb47..6752be7d58 100644 --- a/packages/dashboard-api/go.mod +++ b/packages/dashboard-api/go.mod @@ -80,6 +80,7 @@ require ( github.com/goccy/go-yaml v1.19.2 // indirect github.com/golang-jwt/jwt/v5 v5.3.1 // indirect github.com/gorilla/mux v1.8.1 // indirect + github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 // indirect github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/hashicorp/go-cleanhttp v0.5.2 // indirect @@ -92,6 +93,14 @@ require ( github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.18.5 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/launchdarkly/ccache v1.1.0 // indirect + github.com/launchdarkly/eventsource v1.10.0 // indirect + github.com/launchdarkly/go-jsonstream/v3 v3.1.0 // indirect + github.com/launchdarkly/go-sdk-common/v3 v3.3.0 // indirect + 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/leodido/go-urn v1.4.0 // indirect github.com/lib/pq v1.11.2 // indirect github.com/lufia/plan9stats v0.0.0-20240909124753-873cd0166683 // indirect @@ -115,6 +124,7 @@ require ( github.com/oasdiff/yaml3 v0.0.12 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/paulmach/orb v0.11.1 // indirect github.com/pelletier/go-toml/v2 v2.3.1 // indirect github.com/perimeterx/marshmallow v1.1.5 // indirect @@ -160,6 +170,7 @@ require ( go.uber.org/multierr v1.11.0 // indirect golang.org/x/arch v0.25.0 // indirect golang.org/x/crypto v0.51.0 // indirect + golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect golang.org/x/mod v0.36.0 // indirect golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect diff --git a/packages/dashboard-api/go.sum b/packages/dashboard-api/go.sum index 5083c685a6..e17accad4d 100644 --- a/packages/dashboard-api/go.sum +++ b/packages/dashboard-api/go.sum @@ -141,6 +141,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= +github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 h1:+ngKgrYPPJrOjhax5N+uePQ0Fh1Z7PheYoUI/0nzkPA= +github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA= github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2 h1:sGm2vDRFUrQJO/Veii4h4zG2vvqG6uWNkBHSTqXOZk0= github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2/go.mod h1:wd1YpapPLivG6nQgbf7ZkG1hhSOXDhhn4MLTknx2aAc= github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 h1:HWRh5R2+9EifMyIHV7ZV+MIZqgz+PMpZ14Jynv3O2Zs= @@ -168,6 +170,8 @@ github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFF github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/juju/gnuflag v0.0.0-20171113085948-2ce1bb71843d/go.mod h1:2PavIy+JPciBPrBUjwbNvtwB6RQlve+hkpll6QSNmOE= +github.com/karlseguin/expect v1.0.2-0.20190806010014-778a5f0c6003 h1:vJ0Snvo+SLMY72r5J4sEfkuE7AFbixEP2qRbEcum/wA= +github.com/karlseguin/expect v1.0.2-0.20190806010014-778a5f0c6003/go.mod h1:zNBxMY8P21owkeogJELCLeHIt+voOSduHYTFUbwRAV8= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk= @@ -182,6 +186,24 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/launchdarkly/ccache v1.1.0 h1:voD1M+ZJXR3MREOKtBwgTF9hYHl1jg+vFKS/+VAkR2k= +github.com/launchdarkly/ccache v1.1.0/go.mod h1:TlxzrlnzvYeXiLHmesMuvoZetu4Z97cV1SsdqqBJi1Q= +github.com/launchdarkly/eventsource v1.10.0 h1:H9Tp6AfGu/G2qzBJC26iperrvwhzdbiA/gx7qE2nDFI= +github.com/launchdarkly/eventsource v1.10.0/go.mod h1:J3oa50bPvJesZqNAJtb5btSIo5N6roDWhiAS3IpsKck= +github.com/launchdarkly/go-jsonstream/v3 v3.1.0 h1:U/7/LplZO72XefBQ+FzHf6o4FwLHVqBE+4V58Ornu/E= +github.com/launchdarkly/go-jsonstream/v3 v3.1.0/go.mod h1:2Pt4BR5AwWgsuVTCcIpB6Os04JFIKWfoA+7faKkZB5E= +github.com/launchdarkly/go-sdk-common/v3 v3.3.0 h1:kkf78wcKX+DOXzNjG29i+py/P+XMIw8/mXS7eEWGQwU= +github.com/launchdarkly/go-sdk-common/v3 v3.3.0/go.mod h1:mXFmDGEh4ydK3QilRhrAyKuf9v44VZQWnINyhqbbOd0= +github.com/launchdarkly/go-sdk-events/v3 v3.5.0 h1:Yav8Thm70dZbO8U1foYwZPf3w60n/lNBRaYeeNM/qg4= +github.com/launchdarkly/go-sdk-events/v3 v3.5.0/go.mod h1:oepYWQ2RvvjfL2WxkE1uJJIuRsIMOP4WIVgUpXRPcNI= +github.com/launchdarkly/go-semver v1.0.3 h1:agIy/RN3SqeQDIfKkl+oFslEdeIs7pgsJBs3CdCcGQM= +github.com/launchdarkly/go-semver v1.0.3/go.mod h1:xFmMwXba5Mb+3h72Z+VeSs9ahCvKo2QFUTHRNHVqR28= +github.com/launchdarkly/go-server-sdk-evaluation/v3 v3.0.1 h1:rTgcYAFraGFj7sBMB2b7JCYCm0b9kph4FaMX02t4osQ= +github.com/launchdarkly/go-server-sdk-evaluation/v3 v3.0.1/go.mod h1:fPS5d+zOsgFnMunj+Ki6jjlZtFvo4h9iNbtNXxzYn58= +github.com/launchdarkly/go-server-sdk/v7 v7.13.0 h1:ajiZOPBwmWVFFgP+EMdy3oS1Xl9wNDlEd/7Zn/0I2JU= +github.com/launchdarkly/go-server-sdk/v7 v7.13.0/go.mod h1:6krbDWp417H7lIg+3ehh/A/AW5xwHtiUFg06fvNYHAk= +github.com/launchdarkly/go-test-helpers/v3 v3.1.0 h1:E3bxJMzMoA+cJSF3xxtk2/chr1zshl1ZWa0/oR+8bvg= +github.com/launchdarkly/go-test-helpers/v3 v3.1.0/go.mod h1:Ake5+hZFS/DmIGKx/cizhn5W9pGA7pplcR7xCxWiLIo= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/lib/pq v1.11.2 h1:x6gxUeu39V0BHZiugWe8LXZYZ+Utk7hSJGThs8sdzfs= @@ -243,6 +265,8 @@ github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJw github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= github.com/ory/client-go v1.22.42 h1:uH3IWR1RjP9XCcekTD+SFGp6sZLEwvVTEH0DDqT9Rm4= github.com/ory/client-go v1.22.42/go.mod h1:G1f+5+m/PJVvl40bsRn0QuyVIcXe7EHiWeM7iWpIDjw= +github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= +github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= github.com/paulmach/orb v0.11.1 h1:3koVegMC4X/WeiXYz9iswopaTwMem53NzTJuTF20JzU= github.com/paulmach/orb v0.11.1/go.mod h1:5mULz1xQfs3bmQm63QEJA6lNGujuRafwA5S/EnuLaLU= github.com/paulmach/protoscan v0.2.1/go.mod h1:SpcSwydNLrxUGSDvXvO0P7g7AuhJ7lcKfDlhJCDw2gY= @@ -317,6 +341,8 @@ github.com/ugorji/go/codec v1.3.1 h1:waO7eEiFDwidsBN6agj1vJQ4AG7lh2yqXyOXqhgQuyY github.com/ugorji/go/codec v1.3.1/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4= github.com/woodsbury/decimal128 v1.4.0 h1:xJATj7lLu4f2oObouMt2tgGiElE5gO6mSWUjQsBgUlc= github.com/woodsbury/decimal128 v1.4.0/go.mod h1:BP46FUrVjVhdTbKT+XuQh2xfQaGki9LMIRJSFuh6THU= +github.com/wsxiaoys/terminal v0.0.0-20160513160801-0940f3fc43a0 h1:3UeQBvD0TFrlVjOeLOBz+CPAI8dnbqNSVwUwRrkp7vQ= +github.com/wsxiaoys/terminal v0.0.0-20160513160801-0940f3fc43a0/go.mod h1:IXCdmsXIht47RaVFLEdVnh1t+pgYtTAhQGj73kz+2DM= github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI= github.com/xdg-go/scram v1.1.1/go.mod h1:RaEWvsqvNKKvBPvcKeFjrG2cJqOkHTiyTpzz23ni57g= github.com/xdg-go/stringprep v1.0.3/go.mod h1:W3f5j4i+9rC0kuIEJL0ky1VpHXQU3ocBgklLGvcBnW8= diff --git a/packages/shared/pkg/utils/switching/switching_test.go b/packages/shared/pkg/utils/switching/switching_test.go index e68650ac24..6fed95ea65 100644 --- a/packages/shared/pkg/utils/switching/switching_test.go +++ b/packages/shared/pkg/utils/switching/switching_test.go @@ -8,12 +8,13 @@ import ( "sync/atomic" "testing" - "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/launchdarkly/go-sdk-common/v3/ldvalue" "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/metric/noop" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" ) type fakeResource struct { @@ -303,7 +304,7 @@ func TestSwitcher_WithMeter(t *testing.T) { ff, source := newTestFeatureFlags(t) flag := featureflags.NewStringFlag("test", "") - meter := noop.NewMeterProvider().Meter("test") + meter := noop.NewMeterProvider().Meter("github.com/e2b-dev/infra/packages/shared/pkg/utils/switching") s, err := New[*fakeResource](t.Context(), ff, flag, "default", nil, func(dsn string) (*fakeResource, error) { return &fakeResource{id: dsn}, nil }, WithMeter[*fakeResource](meter))