[EN-1048] fix(api): heal sandbox missing from eviction index

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

@@ -0,0 +1,195 @@
+package redis
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+const (
+	healInterval = 5 * time.Minute
+
+	// healGracePeriod skips recently started sandboxes:
+	// this prevents the healer from clearing in-flight Add/Remove
+	healGracePeriod = time.Minute
+
+	// healScanBatchSize bounds per-command work (SSCAN page, MGET keys,
+	// ZMSCORE members, ZADD members) so teams with many sandboxes can't
+	// produce huge single commands/replies that stall Redis or explode service memory
+	healScanBatchSize = 256
+)
+
+// startHealer restores "sandbox key exists => expiration index member exists".
+// A sandbox missing from the global expiration ZSET is never seen by the evictor
+// and would otherwise live forever.
+//
+// Runs on every API pod with jitter; ZADD NX makes concurrent passes idempotent and harmless.
+func (s *Storage) startHealer(ctx context.Context) {
+	timer := time.NewTimer(jitterBackoff(healInterval))
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-timer.C:
+			healed, err := s.healExpirationIndex(ctx)
+			if err != nil {
+				logger.L().Warn(ctx, "Expiration index heal pass failed", zap.Error(err))
+			}
+			if healed > 0 {
+				logger.L().Warn(ctx, "Healed sandboxes missing from expiration index", zap.Int("count", healed))
+			}
+
+			timer.Reset(jitterBackoff(healInterval))
+		}
+	}
+}
+
+// healExpirationIndex scans all stored sandboxes and re-adds expiration index
+// members missing for live sandbox keys. Returns the number of healed members.
+func (s *Storage) healExpirationIndex(ctx context.Context) (int, error) {
+	teams, err := s.redisClient.ZRange(ctx, globalTeamsSet, 0, -1).Result()
+	if err != nil {
+		return 0, fmt.Errorf("failed to list teams from global index: %w", err)
+	}
+
+	healed := 0
+	for _, teamID := range teams {
+		n, err := s.healTeamExpirationIndex(ctx, teamID)
+		if err != nil {
+			// Isolate per-team failures; other teams still get healed.
+			logger.L().Warn(ctx, "Failed to heal team expiration index", zap.Error(err), zap.String("team_id", teamID))
+
+			continue
+		}
+
+		healed += n
+	}
+
+	return healed, nil
+}
+
+func (s *Storage) healTeamExpirationIndex(ctx context.Context, teamID string) (int, error) {
+	healed := 0
+	var cursor uint64
+
+	for {
+		sandboxIDs, next, err := s.redisClient.SScan(ctx, GetSandboxStorageTeamIndexKey(teamID), cursor, "", healScanBatchSize).Result()
+		if err != nil {
+			return healed, fmt.Errorf("failed to scan team index: %w", err)
+		}
+
+		// SSCAN COUNT is a hint, not a cap: split oversized pages so
+		// downstream commands stay bounded.
+		for start := 0; start < len(sandboxIDs); start += healScanBatchSize {
+			end := min(start+healScanBatchSize, len(sandboxIDs))
+
+			n, err := s.healSandboxBatch(ctx, teamID, sandboxIDs[start:end])
+			if err != nil {
+				return healed, err
+			}
+
+			healed += n
+		}
+
+		cursor = next
+		if cursor == 0 {
+			return healed, nil
+		}
+	}
+}
+
+// healSandboxBatch re-adds missing expiration index members for one bounded
+// batch of sandbox IDs. Returns the number of healed members.
+func (s *Storage) healSandboxBatch(ctx context.Context, teamID string, sandboxIDs []string) (int, error) {
+	if len(sandboxIDs) == 0 {
+		return 0, nil
+	}
+
+	// Per-team MGET: all keys share the team hash tag (cluster slot safe).
+	keys := utils.Map(sandboxIDs, func(id string) string { return getSandboxKey(teamID, id) })
+	vals, err := s.redisClient.MGet(ctx, keys...).Result()
+	if err != nil {
+		return 0, fmt.Errorf("MGET failed: %w", err)
+	}
+
+	now := time.Now()
+	type candidate struct {
+		member       string
+		legacyMember string
+		score        float64
+	}
+	var candidates []candidate
+	for _, raw := range vals {
+		str, ok := raw.(string)
+		if !ok {
+			continue // stale team index entry; TeamItems tolerates these too
+		}
+
+		var sbx sandboxtypes.Sandbox
+		if err := json.Unmarshal([]byte(str), &sbx); err != nil {
+			continue
+		}
+
+		if now.Sub(sbx.StartTime) < healGracePeriod {
+			continue
+		}
+
+		candidates = append(candidates, candidate{
+			member:       sandboxExpirationMember(sbx),
+			legacyMember: legacyExpirationMember(teamID, sbx.SandboxID),
+			score:        float64(sbx.EndTime.UnixMilli()),
+		})
+	}
+	if len(candidates) == 0 {
+		return 0, nil
+	}
+
+	// A sandbox is indexed if either its execution-scoped member or the
+	// legacy member (written by old pods during rolling deploys) exists.
+	// ZMSCORE returns 0 for absent members; legitimate scores are unix
+	// milliseconds and can never be 0.
+	members := make([]string, 0, len(candidates)*2)
+	for _, c := range candidates {
+		members = append(members, c.member, c.legacyMember)
+	}
+
+	scores, err := s.redisClient.ZMScore(ctx, globalExpirationSet, members...).Result()
+	if err != nil {
+		return 0, fmt.Errorf("ZMSCORE failed: %w", err)
+	}
+
+	var missing []redis.Z
+	for i, c := range candidates {
+		if scores[2*i] != 0 || scores[2*i+1] != 0 {
+			continue
+		}
+
+		missing = append(missing, redis.Z{Score: c.score, Member: c.member})
+	}
+	if len(missing) == 0 {
+		return 0, nil
+	}
+
+	// NX: only fill holes. A concurrent Add/Update owns the member's score.
+	// A TOCTOU with a concurrent Remove can only plant an orphan member,
+	// which ExpiredItems sweeps once its score passes — garbage, never a
+	// false eviction (eviction re-checks the stored JSON and re-validates
+	// expiry under the lock in StartRemoving).
+	if err := s.redisClient.ZAddNX(ctx, globalExpirationSet, missing...).Err(); err != nil {
+		return 0, fmt.Errorf("ZADD NX failed: %w", err)
+	}
+
+	s.metrics.indexHealed.Add(ctx, int64(len(missing)))
+
+	return len(missing), nil
+}

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

