diff --git a/spindle/engines/microvm/engine.go b/spindle/engines/microvm/engine.go index 9c599a3e..91061656 100644 --- a/spindle/engines/microvm/engine.go +++ b/spindle/engines/microvm/engine.go @@ -58,6 +58,7 @@ type Engine struct { cgroupParent *CgroupParent budget Resources maxWorkflow Resources + cids *cidRegistry cleanupMu sync.Mutex cleanup map[string][]cleanupFunc @@ -103,6 +104,7 @@ func New(ctx context.Context, cfg *config.Config, d *db.DB, quotaStore quota.Res cgroupParent: cgroupParent, budget: budget, maxWorkflow: max, + cids: newCIDRegistry(), cleanup: make(map[string][]cleanupFunc), debug: make(map[string]debugTarget), } @@ -318,10 +320,23 @@ func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *m return fmt.Errorf("workflow state is not initialized") } - cid, err := AllocateCID() + cid, err := e.cids.allocate() if err != nil { return err } + state.CID = cid + setupDone := false + defer func() { + if setupDone { + return + } + if detail := VMCrashLog(state.VM); detail != "" { + l.Error("microVM setup failed", "detail", detail) + } + if err := e.cleanupState(context.Background(), wid, state); err != nil { + l.Error("failed to cleanup failed setup", "error", err) + } + }() agent, err := e.ensureAgentHub() if err != nil { return err @@ -342,19 +357,6 @@ func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *m } state.WorkDir = workDir - setupDone := false - defer func() { - if setupDone { - return - } - if detail := VMCrashLog(state.VM); detail != "" { - l.Error("microVM setup failed", "detail", detail) - } - if err := e.cleanupState(context.Background(), wid, state); err != nil { - l.Error("failed to cleanup failed setup", "error", err) - } - }() - upstreams, err := BuildCacheUpstreams(e.cfg.NixCache.ReadURLs, state.CacheReadURLs) if err != nil { return err @@ -435,7 +437,6 @@ func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *m return err } state.Agent = agentSession - state.CID = cid wf.Data = state e.registerCleanup(wid, func(ctx context.Context) error { diff --git a/spindle/engines/microvm/vm.go b/spindle/engines/microvm/vm.go index d3902af1..85e77bbf 100644 --- a/spindle/engines/microvm/vm.go +++ b/spindle/engines/microvm/vm.go @@ -30,6 +30,7 @@ import ( const ( minGuestCID = 3 + maxGuestCID = minGuestCID + 59999 vmCrashLogTailBytes = 8192 ) @@ -41,6 +42,43 @@ func AllocateCID() (uint32, error) { return minGuestCID + binary.BigEndian.Uint32(data[:])%60000, nil } +type cidRegistry struct { + mu sync.Mutex + next uint32 + inUse map[uint32]struct{} +} + +func newCIDRegistry() *cidRegistry { + return &cidRegistry{next: minGuestCID, inUse: make(map[uint32]struct{})} +} + +func (r *cidRegistry) allocate() (uint32, error) { + r.mu.Lock() + defer r.mu.Unlock() + for range maxGuestCID - minGuestCID + 1 { + cid := r.next + r.next++ + if r.next > maxGuestCID { + r.next = minGuestCID + } + if _, exists := r.inUse[cid]; exists { + continue + } + r.inUse[cid] = struct{}{} + return cid, nil + } + return 0, errors.New("no guest CIDs available") +} + +func (r *cidRegistry) release(cid uint32) { + if cid == 0 { + return + } + r.mu.Lock() + delete(r.inUse, cid) + r.mu.Unlock() +} + func prepareWorkDir(workDir string) error { if workDir == "" { return fmt.Errorf("microvm work directory is required") @@ -213,6 +251,8 @@ func (e *Engine) cleanupState(ctx context.Context, wid models.WorkflowId, state err = errors.Join(err, closeIO(&state.UploadCache)) err = errors.Join(err, closeIO(&state.DNSProxy)) err = errors.Join(err, removeWorkDir(state)) + e.cids.release(state.CID) + state.CID = 0 return err } diff --git a/spindle/engines/microvm/vm_test.go b/spindle/engines/microvm/vm_test.go index fe9f490b..24530f9b 100644 --- a/spindle/engines/microvm/vm_test.go +++ b/spindle/engines/microvm/vm_test.go @@ -174,3 +174,29 @@ func TestCaptureResourceUsageMarksPartialVolumeSnapshotUnavailable(t *testing.T) t.Fatal("partial volume snapshot lost successful measurements") } } + +func TestCIDRegistryAvoidsCollisionsAndReusesReleasedCID(t *testing.T) { + registry := newCIDRegistry() + first, err := registry.allocate() + if err != nil { + t.Fatal(err) + } + second, err := registry.allocate() + if err != nil { + t.Fatal(err) + } + if first == second { + t.Fatalf("duplicate allocated CID %d", first) + } + registry.release(first) + registry.mu.Lock() + registry.next = first + registry.mu.Unlock() + reused, err := registry.allocate() + if err != nil { + t.Fatal(err) + } + if reused != first { + t.Fatalf("released CID: got %d, want %d", reused, first) + } +}