Skip to content

Commit 057f20c

Browse files
wj-e2bwj-e2b
authored
feat(orchestrator): track sandbox lifecycles (#2998)
Track sandbox lifecycle entries in the sandbox map so shutdown can observe live sandboxes, and mark sandboxes stopped on Close so the map reflects reality even when cleanup partially fails.
1 parent cb6aa0b commit 057f20c

3 files changed

Lines changed: 278 additions & 16 deletions

File tree

‎packages/orchestrator/pkg/sandbox/map.go‎

Lines changed: 96 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,26 +25,51 @@ type MapSubscriber interface {
2525
OnNetworkRelease(ctx context.Context, sbx *Sandbox)
2626
}
2727

28-
// Map holds sandboxes that are live (running) together with a IP-to-sandbox index
29-
// The two maps are managed independently.
28+
// Map tracks sandboxes in three indexes, managed independently:
3029
//
31-
// AssignNetwork/NetworkReleased manage the IP map,
32-
// MarkRunning/MarkStopping manage the live set.
30+
// - live: keyed by sandboxID, holds the current routable lifecycle per
31+
// sandbox from MarkRunning until MarkStopping. It serves the API/proxy
32+
// lookup paths (Get, Items, Count).
33+
// - lifecycles: keyed by sandboxID/lifecycleID, holds every lifecycle whose
34+
// cleanup is still outstanding, from MarkRunning until MarkStopped in
35+
// Close. During checkpoint/resume an old lifecycle can still be cleaning
36+
// up while a new lifecycle with the same sandboxID is already live, so a
37+
// sandboxID can map to multiple lifecycle entries. Shutdown uses this set
38+
// (WaitLifecycles, LifecycleItems) to wait for cleanup to finish, not
39+
// just for sandboxes to stop being routable.
40+
// - network: an IP-to-sandbox index managed by AssignNetwork and
41+
// NetworkReleased, serving GetByHostPort lookups.
42+
//
43+
// Invariant: live is a subset of lifecycles; MarkRunning inserts into both.
44+
// The live and lifecycles maps could later be merged into a single registry
45+
// keyed by sandboxID/lifecycleID with a running/stopping state per entry;
46+
// they are kept separate for now to stay close to the pre-existing live-map
47+
// shape.
3348
type Map struct {
34-
live *smap.Map[*Sandbox]
35-
network *smap.Map[*Sandbox]
49+
live *smap.Map[*Sandbox]
50+
lifecycles *smap.Map[*Sandbox]
51+
network *smap.Map[*Sandbox]
52+
53+
lifecycleMu sync.Mutex
54+
lifecycleChanged chan struct{}
3655

3756
subs []MapSubscriber
3857
subsLock sync.RWMutex
3958
}
4059

4160
func NewSandboxesMap() *Map {
4261
return &Map{
43-
live: smap.New[*Sandbox](),
44-
network: smap.New[*Sandbox](),
62+
live: smap.New[*Sandbox](),
63+
lifecycles: smap.New[*Sandbox](),
64+
network: smap.New[*Sandbox](),
65+
lifecycleChanged: make(chan struct{}),
4566
}
4667
}
4768

69+
func sandboxLifecycleKey(sandboxID, lifecycleID string) string {
70+
return fmt.Sprintf("%s/%s", sandboxID, lifecycleID)
71+
}
72+
4873
func (m *Map) Subscribe(subscriber MapSubscriber) {
4974
m.subsLock.Lock()
5075
defer m.subsLock.Unlock()
@@ -73,6 +98,36 @@ func (m *Map) Get(sandboxID string) (*Sandbox, bool) {
7398
return m.live.Get(sandboxID)
7499
}
75100

101+
func (m *Map) LifecycleItems() []*Sandbox {
102+
items := m.lifecycles.Items()
103+
sandboxes := make([]*Sandbox, 0, len(items))
104+
for _, sbx := range items {
105+
sandboxes = append(sandboxes, sbx)
106+
}
107+
108+
return sandboxes
109+
}
110+
111+
func (m *Map) WaitLifecycles(ctx context.Context) error {
112+
for {
113+
m.lifecycleMu.Lock()
114+
if m.lifecycles.Count() == 0 {
115+
m.lifecycleMu.Unlock()
116+
117+
return nil
118+
}
119+
120+
changed := m.lifecycleChanged
121+
m.lifecycleMu.Unlock()
122+
123+
select {
124+
case <-ctx.Done():
125+
return fmt.Errorf("waiting for sandbox lifecycle cleanup: %w", ctx.Err())
126+
case <-changed:
127+
}
128+
}
129+
}
130+
76131
// GetByHostPort looks up a sandbox by its host IP address parsed from hostPort.
77132
func (m *Map) GetByHostPort(hostPort string) (*Sandbox, error) {
78133
reqIP, _, err := net.SplitHostPort(hostPort)
@@ -100,12 +155,27 @@ func (m *Map) AssignNetwork(ctx context.Context, sbx *Sandbox) {
100155
)
101156
}
102157

158+
func (m *Map) trackLifecycle(ctx context.Context, sbx *Sandbox) {
159+
m.lifecycleMu.Lock()
160+
m.lifecycles.Insert(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID), sbx)
161+
m.notifyLifecycleChangeLocked()
162+
m.lifecycleMu.Unlock()
163+
164+
logger.L().Info(ctx, "sandbox lifecycle tracked",
165+
logger.WithSandboxID(sbx.Runtime.SandboxID),
166+
logger.WithLifecycleID(sbx.LifecycleID),
167+
logger.WithSandboxIP(sbx.Slot.HostIPString()),
168+
)
169+
}
170+
103171
// MarkRunning makes the sandbox visible to Get/Items/Count and notifies OnInsert subscribers.
104172
func (m *Map) MarkRunning(ctx context.Context, sbx *Sandbox) {
105173
if !m.live.InsertIfAbsent(sbx.Runtime.SandboxID, sbx) {
106174
return
107175
}
108176

177+
m.trackLifecycle(ctx, sbx)
178+
109179
m.trigger(ctx, func(ctx context.Context, s MapSubscriber) {
110180
s.OnInsert(ctx, sbx)
111181
})
@@ -150,6 +220,24 @@ func (m *Map) MarkStopping(ctx context.Context, sandboxID, lifecycleID string) b
150220
return stopped
151221
}
152222

223+
func (m *Map) MarkStopped(ctx context.Context, sbx *Sandbox) {
224+
m.lifecycleMu.Lock()
225+
m.lifecycles.Remove(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID))
226+
m.notifyLifecycleChangeLocked()
227+
m.lifecycleMu.Unlock()
228+
229+
logger.L().Info(ctx, "sandbox lifecycle stopped",
230+
logger.WithSandboxID(sbx.Runtime.SandboxID),
231+
logger.WithLifecycleID(sbx.LifecycleID),
232+
logger.WithSandboxIP(sbx.Slot.HostIPString()),
233+
)
234+
}
235+
236+
func (m *Map) notifyLifecycleChangeLocked() {
237+
close(m.lifecycleChanged)
238+
m.lifecycleChanged = make(chan struct{})
239+
}
240+
153241
// NetworkReleased unregisters a sandbox's IP and notifies OnNetworkRelease
154242
// subscribers after a successful removal.
155243
func (m *Map) NetworkReleased(ctx context.Context, ip string) {

‎packages/orchestrator/pkg/sandbox/map_test.go‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
//go:build linux
2+
3+
package sandbox
4+
5+
import (
6+
"context"
7+
"sync"
8+
"testing"
9+
"time"
10+
11+
"github.com/stretchr/testify/require"
12+
13+
"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/network"
14+
)
15+
16+
func TestMapMarkRunningTracksLifecycle(t *testing.T) {
17+
t.Parallel()
18+
19+
sandboxes := NewSandboxesMap()
20+
sbx := testMapSandbox(t, "lifecycle-1")
21+
22+
sandboxes.MarkRunning(t.Context(), sbx)
23+
require.Len(t, sandboxes.Items(), 1)
24+
require.Len(t, sandboxes.LifecycleItems(), 1)
25+
}
26+
27+
func TestMapLifecycleItemsRemainAfterMarkStopping(t *testing.T) {
28+
t.Parallel()
29+
30+
sandboxes := NewSandboxesMap()
31+
sbx := testMapSandbox(t, "lifecycle-1")
32+
33+
sandboxes.MarkRunning(t.Context(), sbx)
34+
require.Len(t, sandboxes.Items(), 1)
35+
require.Len(t, sandboxes.LifecycleItems(), 1)
36+
37+
marked := sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID)
38+
require.True(t, marked)
39+
require.Empty(t, sandboxes.Items())
40+
require.Len(t, sandboxes.LifecycleItems(), 1)
41+
42+
sandboxes.MarkStopped(t.Context(), sbx)
43+
require.Empty(t, sandboxes.LifecycleItems())
44+
}
45+
46+
func TestSandboxCloseMarksLifecycleStopped(t *testing.T) {
47+
t.Parallel()
48+
49+
sandboxes := NewSandboxesMap()
50+
sbx := testMapSandbox(t, "lifecycle-1")
51+
sbx.cleanup = NewCleanup()
52+
sbx.sandboxes = sandboxes
53+
54+
sandboxes.MarkRunning(t.Context(), sbx)
55+
require.Len(t, sandboxes.LifecycleItems(), 1)
56+
57+
require.NoError(t, sbx.Close(t.Context()))
58+
require.Empty(t, sandboxes.LifecycleItems())
59+
}
60+
61+
func TestMapLifecycleItemsAllowDuplicateSandboxIDs(t *testing.T) {
62+
t.Parallel()
63+
64+
sandboxes := NewSandboxesMap()
65+
oldSbx := testMapSandbox(t, "lifecycle-old")
66+
newSbx := testMapSandbox(t, "lifecycle-new")
67+
68+
sandboxes.MarkRunning(t.Context(), oldSbx)
69+
require.True(t, sandboxes.MarkStopping(t.Context(), oldSbx.Runtime.SandboxID, oldSbx.LifecycleID))
70+
sandboxes.MarkRunning(t.Context(), newSbx)
71+
72+
require.Len(t, sandboxes.LifecycleItems(), 2)
73+
}
74+
75+
func TestMapWaitLifecyclesReturnsWhenEmpty(t *testing.T) {
76+
t.Parallel()
77+
78+
sandboxes := NewSandboxesMap()
79+
80+
require.NoError(t, sandboxes.WaitLifecycles(t.Context()))
81+
}
82+
83+
func TestMapWaitLifecyclesWaitsUntilStopped(t *testing.T) {
84+
t.Parallel()
85+
86+
sandboxes := NewSandboxesMap()
87+
sbx := testMapSandbox(t, "lifecycle-1")
88+
sandboxes.MarkRunning(t.Context(), sbx)
89+
90+
waitCtx, cancel := context.WithTimeout(t.Context(), time.Second)
91+
defer cancel()
92+
93+
done := make(chan error, 1)
94+
go func() {
95+
done <- sandboxes.WaitLifecycles(waitCtx)
96+
}()
97+
98+
select {
99+
case err := <-done:
100+
require.Failf(t, "WaitLifecycles returned before lifecycle stopped", "err: %v", err)
101+
case <-time.After(25 * time.Millisecond):
102+
}
103+
104+
sandboxes.MarkStopped(t.Context(), sbx)
105+
require.NoError(t, <-done)
106+
}
107+
108+
func TestMapWaitLifecyclesReturnsContextError(t *testing.T) {
109+
t.Parallel()
110+
111+
sandboxes := NewSandboxesMap()
112+
sbx := testMapSandbox(t, "lifecycle-1")
113+
sandboxes.MarkRunning(t.Context(), sbx)
114+
115+
waitCtx, cancel := context.WithCancel(t.Context())
116+
cancel()
117+
118+
require.ErrorIs(t, sandboxes.WaitLifecycles(waitCtx), context.Canceled)
119+
}
120+
121+
func TestMapConcurrentMarkStoppingAndStoppedDoesNotResurrectLifecycle(t *testing.T) {
122+
t.Parallel()
123+
124+
for range 1000 {
125+
sandboxes := NewSandboxesMap()
126+
sbx := testMapSandbox(t, "lifecycle-1")
127+
sandboxes.MarkRunning(t.Context(), sbx)
128+
129+
start := make(chan struct{})
130+
var wg sync.WaitGroup
131+
wg.Add(2)
132+
133+
go func() {
134+
defer wg.Done()
135+
<-start
136+
sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID)
137+
}()
138+
139+
go func() {
140+
defer wg.Done()
141+
<-start
142+
sandboxes.MarkStopped(t.Context(), sbx)
143+
}()
144+
145+
close(start)
146+
wg.Wait()
147+
148+
require.Empty(t, sandboxes.LifecycleItems())
149+
}
150+
}
151+
152+
func testMapSandbox(t *testing.T, lifecycleID string) *Sandbox {
153+
t.Helper()
154+
155+
slot, err := network.NewSlot("test", 1, network.Config{}, network.NoopEgressProxy{})
156+
require.NoError(t, err)
157+
158+
return &Sandbox{
159+
LifecycleID: lifecycleID,
160+
Metadata: &Metadata{
161+
Config: NewConfig(Config{}),
162+
Runtime: RuntimeMetadata{SandboxID: "sandbox-1"},
163+
},
164+
Resources: &Resources{Slot: slot},
165+
}
166+
}

