Skip to content

Commit e940ec7

Browse files
djeebusclaude
authored andcommitted
fix(orchestrator): pin template cache entries backing live sandboxes
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> GitOrigin-RevId: ccdd930bdeac0c785c75cec8e4fffc6d9ceaf3a8
1 parent 0c21aa2 commit e940ec7

10 files changed

Lines changed: 2188 additions & 49 deletions

File tree

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

Lines changed: 709 additions & 37 deletions
Large diffs are not rendered by default.

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

Lines changed: 1012 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,19 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
290290
}
291291

292292
func (t *storageTemplate) Close(ctx context.Context) error {
293-
return closeTemplate(ctx, t)
293+
err := closeTemplate(ctx, t)
294+
295+
// closeTemplate only removes the files it holds handles for, which leaves the
296+
// metafile and the directory itself behind; nothing else reclaims them, since
297+
// the startup sweep covers DefaultCacheDir and these live under
298+
// TemplateCacheDir. The directory is private to this instance — CachePaths
299+
// mints a fresh identifier per template — so removing it cannot touch another
300+
// instance's files.
301+
if pathsErr := t.paths.Close(); pathsErr != nil {
302+
err = errors.Join(err, fmt.Errorf("failed to remove template cache dir: %w", pathsErr))
303+
}
304+
305+
return err
294306
}
295307

296308
func (t *storageTemplate) Files() storage.CachePaths {
@@ -340,6 +352,72 @@ func (t *storageTemplate) SchedulingMetadata(ctx context.Context) *orchestrator.
340352
return scheduling.FromHeaders(rh.Metadata.BuildId, mh, rh, 0)
341353
}
342354

