microVM: restore and snapshot no longer run under the actor's CPU limit

cmd/ateworker-microvm/checkpoint.go

@@ -172,6 +172,17 @@ func (s *AteWorkerService) CheckpointWorkload(ctx context.Context, req *ateworke
 	if pauseErr != nil {
 		return nil, fmt.Errorf("while pausing guest: %w", pauseErr)
 	}
+	// The paused guest runs none of the actor's work, so cloud-hypervisor writes
+	// the snapshot at the lifecycle CPU limit. A failed checkpoint gets the
+	// actor's own limit back.
+	if ra != nil {
+		raiseActorCPU(ctx, actorUID, ra.milliCPU)
+		defer func() {
+			if err != nil {
+				err = errors.Join(err, capActorCPU(actorUID, ra.milliCPU))
+			}
+		}()
+	}
 
 	checkpointDir := actorDirs.GetCheckpointDir()
 	// Start from a clean dir so CH's snapshot files are the only contents.
@@ -468,6 +479,7 @@ func (s *AteWorkerService) teardownActor(ctx context.Context, id string, actorDi
 
 		// Kill the CH process ateworker launched.
 		if ra.chProc != nil {
+			raiseActorCPU(ctx, id, ra.milliCPU)
 			ra.chProc.Stop()
 		}
 		// Kill the virtiofsd (after CH, its only client).

cmd/ateworker-microvm/hosted.go

@@ -119,6 +119,32 @@ func (s *AteWorkerService) cleanupSandboxState(ctx context.Context, actorUID str
 	kata.CleanupSandboxState(ctx, actorUID)
 }
 
+// openActorLeaf opens the leaf the actor's virtiofsd and cloud-hypervisor run
+// in, at the actor's lifecycle CPU limit (see cgroup.LifecycleMilliCPU) until
+// capActorCPU.
+func (s *AteWorkerService) openActorLeaf(p actorBootParams) (*cgroup.ActorLeaf, error) {
+	return cgroup.OpenActorLeaf(p.actorUID, cgroup.LifecycleMilliCPU(p.size.MilliCPU), s.leafMemoryHigh(p.size.MemoryBytes))
+}
+
+// capActorCPU puts the actor's own CPU limit on its leaf. The worker does this
+// before the wakeup probe: from then on the guest runs the actor's own work, and
+// the actor decides when its probe passes.
+func capActorCPU(actorUID string, milliCPU int64) error {
+	if err := cgroup.SetActorLeafCPU(actorUID, milliCPU); err != nil {
+		return fmt.Errorf("while applying the actor's CPU limit: %w", err)
+	}
+	return nil
+}
+
+// raiseActorCPU puts the lifecycle CPU limit back on the actor's leaf, before
+// the worker snapshots or stops a guest that runs none of the actor's work.
+// cloud-hypervisor frees guest RAM on its way out, charged to the leaf.
+func raiseActorCPU(ctx context.Context, actorUID string, milliCPU int64) {
+	if err := cgroup.SetActorLeafCPU(actorUID, cgroup.LifecycleMilliCPU(milliCPU)); err != nil {
+		slog.WarnContext(ctx, "Working under the actor's CPU limit", slog.String("id", actorUID), slog.Any("err", err))
+	}
+}
+
 // unhostActor removes an actor, its network, and its cgroup. Repeated calls are
 // safe. The caller must stop the VM first.
 func (s *AteWorkerService) unhostActor(ctx context.Context, actorUID string) error {

cmd/ateworker-microvm/restore.go

@@ -33,7 +33,6 @@ import (
 	"go.ate.dev/substrate/internal/workerstats"
 
 	"go.ate.dev/substrate/internal/ateworker/actornet"
-	"go.ate.dev/substrate/internal/ateworker/cgroup"
 
 	"go.ate.dev/substrate/cmd/ateworker-microvm/internal/ch"
 	"go.ate.dev/substrate/cmd/ateworker-microvm/internal/kata"
@@ -343,7 +342,7 @@ func (s *AteWorkerService) restoreMemoryFidelity(ctx context.Context, p actorBoo
 		return untarErr
 	}
 	tUpper := time.Now()
-	leaf, err := cgroup.OpenActorLeaf(actorUID, p.size.MilliCPU, s.leafMemoryHigh(p.size.MemoryBytes))
+	leaf, err := s.openActorLeaf(p)
 	if err != nil {
 		return err
 	}
@@ -401,6 +400,7 @@ func (s *AteWorkerService) restoreMemoryFidelity(ctx context.Context, p actorBoo
 	}
 	defer func() {
 		if retErr != nil {
+			raiseActorCPU(ctx, actorUID, p.size.MilliCPU)
 			chProc.Stop()
 		}
 	}()
@@ -480,6 +480,9 @@ func (s *AteWorkerService) restoreMemoryFidelity(ctx context.Context, p actorBoo
 	}
 	tReseeded := time.Now()
 
+	if err := capActorCPU(actorUID, p.size.MilliCPU); err != nil {
+		return err
+	}
 	// Block until every wakeup-probe-enabled container reports 200.
 	if err := wakeupprobe.WaitAll(ctx, containers, actornet.ActorVethIP, wakeupprobe.DialFunc(s.sandboxDialer(actorUID))); err != nil {
 		return fmt.Errorf("while waiting for container wakeup probe: %w", err)
@@ -524,7 +527,7 @@ func (s *AteWorkerService) restoreMemoryFidelity(ctx context.Context, p actorBoo
 
 	ra := &runningActor{
 		chProc: chProc, vfsdProc: vfsdProc,
-		apiSocket: apiSocket, memoryBytes: p.size.MemoryBytes, baseID: srcID, restoreSourceDir: restoreDir,
+		apiSocket: apiSocket, memoryBytes: p.size.MemoryBytes, milliCPU: p.size.MilliCPU, baseID: srcID, restoreSourceDir: restoreDir,
 		preserveRestoreSource:   preserveRestoreDir,
 		snapshotIsSelfContained: memMode == ch.MemRestoreEager,
 		// Signaling an id the agent does not know fails the whole graceful

cmd/ateworker-microvm/run.go

@@ -75,6 +75,8 @@ type runningActor struct {
 	stopping atomic.Bool
 	// memoryBytes is the actor's memory limit, what it books on the worker.
 	memoryBytes int64
+	// milliCPU is the actor's CPU limit.
+	milliCPU int64
 	// memory kills the actor's cgroup once it is stuck at its memory.high.
 	memory cgroup.MemoryGuard
 
@@ -495,7 +497,7 @@ func (s *AteWorkerService) coldBootActor(ctx context.Context, p actorBootParams)
 	// host upper, mounted into the shared dir) + durable-dir and CSI volumes (if any),
 	// and start the ONE virtiofsd that serves them all. CH connects to it at vm.create
 	// and demand-pages for the actor's lifetime, so ateworker owns the process (killed in teardownActor).
-	leaf, err := cgroup.OpenActorLeaf(actorUID, p.size.MilliCPU, s.leafMemoryHigh(p.size.MemoryBytes))
+	leaf, err := s.openActorLeaf(p)
 	if err != nil {
 		return err
 	}
@@ -524,6 +526,7 @@ func (s *AteWorkerService) coldBootActor(ctx context.Context, p actorBootParams)
 	}
 	defer func() {
 		if retErr != nil {
+			raiseActorCPU(ctx, actorUID, p.size.MilliCPU)
 			chProc.Stop()
 		}
 	}()
@@ -596,6 +599,9 @@ func (s *AteWorkerService) coldBootActor(ctx context.Context, p actorBootParams)
 	}
 	tContainers := time.Now()
 
+	if err := capActorCPU(actorUID, p.size.MilliCPU); err != nil {
+		return err
+	}
 	// Block until every wakeup-probe-enabled container reports 200.
 	if err := wakeupprobe.WaitAll(ctx, containers, actornet.ActorVethIP, wakeupprobe.DialFunc(s.sandboxDialer(actorUID))); err != nil {
 		return fmt.Errorf("while waiting for container wakeup probe: %w", err)
@@ -610,7 +616,7 @@ func (s *AteWorkerService) coldBootActor(ctx context.Context, p actorBootParams)
 		slog.Duration("wakeup_probe", time.Since(tContainers)),
 		slog.Duration("since_boot", time.Since(tBooted)))
 
-	ra := &runningActor{chProc: chProc, vfsdProc: vfsdProc, apiSocket: apiSocket, memoryBytes: p.size.MemoryBytes, baseID: actorUID, guestAgent: ac, workloadIDs: workloadIDs(ctrs)}
+	ra := &runningActor{chProc: chProc, vfsdProc: vfsdProc, apiSocket: apiSocket, memoryBytes: p.size.MemoryBytes, milliCPU: p.size.MilliCPU, baseID: actorUID, guestAgent: ac, workloadIDs: workloadIDs(ctrs)}
 	if err := s.tunnel.ActivateIngress(p.attribution(), s.sandboxDialer(p.actorUID)); err != nil {
 		return err
 	}

docs/reference/sandbox-classes.md

@@ -29,7 +29,7 @@ An [`ActorTemplate`](api-guide.md#2-actortemplate-the-workload-blueprint) select
 An actor's CPU size comes from `resources.limits.cpuMillicores` on its `ActorTemplate`:
 
 - **gVisor (`gvisor`):** `ateworker-gvisor` sets the sandbox's cgroup v2 `cpu.max` quota from `cpuMillicores` and passes `--cpu-num-from-quota` to `runsc`. `runsc` rounds the quota up to whole CPUs, with a minimum of 2, and never reports more CPUs than the sandbox can use. An actor with `cpuMillicores: 500` or `1000` sees 2 CPUs, and the host kernel throttles it to its quota.
-- **MicroVM (`microvm`):** `ateworker-microvm` sets Cloud Hypervisor `BootVcpus` and `MaxVcpus` to `ceil(cpuMillicores / 1000)` with a minimum of 1. It also caps the VMM and `virtiofsd` at `cpuMillicores` with a host cgroup v2 `cpu.max`. An actor with `cpuMillicores: 500` boots with 1 vCPU and gets half a CPU of host time.
+- **MicroVM (`microvm`):** `ateworker-microvm` sets Cloud Hypervisor `BootVcpus` and `MaxVcpus` to `ceil(cpuMillicores / 1000)` with a minimum of 1. It also caps the VMM and `virtiofsd` at `cpuMillicores` with a host cgroup v2 `cpu.max`. An actor with `cpuMillicores: 500` boots with 1 vCPU and gets half a CPU of host time. While the worker boots, restores, snapshots or stops the actor, the cap is four times `cpuMillicores`, and at least one CPU, like gVisor's default `--sandbox-lifecycle-cpu-policy`. The actor's own cap applies from its wakeup probe on.
 
 ## Memory Sizing
 

internal/ateworker/cgroup/actor.go

@@ -43,6 +43,17 @@ type ActorLeaf struct {
 	dir *os.File
 }
 
+// LifecycleMilliCPU is the CPU limit for an actor's leaf while the worker
+// boots, restores, snapshots or tears down the actor: four times milliCPU, and
+// at least one CPU. It matches gVisor's default --sandbox-lifecycle-cpu-policy.
+// An unset milliCPU stays unset.
+func LifecycleMilliCPU(milliCPU int64) int64 {
+	if milliCPU <= 0 {
+		return 0
+	}
+	return max(4*milliCPU, 1000)
+}
+
 // OpenActorLeaf creates the actor's leaf under the delegated scope, or reuses a
 // leftover one, caps it at milliCPU and sets memory.high to memoryHigh bytes
 // when those are set. Close it once its processes are started.
@@ -77,6 +88,20 @@ func openActorLeaf(root, actorUID string, milliCPU, memoryHigh int64) (*ActorLea
 	return &ActorLeaf{dir: dir}, nil
 }
 
+// SetActorLeafCPU caps the actor's leaf at milliCPU, or lifts the cap when
+// milliCPU is unset. A node that did not delegate the cpu controller has
+// nothing written.
+func SetActorLeafCPU(actorUID string, milliCPU int64) error {
+	path, err := actorLeafPath(Root, actorUID)
+	if err != nil {
+		return err
+	}
+	if err := writeExisting(filepath.Join(path, "cpu.max"), cpuMax(milliCPU)); err != nil && !errors.Is(err, fs.ErrNotExist) {
+		return fmt.Errorf("while limiting cgroup %q: %w", path, err)
+	}
+	return nil
+}
+
 // writeExisting writes a cgroup interface file without creating it: cgroupfs
 // refuses creation with EACCES, which would hide a file that is simply absent.
 func writeExisting(path, value string) error {

internal/ateworker/cgroup/actor_test.go

@@ -39,6 +39,20 @@ func TestCPUMax(t *testing.T) {
 	}
 }
 
+func TestLifecycleMilliCPU(t *testing.T) {
+	for _, tc := range []struct{ milli, want int64 }{
+		{0, 0},
+		{100, 1000},
+		{250, 1000},
+		{500, 2000},
+		{2000, 8000},
+	} {
+		if got := LifecycleMilliCPU(tc.milli); got != tc.want {
+			t.Errorf("LifecycleMilliCPU(%d) = %d, want %d", tc.milli, got, tc.want)
+		}
+	}
+}
+
 // A leaf left by an earlier incarnation is reused, and its old cap replaced.
 func TestOpenActorLeafReusesAndResetsALeftover(t *testing.T) {
 	root := t.TempDir()