‎packages/orchestrator/pkg/sandbox/sandbox.go‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,8 @@ type Sandbox struct {
247247
files *storage.SandboxFiles
248248
cleanup *Cleanup
249249

250+
sandboxes *Map
251+
250252
featureFlags *featureflags.Client
251253

252254
process *fc.Process
@@ -500,10 +502,11 @@ func (f *Factory) CreateSandbox(
500502
Metadata: metadata,
501503
cgroupHandle: cgroupHandle,
502504

503-
Template: template,
504-
config: f.config,
505-
files: sandboxFiles,
506-
process: fcHandle,
505+
Template: template,
506+
config: f.config,
507+
files: sandboxFiles,
508+
process: fcHandle,
509+
sandboxes: f.Sandboxes,
507510

508511
cleanup: cleanup,
509512
featureFlags: f.featureFlags,
@@ -845,10 +848,11 @@ func (f *Factory) ResumeSandbox(
845848
Metadata: metadata,
846849
cgroupHandle: cgroupHandle,
847850

848-
Template: t,
849-
config: f.config,
850-
files: sandboxFiles,
851-
process: fcHandle,
851+
Template: t,
852+
config: f.config,
853+
files: sandboxFiles,
854+
process: fcHandle,
855+
sandboxes: f.Sandboxes,
852856

853857
cleanup: cleanup,
854858
featureFlags: f.featureFlags,
@@ -988,6 +992,10 @@ func (s *Sandbox) Wait(ctx context.Context) error {
988992

989993
func (s *Sandbox) Close(ctx context.Context) error {
990994
err := s.cleanup.Run(ctx)
995+
if s.sandboxes != nil {
996+
s.sandboxes.MarkStopped(context.WithoutCancel(ctx), s)
997+
}
998+
991999
if err != nil {
9921000
return fmt.Errorf("failed to cleanup sandbox: %w", err)
9931001
}

0 commit comments

Comments
 (0)