diff --git a/packages/api/internal/orchestrator/cache.go b/packages/api/internal/orchestrator/cache.go index 8ed3e5d475..b7aa4290be 100644 --- a/packages/api/internal/orchestrator/cache.go +++ b/packages/api/internal/orchestrator/cache.go @@ -11,6 +11,7 @@ import ( "go.opentelemetry.io/otel" "go.uber.org/zap" + "github.com/e2b-dev/infra/packages/api/internal/orchestrator/deadnode" "github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager" "github.com/e2b-dev/infra/packages/api/internal/sandbox" "github.com/e2b-dev/infra/packages/shared/pkg/logger" @@ -84,6 +85,9 @@ func (o *Orchestrator) syncNodes(ctx context.Context, store *sandbox.Store, skip // Wait for nodes discovery to finish wg.Wait() + // Discovery succeeded (a listNomadNodes failure returns early above) + o.lastDiscoverySync.Store(new(time.Now())) + // Sync state of all nodes currently in the pool ctx, syncNodesSpan := tracer.Start(ctx, "keep-in-sync-existing") defer syncNodesSpan.End() @@ -114,6 +118,15 @@ func (o *Orchestrator) syncNodes(ctx context.Context, store *sandbox.Store, skip } o.deregisterNode(n) + + return + } + + if _, unreachable := n.UnreachableSince(); !unreachable { + ref := deadnode.NodeRef{ClusterID: n.ClusterID.String(), NodeID: n.ID} + if err := deadnode.RecordNodeSeen(ctx, o.redisClient, ref, time.Now()); err != nil { + logger.L().Warn(ctx, "Failed to record node liveness", zap.Error(err), logger.WithNodeID(n.ID)) + } } }) } diff --git a/packages/api/internal/orchestrator/deadnode/liveness.go b/packages/api/internal/orchestrator/deadnode/liveness.go new file mode 100644 index 0000000000..eb4703915f --- /dev/null +++ b/packages/api/internal/orchestrator/deadnode/liveness.go @@ -0,0 +1,150 @@ +package deadnode + +import ( + "context" + "errors" + "fmt" + "strconv" + "time" + + "github.com/redis/go-redis/v9" + "go.uber.org/zap" + + "github.com/e2b-dev/infra/packages/shared/pkg/logger" +) + +// Node liveness is a cross-replica signal stored in Redis: every API replica +// refreshes a per-node "last seen" key after each successful sync cycle with +// that node. The dead-node sweep only considers a node dead when NO replica +// has seen it for the grace period — a single replica with a broken network +// view can never purge sandboxes of a node other replicas still reach. +// +// For nodes without a last-seen key (fresh feature rollout, or the node died +// before any replica ever synced it), a first-missing marker is written +// once (SET NX) so the grace period ages from the first observation and is +// shared across replicas and their restarts. +const ( + nodeLastSeenKeyPrefix = "node:last-seen:" + nodeFirstMissingKeyPrefix = "node:first-missing:" + + // nodeLivenessKeyTTL self-cleans keys of nodes that were removed from the + // fleet permanently. + nodeLivenessKeyTTL = 24 * time.Hour +) + +// NodeRef identifies a node across clusters. Comparable — used as a map key. +type NodeRef struct { + ClusterID string + NodeID string +} + +// nodeLiveness is the cross-replica view of one node. +type nodeLiveness struct { + // lastSeen is the last time any replica completed a successful sync with + // the node. Zero when no replica has ever recorded seeing it. + lastSeen time.Time + // firstMissing is when a replica first observed the node having no + // last-seen key. Only meaningful when lastSeen is zero. + firstMissing time.Time +} + +func nodeLastSeenKey(ref NodeRef) string { + return nodeLastSeenKeyPrefix + ref.ClusterID + ":" + ref.NodeID +} + +func nodeFirstMissingKey(ref NodeRef) string { + return nodeFirstMissingKeyPrefix + ref.ClusterID + ":" + ref.NodeID +} + +// RecordNodeSeen refreshes the node's last-seen key and clears any +// first-missing marker. Called by every replica after a successful sync +// cycle unconditionally +func RecordNodeSeen(ctx context.Context, client redis.UniversalClient, ref NodeRef, now time.Time) error { + pipe := client.Pipeline() + pipe.Set(ctx, nodeLastSeenKey(ref), strconv.FormatInt(now.Unix(), 10), nodeLivenessKeyTTL) + pipe.Del(ctx, nodeFirstMissingKey(ref)) + _, err := pipe.Exec(ctx) + if err != nil { + return fmt.Errorf("failed to record node seen: %w", err) + } + + return nil +} + +// fetchNodeLiveness returns the liveness for the given nodes. For nodes +// without a last-seen key it plants (SET NX) and reads back the shared +// first-missing marker so the missing-grace ages consistently across replicas +func fetchNodeLiveness(ctx context.Context, client redis.UniversalClient, refs []NodeRef, now time.Time) (map[NodeRef]nodeLiveness, error) { + if len(refs) == 0 { + return map[NodeRef]nodeLiveness{}, nil + } + + pipe := client.Pipeline() + lastSeenCmds := make([]*redis.StringCmd, len(refs)) + for i, ref := range refs { + lastSeenCmds[i] = pipe.Get(ctx, nodeLastSeenKey(ref)) + } + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, fmt.Errorf("failed to fetch node last-seen keys: %w", err) + } + + out := make(map[NodeRef]nodeLiveness, len(refs)) + var missing []NodeRef + for i, ref := range refs { + ts, err := parseUnixSeconds(lastSeenCmds[i]) + if err != nil { + missing = append(missing, ref) + + continue + } + + out[ref] = nodeLiveness{lastSeen: ts} + } + + if len(missing) == 0 { + return out, nil + } + + // Plant the shared first-missing marker (first writer wins) and read the + // authoritative value back in one pipeline. + markerPipe := client.Pipeline() + markerCmds := make([]*redis.StringCmd, len(missing)) + for i, ref := range missing { + markerPipe.SetNX(ctx, nodeFirstMissingKey(ref), strconv.FormatInt(now.Unix(), 10), nodeLivenessKeyTTL) + markerCmds[i] = markerPipe.Get(ctx, nodeFirstMissingKey(ref)) + } + if _, err := markerPipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, fmt.Errorf("failed to plant node first-missing markers: %w", err) + } + + for i, ref := range missing { + ts, err := parseUnixSeconds(markerCmds[i]) + if err != nil { + // Leave the node without data; the sweep skips nodes it has no evidence about + logger.L().Warn(ctx, "Failed to read node first-missing marker", + zap.Error(err), + logger.WithNodeID(ref.NodeID), + ) + + continue + } + + out[ref] = nodeLiveness{firstMissing: ts} + } + + return out, nil +} + +func parseUnixSeconds(cmd *redis.StringCmd) (time.Time, error) { + raw, err := cmd.Result() + if err != nil { + return time.Time{}, err + } + + sec, err := strconv.ParseInt(raw, 10, 64) + if err != nil { + return time.Time{}, fmt.Errorf("invalid unix timestamp %q: %w", raw, err) + } + + return time.Unix(sec, 0), nil +} diff --git a/packages/api/internal/orchestrator/deadnode/liveness_test.go b/packages/api/internal/orchestrator/deadnode/liveness_test.go new file mode 100644 index 0000000000..0fc085c2d7 --- /dev/null +++ b/packages/api/internal/orchestrator/deadnode/liveness_test.go @@ -0,0 +1,100 @@ +package deadnode + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" +) + +func TestNodeLiveness_RecordAndFetch(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + ref := NodeRef{ClusterID: "cluster-a", NodeID: "node-1"} + now := time.Now() + + require.NoError(t, RecordNodeSeen(t.Context(), client, ref, now)) + + liveness, err := fetchNodeLiveness(t.Context(), client, []NodeRef{ref}, now) + require.NoError(t, err) + + nl, ok := liveness[ref] + require.True(t, ok) + assert.WithinDuration(t, now, nl.lastSeen, time.Second) + assert.True(t, nl.firstMissing.IsZero()) +} + +func TestNodeLiveness_MissingMarkerIsFirstWriteWins(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + ref := NodeRef{ClusterID: "cluster-a", NodeID: "node-1"} + first := time.Now().Add(-time.Minute) + + // First observer (e.g. replica A) plants the marker. + livenessA, err := fetchNodeLiveness(t.Context(), client, []NodeRef{ref}, first) + require.NoError(t, err) + require.WithinDuration(t, first, livenessA[ref].firstMissing, time.Second) + + // A later observer (replica B) must read A's marker, not plant its own — + // the grace period ages from the FIRST observation across all replicas. + livenessB, err := fetchNodeLiveness(t.Context(), client, []NodeRef{ref}, time.Now()) + require.NoError(t, err) + assert.Equal(t, livenessA[ref].firstMissing.Unix(), livenessB[ref].firstMissing.Unix()) +} + +func TestNodeLiveness_RecordSeenClearsMissingMarker(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + ref := NodeRef{ClusterID: "cluster-a", NodeID: "node-1"} + past := time.Now().Add(-time.Hour) + + // Marker planted long ago. + _, err := fetchNodeLiveness(t.Context(), client, []NodeRef{ref}, past) + require.NoError(t, err) + + // Node seen: last-seen written, marker cleared. + require.NoError(t, RecordNodeSeen(t.Context(), client, ref, time.Now())) + + // Simulate the last-seen key later expiring (node gone again long after): + // the next missing observation must start a FRESH grace period, not + // resurrect the old marker. + require.NoError(t, client.Del(t.Context(), nodeLastSeenKey(ref)).Err()) + + now := time.Now() + liveness, err := fetchNodeLiveness(t.Context(), client, []NodeRef{ref}, now) + require.NoError(t, err) + assert.WithinDuration(t, now, liveness[ref].firstMissing, time.Second, "old marker must not survive a successful sync") +} + +func TestNodeLiveness_MixedBatch(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + seen := NodeRef{ClusterID: "cluster-a", NodeID: "node-seen"} + missing := NodeRef{ClusterID: "cluster-a", NodeID: "node-missing"} + now := time.Now() + + require.NoError(t, RecordNodeSeen(t.Context(), client, seen, now)) + + liveness, err := fetchNodeLiveness(t.Context(), client, []NodeRef{seen, missing}, now) + require.NoError(t, err) + require.Len(t, liveness, 2) + assert.False(t, liveness[seen].lastSeen.IsZero()) + assert.False(t, liveness[missing].firstMissing.IsZero()) +} + +func TestNodeLiveness_EmptyRefs(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + + liveness, err := fetchNodeLiveness(t.Context(), client, nil, time.Now()) + require.NoError(t, err) + assert.Empty(t, liveness) +} diff --git a/packages/api/internal/orchestrator/deadnode/sweep.go b/packages/api/internal/orchestrator/deadnode/sweep.go new file mode 100644 index 0000000000..03e10ac4ab --- /dev/null +++ b/packages/api/internal/orchestrator/deadnode/sweep.go @@ -0,0 +1,237 @@ +package deadnode + +import ( + "context" + "sync" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel" + "go.uber.org/zap" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + "github.com/e2b-dev/infra/packages/shared/pkg/logger" +) + +var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/api/internal/orchestrator/deadnode") + +const ( + // sweepInterval is the tick interval. + sweepInterval = time.Minute + + // backstopTicks forces a full scan every Nth tick + backstopTicks = 5 + + // gracePeriod is how long a node must go unseen by EVERY API + // replica before its sandboxes are purged from the store. Liveness is a + // cross-replica signal, so a single replica with a broken network + // cannot trigger a purge + gracePeriod = 5 * time.Minute + + // killTimeout bounds a single sandbox removal so one stuck kill + // cannot stall the whole sweep. + killTimeout = 30 * time.Second + + // killConcurrency bounds parallel removals per sweep cycle. + killConcurrency = 32 + + // discoveryFreshnessMax is how stale the last successful node discovery may + // be before the sweep refuses to run + discoveryFreshnessMax = 30 * time.Second +) + +// Sweeper removes store records for sandboxes whose node crashed. +// A node is considered dead only when no API replica has completed a +// successful sync with it for gracePeriod +type Sweeper struct { + flags *featureflags.Client + redisClient redis.UniversalClient + + lastDiscoverySync func() time.Time + listSandboxes func(ctx context.Context) ([]sandbox.Sandbox, error) + // anyNodeUnreachablePast reports whether any pool node has been locally + // unreachable for at least the given duration + anyNodeUnreachablePast func(d time.Duration) bool + removeSandbox func(ctx context.Context, teamID uuid.UUID, sandboxID string, opts sandbox.RemoveOpts) error + + tick int +} + +func New( + flags *featureflags.Client, + redisClient redis.UniversalClient, + listSandboxes func(ctx context.Context) ([]sandbox.Sandbox, error), + lastDiscoverySync func() time.Time, + anyNodeUnreachablePast func(d time.Duration) bool, + removeSandbox func(ctx context.Context, teamID uuid.UUID, sandboxID string, opts sandbox.RemoveOpts) error, +) *Sweeper { + return &Sweeper{ + flags: flags, + redisClient: redisClient, + listSandboxes: listSandboxes, + lastDiscoverySync: lastDiscoverySync, + anyNodeUnreachablePast: anyNodeUnreachablePast, + removeSandbox: removeSandbox, + } +} + +// Start runs the sweep loop until the context is cancelled. +func (s *Sweeper) Start(ctx context.Context) { + ticker := time.NewTicker(sweepInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + logger.L().Info(ctx, "Stopping dead node sweep") + + return + case <-ticker.C: + s.sweep(ctx) + } + } +} + +func (s *Sweeper) sweep(ctx context.Context) { + if !s.flags.BoolFlag(ctx, featureflags.DeadNodeSweepEnabledFlag) { + return + } + + last := s.lastDiscoverySync() + now := time.Now() + if last.IsZero() || now.Sub(last) > discoveryFreshnessMax { + logger.L().Warn(ctx, "Dead node sweep: node discovery is stale, skipping cycle", + logger.Time("last_discovery_sync", last), + ) + + return + } + + // Skip the store scan unless a purge candidate might exist. + // Nodes absent from this replica's pool entirely are only discoverable by scanning, + // hence the periodic unconditional backstop tick. + s.tick++ + backstop := s.tick%backstopTicks == 0 + if !backstop && !s.anyNodeUnreachablePast(gracePeriod) { + return + } + + ctx, span := tracer.Start(ctx, "dead-node-sweep") + defer span.End() + + sandboxes, err := s.listSandboxes(ctx) + if err != nil { + logger.L().Error(ctx, "Dead node sweep: failed to list stored sandboxes, skipping cycle", zap.Error(err)) + + return + } + if len(sandboxes) == 0 { + return + } + + refs := uniqueNodeRefs(sandboxes) + liveness, err := fetchNodeLiveness(ctx, s.redisClient, refs, time.Now()) + if err != nil { + logger.L().Error(ctx, "Dead node sweep: failed to fetch node liveness, skipping cycle", zap.Error(err)) + + return + } + + toKill := collectDeadNodeSandboxes(now, sandboxes, liveness) + if len(toKill) == 0 { + return + } + + logger.L().Warn(ctx, "Dead node sweep: removing sandboxes whose node is gone", + zap.Int("sandbox_count", len(toKill)), + ) + + sem := make(chan struct{}, killConcurrency) + var wg sync.WaitGroup + for _, sbx := range toKill { + sem <- struct{}{} + wg.Add(1) + go func(sbx sandbox.Sandbox) { + defer wg.Done() + defer func() { <-sem }() + + s.kill(ctx, sbx) + }(sbx) + } + wg.Wait() +} + +func (s *Sweeper) kill(ctx context.Context, sbx sandbox.Sandbox) { + killCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), killTimeout) + defer cancel() + + // nil covers the expected outcomes too (already removed elsewhere, node + // kill RPC failing after store cleanup) — the injected removeSandbox owns + // that translation. + err := s.removeSandbox(killCtx, sbx.TeamID, sbx.SandboxID, sandbox.RemoveOpts{ + Action: sandbox.StateActionKill, + Reason: sandbox.KillReasonNodeGone, + }) + if err != nil { + logger.L().Error(ctx, "Dead node sweep: failed to remove sandbox", + zap.Error(err), + logger.WithSandboxID(sbx.SandboxID), + logger.WithNodeID(sbx.NodeID), + ) + } +} + +func uniqueNodeRefs(sandboxes []sandbox.Sandbox) []NodeRef { + seen := make(map[NodeRef]struct{}) + var refs []NodeRef + for _, sbx := range sandboxes { + ref := NodeRef{ClusterID: sbx.ClusterID.String(), NodeID: sbx.NodeID} + if _, ok := seen[ref]; ok { + continue + } + + seen[ref] = struct{}{} + refs = append(refs, ref) + } + + return refs +} + +// collectDeadNodeSandboxes returns the sandboxes whose node no replica has +// seen for gracePeriod: +func collectDeadNodeSandboxes( + now time.Time, + sandboxes []sandbox.Sandbox, + liveness map[NodeRef]nodeLiveness, +) []sandbox.Sandbox { + var toKill []sandbox.Sandbox + + for _, sbx := range sandboxes { + ref := NodeRef{ClusterID: sbx.ClusterID.String(), NodeID: sbx.NodeID} + + evidence, ok := liveness[ref] + if !ok { + // No evidence about this node: never destructive on missing data. + continue + } + + switch { + case !evidence.lastSeen.IsZero(): + if now.Sub(evidence.lastSeen) < gracePeriod { + continue + } + case !evidence.firstMissing.IsZero(): + if now.Sub(evidence.firstMissing) < gracePeriod { + continue + } + default: + continue + } + + toKill = append(toKill, sbx) + } + + return toKill +} diff --git a/packages/api/internal/orchestrator/deadnode/sweep_test.go b/packages/api/internal/orchestrator/deadnode/sweep_test.go new file mode 100644 index 0000000000..999eba73f2 --- /dev/null +++ b/packages/api/internal/orchestrator/deadnode/sweep_test.go @@ -0,0 +1,385 @@ +package deadnode + +import ( + "context" + "errors" + "strconv" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" +) + +func deadNodeTestSandbox(nodeID string, clusterID uuid.UUID, startedAgo time.Duration, now time.Time) sandbox.Sandbox { + return sandbox.Sandbox{ + SandboxID: "sbx-" + uuid.NewString(), + TeamID: uuid.New(), + NodeID: nodeID, + ClusterID: clusterID, + StartTime: now.Add(-startedAgo), + State: sandbox.StateRunning, + } +} + +func refFor(clusterID uuid.UUID, nodeID string) NodeRef { + return NodeRef{ClusterID: clusterID.String(), NodeID: nodeID} +} + +// seenAgo builds a liveness entry for a node last seen the given duration ago. +func seenAgo(now time.Time, ago time.Duration) nodeLiveness { + return nodeLiveness{lastSeen: now.Add(-ago)} +} + +// missingFor builds a liveness entry for a never-seen node whose shared +// first-missing marker was planted the given duration ago. +func missingFor(now time.Time, ago time.Duration) nodeLiveness { + return nodeLiveness{firstMissing: now.Add(-ago)} +} + +func TestCollectDeadNodeSandboxes_RecentlySeenNodeSkipped(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterID := uuid.New() + + sbx := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + liveness := map[NodeRef]nodeLiveness{ + refFor(clusterID, "node-1"): seenAgo(now, gracePeriod-time.Second), + } + + toKill := collectDeadNodeSandboxes(now, []sandbox.Sandbox{sbx}, liveness) + + assert.Empty(t, toKill, "a node seen within the grace period must not be swept") +} + +func TestCollectDeadNodeSandboxes_StaleLastSeenKills(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterID := uuid.New() + + dead1 := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + dead2 := deadNodeTestSandbox("node-1", clusterID, 2*time.Hour, now) + alive := deadNodeTestSandbox("node-2", clusterID, time.Hour, now) + liveness := map[NodeRef]nodeLiveness{ + refFor(clusterID, "node-1"): seenAgo(now, gracePeriod+time.Second), + refFor(clusterID, "node-2"): seenAgo(now, time.Second), + } + + toKill := collectDeadNodeSandboxes(now, []sandbox.Sandbox{dead1, dead2, alive}, liveness) + + require.Len(t, toKill, 2) + assert.ElementsMatch(t, []string{dead1.SandboxID, dead2.SandboxID}, []string{toKill[0].SandboxID, toKill[1].SandboxID}) +} + +func TestCollectDeadNodeSandboxes_NeverSeenNodeAgesViaMarker(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterID := uuid.New() + + sbx := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + + // Marker younger than grace: only tracks. + young := map[NodeRef]nodeLiveness{refFor(clusterID, "node-1"): missingFor(now, gracePeriod-time.Second)} + assert.Empty(t, collectDeadNodeSandboxes(now, []sandbox.Sandbox{sbx}, young)) + + // Marker past grace: kills. + ripe := map[NodeRef]nodeLiveness{refFor(clusterID, "node-1"): missingFor(now, gracePeriod+time.Second)} + toKill := collectDeadNodeSandboxes(now, []sandbox.Sandbox{sbx}, ripe) + require.Len(t, toKill, 1) + assert.Equal(t, sbx.SandboxID, toKill[0].SandboxID) +} + +func TestCollectDeadNodeSandboxes_NoEvidenceSkipped(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterID := uuid.New() + + sbx := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + + // Node absent from the liveness map entirely. + assert.Empty(t, collectDeadNodeSandboxes(now, []sandbox.Sandbox{sbx}, map[NodeRef]nodeLiveness{})) + + // Node present but with a zero-value entry. + empty := map[NodeRef]nodeLiveness{refFor(clusterID, "node-1"): {}} + assert.Empty(t, collectDeadNodeSandboxes(now, []sandbox.Sandbox{sbx}, empty)) +} + +func TestCollectDeadNodeSandboxes_SameNodeIDDifferentClusters(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterA := uuid.New() + clusterB := uuid.New() + + deadClusterSbx := deadNodeTestSandbox("node-1", clusterA, time.Hour, now) + freshClusterSbx := deadNodeTestSandbox("node-1", clusterB, time.Hour, now) + liveness := map[NodeRef]nodeLiveness{ + refFor(clusterA, "node-1"): seenAgo(now, gracePeriod+time.Second), + refFor(clusterB, "node-1"): seenAgo(now, time.Second), + } + + toKill := collectDeadNodeSandboxes(now, []sandbox.Sandbox{deadClusterSbx, freshClusterSbx}, liveness) + + require.Len(t, toKill, 1) + assert.Equal(t, deadClusterSbx.SandboxID, toKill[0].SandboxID) +} + +func TestUniqueNodeRefs(t *testing.T) { + t.Parallel() + + now := time.Now() + clusterA := uuid.New() + clusterB := uuid.New() + + refs := uniqueNodeRefs([]sandbox.Sandbox{ + deadNodeTestSandbox("node-1", clusterA, time.Hour, now), + deadNodeTestSandbox("node-1", clusterA, time.Hour, now), + deadNodeTestSandbox("node-1", clusterB, time.Hour, now), + deadNodeTestSandbox("node-2", clusterA, time.Hour, now), + }) + + assert.ElementsMatch(t, []NodeRef{ + refFor(clusterA, "node-1"), + refFor(clusterB, "node-1"), + refFor(clusterA, "node-2"), + }, refs) +} + +// --- sweeper loop tests --- + +type fakeRemover struct { + mu sync.Mutex + calls []string + opts map[string]sandbox.RemoveOpts + returns map[string]error // sandboxID -> error +} + +func (f *fakeRemover) remove(_ context.Context, _ uuid.UUID, sandboxID string, opts sandbox.RemoveOpts) error { + f.mu.Lock() + defer f.mu.Unlock() + + f.calls = append(f.calls, sandboxID) + if f.opts == nil { + f.opts = map[string]sandbox.RemoveOpts{} + } + f.opts[sandboxID] = opts + + return f.returns[sandboxID] +} + +type countingLister struct { + mu sync.Mutex + calls int + sbxs []sandbox.Sandbox + err error +} + +func (l *countingLister) list(context.Context) ([]sandbox.Sandbox, error) { + l.mu.Lock() + defer l.mu.Unlock() + l.calls++ + + return l.sbxs, l.err +} + +func seedLastSeen(t *testing.T, client redis.UniversalClient, ref NodeRef, ts time.Time) { + t.Helper() + + require.NoError(t, client.Set(t.Context(), nodeLastSeenKey(ref), strconv.FormatInt(ts.Unix(), 10), nodeLivenessKeyTTL).Err()) +} + +func newTestSweeper(t *testing.T, enabled bool, client redis.UniversalClient, lister *countingLister, remover *fakeRemover) *Sweeper { + t.Helper() + + td := ldtestdata.DataSource() + td.Update(td.Flag(featureflags.DeadNodeSweepEnabledFlag.Key()).VariationForAll(enabled)) + ff, err := featureflags.NewClientWithDatasource(td) + require.NoError(t, err) + t.Cleanup(func() { _ = ff.Close(context.Background()) }) + + return New( + ff, + client, + lister.list, + time.Now, // discovery fresh by default + // Pre-gate open by default so decision tests exercise the full path. + func(time.Duration) bool { return true }, + remover.remove, + ) +} + +func TestSweep_FlagDisabledDoesNothing(t *testing.T) { + t.Parallel() + + remover := &fakeRemover{} + lister := &countingLister{sbxs: []sandbox.Sandbox{deadNodeTestSandbox("node-1", uuid.New(), time.Hour, time.Now())}} + + // nil redis client: a disabled sweep must return before touching anything. + s := newTestSweeper(t, false, nil, lister, remover) + s.sweep(context.Background()) + + assert.Empty(t, remover.calls, "disabled sweep must never call removeSandbox") + assert.Zero(t, lister.calls, "disabled sweep must not scan the store") +} + +func TestSweep_StaleDiscoverySkipsCycle(t *testing.T) { + t.Parallel() + + remover := &fakeRemover{} + lister := &countingLister{sbxs: []sandbox.Sandbox{deadNodeTestSandbox("node-1", uuid.New(), time.Hour, time.Now())}} + + s := newTestSweeper(t, true, nil, lister, remover) + + for name, last := range map[string]func() time.Time{ + "never synced": func() time.Time { return time.Time{} }, + "stale sync": func() time.Time { return time.Now().Add(-discoveryFreshnessMax - time.Second) }, + } { + s.lastDiscoverySync = last + s.sweep(context.Background()) + + assert.Empty(t, remover.calls, "%s: sweep must not act on stale observations", name) + assert.Zero(t, lister.calls, "%s: sweep must not scan on stale observations", name) + } +} + +func TestSweep_PreGateSkipsScanWhenIdle(t *testing.T) { + t.Parallel() + + remover := &fakeRemover{} + lister := &countingLister{} + s := newTestSweeper(t, true, nil, lister, remover) + s.anyNodeUnreachablePast = func(time.Duration) bool { return false } + + // Ticks 1..backstopTicks-1: no candidates -> no store scan. + for range backstopTicks - 1 { + s.sweep(context.Background()) + } + assert.Zero(t, lister.calls, "idle ticks must not scan the store") + + // Backstop tick scans unconditionally. + s.sweep(context.Background()) + assert.Equal(t, 1, lister.calls, "backstop tick must scan") +} + +func TestSweep_UnreachableCandidateTriggersScan(t *testing.T) { + t.Parallel() + + remover := &fakeRemover{} + lister := &countingLister{} + s := newTestSweeper(t, true, nil, lister, remover) + + s.sweep(context.Background()) // tick 1, not a backstop tick; pre-gate open + + assert.Equal(t, 1, lister.calls, "a local unreachable candidate must trigger a scan without waiting for the backstop") +} + +func TestSweep_ListErrorSkipsCycle(t *testing.T) { + t.Parallel() + + remover := &fakeRemover{} + lister := &countingLister{err: errors.New("redis down")} + s := newTestSweeper(t, true, nil, lister, remover) + + s.sweep(context.Background()) + + assert.Empty(t, remover.calls) +} + +func TestSweep_LivenessErrorSkipsCycle(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + clusterID := uuid.New() + sbx := deadNodeTestSandbox("node-1", clusterID, time.Hour, time.Now()) + remover := &fakeRemover{} + lister := &countingLister{sbxs: []sandbox.Sandbox{sbx}} + + s := newTestSweeper(t, true, client, lister, remover) + // Break the liveness fetch: closed client errors on every command. + require.NoError(t, client.Close()) + + s.sweep(context.Background()) + + assert.Empty(t, remover.calls, "no liveness evidence must never purge") +} + +func TestSweep_KillsWithNodeGoneReason(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + clusterID := uuid.New() + now := time.Now() + dead := deadNodeTestSandbox("node-dead", clusterID, time.Hour, now) + alive := deadNodeTestSandbox("node-alive", clusterID, time.Hour, now) + + seedLastSeen(t, client, refFor(clusterID, "node-dead"), now.Add(-gracePeriod-time.Second)) + seedLastSeen(t, client, refFor(clusterID, "node-alive"), now) + + remover := &fakeRemover{} + lister := &countingLister{sbxs: []sandbox.Sandbox{dead, alive}} + + s := newTestSweeper(t, true, client, lister, remover) + s.sweep(context.Background()) + + require.Equal(t, []string{dead.SandboxID}, remover.calls) + opts := remover.opts[dead.SandboxID] + assert.Equal(t, sandbox.StateActionKill, opts.Action) + assert.Equal(t, sandbox.KillReasonNodeGone, opts.Reason) + assert.False(t, opts.Eviction) +} + +func TestSweep_NeverSeenNodePurgedOnlyAfterMarkerGrace(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + clusterID := uuid.New() + now := time.Now() + sbx := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + + remover := &fakeRemover{} + lister := &countingLister{sbxs: []sandbox.Sandbox{sbx}} + + // First sweep observes the never-seen node: plants the marker, kills nothing. + s := newTestSweeper(t, true, client, lister, remover) + s.sweep(context.Background()) + assert.Empty(t, remover.calls, "first observation must only plant the marker") + + // Backdate the shared marker past the grace period: next sweep purges. + require.NoError(t, client.Set(t.Context(), nodeFirstMissingKey(refFor(clusterID, "node-1")), + strconv.FormatInt(now.Add(-gracePeriod-time.Second).Unix(), 10), nodeLivenessKeyTTL).Err()) + + s.sweep(context.Background()) + assert.Equal(t, []string{sbx.SandboxID}, remover.calls) +} + +func TestSweep_OneRemoveErrorDoesNotStopOthers(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + clusterID := uuid.New() + now := time.Now() + a := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + b := deadNodeTestSandbox("node-1", clusterID, time.Hour, now) + + seedLastSeen(t, client, refFor(clusterID, "node-1"), now.Add(-gracePeriod-time.Second)) + + remover := &fakeRemover{returns: map[string]error{a.SandboxID: errors.New("boom")}} + lister := &countingLister{sbxs: []sandbox.Sandbox{a, b}} + + s := newTestSweeper(t, true, client, lister, remover) + s.sweep(context.Background()) + + assert.ElementsMatch(t, []string{a.SandboxID, b.SandboxID}, remover.calls) +} diff --git a/packages/api/internal/orchestrator/delete_instance.go b/packages/api/internal/orchestrator/delete_instance.go index d7817980e1..d83ceaf2a8 100644 --- a/packages/api/internal/orchestrator/delete_instance.go +++ b/packages/api/internal/orchestrator/delete_instance.go @@ -135,6 +135,16 @@ func (o *Orchestrator) removeSandboxFromNode( node := o.getOrConnectNode(ctx, sbx.ClusterID, sbx.NodeID) if node == nil { + // For a node-gone kill, this is expected + if stateAction == sandbox.StateActionKill && reason == sandbox.KillReasonNodeGone { + logger.L().Warn(ctx, "Node gone, removed sandbox from store without node kill", + logger.WithSandboxID(sbx.SandboxID), + logger.WithNodeID(sbx.NodeID), + ) + + return nil + } + fields := []zap.Field{ logger.WithNodeID(sbx.NodeID), } @@ -192,7 +202,18 @@ func (o *Orchestrator) removeSandboxFromNode( return nil case sandbox.StateActionKill: - return o.killSandboxOnNode(ctx, node, sbx, reason) + err := o.killSandboxOnNode(ctx, node, sbx, reason) + if err != nil && reason == sandbox.KillReasonNodeGone { + logger.L().Warn(ctx, "Node unreachable, removed sandbox from store without node kill", + logger.WithSandboxID(sbx.SandboxID), + logger.WithNodeID(sbx.NodeID), + zap.Error(err), + ) + + return nil + } + + return err } return nil diff --git a/packages/api/internal/orchestrator/nodemanager/node.go b/packages/api/internal/orchestrator/nodemanager/node.go index 6f8ad11bdd..5be9e2db64 100644 --- a/packages/api/internal/orchestrator/nodemanager/node.go +++ b/packages/api/internal/orchestrator/nodemanager/node.go @@ -43,6 +43,10 @@ type Node struct { client *clusters.GRPCClient status StatusInfo + // unreachableSince is the time of the first local observation of a full + // sync-cycle failure against this node + unreachableSince time.Time + metrics Metrics metricsMu sync.RWMutex diff --git a/packages/api/internal/orchestrator/nodemanager/status.go b/packages/api/internal/orchestrator/nodemanager/status.go index 32350a0c6f..b2a6d39229 100644 --- a/packages/api/internal/orchestrator/nodemanager/status.go +++ b/packages/api/internal/orchestrator/nodemanager/status.go @@ -64,6 +64,10 @@ func (n *Node) markUnhealthyLocal(ctx context.Context) { n.mutex.Lock() defer n.mutex.Unlock() + if n.unreachableSince.IsZero() { + n.unreachableSince = time.Now() + } + if n.status.Status == api.NodeStatusUnhealthy { return } @@ -72,6 +76,21 @@ func (n *Node) markUnhealthyLocal(ctx context.Context) { n.status = StatusInfo{Status: api.NodeStatusUnhealthy, ChangedAt: time.Now()} } +// markReachable records a successful sync cycle, clearing the unreachable observation +func (n *Node) markReachable() { + n.mutex.Lock() + defer n.mutex.Unlock() + + n.unreachableSince = time.Time{} +} + +func (n *Node) UnreachableSince() (time.Time, bool) { + n.mutex.RLock() + defer n.mutex.RUnlock() + + return n.unreachableSince, !n.unreachableSince.IsZero() +} + func (n *Node) SendStatusChange(ctx context.Context, s api.NodeStatus) error { nodeStatus, ok := ApiNodeToOrchestratorStateMapper[s] if !ok { diff --git a/packages/api/internal/orchestrator/nodemanager/status_test.go b/packages/api/internal/orchestrator/nodemanager/status_test.go new file mode 100644 index 0000000000..c45afb47b3 --- /dev/null +++ b/packages/api/internal/orchestrator/nodemanager/status_test.go @@ -0,0 +1,74 @@ +package nodemanager + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/api" +) + +func TestUnreachableSince_DefaultReachable(t *testing.T) { + t.Parallel() + + n := &Node{} + + _, unreachable := n.UnreachableSince() + assert.False(t, unreachable) +} + +func TestMarkUnhealthyLocal_SetsUnreachableSinceOnce(t *testing.T) { + t.Parallel() + + n := &Node{} + ctx := context.Background() + + n.markUnhealthyLocal(ctx) + first, unreachable := n.UnreachableSince() + require.True(t, unreachable) + + // A repeated local unhealthy observation must preserve the first + // timestamp so the unreachable duration keeps accumulating. + time.Sleep(5 * time.Millisecond) + n.markUnhealthyLocal(ctx) + second, unreachable := n.UnreachableSince() + require.True(t, unreachable) + assert.Equal(t, first, second) + + assert.Equal(t, api.NodeStatusUnhealthy, n.Status()) +} + +func TestMarkReachable_ClearsUnreachableSince(t *testing.T) { + t.Parallel() + + n := &Node{} + ctx := context.Background() + + n.markUnhealthyLocal(ctx) + n.markReachable() + + _, unreachable := n.UnreachableSince() + assert.False(t, unreachable) + + // The next unreachable observation starts a fresh clock. + n.markUnhealthyLocal(ctx) + since, unreachable := n.UnreachableSince() + require.True(t, unreachable) + assert.WithinDuration(t, time.Now(), since, time.Second) +} + +// TestSelfReportedUnhealthyIsNotUnreachable pins the distinction the dead-node +// sweep relies on: a node whose orchestrator self-reports Unhealthy via a +// successful sync is responsive and must never be considered unreachable. +func TestSelfReportedUnhealthyIsNotUnreachable(t *testing.T) { + t.Parallel() + + n := &Node{} + n.setStatus(context.Background(), api.NodeStatusUnhealthy, time.Now().Add(-time.Hour)) + + _, unreachable := n.UnreachableSince() + assert.False(t, unreachable) +} diff --git a/packages/api/internal/orchestrator/nodemanager/sync.go b/packages/api/internal/orchestrator/nodemanager/sync.go index db86ff581d..5660a3394b 100644 --- a/packages/api/internal/orchestrator/nodemanager/sync.go +++ b/packages/api/internal/orchestrator/nodemanager/sync.go @@ -60,6 +60,9 @@ func (n *Node) Sync(ctx context.Context, store *sandbox.Store) { store.Reconcile(ctx, activeInstances, n.ID) + // Full sync cycle succeeded: the node is reachable. + n.markReachable() + syncRetrySuccess = true break diff --git a/packages/api/internal/orchestrator/orchestrator.go b/packages/api/internal/orchestrator/orchestrator.go index 2200b302fb..ced5ee01e7 100644 --- a/packages/api/internal/orchestrator/orchestrator.go +++ b/packages/api/internal/orchestrator/orchestrator.go @@ -5,8 +5,10 @@ import ( "errors" "fmt" "net/http" + "sync/atomic" "time" + "github.com/google/uuid" "github.com/redis/go-redis/v9" "go.opentelemetry.io/otel/metric" "go.uber.org/zap" @@ -16,6 +18,7 @@ import ( "github.com/e2b-dev/infra/packages/api/internal/cfg" "github.com/e2b-dev/infra/packages/api/internal/clusters" "github.com/e2b-dev/infra/packages/api/internal/metrics" + "github.com/e2b-dev/infra/packages/api/internal/orchestrator/deadnode" "github.com/e2b-dev/infra/packages/api/internal/orchestrator/discovery" "github.com/e2b-dev/infra/packages/api/internal/orchestrator/evictor" "github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager" @@ -45,6 +48,7 @@ type SnapshotCacheInvalidator interface { type Orchestrator struct { httpClient *http.Client nodeDiscovery discovery.Discovery + redisClient redis.UniversalClient sandboxStore *sandbox.Store nodes *smap.Map[*nodemanager.Node] placementAlgorithm *placement.BestOfK @@ -81,6 +85,15 @@ type Orchestrator struct { // same string, and nesting Do calls for the same key on the same Group would // block forever. discoveryGroup singleflight.Group + + // lastDiscoverySync is the time of the last syncNodes cycle that + // successfully completed the node discovery phase. Nil until the first + // success. The dead-node sweep refuses to act on observations older + // than this: "node not seen" is only meaningful when nodes are actually + // being synced — otherwise a discovery outage (or a freshly booted + // replica that cannot reach service discovery) would make every node + // look dead. + lastDiscoverySync atomic.Pointer[time.Time] } func New( @@ -138,6 +151,7 @@ func New( analytics: analyticsInstance, posthogClient: posthogClient, nodeDiscovery: nodeDiscovery, + redisClient: redisClient, nodes: smap.New[*nodemanager.Node](), placementAlgorithm: bestOfKAlgorithm, featureFlagsClient: featureFlags, @@ -194,6 +208,40 @@ func New( go o.startStatusLogging(ctx) go o.updateBestOfKConfig(ctx) + // Purge sandboxes from the store when their node crashed and never came back. + deadNodeSweeper := deadnode.New( + featureFlags, + redisClient, + o.sandboxStore.AllRunningItems, + func() time.Time { + if ts := o.lastDiscoverySync.Load(); ts != nil { + return *ts + } + + return time.Time{} + }, + func(d time.Duration) bool { + now := time.Now() + for _, n := range o.nodes.Items() { + if since, ok := n.UnreachableSince(); ok && now.Sub(since) >= d { + return true + } + } + + return false + }, + func(ctx context.Context, teamID uuid.UUID, sandboxID string, opts sandbox.RemoveOpts) error { + err := o.RemoveSandbox(ctx, teamID, sandboxID, opts) + if errors.Is(err, ErrSandboxNotFound) { + // Another replica (or the user) removed it concurrently. + return nil + } + + return err + }, + ) + go deadNodeSweeper.Start(ctx) + return &o, nil } diff --git a/packages/api/internal/sandbox/aliases.go b/packages/api/internal/sandbox/aliases.go index 258fa2fb32..d525b35f11 100644 --- a/packages/api/internal/sandbox/aliases.go +++ b/packages/api/internal/sandbox/aliases.go @@ -44,6 +44,7 @@ const ( KillReasonAdmin = sandboxtypes.KillReasonAdmin KillReasonOrphaned = sandboxtypes.KillReasonOrphaned KillReasonBaseTemplateMissing = sandboxtypes.KillReasonBaseTemplateMissing + KillReasonNodeGone = sandboxtypes.KillReasonNodeGone ) // Errors and pre-defined state actions / transition tables. diff --git a/packages/api/internal/sandbox/sandboxtypes/states.go b/packages/api/internal/sandbox/sandboxtypes/states.go index 4b462d2d55..c6abaf4481 100644 --- a/packages/api/internal/sandbox/sandboxtypes/states.go +++ b/packages/api/internal/sandbox/sandboxtypes/states.go @@ -55,6 +55,9 @@ const ( KillReasonAdmin KillReason = "admin" KillReasonOrphaned KillReason = "orphaned" KillReasonBaseTemplateMissing KillReason = "base_template_missing" + // KillReasonNodeGone marks sandboxes removed from the store because the + // node they were running on disappeared from the pool (crashed or was removed) + KillReasonNodeGone KillReason = "node_gone" ) // String returns the reason as a string, normalizing the empty value to diff --git a/packages/api/internal/sandbox/sandboxtypes/states_test.go b/packages/api/internal/sandbox/sandboxtypes/states_test.go index 2681efe229..157992da82 100644 --- a/packages/api/internal/sandbox/sandboxtypes/states_test.go +++ b/packages/api/internal/sandbox/sandboxtypes/states_test.go @@ -12,6 +12,7 @@ func TestKillReasonValues(t *testing.T) { KillReasonAdmin: "admin", KillReasonOrphaned: "orphaned", KillReasonBaseTemplateMissing: "base_template_missing", + KillReasonNodeGone: "node_gone", } { if string(reason) != want { t.Errorf("KillReason = %q, want %q", string(reason), want) @@ -62,6 +63,11 @@ func TestKillReasonString(t *testing.T) { in: KillReasonBaseTemplateMissing, want: "base_template_missing", }, + { + name: "node gone", + in: KillReasonNodeGone, + want: "node_gone", + }, } for _, tt := range tests { diff --git a/packages/api/internal/sandbox/sandboxtypes/storage.go b/packages/api/internal/sandbox/sandboxtypes/storage.go index f01f93e1ce..a99ecf2c39 100644 --- a/packages/api/internal/sandbox/sandboxtypes/storage.go +++ b/packages/api/internal/sandbox/sandboxtypes/storage.go @@ -11,6 +11,8 @@ const ( ) // Storage is the persistence interface implemented by the redis backend. +// +//nolint:interfacebloat type Storage interface { Add(ctx context.Context, sandbox Sandbox) error Get(ctx context.Context, teamID uuid.UUID, sandboxID string) (Sandbox, error) @@ -19,6 +21,7 @@ type Storage interface { TeamItems(ctx context.Context, teamID uuid.UUID, states []State) ([]Sandbox, error) ExpiredItems(ctx context.Context) ([]Sandbox, error) TeamsWithSandboxCount(ctx context.Context) (map[uuid.UUID]int64, error) + AllRunningItems(ctx context.Context) ([]Sandbox, error) Update(ctx context.Context, teamID uuid.UUID, sandboxID string, updateFunc func(sandbox Sandbox) (Sandbox, error)) (Sandbox, error) StartRemoving(ctx context.Context, teamID uuid.UUID, sandboxID string, opts RemoveOpts) (Sandbox, bool, func(context.Context, error), error) diff --git a/packages/api/internal/sandbox/storage/redis/expiration_index_test.go b/packages/api/internal/sandbox/storage/redis/expiration_index_test.go index b6d7f502de..f5271adcd0 100644 --- a/packages/api/internal/sandbox/storage/redis/expiration_index_test.go +++ b/packages/api/internal/sandbox/storage/redis/expiration_index_test.go @@ -304,15 +304,15 @@ func TestHeal_SkipsYoungSandbox(t *testing.T) { } // TestHeal_ManySandboxesMultipleBatches exercises the SSCAN pagination in -// healTeamExpirationIndex: more sandboxes than healScanBatchSize, with holes -// spread across the whole ID range. +// the shared sandbox scanner: more sandboxes than sandboxScanBatchSize, 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 + total := sandboxScanBatchSize*2 + 37 // force >2 batches with a partial tail pipe := client.Pipeline() var missingMembers []string @@ -334,10 +334,11 @@ func TestHeal_ManySandboxesMultipleBatches(t *testing.T) { missingMembers = append(missingMembers, member) } } + pipe.ZAdd(t.Context(), globalTeamsSet, redis.Z{Score: float64(time.Now().Unix()), Member: teamID.String()}) _, err := pipe.Exec(t.Context()) require.NoError(t, err) - healed, err := storage.healTeamExpirationIndex(t.Context(), teamID.String()) + healed, err := storage.healExpirationIndex(t.Context()) require.NoError(t, err) require.Equal(t, len(missingMembers), healed) @@ -348,7 +349,7 @@ func TestHeal_ManySandboxesMultipleBatches(t *testing.T) { } // Idempotent: a second pass heals nothing. - healed, err = storage.healTeamExpirationIndex(t.Context(), teamID.String()) + healed, err = storage.healExpirationIndex(t.Context()) require.NoError(t, err) require.Zero(t, healed) } diff --git a/packages/api/internal/sandbox/storage/redis/heal.go b/packages/api/internal/sandbox/storage/redis/heal.go index 9a17506d21..e8ef5b0770 100644 --- a/packages/api/internal/sandbox/storage/redis/heal.go +++ b/packages/api/internal/sandbox/storage/redis/heal.go @@ -2,7 +2,6 @@ package redis import ( "context" - "encoding/json" "fmt" "time" @@ -12,7 +11,6 @@ import ( "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" - "github.com/e2b-dev/infra/packages/shared/pkg/utils" ) const ( @@ -21,11 +19,6 @@ const ( // 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". @@ -67,94 +60,42 @@ func (s *Storage) healerEnabled(ctx context.Context) bool { // healExpirationIndex scans all stored sandboxes and re-adds expiration index // members missing for live sandbox keys. Returns the number of healed members. +// Per-team failures are isolated by the shared scanner; other teams still get +// healed. func (s *Storage) healExpirationIndex(ctx context.Context) (int, error) { // Re-evaluated every tick: acts as a kill switch without redeploy. if !s.healerEnabled(ctx) { return 0, nil } - 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) + err := s.forEachSandboxBatch(ctx, func(_ string, batch []sandboxtypes.Sandbox) error { + n, err := s.healSandboxBatch(ctx, batch) 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 + return err } healed += n + + return nil + }) + if err != nil { + return healed, err } 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) - } - +// batch of sandbox records. Returns the number of healed members. +func (s *Storage) healSandboxBatch(ctx context.Context, batch []sandboxtypes.Sandbox) (int, error) { now := time.Now() type candidate struct { member 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 - } - + for _, sbx := range batch { if now.Sub(sbx.StartTime) < healGracePeriod { continue } diff --git a/packages/api/internal/sandbox/storage/redis/scan.go b/packages/api/internal/sandbox/storage/redis/scan.go new file mode 100644 index 0000000000..96c06ce7f6 --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/scan.go @@ -0,0 +1,136 @@ +package redis + +import ( + "context" + "encoding/json" + "fmt" + + "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" +) + +// sandboxScanBatchSize bounds per-command work (SSCAN page, MGET keys) so +// teams with many sandboxes can't produce huge single commands/replies that +// stall Redis or explode service memory. +const sandboxScanBatchSize = 256 + +// forEachSandboxBatch scans every team's sandbox index and invokes fn with +// bounded batches of unmarshalled records (all states, stale index entries +// skipped). Per-team failures — including errors returned by fn — are +// isolated: the failing team is skipped and the remaining teams are still +// scanned. Only a failure to list the teams themselves is returned. +func (s *Storage) forEachSandboxBatch(ctx context.Context, fn func(teamID string, batch []sandboxtypes.Sandbox) error) error { + teams, err := s.redisClient.ZRange(ctx, globalTeamsSet, 0, -1).Result() + if err != nil { + return fmt.Errorf("failed to list teams from global index: %w", err) + } + + for _, teamID := range teams { + if err := s.forEachTeamSandboxBatch(ctx, teamID, fn); err != nil { + logger.L().Warn(ctx, "Failed to scan team sandboxes, skipping team", + zap.Error(err), + logger.WithTeamID(teamID), + ) + + continue + } + } + + return nil +} + +// forEachTeamSandboxBatch pages through one team's sandbox index with SSCAN +// and feeds fn bounded batches fetched via MGET. +func (s *Storage) forEachTeamSandboxBatch(ctx context.Context, teamID string, fn func(teamID string, batch []sandboxtypes.Sandbox) error) error { + var cursor uint64 + + for { + sandboxIDs, next, err := s.redisClient.SScan(ctx, GetSandboxStorageTeamIndexKey(teamID), cursor, "", sandboxScanBatchSize).Result() + if err != nil { + return fmt.Errorf("failed to scan team index: %w", err) + } + + // SSCAN COUNT is a hint, not a cap: split oversized pages so + // downstream MGET commands stay bounded. + for start := 0; start < len(sandboxIDs); start += sandboxScanBatchSize { + end := min(start+sandboxScanBatchSize, len(sandboxIDs)) + + batch, err := s.fetchSandboxBatch(ctx, teamID, sandboxIDs[start:end]) + if err != nil { + return err + } + if len(batch) == 0 { + continue + } + + if err := fn(teamID, batch); err != nil { + return err + } + } + + cursor = next + if cursor == 0 { + return nil + } + } +} + +// fetchSandboxBatch MGETs one bounded batch of sandbox records. Stale index +// entries (key already deleted) and unparseable records are skipped. +func (s *Storage) fetchSandboxBatch(ctx context.Context, teamID string, sandboxIDs []string) ([]sandboxtypes.Sandbox, error) { + if len(sandboxIDs) == 0 { + return nil, 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 nil, fmt.Errorf("MGET failed: %w", err) + } + + out := make([]sandboxtypes.Sandbox, 0, len(vals)) + 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 { + logger.L().Error(ctx, "Failed to unmarshal sandbox during scan", zap.Error(err)) + + continue + } + + out = append(out, sbx) + } + + return out, nil +} + +// AllRunningItems returns every running sandbox in the store across all +// teams, scanned in bounded batches. Used by the dead-node sweep. +func (s *Storage) AllRunningItems(ctx context.Context) ([]sandboxtypes.Sandbox, error) { + var out []sandboxtypes.Sandbox + + err := s.forEachSandboxBatch(ctx, func(_ string, batch []sandboxtypes.Sandbox) error { + for _, sbx := range batch { + if sbx.State != sandboxtypes.StateRunning { + continue + } + + out = append(out, sbx) + } + + return nil + }) + if err != nil { + return nil, err + } + + return out, nil +} diff --git a/packages/api/internal/sandbox/storage/redis/scan_test.go b/packages/api/internal/sandbox/storage/redis/scan_test.go new file mode 100644 index 0000000000..e080fa86dc --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/scan_test.go @@ -0,0 +1,93 @@ +package redis + +import ( + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox/sandboxtypes" +) + +func TestAllRunningItems_Empty(t *testing.T) { + t.Parallel() + + storage, _ := setupTestStorage(t) + + items, err := storage.AllRunningItems(t.Context()) + require.NoError(t, err) + assert.Empty(t, items) +} + +func TestAllRunningItems_MultipleBatchesAndStateFilter(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamA := uuid.New() + teamB := uuid.New() + + // Force >2 SSCAN/MGET batches with a partial tail for team A. + totalA := sandboxScanBatchSize*2 + 37 + seedTeamSandboxes(t, client, teamA, totalA, sandboxtypes.StateRunning) + + // Team B: a few running plus non-running records that must be filtered. + seedTeamSandboxes(t, client, teamB, 5, sandboxtypes.StateRunning) + seedTeamSandboxes(t, client, teamB, 3, sandboxtypes.StateKilling) + + items, err := storage.AllRunningItems(t.Context()) + require.NoError(t, err) + + assert.Len(t, items, totalA+5, "must return all running sandboxes across teams and skip non-running states") + + perTeam := map[uuid.UUID]int{} + for _, sbx := range items { + assert.Equal(t, sandboxtypes.StateRunning, sbx.State) + perTeam[sbx.TeamID]++ + } + assert.Equal(t, totalA, perTeam[teamA]) + assert.Equal(t, 5, perTeam[teamB]) +} + +func TestAllRunningItems_ToleratesStaleIndexEntries(t *testing.T) { + t.Parallel() + + storage, client := setupTestStorage(t) + + teamID := uuid.New() + seedTeamSandboxes(t, client, teamID, 2, sandboxtypes.StateRunning) + + // Stale index entry: ID in the team set without a sandbox key. + require.NoError(t, client.SAdd(t.Context(), GetSandboxStorageTeamIndexKey(teamID.String()), "sbx-deleted").Err()) + + items, err := storage.AllRunningItems(t.Context()) + require.NoError(t, err) + assert.Len(t, items, 2) +} + +// seedTeamSandboxes writes sandbox records + index entries directly (bypassing +// Storage.Add) and registers the team in the global teams index. +func seedTeamSandboxes(t *testing.T, client redis.UniversalClient, teamID uuid.UUID, count int, state sandboxtypes.State) { + t.Helper() + + pipe := client.Pipeline() + for i := range count { + sbx := makeIndexedSandbox(teamID, fmt.Sprintf("sbx-%s-%04d", state, i), uuid.NewString(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour)) + sbx.State = state + sbx.NodeID = "test-node" + + 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) + } + pipe.ZAdd(t.Context(), globalTeamsSet, redis.Z{Score: float64(time.Now().Unix()), Member: teamID.String()}) + _, err := pipe.Exec(t.Context()) + require.NoError(t, err) +} diff --git a/packages/api/internal/sandbox/store.go b/packages/api/internal/sandbox/store.go index 2a69e5eb7b..c1f6cb51b3 100644 --- a/packages/api/internal/sandbox/store.go +++ b/packages/api/internal/sandbox/store.go @@ -125,6 +125,11 @@ func (s *Store) TeamsWithSandboxes(ctx context.Context) (map[uuid.UUID]int64, er return s.storage.TeamsWithSandboxCount(ctx) } +// AllRunningItems returns every running sandbox across all teams. +func (s *Store) AllRunningItems(ctx context.Context) ([]Sandbox, error) { + return s.storage.AllRunningItems(ctx) +} + func (s *Store) Update(ctx context.Context, teamID uuid.UUID, sandboxID string, updateFunc func(sandbox Sandbox) (Sandbox, error)) (Sandbox, error) { return s.storage.Update(ctx, teamID, sandboxID, updateFunc) } diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 2373a115be..affb1d452f 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -219,6 +219,13 @@ var ( // On by default; acts as a kill switch if a heal pass misbehaves. ExpirationIndexHealerFlag = NewBoolFlag("expiration-index-healer", true) + // DeadNodeSweepEnabledFlag gates the API's dead-node sweep, which purges + // sandbox records from the store when the node they ran on has been gone + // from the pool past a grace period. Checked on every sweep tick, so it can + // be toggled without a redeploy. Off by default: this is a destructive + // background loop and the flag doubles as its fleet-wide kill switch. + DeadNodeSweepEnabledFlag = NewBoolFlag("dead-node-sweep-enabled", 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