@@ -36,13 +36,18 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 	}
 
 	// Group by team for per-team MGET (Redis Cluster slot compatibility).
+	type memberRef struct {
+		member      string
+		sandboxID   string
+		executionID string // "" for legacy (execution-less) members
+	}
 	type teamEntry struct {
-		sandboxIDs []string
-		members    []string // original ZSET members, aligned 1:1 with sandboxIDs
+		teamID string
+		refs   []memberRef
 	}
 	teamSandboxes := make(map[string]*teamEntry)
 	for _, member := range expiredMembers {
-		teamID, sandboxID, ok := parseExpirationMember(member)
+		teamID, sandboxID, executionID, ok := parseExpirationMember(member)
 		if !ok {
 			logger.L().Warn(ctx, "Invalid expiration index member", zap.String("member", member))
 
@@ -51,28 +56,28 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 
 		entry, ok := teamSandboxes[teamID]
 		if !ok {
-			entry = &teamEntry{}
+			entry = &teamEntry{teamID: teamID}
 			teamSandboxes[teamID] = entry
 		}
 
-		entry.sandboxIDs = append(entry.sandboxIDs, sandboxID)
-		entry.members = append(entry.members, member)
+		entry.refs = append(entry.refs, memberRef{member: member, sandboxID: sandboxID, executionID: executionID})
 	}
 
 	pipe := s.redisClient.Pipeline()
 	type batchInfo struct {
-		cmd     *redis.SliceCmd
-		members []string // aligned 1:1 with MGET keys
+		cmd    *redis.SliceCmd
+		teamID string
+		refs   []memberRef // aligned 1:1 with MGET keys
 	}
 	var batches []batchInfo
 	for teamID, entry := range teamSandboxes {
-		keys := make([]string, len(entry.sandboxIDs))
-		for i, id := range entry.sandboxIDs {
-			keys[i] = getSandboxKey(teamID, id)
+		keys := make([]string, len(entry.refs))
+		for i, ref := range entry.refs {
+			keys[i] = getSandboxKey(teamID, ref.sandboxID)
 		}
 
 		cmd := pipe.MGet(ctx, keys...)
-		batches = append(batches, batchInfo{cmd: cmd, members: entry.members})
+		batches = append(batches, batchInfo{cmd: cmd, teamID: teamID, refs: entry.refs})
 	}
 
 	_, err = pipe.Exec(ctx)
@@ -83,12 +88,21 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 	// Deserialize and filter; collect stale ZSET members for cleanup.
 	var result []sandboxtypes.Sandbox
 	var staleMembers []any
+	var upgrades []redis.Z   // legacy members re-added in execution-scoped format
+	var upgradedLegacy []any // legacy members to retire once upgrades land
+	var rescores []redis.Z   // live members whose score drifted from EndTime
+	var orphanCount, deadExecutionCount, upgradedCount int64
 
 	for _, batch := range batches {
 		for i, raw := range batch.cmd.Val() {
-			// Sandbox key gone but ZSET entry remains — orphaned.
+			ref := batch.refs[i]
+
+			// Sandbox key gone but ZSET entry remains — orphaned. Members are
+			// execution-scoped, so this removal can never unindex a fresh
+			// execution concurrently re-added under the same sandbox ID.
 			if raw == nil {
-				staleMembers = append(staleMembers, batch.members[i])
+				staleMembers = append(staleMembers, ref.member)
+				orphanCount++
 
 				continue
 			}
@@ -105,10 +119,39 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 				continue
 			}
 
+			// Member names a dead execution; the live execution has its own member.
+			if ref.executionID != "" && ref.executionID != sbx.ExecutionID {
+				staleMembers = append(staleMembers, ref.member)
+				deadExecutionCount++
+
+				continue
+			}
+
+			// Legacy (execution-less) member: upgrade in place so future
+			// removals are execution-exact. ZADD NX first (never a window
+			// with no member for a live key), then retire the legacy member.
+			if ref.executionID == "" && sbx.ExecutionID != "" {
+				upgrades = append(upgrades, redis.Z{
+					Score:  float64(sbx.EndTime.UnixMilli()),
+					Member: expirationMember(batch.teamID, sbx.SandboxID, sbx.ExecutionID),
+				})
+				upgradedLegacy = append(upgradedLegacy, ref.member)
+			}
+
 			// In case that index have failed to be updated
 			if !sbx.IsExpired(now) {
 				logger.L().Debug(ctx, "ExpiredItems: Sandbox marked as expried in index, but state say otherwise", logger.WithSandboxID(sbx.SandboxID), logger.Time("end_time", sbx.EndTime))
 
+				// Re-score the drifted member to the stored EndTime so it
+				// stops occupying the expired scan window on every tick.
+				// XX: never resurrect a member a concurrent Remove deleted.
+				if ref.executionID != "" {
+					rescores = append(rescores, redis.Z{
+						Score:  float64(sbx.EndTime.UnixMilli()),
+						Member: ref.member,
+					})
+				}
+
 				continue
 			}
 
@@ -128,10 +171,40 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 		}
 	}
 
+	// Upgrades before legacy removal: a live key must have a member in the
+	// index at every step. NX keeps a fresher score if one already exists.
+	// Legacy members are retired only after the upgrade landed.
+	if len(upgrades) > 0 {
+		if err := s.redisClient.ZAddNX(ctx, globalExpirationSet, upgrades...).Err(); err != nil {
+			logger.L().Warn(ctx, "Failed to upgrade legacy expiration index entries", zap.Error(err), zap.Int("count", len(upgrades)))
+		} else {
+			staleMembers = append(staleMembers, upgradedLegacy...)
+			upgradedCount = int64(len(upgradedLegacy))
+		}
+	}
+
 	// Remove orphaned ZSET entries so the set doesn't grow unboundedly.
 	if len(staleMembers) > 0 {
 		if err := s.redisClient.ZRem(ctx, globalExpirationSet, staleMembers...).Err(); err != nil {
 			logger.L().Warn(ctx, "Failed to clean up stale expiration index entries", zap.Error(err), zap.Int("count", len(staleMembers)))
+		} else {
+			if orphanCount > 0 {
+				s.metrics.indexSwept.Add(ctx, orphanCount, s.metrics.sweptOrphan)
+			}
+			if deadExecutionCount > 0 {
+				s.metrics.indexSwept.Add(ctx, deadExecutionCount, s.metrics.sweptDeadExecution)
+			}
+			if upgradedCount > 0 {
+				s.metrics.indexSwept.Add(ctx, upgradedCount, s.metrics.sweptLegacyUpgraded)
+			}
+		}
+	}
+
+	if len(rescores) > 0 {
+		if err := s.redisClient.ZAddXX(ctx, globalExpirationSet, rescores...).Err(); err != nil {
+			logger.L().Warn(ctx, "Failed to re-score drifted expiration index entries", zap.Error(err), zap.Int("count", len(rescores)))
+		} else {
+			s.metrics.indexRescored.Add(ctx, int64(len(rescores)))
 		}
 	}