perf(api): add per-allocation sandbox cache to eliminate TeamItems MGET bandwidth

packages/api/internal/sandbox/storage/redis/main.go

@@ -39,8 +39,30 @@ type Storage struct {
 	subManager   *subscriptionManager
 	publisher    *publisher
 	featureFlags *featureflags.Client
+	cacheForced  bool // set only in tests via forceCacheEnabled
 
-	metrics expirationIndexMetrics
+	metrics      expirationIndexMetrics
+	cacheMetrics sandboxCacheMetrics
+}
+
+// sandboxCacheMetrics tracks per-allocation TeamItems cache effectiveness.
+type sandboxCacheMetrics struct {
+	hits   metric.Int64Counter
+	misses metric.Int64Counter
+}
+
+func newSandboxCacheMetrics(meter metric.Meter) (sandboxCacheMetrics, error) {
+	hits, err := telemetry.GetCounter(meter, telemetry.ApiRedisStorageSandboxCacheHits)
+	if err != nil {
+		return sandboxCacheMetrics{}, fmt.Errorf("sandbox cache hits counter: %w", err)
+	}
+
+	misses, err := telemetry.GetCounter(meter, telemetry.ApiRedisStorageSandboxCacheMisses)
+	if err != nil {
+		return sandboxCacheMetrics{}, fmt.Errorf("sandbox cache misses counter: %w", err)
+	}
+
+	return sandboxCacheMetrics{hits: hits, misses: misses}, nil
 }
 
 // expirationIndexMetrics observes global expiration index consistency.
@@ -103,13 +125,19 @@ func NewStorage(
 		return nil, fmt.Errorf("failed to create expiration index metrics: %w", err)
 	}
 
+	cacheMetrics, err := newSandboxCacheMetrics(meter)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create sandbox cache metrics: %w", err)
+	}
+
 	return &Storage{
 		redisClient:  redisClient,
 		locker:       newStorageLocker(redisClient, subManager, pub),
 		subManager:   subManager,
 		publisher:    pub,
 		featureFlags: featureFlags,
 		metrics:      metrics,
+		cacheMetrics: cacheMetrics,
 	}, nil
 }
 

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,79 @@ 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 {
+			s.cacheMetrics.hits.Add(ctx, 1)
+			return sandboxes, nil
+		}
+
+		// Cache cold for this team: fall through to Redis, then warm the cache.
+		s.cacheMetrics.misses.Add(ctx, 1)
+		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 +213,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 +285,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 +354,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)
+}

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)
+}

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
+}