From 23a9eede4e57f98b151f8308139613b43ab5d889 Mon Sep 17 00:00:00 2001 From: AdaAibaby Date: Sat, 22 Aug 2026 17:57:53 +0800 Subject: [PATCH 1/2] perf(api): add per-allocation sandbox cache to eliminate TeamItems MGET bandwidth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TeamItems issues SMEMBERS + MGET on every call. With teams of 9 000+ sandboxes and multiple API allocations the read bandwidth scales as O(allocations × team_size), saturating the API allocation NIC. This PR introduces a per-allocation in-process cache for sandbox state backed by the existing pub/sub infrastructure (publisher + subscriptionManager, introduced in #2099 / #2668). The cache eliminates the O(allocations × team_size) multiplier: each allocation maintains a local snapshot and TeamItems reads from memory after the first cold-fetch. Design - sandbox_event.go: sandboxEvent JSON type published alongside existing plain routing-key strings on globalStorageNotifyChannel. JSON prefix '{' is an unambiguous discriminator from routing keys ('sandbox:...', 'lock:...'). - sandbox_cache.go: sandboxCache keyed by sandbox ID, indexed by team, with warm/cold state per team. Thread-safe via sync.RWMutex. - publisher.go: publishSandboxEvent marshals and enqueues events on the existing 32-worker publish pool. - subscription_manager.go: dispatch detects JSON events and applies them to the embedded sandboxCache before routing-key fan-out. - operations.go: Add/Update/Remove broadcast events after each Redis write; TeamItems checks the cache (warm-path) or falls back to SMEMBERS+MGET and warms the team on cold-start. Gated by SandboxTeamItemsCacheFlag (default false). - featureflags/flags.go: SandboxTeamItemsCacheFlag for safe rollout. Tests - sandbox_cache_test.go: unit tests for apply/evict/warmTeam/getTeam, state filtering, team isolation, stale-entry eviction, event marshal/unmarshal, routing-key disambiguation. - team_items_test.go: extended with 9 new integration tests (real Redis via testcontainers) covering cold-start warming, empty-team warming, event-driven add/remove reflection, end-to-end Add/Remove/Update broadcast, team isolation, and flag-off behaviour. Closes #3593 --- .../internal/sandbox/storage/redis/main.go | 1 + .../sandbox/storage/redis/operations.go | 91 +++++- .../sandbox/storage/redis/publisher.go | 12 + .../sandbox/storage/redis/sandbox_cache.go | 156 ++++++++++ .../storage/redis/sandbox_cache_test.go | 216 ++++++++++++++ .../sandbox/storage/redis/sandbox_event.go | 54 ++++ .../storage/redis/subscription_manager.go | 24 +- .../sandbox/storage/redis/team_items_test.go | 267 ++++++++++++++++++ packages/shared/pkg/featureflags/flags.go | 7 + 9 files changed, 820 insertions(+), 8 deletions(-) create mode 100644 packages/api/internal/sandbox/storage/redis/sandbox_cache.go create mode 100644 packages/api/internal/sandbox/storage/redis/sandbox_cache_test.go create mode 100644 packages/api/internal/sandbox/storage/redis/sandbox_event.go diff --git a/packages/api/internal/sandbox/storage/redis/main.go b/packages/api/internal/sandbox/storage/redis/main.go index e0cdebb437..24d3ff0794 100644 --- a/packages/api/internal/sandbox/storage/redis/main.go +++ b/packages/api/internal/sandbox/storage/redis/main.go @@ -39,6 +39,7 @@ type Storage struct { subManager *subscriptionManager publisher *publisher featureFlags *featureflags.Client + cacheForced bool // set only in tests via forceCacheEnabled metrics expirationIndexMetrics } diff --git a/packages/api/internal/sandbox/storage/redis/operations.go b/packages/api/internal/sandbox/storage/redis/operations.go index 7fa0574b69..d69325f128 100644 --- a/packages/api/internal/sandbox/storage/redis/operations.go +++ b/packages/api/internal/sandbox/storage/redis/operations.go @@ -13,6 +13,7 @@ import ( "go.uber.org/zap" "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/e2b-dev/infra/packages/shared/pkg/logger" redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" ) @@ -52,6 +53,9 @@ func (s *Storage) Add(ctx context.Context, sbx sandboxtypes.Sandbox) error { logger.L().Warn(ctx, "failed to add team to global teams index", zap.Error(err), logger.WithSandboxID(sbx.SandboxID)) } + // Broadcast to all allocations so their caches stay consistent with Redis. + s.publisher.publishSandboxEvent(ctx, sandboxEvent{Op: sandboxEventOpAdd, Sandbox: &sbx}) + return nil } @@ -117,12 +121,77 @@ func (s *Storage) Remove(ctx context.Context, teamID uuid.UUID, sandboxID string } } + // Evict from all allocations' caches. + s.publisher.publishSandboxEvent(ctx, sandboxEvent{ + Op: sandboxEventOpRemove, + SandboxID: sandboxID, + TeamID: teamID.String(), + }) + return nil } -// TeamItems retrieves sandboxes for a specific team, filtered by states and options +// TeamItems retrieves sandboxes for a specific team, filtered by states. +// +// When the sandbox-team-items-cache feature flag is enabled the result is +// served from the per-allocation in-process cache after the first call for a +// team (cold-fetch path). Subsequent calls are zero-Redis-read. +// +// The cache is kept consistent by sandbox state-change events published on the +// shared pub/sub channel by Add, Update, and Remove. Dropped events cause +// temporary staleness; the cold-fetch on startup recovers full consistency. func (s *Storage) TeamItems(ctx context.Context, teamID uuid.UUID, states []sandboxtypes.State) ([]sandboxtypes.Sandbox, error) { - // Get sandbox IDs from team index + if s.cacheEnabled(ctx) { + if sandboxes, ok := s.subManager.cache.getTeam(teamID, states); ok { + return sandboxes, nil + } + + // Cache cold for this team: fall through to Redis, then warm the cache. + return s.teamItemsFromRedisAndWarm(ctx, teamID, states) + } + + return s.teamItemsFromRedis(ctx, teamID, states) +} + +// teamItemsFromRedisAndWarm fetches from Redis, warms the cache for teamID, +// then returns the state-filtered slice. +func (s *Storage) teamItemsFromRedisAndWarm(ctx context.Context, teamID uuid.UUID, states []sandboxtypes.State) ([]sandboxtypes.Sandbox, error) { + teamKey := GetSandboxStorageTeamIndexKey(teamID.String()) + sandboxIDs, err := s.redisClient.SMembers(ctx, teamKey).Result() + if err != nil { + return nil, fmt.Errorf("failed to get sandbox IDs from team index: %w", err) + } + + if len(sandboxIDs) == 0 { + s.subManager.cache.warmTeam(teamID.String(), nil) + + return []sandboxtypes.Sandbox{}, nil + } + + fetched, err := s.fetchSandboxBatch(ctx, teamID.String(), sandboxIDs) + if err != nil { + return nil, err + } + + // Warm the cache with all sandboxes (unfiltered) so subsequent calls for + // different state filters can still be served from memory. + s.subManager.cache.warmTeam(teamID.String(), fetched) + + var sandboxes []sandboxtypes.Sandbox + for _, sbx := range fetched { + if len(states) > 0 && !slices.Contains(states, sbx.State) { + continue + } + + sandboxes = append(sandboxes, sbx) + } + + return sandboxes, nil +} + +// teamItemsFromRedis is the original Redis-only path, used when the cache +// feature flag is disabled. +func (s *Storage) teamItemsFromRedis(ctx context.Context, teamID uuid.UUID, states []sandboxtypes.State) ([]sandboxtypes.Sandbox, error) { teamKey := GetSandboxStorageTeamIndexKey(teamID.String()) sandboxIDs, err := s.redisClient.SMembers(ctx, teamKey).Result() if err != nil { @@ -142,7 +211,6 @@ func (s *Storage) TeamItems(ctx context.Context, teamID uuid.UUID, states []sand return nil, err } - // Filter by state if states are specified var sandboxes []sandboxtypes.Sandbox for _, sbx := range fetched { if len(states) > 0 && !slices.Contains(states, sbx.State) { @@ -215,6 +283,9 @@ func (s *Storage) Update(ctx context.Context, teamID uuid.UUID, sandboxID string } } + // Broadcast the updated state to all allocations. + s.publisher.publishSandboxEvent(ctx, sandboxEvent{Op: sandboxEventOpUpdate, Sandbox: &updatedSbx}) + return updatedSbx, nil } @@ -281,3 +352,17 @@ func (s *Storage) TeamsWithSandboxCount(ctx context.Context) (map[uuid.UUID]int6 return teams, nil } + +// cacheEnabled reports whether the per-allocation sandbox cache is active. +// Falls back to the flag's default when the feature-flag client is unavailable. +func (s *Storage) cacheEnabled(ctx context.Context) bool { + if s.cacheForced { + return true + } + + if s.featureFlags == nil { + return featureflags.SandboxTeamItemsCacheFlag.Fallback() + } + + return s.featureFlags.BoolFlag(ctx, featureflags.SandboxTeamItemsCacheFlag) +} diff --git a/packages/api/internal/sandbox/storage/redis/publisher.go b/packages/api/internal/sandbox/storage/redis/publisher.go index 7ab42f3afa..781741cd16 100644 --- a/packages/api/internal/sandbox/storage/redis/publisher.go +++ b/packages/api/internal/sandbox/storage/redis/publisher.go @@ -283,3 +283,15 @@ func (p *publisher) close(ctx context.Context) { }) <-p.done } + +// publishSandboxEvent marshals evt to JSON and enqueues it on the shared +// pub/sub channel, alongside the existing plain routing-key messages. +// Never blocks; drops silently under backpressure (same policy as Publish). +func (p *publisher) publishSandboxEvent(ctx context.Context, evt sandboxEvent) { + data, err := marshalSandboxEvent(evt) + if err != nil { + return + } + + p.Publish(ctx, data) +} diff --git a/packages/api/internal/sandbox/storage/redis/sandbox_cache.go b/packages/api/internal/sandbox/storage/redis/sandbox_cache.go new file mode 100644 index 0000000000..6d0443d571 --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/sandbox_cache.go @@ -0,0 +1,156 @@ +package redis + +import ( + "slices" + "sync" + + "github.com/google/uuid" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" +) + +// sandboxCache is an in-process read-through cache for sandbox state, keyed by +// sandbox ID and indexed by team. It is populated via sandbox state-change +// events received over the shared pub/sub channel and by cold-fetch on the +// first TeamItems call for a team. +// +// Consistency model: the cache is eventually consistent with Redis. Dropped +// pub/sub events cause temporary staleness until the next TeamItems cold-fetch +// for that team. Callers that require strong consistency (e.g. ExpiredItems, +// Reconcile) must bypass the cache and query Redis directly. +type sandboxCache struct { + mu sync.RWMutex + byID map[string]sandboxtypes.Sandbox // sandboxID → sandbox + byTeam map[string]map[string]struct{} // teamID → set of sandboxIDs + warm map[string]struct{} // teamIDs whose cache is fully warm +} + +func newSandboxCache() *sandboxCache { + return &sandboxCache{ + byID: make(map[string]sandboxtypes.Sandbox), + byTeam: make(map[string]map[string]struct{}), + warm: make(map[string]struct{}), + } +} + +// apply updates the cache from a sandbox state-change event. +func (c *sandboxCache) apply(evt sandboxEvent) { + c.mu.Lock() + defer c.mu.Unlock() + + switch evt.Op { + case sandboxEventOpAdd, sandboxEventOpUpdate: + if evt.Sandbox == nil { + return + } + + sbx := *evt.Sandbox + teamID := sbx.TeamID.String() + c.byID[sbx.SandboxID] = sbx + + if c.byTeam[teamID] == nil { + c.byTeam[teamID] = make(map[string]struct{}) + } + + c.byTeam[teamID][sbx.SandboxID] = struct{}{} + + case sandboxEventOpRemove: + c.evictLocked(evt.TeamID, evt.SandboxID) + } +} + +// evictLocked removes a sandbox entry. Must be called with c.mu held. +func (c *sandboxCache) evictLocked(teamID, sandboxID string) { + delete(c.byID, sandboxID) + + if ids, ok := c.byTeam[teamID]; ok { + delete(ids, sandboxID) + if len(ids) == 0 { + delete(c.byTeam, teamID) + } + } +} + +// warmTeam atomically replaces the cached snapshot for teamID with the +// freshly-fetched set of sandboxes, then marks the team as warm. +// +// Sandboxes already in the cache for this team but absent from the fresh fetch +// are evicted (they were removed from Redis while the cache was cold). +func (c *sandboxCache) warmTeam(teamID string, sandboxes []sandboxtypes.Sandbox) { + c.mu.Lock() + defer c.mu.Unlock() + + // Build lookup of the fresh set to detect removals. + fresh := make(map[string]struct{}, len(sandboxes)) + for _, sbx := range sandboxes { + fresh[sbx.SandboxID] = struct{}{} + } + + // Evict stale entries that are no longer present in Redis. + if existing, ok := c.byTeam[teamID]; ok { + for sid := range existing { + if _, seen := fresh[sid]; !seen { + c.evictLocked(teamID, sid) + } + } + } + + // Insert or overwrite with the fresh state. + for _, sbx := range sandboxes { + c.byID[sbx.SandboxID] = sbx + + if c.byTeam[teamID] == nil { + c.byTeam[teamID] = make(map[string]struct{}) + } + + c.byTeam[teamID][sbx.SandboxID] = struct{}{} + } + + c.warm[teamID] = struct{}{} +} + +// getTeam returns the cached sandboxes for teamID filtered by states. +// Returns (nil, false) when the team has not been warmed yet; the caller +// should fall back to a Redis cold-fetch. +func (c *sandboxCache) getTeam(teamID uuid.UUID, states []sandboxtypes.State) ([]sandboxtypes.Sandbox, bool) { + tid := teamID.String() + + c.mu.RLock() + defer c.mu.RUnlock() + + if _, ok := c.warm[tid]; !ok { + return nil, false + } + + ids, ok := c.byTeam[tid] + if !ok { + return []sandboxtypes.Sandbox{}, true + } + + result := make([]sandboxtypes.Sandbox, 0, len(ids)) + for id := range ids { + sbx, found := c.byID[id] + if !found { + continue + } + + if len(states) > 0 && !slices.Contains(states, sbx.State) { + continue + } + + result = append(result, sbx) + } + + return result, true +} + +// isWarm reports whether the cache holds a complete snapshot for teamID. +// Used only in tests to assert cold-start and warm-path behaviour. +func (c *sandboxCache) isWarm(teamID string) bool { + c.mu.RLock() + defer c.mu.RUnlock() + + _, ok := c.warm[teamID] + + return ok +} diff --git a/packages/api/internal/sandbox/storage/redis/sandbox_cache_test.go b/packages/api/internal/sandbox/storage/redis/sandbox_cache_test.go new file mode 100644 index 0000000000..0e3de2f7ff --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/sandbox_cache_test.go @@ -0,0 +1,216 @@ +package redis + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" +) + +func makeCacheSandbox(teamID uuid.UUID, sandboxID string, state sandboxtypes.State) sandboxtypes.Sandbox { + return sandboxtypes.Sandbox{ + SandboxID: sandboxID, + TeamID: teamID, + ExecutionID: uuid.NewString(), + StartTime: time.Now().Add(-time.Hour), + EndTime: time.Now().Add(time.Hour), + State: state, + } +} + +// TestSandboxCache_AddAndGet verifies that an add event populates the cache +// and getTeam returns it after warmTeam is called. +func TestSandboxCache_AddAndGet(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + sbx := makeCacheSandbox(teamID, "sbx-1", sandboxtypes.StateRunning) + + c.warmTeam(teamID.String(), []sandboxtypes.Sandbox{sbx}) + + got, ok := c.getTeam(teamID, nil) + require.True(t, ok, "cache should be warm after warmTeam") + require.Len(t, got, 1) + assert.Equal(t, sbx.SandboxID, got[0].SandboxID) +} + +// TestSandboxCache_ApplyAdd verifies that an add event populates the cache +// even without an explicit warmTeam call, once the team is already warm. +func TestSandboxCache_ApplyAdd(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + + // Warm with empty set. + c.warmTeam(teamID.String(), nil) + + sbx := makeCacheSandbox(teamID, "sbx-2", sandboxtypes.StateRunning) + c.apply(sandboxEvent{Op: sandboxEventOpAdd, Sandbox: &sbx}) + + got, ok := c.getTeam(teamID, nil) + require.True(t, ok) + require.Len(t, got, 1) + assert.Equal(t, "sbx-2", got[0].SandboxID) +} + +// TestSandboxCache_ApplyUpdate verifies that an update event replaces the +// existing sandbox entry in the cache. +func TestSandboxCache_ApplyUpdate(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + sbx := makeCacheSandbox(teamID, "sbx-3", sandboxtypes.StateRunning) + c.warmTeam(teamID.String(), []sandboxtypes.Sandbox{sbx}) + + sbx.State = sandboxtypes.StatePausing + c.apply(sandboxEvent{Op: sandboxEventOpUpdate, Sandbox: &sbx}) + + got, ok := c.getTeam(teamID, nil) + require.True(t, ok) + require.Len(t, got, 1) + assert.Equal(t, sandboxtypes.StatePausing, got[0].State) +} + +// TestSandboxCache_ApplyRemove verifies that a remove event evicts the entry. +func TestSandboxCache_ApplyRemove(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + sbx := makeCacheSandbox(teamID, "sbx-4", sandboxtypes.StateRunning) + c.warmTeam(teamID.String(), []sandboxtypes.Sandbox{sbx}) + + c.apply(sandboxEvent{Op: sandboxEventOpRemove, SandboxID: sbx.SandboxID, TeamID: teamID.String()}) + + got, ok := c.getTeam(teamID, nil) + require.True(t, ok, "team should still be warm after remove") + assert.Empty(t, got, "removed sandbox should not appear") +} + +// TestSandboxCache_ColdTeamReturnsFalse verifies that getTeam returns false +// for a team that has never been warmed. +func TestSandboxCache_ColdTeamReturnsFalse(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + _, ok := c.getTeam(uuid.New(), nil) + assert.False(t, ok, "unwarmed team should return false") +} + +// TestSandboxCache_FiltersByState verifies state filtering in getTeam. +func TestSandboxCache_FiltersByState(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + + sandboxes := []sandboxtypes.Sandbox{ + makeCacheSandbox(teamID, "sbx-running", sandboxtypes.StateRunning), + makeCacheSandbox(teamID, "sbx-killing", sandboxtypes.StateKilling), + makeCacheSandbox(teamID, "sbx-pausing", sandboxtypes.StatePausing), + } + c.warmTeam(teamID.String(), sandboxes) + + running, ok := c.getTeam(teamID, []sandboxtypes.State{sandboxtypes.StateRunning}) + require.True(t, ok) + require.Len(t, running, 1) + assert.Equal(t, "sbx-running", running[0].SandboxID) + + multi, ok := c.getTeam(teamID, []sandboxtypes.State{sandboxtypes.StateRunning, sandboxtypes.StatePausing}) + require.True(t, ok) + assert.Len(t, multi, 2) + + all, ok := c.getTeam(teamID, nil) + require.True(t, ok) + assert.Len(t, all, 3) +} + +// TestSandboxCache_TeamIsolation verifies that events for one team do not +// appear in another team's results. +func TestSandboxCache_TeamIsolation(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamA, teamB := uuid.New(), uuid.New() + + sbxA := makeCacheSandbox(teamA, "sbx-a", sandboxtypes.StateRunning) + sbxB := makeCacheSandbox(teamB, "sbx-b", sandboxtypes.StateRunning) + + c.warmTeam(teamA.String(), []sandboxtypes.Sandbox{sbxA}) + c.warmTeam(teamB.String(), []sandboxtypes.Sandbox{sbxB}) + + gotA, ok := c.getTeam(teamA, nil) + require.True(t, ok) + require.Len(t, gotA, 1) + assert.Equal(t, "sbx-a", gotA[0].SandboxID) + + gotB, ok := c.getTeam(teamB, nil) + require.True(t, ok) + require.Len(t, gotB, 1) + assert.Equal(t, "sbx-b", gotB[0].SandboxID) +} + +// TestSandboxCache_WarmTeamEvictsStaleSandboxes verifies that re-warming a +// team removes sandboxes that are no longer present in the fresh Redis fetch. +func TestSandboxCache_WarmTeamEvictsStaleSandboxes(t *testing.T) { + t.Parallel() + + c := newSandboxCache() + teamID := uuid.New() + + old := makeCacheSandbox(teamID, "sbx-old", sandboxtypes.StateRunning) + fresh := makeCacheSandbox(teamID, "sbx-fresh", sandboxtypes.StateRunning) + + c.warmTeam(teamID.String(), []sandboxtypes.Sandbox{old}) + + // Re-warm with only the fresh sandbox — old should be evicted. + c.warmTeam(teamID.String(), []sandboxtypes.Sandbox{fresh}) + + got, ok := c.getTeam(teamID, nil) + require.True(t, ok) + require.Len(t, got, 1) + assert.Equal(t, "sbx-fresh", got[0].SandboxID) +} + +// TestSandboxEvent_Roundtrip verifies JSON marshal/unmarshal of sandboxEvent. +func TestSandboxEvent_Roundtrip(t *testing.T) { + t.Parallel() + + teamID := uuid.New() + sbx := makeCacheSandbox(teamID, "sbx-roundtrip", sandboxtypes.StateRunning) + + evt := sandboxEvent{Op: sandboxEventOpAdd, Sandbox: &sbx} + payload, err := marshalSandboxEvent(evt) + require.NoError(t, err) + assert.True(t, isSandboxEvent(payload)) + + got, ok := parseSandboxEvent(payload) + require.True(t, ok) + assert.Equal(t, sandboxEventOpAdd, got.Op) + require.NotNil(t, got.Sandbox) + assert.Equal(t, sbx.SandboxID, got.Sandbox.SandboxID) +} + +// TestSandboxEvent_RoutingKeyIsNotEvent ensures existing routing-key strings +// are not misidentified as sandbox events. +func TestSandboxEvent_RoutingKeyIsNotEvent(t *testing.T) { + t.Parallel() + + routingKeys := []string{ + "sandbox:storage:team-abc:transition:sbx-1:txn-1:notify", + "lock:sandbox:storage:team-abc:sandboxes:sbx-1:notify", + "", + "plain-string", + } + + for _, key := range routingKeys { + assert.False(t, isSandboxEvent(key), "routing key %q should not be detected as event", key) + } +} diff --git a/packages/api/internal/sandbox/storage/redis/sandbox_event.go b/packages/api/internal/sandbox/storage/redis/sandbox_event.go new file mode 100644 index 0000000000..43acc980e4 --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/sandbox_event.go @@ -0,0 +1,54 @@ +package redis + +import ( + "encoding/json" + "strings" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" +) + +const ( + sandboxEventOpAdd = "add" + sandboxEventOpUpdate = "update" + sandboxEventOpRemove = "remove" +) + +// sandboxEvent is published on globalStorageNotifyChannel alongside existing +// plain-string routing keys. It carries the full sandbox payload so consumers +// can update their local cache without a follow-up Redis GET. +// +// Disambiguation: routing keys always begin with "sandbox:storage:" or "lock:", +// so a JSON object prefix "{" is an unambiguous discriminator. +type sandboxEvent struct { + Op string `json:"op"` + Sandbox *sandboxtypes.Sandbox `json:"sandbox,omitempty"` + SandboxID string `json:"sandbox_id,omitempty"` + TeamID string `json:"team_id,omitempty"` +} + +// isSandboxEvent reports whether a pub/sub payload is a sandboxEvent rather +// than a plain routing-key string. +func isSandboxEvent(payload string) bool { + return strings.HasPrefix(payload, "{") +} + +// parseSandboxEvent deserializes a sandboxEvent from a pub/sub payload. +// Returns false if the payload is not a valid sandboxEvent. +func parseSandboxEvent(payload string) (sandboxEvent, bool) { + var evt sandboxEvent + if err := json.Unmarshal([]byte(payload), &evt); err != nil || evt.Op == "" { + return sandboxEvent{}, false + } + + return evt, true +} + +// marshalSandboxEvent serializes a sandboxEvent to a JSON string for publishing. +func marshalSandboxEvent(evt sandboxEvent) (string, error) { + data, err := json.Marshal(evt) + if err != nil { + return "", err + } + + return string(data), nil +} diff --git a/packages/api/internal/sandbox/storage/redis/subscription_manager.go b/packages/api/internal/sandbox/storage/redis/subscription_manager.go index 3431bc2038..5bf61ff437 100644 --- a/packages/api/internal/sandbox/storage/redis/subscription_manager.go +++ b/packages/api/internal/sandbox/storage/redis/subscription_manager.go @@ -7,12 +7,16 @@ import ( "github.com/redis/go-redis/v9" ) -// subscriptionManager maintains a Redis PubSub connection and -// fans out storage notifications to registered in-process waiters. +// subscriptionManager maintains a Redis PubSub connection and fans out storage +// notifications to registered in-process waiters. It also maintains a +// sandboxCache that is updated when sandbox state-change events arrive on the +// channel alongside the existing routing-key messages. type subscriptionManager struct { mu sync.RWMutex waiters map[string]map[chan struct{}]struct{} // routingKey → registered waiters + cache *sandboxCache + redisClient redis.UniversalClient channel string stop chan struct{} @@ -22,6 +26,7 @@ type subscriptionManager struct { func newSubscriptionManager(redisClient redis.UniversalClient, channel string) *subscriptionManager { return &subscriptionManager{ waiters: make(map[string]map[chan struct{}]struct{}), + cache: newSandboxCache(), redisClient: redisClient, channel: channel, stop: make(chan struct{}), @@ -88,12 +93,21 @@ func (m *subscriptionManager) subscribe(routingKey string) (<-chan struct{}, fun return channel, cleanup } -// dispatch signals all waiters registered for the given routing key. -func (m *subscriptionManager) dispatch(routingKey string) { +// dispatch processes an incoming pub/sub payload. Payloads that parse as +// sandboxEvents are applied to the local cache; all other payloads are treated +// as plain routing keys and fan out to registered waiters. +func (m *subscriptionManager) dispatch(payload string) { + if isSandboxEvent(payload) { + if evt, ok := parseSandboxEvent(payload); ok { + m.cache.apply(evt) + return + } + } + m.mu.RLock() defer m.mu.RUnlock() - for waiter := range m.waiters[routingKey] { + for waiter := range m.waiters[payload] { select { case waiter <- struct{}{}: default: diff --git a/packages/api/internal/sandbox/storage/redis/team_items_test.go b/packages/api/internal/sandbox/storage/redis/team_items_test.go index 9089f15208..f567113e61 100644 --- a/packages/api/internal/sandbox/storage/redis/team_items_test.go +++ b/packages/api/internal/sandbox/storage/redis/team_items_test.go @@ -1,6 +1,7 @@ package redis import ( + "context" "encoding/json" "testing" "time" @@ -48,6 +49,18 @@ func sandboxIDsOf(sbxs []sandboxtypes.Sandbox) []string { return ids } +// enableCache forces the sandbox cache on for the duration of the test. +// It sets the internal cacheForced field, which bypasses the LaunchDarkly +// feature flag so tests work without a real feature-flag client. +func enableCache(t *testing.T, storage *Storage) { + t.Helper() + + storage.cacheForced = true + t.Cleanup(func() { storage.cacheForced = false }) +} + +// ── Redis-only path (cache disabled) ───────────────────────────────────────── + func TestTeamItems_FiltersByState(t *testing.T) { t.Parallel() @@ -146,3 +159,257 @@ func TestTeamItems_IsScopedToOneTeam(t *testing.T) { require.NoError(t, err) assert.Equal(t, []string{"sbx-a"}, sandboxIDsOf(items)) } + +func TestTeamItems_BatchesLargeTeams(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + teamID := uuid.New() + + const n = 513 // 2×256+1: exercises all three MGET batches + for i := range n { + sbx := seedTeamSandbox(t, teamID, uuid.NewString(), sandboxtypes.StateRunning) + sbx.SandboxID = sbx.SandboxID + "-" + string(rune('a'+i%26)) + writeTeamSandbox(t, storage, sbx) + } + + items, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Len(t, items, n) +} + +// ── Cache path ──────────────────────────────────────────────────────────────── + +// TestTeamItems_CacheColdStartWarms verifies that the first TeamItems call +// populates the cache and subsequent calls hit memory. +func TestTeamItems_CacheColdStartWarms(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + sbx := seedTeamSandbox(t, teamID, "sbx-warm", sandboxtypes.StateRunning) + writeTeamSandbox(t, storage, sbx) + + // First call: cache is cold → falls back to Redis and warms. + items, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Equal(t, []string{"sbx-warm"}, sandboxIDsOf(items)) + assert.True(t, storage.subManager.cache.isWarm(teamID.String()), "cache should be warm after cold-fetch") + + // Second call: cache is warm → zero Redis reads (no way to count Redis + // commands in unit tests, but we verify the result is still correct). + items2, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Equal(t, []string{"sbx-warm"}, sandboxIDsOf(items2)) +} + +// TestTeamItems_CacheEmptyTeamIsWarm verifies that a team with no sandboxes +// is marked warm after the cold-fetch so it does not hit Redis on every call. +func TestTeamItems_CacheEmptyTeamIsWarm(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + items, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Empty(t, items) + assert.True(t, storage.subManager.cache.isWarm(teamID.String())) +} + +// TestTeamItems_CacheReflectsAddEvent verifies that after the cache is warm, +// an add event (delivered via pub/sub) makes the new sandbox visible without +// a Redis round-trip. +func TestTeamItems_CacheReflectsAddEvent(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + // Warm with empty team. + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + // Simulate an add event arriving from another allocation. + sbx := seedTeamSandbox(t, teamID, "sbx-via-event", sandboxtypes.StateRunning) + storage.subManager.cache.apply(sandboxEvent{Op: sandboxEventOpAdd, Sandbox: &sbx}) + + items, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Equal(t, []string{"sbx-via-event"}, sandboxIDsOf(items)) +} + +// TestTeamItems_CacheReflectsRemoveEvent verifies that a remove event evicts +// the sandbox from the cache. +func TestTeamItems_CacheReflectsRemoveEvent(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + sbx := seedTeamSandbox(t, teamID, "sbx-to-remove", sandboxtypes.StateRunning) + writeTeamSandbox(t, storage, sbx) + + // Warm the cache. + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + // Simulate a remove event. + storage.subManager.cache.apply(sandboxEvent{ + Op: sandboxEventOpRemove, + SandboxID: sbx.SandboxID, + TeamID: teamID.String(), + }) + + items, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + assert.Empty(t, items) +} + +// TestTeamItems_CacheAddBroadcastedViaStorage verifies the end-to-end path: +// Storage.Add publishes an event that is received by the subscriptionManager +// and applied to the local cache so TeamItems sees the new sandbox. +func TestTeamItems_CacheAddBroadcastedViaStorage(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + // Warm with empty team so the cache is active. + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + // Add via the real Storage.Add path (writes Redis + publishes event). + sbx := makeIndexedSandbox(teamID, "sbx-e2e", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + sbx.State = sandboxtypes.StateRunning + require.NoError(t, storage.Add(t.Context(), sbx)) + + // The event is delivered asynchronously; poll with a short deadline. + require.Eventually(t, func() bool { + items, err := storage.TeamItems(t.Context(), teamID, nil) + if err != nil { + return false + } + return len(items) == 1 && items[0].SandboxID == "sbx-e2e" + }, 2*time.Second, 50*time.Millisecond, "cache should reflect Add within 2s") +} + +// TestTeamItems_CacheRemoveBroadcastedViaStorage verifies that Storage.Remove +// publishes an event that evicts the sandbox from the cache. +func TestTeamItems_CacheRemoveBroadcastedViaStorage(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + sbx := makeIndexedSandbox(teamID, "sbx-e2e-rm", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + sbx.State = sandboxtypes.StateRunning + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Warm the cache. + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + // Remove via the real path. + require.NoError(t, storage.Remove(context.WithoutCancel(t.Context()), teamID, sbx.SandboxID)) + + require.Eventually(t, func() bool { + items, err := storage.TeamItems(t.Context(), teamID, nil) + return err == nil && len(items) == 0 + }, 2*time.Second, 50*time.Millisecond, "cache should reflect Remove within 2s") +} + +// TestTeamItems_CacheUpdateBroadcastedViaStorage verifies that Storage.Update +// publishes an event that refreshes the cached sandbox state. +func TestTeamItems_CacheUpdateBroadcastedViaStorage(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamID := uuid.New() + + sbx := makeIndexedSandbox(teamID, "sbx-e2e-upd", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + sbx.State = sandboxtypes.StateRunning + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Warm the cache. + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + // Update state via the real path. + _, err = storage.Update(t.Context(), teamID, sbx.SandboxID, func(s sandboxtypes.Sandbox) (sandboxtypes.Sandbox, error) { + s.State = sandboxtypes.StatePausing + + return s, nil + }) + require.NoError(t, err) + + require.Eventually(t, func() bool { + items, err := storage.TeamItems(t.Context(), teamID, nil) + if err != nil || len(items) != 1 { + return false + } + + return items[0].State == sandboxtypes.StatePausing + }, 2*time.Second, 50*time.Millisecond, "cache should reflect Update within 2s") +} + +// TestTeamItems_CacheIsolatesTeams verifies that events for one team cannot +// affect another team's results when the cache is enabled. +func TestTeamItems_CacheIsolatesTeams(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + teamA, teamB := uuid.New(), uuid.New() + + sbxA := seedTeamSandbox(t, teamA, "sbx-a", sandboxtypes.StateRunning) + sbxB := seedTeamSandbox(t, teamB, "sbx-b", sandboxtypes.StateRunning) + writeTeamSandbox(t, storage, sbxA) + writeTeamSandbox(t, storage, sbxB) + + // Warm both teams. + _, err := storage.TeamItems(t.Context(), teamA, nil) + require.NoError(t, err) + _, err = storage.TeamItems(t.Context(), teamB, nil) + require.NoError(t, err) + + // Remove sbx-a; sbx-b must remain. + storage.subManager.cache.apply(sandboxEvent{ + Op: sandboxEventOpRemove, SandboxID: sbxA.SandboxID, TeamID: teamA.String(), + }) + + gotA, err := storage.TeamItems(t.Context(), teamA, nil) + require.NoError(t, err) + assert.Empty(t, gotA) + + gotB, err := storage.TeamItems(t.Context(), teamB, nil) + require.NoError(t, err) + assert.Equal(t, []string{"sbx-b"}, sandboxIDsOf(gotB)) +} + +// TestTeamItems_CacheDisabledByDefault verifies that with no flag override +// the cache is not consulted (team is never warmed via TeamItems alone). +func TestTeamItems_CacheDisabledByDefault(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + // No enableCache call — flag defaults to false. + teamID := uuid.New() + + sbx := seedTeamSandbox(t, teamID, "sbx-nocache", sandboxtypes.StateRunning) + writeTeamSandbox(t, storage, sbx) + + _, err := storage.TeamItems(t.Context(), teamID, nil) + require.NoError(t, err) + + assert.False(t, storage.subManager.cache.isWarm(teamID.String()), + "cache should NOT be warmed when flag is off") +} diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 1b7935044e..501c374153 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -351,6 +351,13 @@ var ( // On by default; acts as a kill switch if a heal pass misbehaves. ExpirationIndexHealerFlag = NewBoolFlag("expiration-index-healer", true) + // SandboxTeamItemsCacheFlag enables the per-allocation in-process sandbox cache + // for TeamItems. When enabled, TeamItems is served from memory after the first + // cold-fetch; subsequent mutations broadcast events over the existing pub/sub + // channel so every allocation’s cache stays consistent with Redis. + // Defaults to false for safe rollout; flip to true once observed in staging. + SandboxTeamItemsCacheFlag = NewBoolFlag("sandbox-team-items-cache", false) + // DisableE2BAccessTokenProvisioningFlag stops POST /access-tokens from issuing // new E2B access tokens (sk_e2b_) once enabled. E2B_ACCESS_TOKEN is deprecated // in favor of E2B_API_KEY; the CLI now authenticates via Hydra JWTs. Off by From 62eac6c9ddcda69bb8d1b5f0f50307dd079ff357 Mon Sep 17 00:00:00 2001 From: AdaAibaby Date: Sat, 22 Aug 2026 18:52:16 +0800 Subject: [PATCH 2/2] fix(sandbox-cache): broadcast state transition event from StartRemoving MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before this change, StartRemoving atomically wrote the new sandbox state (e.g. Running→Killing) to Redis via startTransitionScript, but never called publishSandboxEvent. Every allocation continued to see the old Running state in its local TeamItems cache until the sandbox was eventually deleted and a remove event arrived. Fix: publish a sandboxEvent{update} immediately after the Lua script succeeds. dispatch() applies it via cache.apply, so all allocations reflect the intermediate state (Killing, Pausing, Snapshotting) without any Redis read. The createCallback path needs no separate publish: - Terminal transitions: the subsequent Remove() already broadcasts remove. - Transient transitions: restoreToRunning calls Update() which already broadcasts the restored Running state. Adds two integration tests: - TestStartRemoving_CacheBroadcastsTransitionState: Kill transition visible as Killing in cache before Remove is called. - TestStartRemoving_CacheBroadcastsTransientTransition: Snapshot transition visible as Snapshotting, then restored to Running. Fixes #3595 --- .../sandbox/storage/redis/state_change.go | 5 ++ .../storage/redis/state_change_test.go | 85 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/packages/api/internal/sandbox/storage/redis/state_change.go b/packages/api/internal/sandbox/storage/redis/state_change.go index 072a997836..cd7b065758 100644 --- a/packages/api/internal/sandbox/storage/redis/state_change.go +++ b/packages/api/internal/sandbox/storage/redis/state_change.go @@ -170,6 +170,11 @@ func (s *Storage) StartRemoving(ctx context.Context, teamID uuid.UUID, sandboxID logger.L().Debug(ctx, "Started state transition", logger.WithSandboxID(sandboxID), zap.String("state", string(newState)), zap.String("transitionID", transitionID)) + // Broadcast the state change so all allocations caches reflect the new + // intermediate state immediately. Without this, caches show State=Running + // until the sandbox is eventually deleted and a remove event arrives. + s.publisher.publishSandboxEvent(ctx, sandboxEvent{Op: sandboxEventOpUpdate, Sandbox: &updated}) + return updated, false, s.createCallback(teamID, sandboxID, transitionKey, resultKey, transitionID, opts.Action), nil } diff --git a/packages/api/internal/sandbox/storage/redis/state_change_test.go b/packages/api/internal/sandbox/storage/redis/state_change_test.go index 0c0ce298c3..fe7987a7c2 100644 --- a/packages/api/internal/sandbox/storage/redis/state_change_test.go +++ b/packages/api/internal/sandbox/storage/redis/state_change_test.go @@ -1269,3 +1269,88 @@ func TestStartRemoving_CompletedTransitionAllowsNewTransition(t *testing.T) { callback2(ctx, nil) } + +// TestStartRemoving_CacheBroadcastsTransitionState verifies that StartRemoving +// publishes a sandboxEvent so the per-allocation cache reflects the new +// intermediate state (e.g. Killing) without waiting for the final Remove. +func TestStartRemoving_CacheBroadcastsTransitionState(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + ctx := t.Context() + + sbx := createTestSandbox("sbx-cache-transition") + sbx.State = sandboxtypes.StateRunning + require.NoError(t, storage.Add(ctx, sbx)) + + // Warm the cache so TeamItems is served from memory. + _, err := storage.TeamItems(ctx, sbx.TeamID, nil) + require.NoError(t, err) + + // Sanity: cache shows Running before the transition. + items, err := storage.TeamItems(ctx, sbx.TeamID, []sandboxtypes.State{sandboxtypes.StateRunning}) + require.NoError(t, err) + require.Len(t, items, 1) + + // Start a Kill transition: Lua writes State=Killing to Redis and + // publishSandboxEvent must broadcast the update to the cache. + _, alreadyDone, callback, err := storage.StartRemoving(ctx, sbx.TeamID, sbx.SandboxID, sandboxtypes.RemoveOpts{ + Action: sandboxtypes.StateActionKill, + }) + require.NoError(t, err) + require.False(t, alreadyDone) + require.NotNil(t, callback) + defer callback(ctx, nil) + + // The event is delivered asynchronously via pub/sub; poll with a short deadline. + require.Eventually(t, func() bool { + killing, err := storage.TeamItems(ctx, sbx.TeamID, []sandboxtypes.State{sandboxtypes.StateKilling}) + return err == nil && len(killing) == 1 + }, 2*time.Second, 50*time.Millisecond, + "cache should reflect Killing state after StartRemoving without waiting for Remove") + + // TeamItems filtered to Running must now return nothing. + running, err := storage.TeamItems(ctx, sbx.TeamID, []sandboxtypes.State{sandboxtypes.StateRunning}) + require.NoError(t, err) + require.Empty(t, running, "cache must not show Running once transition is in flight") +} + +// TestStartRemoving_CacheBroadcastsTransientTransition verifies that a +// transient (snapshot) transition is visible in the cache as Snapshotting, and +// that the subsequent restore-to-Running is equally reflected. +func TestStartRemoving_CacheBroadcastsTransientTransition(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + enableCache(t, storage) + ctx := t.Context() + + sbx := createTestSandbox("sbx-cache-transient") + sbx.State = sandboxtypes.StateRunning + require.NoError(t, storage.Add(ctx, sbx)) + + _, err := storage.TeamItems(ctx, sbx.TeamID, nil) + require.NoError(t, err) + + // Start a Snapshot transition (TransitionTransient: restores to Running on success). + _, _, callback, err := storage.StartRemoving(ctx, sbx.TeamID, sbx.SandboxID, sandboxtypes.RemoveOpts{ + Action: sandboxtypes.StateActionSnapshot, + }) + require.NoError(t, err) + require.NotNil(t, callback) + + // Cache should show Snapshotting. + require.Eventually(t, func() bool { + snapshotting, err := storage.TeamItems(ctx, sbx.TeamID, []sandboxtypes.State{sandboxtypes.StateSnapshotting}) + return err == nil && len(snapshotting) == 1 + }, 2*time.Second, 50*time.Millisecond, "cache should reflect Snapshotting state") + + // Signal success: restoreToRunning calls Update which re-broadcasts Running. + callback(ctx, nil) + + require.Eventually(t, func() bool { + running, err := storage.TeamItems(ctx, sbx.TeamID, []sandboxtypes.State{sandboxtypes.StateRunning}) + return err == nil && len(running) == 1 + }, 2*time.Second, 50*time.Millisecond, "cache should reflect Running after transient transition completes") +}