diff --git a/packages/api/internal/sandbox/storage/redis/expiration_index_test.go b/packages/api/internal/sandbox/storage/redis/expiration_index_test.go new file mode 100644 index 0000000000..ccfc718115 --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/expiration_index_test.go @@ -0,0 +1,415 @@ +package redis + +import ( + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" +) + +// makeIndexedSandbox builds a running sandbox fixture with explicit lifecycle +// timestamps for expiration index tests. +func makeIndexedSandbox(teamID uuid.UUID, sandboxID, executionID string, startTime, endTime time.Time) sandboxtypes.Sandbox { + return sandboxtypes.Sandbox{ + SandboxID: sandboxID, + TemplateID: "test-template", + ClientID: "test-client", + TeamID: teamID, + ExecutionID: executionID, + StartTime: startTime, + EndTime: endTime, + MaxInstanceLength: 24 * time.Hour, + State: sandboxtypes.StateRunning, + } +} + +func requireMemberScore(t *testing.T, client redis.UniversalClient, member string, wantScore float64) { + t.Helper() + + score, err := client.ZScore(t.Context(), globalExpirationSet, member).Result() + require.NoError(t, err) + require.InDelta(t, wantScore, score, 0.5) +} + +func requireMemberAbsent(t *testing.T, client redis.UniversalClient, member string) { + t.Helper() + + err := client.ZScore(t.Context(), globalExpirationSet, member).Err() + require.ErrorIs(t, err, redis.Nil) +} + +func TestParseExpirationMember(t *testing.T) { + t.Parallel() + + teamID := uuid.NewString() + execID := uuid.NewString() + + tests := []struct { + name string + member string + wantTeam string + wantSbx string + wantExec string + wantParsed bool + }{ + {"execution scoped", teamID + ":sbx-1:" + execID, teamID, "sbx-1", execID, true}, + {"legacy", teamID + ":sbx-1", teamID, "sbx-1", "", true}, + {"legacy with uuid sandbox id", teamID + ":" + execID, teamID, execID, "", true}, + {"no separator", "garbage", "", "", "", false}, + {"empty rest", teamID + ":", "", "", "", false}, + {"non-uuid execution tag", teamID + ":sbx-1:not-a-uuid", "", "", "", false}, + {"too many parts", teamID + ":sbx-1:x:" + execID, "", "", "", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + gotTeam, gotSbx, gotExec, ok := parseExpirationMember(tt.member) + require.Equal(t, tt.wantParsed, ok) + require.Equal(t, tt.wantTeam, gotTeam) + require.Equal(t, tt.wantSbx, gotSbx) + require.Equal(t, tt.wantExec, gotExec) + }) + } +} + +func TestAddRemove_ExecutionScopedMember(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-roundtrip", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + requireMemberScore(t, client, member, float64(sbx.EndTime.UnixMilli())) + + require.NoError(t, storage.Remove(t.Context(), teamID, sbx.SandboxID)) + requireMemberAbsent(t, client, member) + + err := client.Get(t.Context(), getSandboxKey(teamID.String(), sbx.SandboxID)).Err() + require.ErrorIs(t, err, redis.Nil) +} + +// TestRemove_DoesNotUnindexFreshExecution replays Race B: a lockless Add for +// a fresh execution of the same sandbox ID lands its index write while Remove +// of the old execution runs. The fresh execution's member must survive. +func TestRemove_DoesNotUnindexFreshExecution(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + const sandboxID = "sbx-race-b" + + old := makeIndexedSandbox(teamID, sandboxID, uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(-time.Minute)) + require.NoError(t, storage.Add(t.Context(), old)) + + // Simulate the concurrent Add's first phase (index write) for a fresh + // execution landing before Remove's cleanup. + fresh := makeIndexedSandbox(teamID, sandboxID, uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + freshMember := expirationMember(teamID.String(), sandboxID, fresh.ExecutionID) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(fresh.EndTime.UnixMilli()), + Member: freshMember, + }).Err()) + + require.NoError(t, storage.Remove(t.Context(), teamID, sandboxID)) + + // Old execution's member removed, fresh execution's member intact. + requireMemberAbsent(t, client, expirationMember(teamID.String(), sandboxID, old.ExecutionID)) + requireMemberScore(t, client, freshMember, float64(fresh.EndTime.UnixMilli())) +} + +// TestExpiredItems_StaleSweepIsExecutionScoped replays Race A: the evictor's +// stale sweep processes an expired member of a dead execution while a fresh +// execution of the same sandbox ID is live. Only the dead execution's member +// may be removed. +func TestExpiredItems_StaleSweepIsExecutionScoped(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + const sandboxID = "sbx-race-a" + + // Fresh execution: live key, future member. + fresh := makeIndexedSandbox(teamID, sandboxID, uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), fresh)) + + // Dead execution: expired member left behind (e.g. failed Remove cleanup). + deadExecutionID := uuid.NewString() + deadMember := expirationMember(teamID.String(), sandboxID, deadExecutionID) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(time.Now().Add(-time.Minute).UnixMilli()), + Member: deadMember, + }).Err()) + + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Empty(t, items) + + requireMemberAbsent(t, client, deadMember) + requireMemberScore(t, client, expirationMember(teamID.String(), sandboxID, fresh.ExecutionID), float64(fresh.EndTime.UnixMilli())) +} + +func TestExpiredItems_RemovesOrphanMember(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + member := expirationMember(uuid.NewString(), "sbx-orphan", uuid.NewString()) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(time.Now().Add(-time.Minute).UnixMilli()), + Member: member, + }).Err()) + + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Empty(t, items) + + requireMemberAbsent(t, client, member) +} + +func TestExpiredItems_UpgradesLegacyMember(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-legacy", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Rewrite the index entry into the legacy format with a drifted (past) + // score, simulating pre-deploy data. + execMember := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + legacyMember := legacyExpirationMember(teamID.String(), sbx.SandboxID) + require.NoError(t, client.ZRem(t.Context(), globalExpirationSet, execMember).Err()) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(time.Now().Add(-time.Minute).UnixMilli()), + Member: legacyMember, + }).Err()) + + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Empty(t, items) // sandbox is not expired + + requireMemberAbsent(t, client, legacyMember) + requireMemberScore(t, client, execMember, float64(sbx.EndTime.UnixMilli())) +} + +func TestExpiredItems_RescoresDriftedMember(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-drift", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Drift the member's score into the expired window while the stored + // EndTime stays in the future. + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(time.Now().Add(-time.Minute).UnixMilli()), + Member: member, + }).Err()) + + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Empty(t, items) + + requireMemberScore(t, client, member, float64(sbx.EndTime.UnixMilli())) +} + +func TestExpiredItems_ReturnsExpiredRunningSandbox(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-expired", uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(-time.Minute)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, sbx.SandboxID, items[0].SandboxID) + require.Equal(t, sbx.ExecutionID, items[0].ExecutionID) +} + +// TestHeal_RestoresMissingMember reproduces the production incident: a +// sandbox present in team storage but missing from the global expiration +// index is invisible to the evictor and would live forever. +func TestHeal_RestoresMissingMember(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-immortal", uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(-time.Minute)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + require.NoError(t, client.ZRem(t.Context(), globalExpirationSet, member).Err()) + + // Immortal: the evictor cannot see it. + items, err := storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Empty(t, items) + + healed, err := storage.healExpirationIndex(t.Context()) + require.NoError(t, err) + require.Equal(t, 1, healed) + requireMemberScore(t, client, member, float64(sbx.EndTime.UnixMilli())) + + items, err = storage.ExpiredItems(t.Context()) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, sbx.SandboxID, items[0].SandboxID) +} + +func TestHeal_SkipsLegacyIndexedSandbox(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-legacy-indexed", uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Sandbox indexed only via the legacy member (written by an old pod). + require.NoError(t, client.ZRem(t.Context(), globalExpirationSet, expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID)).Err()) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(sbx.EndTime.UnixMilli()), + Member: legacyExpirationMember(teamID.String(), sbx.SandboxID), + }).Err()) + + healed, err := storage.healExpirationIndex(t.Context()) + require.NoError(t, err) + require.Zero(t, healed) +} + +func TestHeal_SkipsYoungSandbox(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-young", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + require.NoError(t, client.ZRem(t.Context(), globalExpirationSet, member).Err()) + + healed, err := storage.healExpirationIndex(t.Context()) + require.NoError(t, err) + require.Zero(t, healed) + requireMemberAbsent(t, client, member) +} + +// TestHeal_ManySandboxesMultipleBatches exercises the SSCAN pagination in +// healTeamExpirationIndex: more sandboxes than healScanBatchSize, with holes +// spread across the whole ID range. +func TestHeal_ManySandboxesMultipleBatches(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + total := healScanBatchSize*2 + 37 // force >2 batches with a partial tail + + pipe := client.Pipeline() + var missingMembers []string + for i := range total { + sbx := makeIndexedSandbox(teamID, fmt.Sprintf("sbx-batch-%04d", i), uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour)) + data, err := json.Marshal(sbx) + require.NoError(t, err) + + pipe.Set(t.Context(), getSandboxKey(teamID.String(), sbx.SandboxID), data, 0) + pipe.SAdd(t.Context(), GetSandboxStorageTeamIndexKey(teamID.String()), sbx.SandboxID) + + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + if i%2 == 0 { + pipe.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(sbx.EndTime.UnixMilli()), + Member: member, + }) + } else { + missingMembers = append(missingMembers, member) + } + } + _, err := pipe.Exec(t.Context()) + require.NoError(t, err) + + healed, err := storage.healTeamExpirationIndex(t.Context(), teamID.String()) + require.NoError(t, err) + require.Equal(t, len(missingMembers), healed) + + scores, err := client.ZMScore(t.Context(), globalExpirationSet, missingMembers...).Result() + require.NoError(t, err) + for i, score := range scores { + require.NotZero(t, score, "member %s not healed", missingMembers[i]) + } + + // Idempotent: a second pass heals nothing. + healed, err = storage.healTeamExpirationIndex(t.Context(), teamID.String()) + require.NoError(t, err) + require.Zero(t, healed) +} + +func TestHeal_DoesNotClobberExistingScore(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-fresh-score", uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + // A concurrent SetTimeout moved the score past what the stored JSON says. + member := expirationMember(teamID.String(), sbx.SandboxID, sbx.ExecutionID) + freshScore := float64(time.Now().Add(2 * time.Hour).UnixMilli()) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{Score: freshScore, Member: member}).Err()) + + healed, err := storage.healExpirationIndex(t.Context()) + require.NoError(t, err) + require.Zero(t, healed) + requireMemberScore(t, client, member, freshScore) +} + +func TestUpdate_RescoresExecutionMemberAndRetiresLegacy(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + sbx := makeIndexedSandbox(teamID, "sbx-update", uuid.NewString(), time.Now(), time.Now().Add(time.Hour)) + require.NoError(t, storage.Add(t.Context(), sbx)) + + // Seed a leftover legacy member (pre-deploy data). + legacyMember := legacyExpirationMember(teamID.String(), sbx.SandboxID) + require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{ + Score: float64(sbx.EndTime.UnixMilli()), + Member: legacyMember, + }).Err()) + + newEndTime := time.Now().Add(2 * time.Hour) + updated, err := storage.Update(t.Context(), teamID, sbx.SandboxID, func(cur sandboxtypes.Sandbox) (sandboxtypes.Sandbox, error) { + cur.EndTime = newEndTime + + return cur, nil + }) + require.NoError(t, err) + + requireMemberScore(t, client, expirationMember(teamID.String(), sbx.SandboxID, updated.ExecutionID), float64(newEndTime.UnixMilli())) + requireMemberAbsent(t, client, legacyMember) +} diff --git a/packages/api/internal/sandbox/storage/redis/heal.go b/packages/api/internal/sandbox/storage/redis/heal.go new file mode 100644 index 0000000000..04fb7f0c38 --- /dev/null +++ b/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 +} diff --git a/packages/api/internal/sandbox/storage/redis/items.go b/packages/api/internal/sandbox/storage/redis/items.go index 50eccb8ecb..527fceb612 100644 --- a/packages/api/internal/sandbox/storage/redis/items.go +++ b/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))) } } diff --git a/packages/api/internal/sandbox/storage/redis/main.go b/packages/api/internal/sandbox/storage/redis/main.go index 57d98a46eb..1b83ed74f9 100644 --- a/packages/api/internal/sandbox/storage/redis/main.go +++ b/packages/api/internal/sandbox/storage/redis/main.go @@ -6,11 +6,13 @@ import ( "time" "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" "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/telemetry" ) const ( @@ -35,6 +37,48 @@ type Storage struct { locker *storageLocker subManager *subscriptionManager publisher *publisher + + metrics expirationIndexMetrics +} + +// expirationIndexMetrics observes global expiration index consistency. +// indexHealed is the primary alert signal: healthy steady state is zero. +type expirationIndexMetrics struct { + indexHealed metric.Int64Counter + indexRescored metric.Int64Counter + + indexSwept metric.Int64Counter + sweptOrphan metric.MeasurementOption + sweptDeadExecution metric.MeasurementOption + sweptLegacyUpgraded metric.MeasurementOption +} + +const sweptReasonAttr = "reason" + +func newExpirationIndexMetrics(meter metric.Meter) (expirationIndexMetrics, error) { + healed, err := telemetry.GetCounter(meter, telemetry.ApiRedisStorageExpirationIndexHealed) + if err != nil { + return expirationIndexMetrics{}, fmt.Errorf("expiration index healed counter: %w", err) + } + + rescored, err := telemetry.GetCounter(meter, telemetry.ApiRedisStorageExpirationIndexRescored) + if err != nil { + return expirationIndexMetrics{}, fmt.Errorf("expiration index rescored counter: %w", err) + } + + swept, err := telemetry.GetCounter(meter, telemetry.ApiRedisStorageExpirationIndexSwept) + if err != nil { + return expirationIndexMetrics{}, fmt.Errorf("expiration index swept counter: %w", err) + } + + return expirationIndexMetrics{ + indexHealed: healed, + indexRescored: rescored, + indexSwept: swept, + sweptOrphan: metric.WithAttributeSet(attribute.NewSet(attribute.String(sweptReasonAttr, "orphan"))), + sweptDeadExecution: metric.WithAttributeSet(attribute.NewSet(attribute.String(sweptReasonAttr, "dead_execution"))), + sweptLegacyUpgraded: metric.WithAttributeSet(attribute.NewSet(attribute.String(sweptReasonAttr, "legacy_upgraded"))), + }, nil } const meterScope = "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis" @@ -51,23 +95,35 @@ func NewStorage( return nil, fmt.Errorf("failed to create publisher: %w", err) } + metrics, err := newExpirationIndexMetrics(meter) + if err != nil { + return nil, fmt.Errorf("failed to create expiration index metrics: %w", err) + } + return &Storage{ redisClient: redisClient, locker: newStorageLocker(redisClient, subManager, pub), subManager: subManager, publisher: pub, + metrics: metrics, }, nil } // Start subscribes to the global PubSub channel and launches the publish -// worker. Blocks until the context is cancelled or Close is called. +// worker and the expiration index healer. Blocks until the context is +// cancelled or Close is called. func (s *Storage) Start(ctx context.Context) { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + pubDone := make(chan struct{}) go func() { defer close(pubDone) s.publisher.run(ctx) }() + go s.startHealer(ctx) + s.subManager.start(ctx) <-pubDone } diff --git a/packages/api/internal/sandbox/storage/redis/operations.go b/packages/api/internal/sandbox/storage/redis/operations.go index 7aee2a745d..aaff4a2341 100644 --- a/packages/api/internal/sandbox/storage/redis/operations.go +++ b/packages/api/internal/sandbox/storage/redis/operations.go @@ -30,9 +30,11 @@ func (s *Storage) Add(ctx context.Context, sbx sandboxtypes.Sandbox) error { // Add to the index before adding to the cache, so there's no possibility of leaking // Index by EndTime so ExpiredItems can use ZRANGEBYSCORE instead of scanning all sandboxes. + // The member is scoped to this execution: concurrent removals of older + // executions of the same sandbox ID can never unindex this one. if err := s.redisClient.ZAdd(ctx, globalExpirationSet, redis.Z{ Score: float64(sbx.EndTime.UnixMilli()), - Member: expirationMember(sbx.TeamID.String(), sbx.SandboxID), + Member: sandboxExpirationMember(sbx), }).Err(); err != nil { return fmt.Errorf("failed to add sandbox to global expiration index: %w", err) } @@ -92,15 +94,30 @@ func (s *Storage) Remove(ctx context.Context, teamID uuid.UUID, sandboxID string } }() - // Execute Lua script for atomic DEL + SREM - err = removeSandboxScript.Run(ctx, s.redisClient, []string{key, teamKey}, sandboxID).Err() - if err != nil { + // Execute Lua script for atomic DEL + SREM; it returns the deleted JSON + // so the expiration-index cleanup below is scoped to the execution we + // actually removed. + raw, err := removeSandboxScript.Run(ctx, s.redisClient, []string{key, teamKey}, sandboxID).Text() + if err != nil && !errors.Is(err, redis.Nil) { return fmt.Errorf("failed to remove sandbox from Redis: %w", err) } // Clean up from the global expiration index. // Do it after the removal to prevent leaking expired sandboxes. - if err := s.redisClient.ZRem(ctx, globalExpirationSet, expirationMember(teamID.String(), sandboxID)).Err(); err != nil { + // Always drop the legacy (execution-less) member; drop the execution + // member only for the execution we deleted. A concurrent lockless Add for + // a newer execution wrote a different member, so it can never be + // unindexed here. If the key was already gone, any leftover execution + // member is swept by ExpiredItems once its score passes. + members := []any{legacyExpirationMember(teamID.String(), sandboxID)} + if raw != "" { + var sbx sandboxtypes.Sandbox + if unmarshalErr := json.Unmarshal([]byte(raw), &sbx); unmarshalErr == nil && sbx.ExecutionID != "" { + members = append(members, expirationMember(teamID.String(), sandboxID, sbx.ExecutionID)) + } + } + + if err := s.redisClient.ZRem(ctx, globalExpirationSet, members...).Err(); err != nil { logger.L().Warn(ctx, "Failed to remove sandbox from global expiration index", zap.Error(err), logger.WithSandboxID(sandboxID)) } @@ -217,10 +234,19 @@ func (s *Storage) Update(ctx context.Context, teamID uuid.UUID, sandboxID string if !updatedSbx.EndTime.Equal(sbx.EndTime) { if err := s.redisClient.ZAdd(ctx, globalExpirationSet, redis.Z{ Score: float64(updatedSbx.EndTime.UnixMilli()), - Member: expirationMember(teamID.String(), sandboxID), + Member: sandboxExpirationMember(updatedSbx), }).Err(); err != nil { return sandboxtypes.Sandbox{}, fmt.Errorf("failed to update sandbox in global expiration index: %w", err) } + + // Migrate-on-write: retire the legacy (execution-less) member so a + // stale-scored duplicate can't shadow the fresh score. Warn-only — + // the lazy upgrade in ExpiredItems covers failures. + if updatedSbx.ExecutionID != "" { + if err := s.redisClient.ZRem(ctx, globalExpirationSet, legacyExpirationMember(teamID.String(), sandboxID)).Err(); err != nil { + logger.L().Warn(ctx, "Failed to remove legacy expiration index member", zap.Error(err), logger.WithSandboxID(sandboxID)) + } + } } return updatedSbx, nil diff --git a/packages/api/internal/sandbox/storage/redis/scripts.go b/packages/api/internal/sandbox/storage/redis/scripts.go index dca28a7fdd..147afd48a7 100644 --- a/packages/api/internal/sandbox/storage/redis/scripts.go +++ b/packages/api/internal/sandbox/storage/redis/scripts.go @@ -15,12 +15,16 @@ var ( `) // removeSandboxScript atomically removes a sandbox and its team index entry. + // It returns the stored JSON (or nil if the key was already gone) so the + // caller knows exactly which execution it removed and can scope the + // expiration-index cleanup to that execution. // KEYS[1] = sandbox key, KEYS[2] = team index key // ARGV[1] = sandbox ID removeSandboxScript = redis.NewScript(` + local data = redis.call('GET', KEYS[1]) redis.call('DEL', KEYS[1]) redis.call('SREM', KEYS[2], ARGV[1]) - return 1 + return data `) // startTransitionScript atomically updates sandbox and sets transition key with UUID. diff --git a/packages/api/internal/sandbox/storage/redis/utils.go b/packages/api/internal/sandbox/storage/redis/utils.go index a1a2dfe997..7570da6b9d 100644 --- a/packages/api/internal/sandbox/storage/redis/utils.go +++ b/packages/api/internal/sandbox/storage/redis/utils.go @@ -3,6 +3,9 @@ package redis import ( "strings" + "github.com/google/uuid" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" ) @@ -25,14 +28,67 @@ var ( globalExpirationSet = redis_utils.CreateKey(sandboxKeyPrefix, "global:expiration") ) -func expirationMember(teamID, sandboxID string) string { +// expirationMember identifies one *execution* (incarnation) of a sandbox in +// the global expiration ZSET. Scoping the member to the execution makes every +// ZREM structurally safe: removing a dead execution's member can never +// unindex a live one, even when a lockless Add for the same sandbox ID races +// a Remove or the evictor's stale sweep. +func expirationMember(teamID, sandboxID, executionID string) string { + return redis_utils.CreateKey(teamID, sandboxID, executionID) +} + +// legacyExpirationMember is the pre-executionID member format still present +// in live data (and written by old pods during a rolling deploy). It drains +// via Remove's dual ZREM and the lazy upgrade in ExpiredItems. +// +// TODO [EN-1602]: remove the legacy member format once the migration is complete +// Then delete: +// - legacyExpirationMember and the ExecutionID == "" fallback in +// sandboxExpirationMember (utils.go) +// - the executionID == "" legacy branch in parseExpirationMember (utils.go) +// - the legacy member in Remove's dual ZREM (operations.go) +// - the migrate-on-write ZRem in Update (operations.go) +// - the legacy upgrade path in ExpiredItems: upgrades/upgradedLegacy and +// the sweptLegacyUpgraded metric attribute (items.go, main.go) +// - the legacyMember half of the ZMSCORE existence check in +// healTeamExpirationIndex (heal.go) +// - the legacy cases in expiration_index_test.go +func legacyExpirationMember(teamID, sandboxID string) string { return redis_utils.CreateKey(teamID, sandboxID) } -func parseExpirationMember(member string) (teamID, sandboxID string, ok bool) { - teamID, sandboxID, ok = strings.Cut(member, ":") +// sandboxExpirationMember returns the expiration index member for a stored +// sandbox, falling back to the legacy format when ExecutionID is absent. +func sandboxExpirationMember(sbx sandboxtypes.Sandbox) string { + if sbx.ExecutionID == "" { + return legacyExpirationMember(sbx.TeamID.String(), sbx.SandboxID) + } + + return expirationMember(sbx.TeamID.String(), sbx.SandboxID, sbx.ExecutionID) +} + +// parseExpirationMember parses both member formats: +// "teamID:sandboxID:executionID" (current) and "teamID:sandboxID" (legacy). +// executionID == "" marks a legacy member. Sandbox IDs never contain ':' +// and execution IDs are always UUIDs. +func parseExpirationMember(member string) (teamID, sandboxID, executionID string, ok bool) { + parts := strings.Split(member, ":") + switch len(parts) { + case 2: // legacy + if parts[1] == "" { + return "", "", "", false + } + + return parts[0], parts[1], "", true + case 3: + if _, err := uuid.Parse(parts[2]); err != nil { + return "", "", "", false + } - return + return parts[0], parts[1], parts[2], true + default: + return "", "", "", false + } } // GetTeamPrefix returns the storage team prefix for external packages (e.g. reservations). diff --git a/packages/shared/pkg/telemetry/meters.go b/packages/shared/pkg/telemetry/meters.go index c80db32e66..418613cdfe 100644 --- a/packages/shared/pkg/telemetry/meters.go +++ b/packages/shared/pkg/telemetry/meters.go @@ -52,6 +52,20 @@ const ( ApiRedisStoragePublisherPublished CounterType = "api.redis_storage.publisher.published" ApiRedisStoragePublisherDropped CounterType = "api.redis_storage.publisher.dropped" + + // ApiRedisStorageExpirationIndexHealed counts sandboxes the healer re-added + // to the global expiration index. Healthy steady state is zero; a sustained + // non-zero rate means expiration index writes are being lost and sandboxes + // would otherwise become invisible to the evictor (immortal). + ApiRedisStorageExpirationIndexHealed CounterType = "api.redis_storage.expiration_index.healed" + // ApiRedisStorageExpirationIndexSwept counts members removed from the + // global expiration index by the evictor scan + // (reason=orphan|dead_execution|legacy_upgraded). + ApiRedisStorageExpirationIndexSwept CounterType = "api.redis_storage.expiration_index.swept" + // ApiRedisStorageExpirationIndexRescored counts live members whose index + // score drifted from the stored EndTime and were re-scored by the evictor + // scan. Sustained non-zero rate means score updates are being lost. + ApiRedisStorageExpirationIndexRescored CounterType = "api.redis_storage.expiration_index.rescored" ) const ( @@ -237,6 +251,10 @@ var counterDesc = map[CounterType]string{ ApiRedisStoragePublisherPublished: "Total Redis PUBLISH calls completed by the storage publisher (result=success|failure)", ApiRedisStoragePublisherDropped: "Total storage notifications dropped before reaching Redis (reason=queue_full|closed)", + + ApiRedisStorageExpirationIndexHealed: "Sandboxes re-added to the global expiration index by the healer; sustained non-zero rate means index writes are being lost", + ApiRedisStorageExpirationIndexSwept: "Members removed from the global expiration index by the evictor scan (reason=orphan|dead_execution|legacy_upgraded)", + ApiRedisStorageExpirationIndexRescored: "Live expiration index members re-scored after drifting from the stored EndTime", } var counterUnits = map[CounterType]string{ @@ -268,6 +286,10 @@ var counterUnits = map[CounterType]string{ ApiRedisStoragePublisherPublished: "{notification}", ApiRedisStoragePublisherDropped: "{notification}", + + ApiRedisStorageExpirationIndexHealed: "{sandbox}", + ApiRedisStorageExpirationIndexSwept: "{member}", + ApiRedisStorageExpirationIndexRescored: "{member}", } var observableCounterDesc = map[ObservableCounterType]string{