355+
// headerFootprint reports the mapping entry count and approximate heap bytes
356+
// held by this template's resolved headers. It never blocks: the headers live
357+
// on the memfile/rootfs devices (the header holders stay unset for templates
358+
// loaded from storage — see SchedulingMetadata), and a device whose SetOnce has
359+
// not resolved yet is skipped. The gauge therefore undercounts still-fetching
360+
// templates rather than stalling the metrics callback on them.
361+
func (t *storageTemplate) headerFootprint() (entries int, bytes int) {
362+
// Deduplicated by the mapping's backing storage rather than by header
363+
// identity, because one allocation reaches this function by two routes. The
364+
// holders below usually resolve to the very headers the devices carry; and
365+
// once a snapshot's upload publishes, the device carries a CloneForUpload
366+
// while the holder still carries the source — two distinct *Header sharing
367+
// one Mapping, since the clone copies the struct and the copy shares its
368+
// slices. Counting that allocation twice would overstate the number this
369+
// gauge exists to make trustworthy, and it would do so as uploads land,
370+
// which reads like retention growth rather than a counting artifact.
371+
seen := make(map[*header.Header]struct{}, 4)
372+
counted := make([]header.Mapping, 0, 4)
373+
374+
add := func(h *header.Header) {
375+
if h == nil {
376+
return
377+
}
378+
if _, dup := seen[h]; dup {
379+
return
380+
}
381+
seen[h] = struct{}{}
382+
383+
for _, m := range counted {
384+
if m.SharesStorageWith(h.Mapping) {
385+
return
386+
}
387+
}
388+
counted = append(counted, h.Mapping)
389+
390+
entries += h.Mapping.Len()
391+
bytes += h.Mapping.ByteSize()
392+
}
393+
394+
if dev, err := t.memfile.Result(); err == nil && dev != nil {
395+
add(dev.Header())
396+
}
397+
398+
if dev, err := t.rootfs.Result(); err == nil && dev != nil {
399+
add(dev.Header())
400+
}
401+
402+
// The holders are not always the headers the devices ended up on, and the
403+
// difference is retained memory. A template built from a provisional memfile
404+
// header keeps that header alive in memfileHeader after SwapHeaderIfCurrent
405+
// has moved the device on to the deduped one, so a paused-and-deduped
406+
// template holds two distinct mappings while the device reports one. Count
407+
// every distinct header the template still references.
408+
for _, holder := range []*utils.SetOnce[*header.Header]{t.memfileHeader, t.rootfsHeader, t.durableMemfileHeader} {
409+
if holder == nil {
410+
continue
411+
}
412+
413+
if h, err := holder.Result(); err == nil {
414+
add(h)
415+
}
416+
}
417+
418+
return entries, bytes
419+
}
420+
343421
func (t *storageTemplate) Memfile(ctx context.Context) (block.ReadonlyDevice, error) {
344422
_, span := tracer.Start(ctx, "storage-template-memfile")
345423
defer span.End()

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

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ package template
44

55
import (
66
"testing"
7+
"time"
78

89
"github.com/google/uuid"
910
"github.com/stretchr/testify/assert"
@@ -80,3 +81,162 @@ func TestStorageTemplate_SchedulingMetadataSkipsPendingMemfile(t *testing.T) {
8081
assert.Empty(t, md.GetMemfileBaseBuildId())
8182
assert.Equal(t, rootfsBase.String(), md.GetRootfsBaseBuildId())
8283
}
84+
85+
// footprintHeader builds a header whose mapping has exactly n entries, using a
86+
// distinct build id per entry so nothing merges.
87+
func footprintHeader(t *testing.T, n int) *header.Header {
88+
t.Helper()
89+
90+
const blockSize = uint64(4096)
91+
92+
maps := make([]header.BuildMap, n)
93+
for i := range maps {
94+
maps[i] = header.BuildMap{
95+
Offset: uint64(i) * blockSize,
96+
Length: blockSize,
97+
BuildId: uuid.New(),
98+
}
99+
}
100+
101+
h, err := header.NewHeader(&header.Metadata{
102+
Version: 3,
103+
BlockSize: blockSize,
104+
Size: uint64(n) * blockSize,
105+
BuildId: uuid.New(),
106+
BaseBuildId: uuid.New(),
107+
}, maps)
108+
require.NoError(t, err)
109+
require.Equal(t, n, h.Mapping.Len())
110+
111+
return h
112+
}
113+
114+
// The gauge exists to size a bound, so it has to see every mapping the template
115+
// keeps alive. After a pause the provisional header stays referenced by
116+
// memfileHeader while the device has already moved on to the deduped one, so a
117+
// gauge that reads only the devices reports one mapping where two are resident
118+
// — understating exactly the footprint it is used to size.
119+
func TestStorageTemplate_HeaderFootprintCountsRetainedHolders(t *testing.T) {
120+
t.Parallel()
121+
122+
provisional := footprintHeader(t, 12)
123+
deduped := footprintHeader(t, 3)
124+
rootfsHdr := footprintHeader(t, 5)
125+
126+
memDev := blockmocks.NewMockReadonlyDevice(t)
127+
memDev.EXPECT().Header().Return(deduped)
128+
rootfsDev := blockmocks.NewMockReadonlyDevice(t)
129+
rootfsDev.EXPECT().Header().Return(rootfsHdr)
130+
131+
tmpl := &storageTemplate{
132+
memfile: utils.NewSetOnce[block.ReadonlyDevice](),
133+
rootfs: utils.NewSetOnce[block.ReadonlyDevice](),
134+
memfileHeader: resolvedHeader(provisional),
135+
rootfsHeader: resolvedHeader(rootfsHdr),
136+
durableMemfileHeader: resolvedHeader(deduped),
137+
}
138+
require.NoError(t, tmpl.memfile.SetValue(memDev))
139+
require.NoError(t, tmpl.rootfs.SetValue(rootfsDev))
140+
141+
entries, bytes := tmpl.headerFootprint()
142+
143+
wantEntries := provisional.Mapping.Len() + deduped.Mapping.Len() + rootfsHdr.Mapping.Len()
144+
wantBytes := provisional.Mapping.ByteSize() + deduped.Mapping.ByteSize() + rootfsHdr.Mapping.ByteSize()
145+
146+
assert.Equal(t, wantEntries, entries, "the provisional header is still resident and must be counted")
147+
assert.Equal(t, wantBytes, bytes)
148+
}
149+
150+
// The holders usually resolve to the same headers the devices carry. Counting
151+
// those twice would overstate the number the sizing decisions are made from.
152+
func TestStorageTemplate_HeaderFootprintCountsEachMappingOnce(t *testing.T) {
153+
t.Parallel()
154+
155+
memHdr := footprintHeader(t, 7)
156+
rootfsHdr := footprintHeader(t, 4)
157+
158+
memDev := blockmocks.NewMockReadonlyDevice(t)
159+
memDev.EXPECT().Header().Return(memHdr)
160+
rootfsDev := blockmocks.NewMockReadonlyDevice(t)
161+
rootfsDev.EXPECT().Header().Return(rootfsHdr)
162+
163+
tmpl := &storageTemplate{
164+
memfile: utils.NewSetOnce[block.ReadonlyDevice](),
165+
rootfs: utils.NewSetOnce[block.ReadonlyDevice](),
166+
memfileHeader: resolvedHeader(memHdr),
167+
rootfsHeader: resolvedHeader(rootfsHdr),
168+
}
169+
require.NoError(t, tmpl.memfile.SetValue(memDev))
170+
require.NoError(t, tmpl.rootfs.SetValue(rootfsDev))
171+
172+
entries, bytes := tmpl.headerFootprint()
173+
174+
assert.Equal(t, memHdr.Mapping.Len()+rootfsHdr.Mapping.Len(), entries)
175+
assert.Equal(t, memHdr.Mapping.ByteSize()+rootfsHdr.Mapping.ByteSize(), bytes)
176+
}
177+
178+
// A template that is still fetching must contribute nothing rather than block
179+
// the metrics collection goroutine on its unresolved futures.
180+
func TestStorageTemplate_HeaderFootprintSkipsUnresolved(t *testing.T) {
181+
t.Parallel()
182+
183+
tmpl := &storageTemplate{
184+
memfile: utils.NewSetOnce[block.ReadonlyDevice](),
185+
rootfs: utils.NewSetOnce[block.ReadonlyDevice](),
186+
memfileHeader: utils.NewSetOnce[*header.Header](),
187+
rootfsHeader: utils.NewSetOnce[*header.Header](),
188+
}
189+
190+
done := make(chan struct{})
191+
go func() {
192+
defer close(done)
193+
194+
entries, bytes := tmpl.headerFootprint()
195+
assert.Zero(t, entries)
196+
assert.Zero(t, bytes)
197+
}()
198+
199+
select {
200+
case <-done:
201+
case <-time.After(5 * time.Second):
202+
t.Fatal("headerFootprint blocked on an unresolved template")
203+
}
204+
}
205+
206+
// One allocation reaches headerFootprint by two routes once a snapshot's upload
207+
// publishes: CloneForUpload copies the Header struct, so the clone shares the
208+
// source's Mapping slices, and publish installs the clone on the device while
209+
// the holder keeps the source. Header identity does not see that, so the gauge
210+
// would step up as uploads land — reading like retention growth rather than a
211+
// counting artifact, on exactly the population a byte budget gets sized from.
212+
func TestStorageTemplate_HeaderFootprintCountsSharedMappingOnce(t *testing.T) {
213+
t.Parallel()
214+
215+
source := footprintHeader(t, 9)
216+
published := source.CloneForUpload(source.Metadata.Version + 1)
217+
require.NotSame(t, source, published, "the clone is a distinct header")
218+
require.True(t, source.Mapping.SharesStorageWith(published.Mapping),
219+
"the clone must share the source's mapping, or this test proves nothing")
220+
221+
rootfsHdr := footprintHeader(t, 4)
222+
223+
memDev := blockmocks.NewMockReadonlyDevice(t)
224+
memDev.EXPECT().Header().Return(published)
225+
rootfsDev := blockmocks.NewMockReadonlyDevice(t)
226+
rootfsDev.EXPECT().Header().Return(rootfsHdr)
227+
228+
tmpl := &storageTemplate{
229+
memfile: utils.NewSetOnce[block.ReadonlyDevice](),
230+
rootfs: utils.NewSetOnce[block.ReadonlyDevice](),
231+
memfileHeader: resolvedHeader(source),
232+
rootfsHeader: resolvedHeader(rootfsHdr),
233+
}
234+
require.NoError(t, tmpl.memfile.SetValue(memDev))
235+
require.NoError(t, tmpl.rootfs.SetValue(rootfsDev))
236+
237+
entries, bytes := tmpl.headerFootprint()
238+
239+
assert.Equal(t, source.Mapping.Len()+rootfsHdr.Mapping.Len(), entries,
240+
"the shared mapping must be counted once, not once per header")
241+
assert.Equal(t, source.Mapping.ByteSize()+rootfsHdr.Mapping.ByteSize(), bytes)
242+
}

‎packages/orchestrator/pkg/server/prefetch_harvest.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ type harvestInstance interface {
108108

109109
// harvestTemplates is the subset of the template cache the harvest uses.
110110
type harvestTemplates interface {
111-
GetTemplate(ctx context.Context, buildID string, isSnapshot, isBuilding bool, opts ...sbxtemplate.GetTemplateOpts) (sbxtemplate.Template, error)
111+
GetTemplatePinned(ctx context.Context, buildID string, isSnapshot, isBuilding bool, opts ...sbxtemplate.GetTemplateOpts) (sbxtemplate.Template, func(), error)
112112
UpdateMetadata(buildID string, meta metadata.Template) error
113113
}
114114

@@ -394,12 +394,18 @@ func (h *prefetchHarvester) resumeMapping(
394394
// pays no cold GCS/NFS fetch, only a local re-fault. isSnapshot=true mirrors
395395
// Checkpoint's resume. The pause artifact carries no Prefetch (SameVersion
396396
// dropped it), so no prefetcher runs and the trace is clean demand faults.
397-
tmpl, err := h.templates.GetTemplate(ctx, buildID, true, false,
397+
// Pinned: the throwaway below is a real Firecracker VM running off this
398+
// template's snapfile, so for its lifetime the template must not be
399+
// evictable — eviction Closes it, and Close removes the snapfile. Released
400+
// last, after the reap deferred below has torn the throwaway down.
401+
tmpl, releaseTemplate, err := h.templates.GetTemplatePinned(ctx, buildID, true, false,
398402
sbxtemplate.GetTemplateOpts{MaxSandboxLengthHours: sbx.Config.MaxSandboxLengthHours})
399403
if err != nil {
400404
return nil, harvestResumeFailed, fmt.Errorf("get template: %w", err)
401405
}
402406

407+
defer releaseTemplate()
408+
403409
// Throwaway identity: distinct SandboxID/ExecutionID from the (being-stopped)
404410
// original so it never collides in the sandbox map. ResumeSandbox registers
405411
// it in the factory's sandbox table (for network assignment and health), but

‎packages/orchestrator/pkg/server/prefetch_harvest_test.go‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,10 +83,21 @@ type fakeHarvestTemplates struct {
8383
updateErr error
8484
updated bool
8585
updatedMeta metadata.Template
86+
87+
// pins and releases must balance: a pin the harvest never releases keeps the
88+
// template resident for the life of the process.
89+
pins int
90+
releases int
8691
}
8792

88-
func (f *fakeHarvestTemplates) GetTemplate(context.Context, string, bool, bool, ...sbxtemplate.GetTemplateOpts) (sbxtemplate.Template, error) {
89-
return nil, f.getErr
93+
func (f *fakeHarvestTemplates) GetTemplatePinned(context.Context, string, bool, bool, ...sbxtemplate.GetTemplateOpts) (sbxtemplate.Template, func(), error) {
94+
if f.getErr != nil {
95+
return nil, func() {}, f.getErr
96+
}
97+
98+
f.pins++
99+
100+
return nil, func() { f.releases++ }, nil
90101
}
91102

92103
func (f *fakeHarvestTemplates) UpdateMetadata(_ string, meta metadata.Template) error {
@@ -204,6 +215,8 @@ func TestHarvestRun_ConsumePersistsAndReaps(t *testing.T) {
204215
require.NotNil(t, p.tmpls.updatedMeta.Prefetch, "persisted metadata must carry the mapping")
205216
require.Equal(t, 2, p.tmpls.updatedMeta.Prefetch.Memory.Count(), "the persisted mapping must hold the harvested blocks")
206217
require.True(t, p.uploadCalled, "remote metadata must be re-uploaded")
218+
require.Equal(t, 1, p.tmpls.pins, "the throwaway runs off the template, so it must be pinned")
219+
require.Equal(t, p.tmpls.pins, p.tmpls.releases, "the template pin must be released")
207220
}
208221

209222
// TestHarvestRun_ThrowawayIdentity checks the throwaway resume runtime: it is
@@ -312,6 +325,8 @@ func TestHarvestRun_CollectErrorStillReaps(t *testing.T) {
312325
require.True(t, p.inst.stopped && p.inst.closed, "throwaway must be reaped even when collection fails")
313326
require.True(t, p.released)
314327
require.False(t, p.tmpls.updated)
328+
require.Equal(t, p.tmpls.pins, p.tmpls.releases,
329+
"a pin leaked on the error path keeps the template resident for the life of the process")
315330
}
316331

317332
// TestHarvestRun_AcquireErrorSkips: if the start slot can't be acquired the

0 commit comments

Comments
 (0)