Handle graceful termination in microvm ateom (#893)

Part of https://github.com/agent-substrate/substrate/issues/23

Trap SIGTERM in the ateom, and forward SIGTERM to all application
containers.
Once all application containers exit, ateom will exit itself.

cmd/ateom-microvm/checkpoint.go

@@ -58,9 +58,18 @@ import (
 // Either way the guest is paused first, which is what makes the tar coherent: the
 // durable share is served write-through, so every completed guest write is already on
 // the host and no further ones can arrive.
+//
+// Allow checkpointing even if the pod is shutting down. This will allow actors
+// (or the harness) to suspend on shutdown.
 func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.CheckpointWorkloadRequest) (*ateompb.CheckpointWorkloadResponse, error) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	s.setActiveRPC(rpcCheckpointWorkload, cancel)
+	defer s.clearActiveRPC()
+
 	if err := s.deactivateActorNetworking(ctx); err != nil {
 		return nil, err
 	}

cmd/ateom-microvm/internal/kata/agentclient.go

@@ -203,6 +203,32 @@ func (a *AgentClient) AddARPNeighbors(ctx context.Context, neighbors []*agentpb.
 	return nil
 }
 
+// SignalProcess sends signal to a process in the guest. Targeting the container's
+// init process (ExecId == ContainerId) delivers the signal to the workload; an
+// empty execID delivers it to ALL processes in the container. Used during
+// graceful shutdown to propagate SIGTERM into the actor. Mirrors
+// grpc.AgentService/SignalProcess (returns google.protobuf.Empty).
+func (a *AgentClient) SignalProcess(ctx context.Context, containerID, execID string, signal uint32) error {
+	req := &agentpb.SignalProcessRequest{ContainerId: containerID, ExecId: execID, Signal: signal}
+	if err := a.client.Call(ctx, "grpc.AgentService", "SignalProcess", req, &emptypb.Empty{}); err != nil {
+		return fmt.Errorf("agent SignalProcess: %w", err)
+	}
+	return nil
+}
+
+// WaitProcess blocks until the identified guest process exits and returns its
+// exit status (mimics waitpid(2)). Used during graceful shutdown to confirm the
+// actor has stopped before ateom tears the VM down. Mirrors
+// grpc.AgentService/WaitProcess.
+func (a *AgentClient) WaitProcess(ctx context.Context, containerID, execID string) (int32, error) {
+	resp := &agentpb.WaitProcessResponse{}
+	req := &agentpb.WaitProcessRequest{ContainerId: containerID, ExecId: execID}
+	if err := a.client.Call(ctx, "grpc.AgentService", "WaitProcess", req, resp); err != nil {
+		return 0, fmt.Errorf("agent WaitProcess: %w", err)
+	}
+	return resp.GetStatus(), nil
+}
+
 // ReadStdout reads up to max bytes from the container process's stdout. It is a
 // unary RPC (NOT a server stream): each call returns whatever bytes the agent has
 // buffered (up to max), so callers loop until it returns an error — the agent

cmd/ateom-microvm/main.go

@@ -31,9 +31,11 @@ import (
 	"net"
 	"net/url"
 	"os"
+	"os/signal"
 	"strings"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
@@ -52,7 +54,9 @@ import (
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"golang.org/x/sys/unix"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/reflection"
+	"google.golang.org/grpc/status"
 )
 
 var (
@@ -238,13 +242,34 @@ func do(ctx context.Context) error {
 	}()
 	slog.InfoContext(ctx, "atunnel egress serving", slog.String("address", *atunnelEgressListenAddress))
 
+	ateomService := NewService(*podUID, *chBinary, *kataConfig, *kataDebug, *vmmMemReserve, interiorNetNS, actorLogger, atunnelIngress, atunnelEgress, atunnelEgressPort, *workerCredentialBundle, *podIdentityTrustBundle, *egressGatewayTrustBundle)
+
 	svr := grpc.NewServer(
 		grpc.StatsHandler(otelgrpc.NewServerHandler()),
 		grpc.UnaryInterceptor(ateinterceptors.InternalServerUnaryInterceptor),
 	)
-	ateompb.RegisterAteomServer(svr, NewService(*podUID, *chBinary, *kataConfig, *kataDebug, *vmmMemReserve, interiorNetNS, actorLogger, atunnelIngress, atunnelEgress, atunnelEgressPort, *workerCredentialBundle, *podIdentityTrustBundle, *egressGatewayTrustBundle))
+	ateompb.RegisterAteomServer(svr, ateomService)
 	reflection.Register(svr)
 
+	// Trap SIGTERM (sent by the kubelet at the start of the pod's termination grace
+	// period) and propagate it into the guest so the actor can save its state and
+	// exit cleanly before the grace period expires. The server deliberately keeps
+	// serving throughout gracefulShutdown: new workload RPCs are rejected with
+	// codes.Unavailable (see rejectIfDraining) while a suspend arriving mid-drain
+	// is still honored, which is what lets an actor checkpoint itself on eviction.
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGTERM)
+	go func() {
+		sig := <-sigCh
+		slog.InfoContext(ctx, "Received signal; beginning graceful shutdown", slog.String("signal", sig.String()))
+		// Use a fresh context: the do() context is torn down on return, but the
+		// shutdown must outlive it until the guest has stopped and the VM is down.
+		ateomService.gracefulShutdown(context.Background())
+		// Only now stop the server, which blocks until any in-flight RPC (notably a
+		// concurrent CheckpointWorkload) has completed, then unblocks svr.Serve below.
+		svr.GracefulStop()
+	}()
+
 	slog.InfoContext(ctx, "ateom-microvm serving", slog.String("socket", sockPath))
 	if err := svr.Serve(lis); err != nil {
 		return fmt.Errorf("while serving: %w", err)
@@ -280,13 +305,74 @@ func ensureSharedPropagation(ctx context.Context, path string) error {
 	return nil
 }
 
+const (
+	rpcRunWorkload        = "RunWorkload"
+	rpcRestoreWorkload    = "RestoreWorkload"
+	rpcCheckpointWorkload = "CheckpointWorkload"
+)
+
+// activeRPCInfo identifies the workload RPC currently holding lock, so graceful
+// shutdown can cancel a boot that would otherwise hold it for minutes.
+type activeRPCInfo struct {
+	name   string
+	cancel context.CancelFunc
+}
+
+// cancelableMutex is a mutex whose acquisition can be abandoned. sync.Mutex has
+// no bounded Lock, and graceful shutdown must not park forever behind an RPC
+// that is wedged: it needs to give up and get on with signaling the guest
+// while the pod's termination grace period still has room.
+type cancelableMutex struct {
+	ch chan struct{}
+}
+
+func newCancelableMutex() *cancelableMutex {
+	ch := make(chan struct{}, 1)
+	ch <- struct{}{}
+	return &cancelableMutex{ch: ch}
+}
+
+func (m *cancelableMutex) Lock() {
+	<-m.ch
+}
+
+func (m *cancelableMutex) Unlock() {
+	m.ch <- struct{}{}
+}
+
+// LockContext acquires the mutex, reporting false if ctx terminates first. On
+// false the mutex is NOT held and must not be unlocked.
+func (m *cancelableMutex) LockContext(ctx context.Context) bool {
+	select {
+	case <-m.ch:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
 // AteomService is the cloud-hypervisor implementation of ateompb.AteomServer.
 type AteomService struct {
 	ateompb.UnimplementedAteomServer
 
 	// lock serializes RPCs; like ateom-gvisor, the run/checkpoint/restore
 	// lifecycle is not safe to drive concurrently.
-	lock sync.Mutex
+	lock *cancelableMutex
+
+	// shuttingDown is set once SIGTERM has been received. While true, new workload
+	// RPCs are rejected with codes.Unavailable so the control plane reschedules.
+	//
+	// Atomic rather than lock-guarded: gracefulShutdown sets it before it tries to
+	// take lock, precisely so an RPC that arrives while it is still waiting is
+	// turned away instead of queueing behind it.
+	shuttingDown atomic.Bool
+
+	// activeRPC is the workload RPC in flight, tracked so gracefulShutdown can
+	// cancel a run or restore rather than wait out its boot. Guarded by
+	// activeRPCMu, which is separate from lock because the whole point is to reach
+	// it while lock is held by the RPC being cancelled.
+	activeRPCMu sync.Mutex
+	activeRPC   *activeRPCInfo
 
 	podUID     string
 	chBinary   string
@@ -374,6 +460,7 @@ var _ ateompb.AteomServer = (*AteomService)(nil)
 // NewService creates a new AteomService.
 func NewService(podUID, chBinary, kataConfig string, kataDebug bool, memReserveMiB int, interiorNetNS netns.NsHandle, actorLogger *actorlog.ActorLogger, atunnelIngress *atunnel.Server, atunnelEgress *atunnel.Egress, atunnelEgressPort uint16, workerCredentialBundlePath, podIdentityTrustBundlePath, egressGatewayTrustBundlePath string) *AteomService {
 	return &AteomService{
+		lock:                         newCancelableMutex(),
 		podUID:                       podUID,
 		chBinary:                     chBinary,
 		kataConfig:                   kataConfig,
@@ -469,3 +556,37 @@ func (s *AteomService) egressRedirectPort(redirectEgress bool) uint16 {
 	}
 	return s.atunnelEgressPort
 }
+
+// rejectIfDraining returns a codes.Unavailable error if ateom has begun graceful
+// shutdown, so the control plane reschedules the actor onto a live worker.
+func (s *AteomService) rejectIfDraining() error {
+	if s.shuttingDown.Load() {
+		return status.Error(codes.Unavailable, "worker draining: not accepting new workloads")
+	}
+	return nil
+}
+
+func (s *AteomService) setActiveRPC(name string, cancel context.CancelFunc) {
+	s.activeRPCMu.Lock()
+	defer s.activeRPCMu.Unlock()
+	s.activeRPC = &activeRPCInfo{name: name, cancel: cancel}
+}
+
+func (s *AteomService) clearActiveRPC() {
+	s.activeRPCMu.Lock()
+	defer s.activeRPCMu.Unlock()
+	s.activeRPC = nil
+}
+
+// cancelActiveRestoreOrRunRPC cancels an in-flight run or restore so it releases
+// lock instead of running its boot to completion. A checkpoint is deliberately
+// left alone: it is the one workload RPC worth finishing during a drain, since
+// it is what saves the actor's state.
+func (s *AteomService) cancelActiveRestoreOrRunRPC() {
+	s.activeRPCMu.Lock()
+	defer s.activeRPCMu.Unlock()
+	if s.activeRPC != nil && (s.activeRPC.name == rpcRestoreWorkload || s.activeRPC.name == rpcRunWorkload) {
+		slog.Info("Cancelling in-progress workload startup RPC due to graceful shutdown", slog.String("rpc", s.activeRPC.name))
+		s.activeRPC.cancel()
+	}
+}

cmd/ateom-microvm/restore.go

@@ -83,6 +83,17 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
+	if err := s.rejectIfDraining(); err != nil {
+		return nil, err
+	}
+
+	// Same as RunWorkload: a restore is a boot, and graceful shutdown cancels it
+	// rather than queueing behind it.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	s.setActiveRPC(rpcRestoreWorkload, cancel)
+	defer s.clearActiveRPC()
+
 	if err := s.deactivateActorNetworking(ctx); err != nil {
 		return nil, err
 	}
@@ -411,6 +422,7 @@ func (s *AteomService) restoreFullScope(ctx context.Context, p actorBootParams,
 		chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd,
 		apiSocket: apiSocket, baseID: srcID, restoreSourceDir: restoreDir,
 		snapshotIsSelfContained: memMode == ch.MemRestoreEager,
+		workloadIDs:             overlayWorkloadIDs(ctrs),
 	}
 
 	// Re-attach stdout/stderr forwarding for each container: the restored guest's

cmd/ateom-microvm/run.go

@@ -94,6 +94,11 @@ type runningActor struct {
 	// post-restore dial), which loses both log forwarding and guest stats for
 	// this activation.
 	guestAgent *kata.AgentClient
+
+	// workloadIDs are the kata overlay-workload container ids (overlayWorkloadID of
+	// each container name) running in the guest. The SIGTERM handler signals and
+	// waits on these to gracefully stop the actor before the VM is torn down.
+	workloadIDs []string
 }
 
 // baseIDFile is a tiny snapshot file (under the checkpoint/restore dir) holding
@@ -143,6 +148,17 @@ const maxActorContainers = 25
 // the _ovl containers, so log forwarding for those lineages keys on this.
 func overlayWorkloadID(name string) string { return name + "_ovl" }
 
+// overlayWorkloadIDs returns the overlay-workload container ids for the actor's
+// containers, in order. Recorded on runningActor so the SIGTERM handler knows
+// which guest workloads to signal and wait on.
+func overlayWorkloadIDs(ctrs []actorContainer) []string {
+	ids := make([]string, 0, len(ctrs))
+	for _, c := range ctrs {
+		ids = append(ids, overlayWorkloadID(c.name))
+	}
+	return ids
+}
+
 // actorContainer is one of the actor's containers prepared for the shared micro-VM:
 // its name (also the kata containerID + the merged rootfs's find-paths subdir), the
 // host OCI bundle rootfs that backs the overlay lower, and its OCI spec. The writable
@@ -222,6 +238,17 @@ func writeGuestResolvConf(rootfs string) error {
 func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkloadRequest) (resp *ateompb.RunWorkloadResponse, retErr error) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
+	if err := s.rejectIfDraining(); err != nil {
+		return nil, err
+	}
+
+	// Register the boot so a SIGTERM arriving mid-cold-boot cancels it rather than
+	// waiting out the whole thing holding lock.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	s.setActiveRPC(rpcRunWorkload, cancel)
+	defer s.clearActiveRPC()
+
 	if err := s.deactivateActorNetworking(ctx); err != nil {
 		return nil, err
 	}
@@ -544,7 +571,7 @@ func (s *AteomService) coldBootActor(ctx context.Context, p actorBootParams) (re
 		return fmt.Errorf("while waiting for container readyz: %w", err)
 	}
 
-	ra := &runningActor{chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd, apiSocket: apiSocket, baseID: actorUID, guestAgent: ac}
+	ra := &runningActor{chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd, apiSocket: apiSocket, baseID: actorUID, guestAgent: ac, workloadIDs: overlayWorkloadIDs(ctrs)}
 	if err := s.activateActorNetworking(p.actorRef.Atespace, p.actorRef.Name, egress); err != nil {
 		return err
 	}

cmd/ateom-microvm/shutdown.go

@@ -0,0 +1,262 @@
+//go:build linux
+
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Graceful termination. The kubelet sends SIGTERM at the start of the pod's
+// termination grace period; main.go traps it and calls gracefulShutdown, which
+// propagates the signal into the guest so the actor can save its state and exit
+// on its own before the pod goes away. Stopping the workloads is all this does —
+// the VM around them is left for the pod's own teardown to reap.
+
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"sync"
+	"syscall"
+	"time"
+
+	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/kata"
+)
+
+const (
+	// workloadGracePeriod is how long a guest workload gets to handle SIGTERM and
+	// exit on its own before ateom escalates to SIGKILL. Matches ateom-gvisor, and
+	// is deliberately shorter than the pod's own termination grace period so the
+	// escalation happens here rather than as a kubelet SIGKILL of ateom itself.
+	workloadGracePeriod = 1 * time.Minute
+
+	// workloadKillTimeout bounds the post-SIGKILL wait. The VM teardown that
+	// follows is what ultimately guarantees the workload is gone, so a wedged
+	// kata-agent must not hold shutdown open past this.
+	workloadKillTimeout = 5 * time.Second
+
+	// signalDeliveryTimeout bounds one SignalProcess round-trip. Delivering a signal
+	// is a local ttrpc call that returns in microseconds; if it has not come back by
+	// now the agent is not answering, and waiting longer will not change that.
+	//
+	// Separate from workloadGracePeriod on purpose. That is the allowance we owe the
+	// actor to save its state, not a budget for a stalled transport to spend, and
+	// sharing one number would both cheat the actor out of part of its grace period
+	// and silently lengthen how long a wedged agent can stall shutdown whenever the
+	// grace period is raised.
+	signalDeliveryTimeout = 10 * time.Second
+)
+
+// gracefulShutdown propagates SIGTERM into every running actor's guest and waits
+// for the workloads to exit, so the caller can exit cleanly. It holds lock only
+// long enough to snapshot the running actors, and releases it before any blocking
+// signaling or waiting, so it never holds it for the whole grace period and a
+// suspend can still land mid-drain.
+func (s *AteomService) gracefulShutdown(ctx context.Context) {
+	// Set this first, before contending for lock, so an RPC that arrives while we
+	// are still waiting is turned away rather than queued behind us.
+	s.shuttingDown.Store(true)
+
+	// Cancel an in-flight run or restore. Waiting for a cold boot to finish only to
+	// SIGTERM the guest it just produced is strictly worse than aborting it.
+	s.cancelActiveRestoreOrRunRPC()
+
+	// Wait for whatever still holds lock — a suspend, a resume — to finish, but
+	// only for the grace period. In the worst case that RPC burns nearly all of it
+	// and then fails, and the stop below spends another grace period on top; that
+	// is bounded well inside the pod's own termination grace period, and is the
+	// price of not truncating an RPC that may be saving the actor's state.
+	lockCtx, lockCancel := context.WithTimeout(ctx, workloadGracePeriod)
+	defer lockCancel()
+	if !s.lock.LockContext(lockCtx) {
+		slog.ErrorContext(ctx, "Failed to acquire lock during graceful shutdown; another RPC is still running")
+		return
+	}
+	// Snapshot by value rather than ranging over s.running directly: we drop the
+	// lock immediately below, and a suspend landing mid-drain deletes from the live
+	// map and writes through the *runningActor it finds there (teardownActor closes
+	// guestAgent and nils the field). Copying the map alone would not help — its
+	// values are pointers into that same mutable state.
+	targets := make([]drainTarget, 0, len(s.running))
+	for id, ra := range s.running {
+		if ra == nil {
+			continue
+		}
+		targets = append(targets, drainTarget{id: id, agent: ra.guestAgent, workloadIDs: ra.workloadIDs})
+	}
+
+	// Release lock so the service can answer new RPCs — notably a suspend arriving
+	// mid-drain — while the stop below waits out the grace period.
+	s.lock.Unlock()
+
+	if len(targets) == 0 {
+		slog.InfoContext(ctx, "No active actor sessions at shutdown; exiting cleanly")
+		return
+	}
+
+	for _, t := range targets {
+		gracefullyStopActor(ctx, t)
+	}
+	slog.InfoContext(ctx, "Shutting down")
+}
+
+// drainTarget is what gracefulShutdown needs from one runningActor, copied out
+// under lock so the drain below never dereferences the shared struct. workloadIDs
+// is set once before the actor is published to s.running and never mutated, so
+// sharing the backing array is safe; guestAgent is the field a concurrent
+// teardownActor writes, and is the reason this snapshot exists.
+//
+// The client the snapshot holds can still be closed under us by that teardown,
+// which is not a problem here: every call on a closed AgentClient fails fast with
+// ttrpc.ErrClosed instead of blocking or faulting, and an actor that has been torn
+// down has no workload left for this path to stop.
+type drainTarget struct {
+	id          string
+	agent       *kata.AgentClient
+	workloadIDs []string
+}
+
+// gracefullyStopActor signals the actor's guest workloads with SIGTERM and waits
+// out the grace period, escalating to SIGKILL.
+func gracefullyStopActor(ctx context.Context, t drainTarget) {
+	id := t.id
+
+	// Obtain a kata-agent client to signal the guest: reuse the log-forwarding
+	// connection if it's open, else dial a fresh one (best-effort). A dial we open
+	// here is closed here; the snapshotted client belongs to the log forwarder.
+	agent := t.agent
+	var dialed *kata.AgentClient
+	if agent == nil {
+		a, err := dialAgentRetry(ctx, kata.VsockSocketPath(id), 15*time.Second)
+		if err != nil {
+			// Without an agent there is no way to reach the guest's processes. They
+			// go down with the VM when the pod's containers are killed.
+			slog.WarnContext(ctx, "Could not dial kata-agent for graceful stop; leaving the workload to pod teardown", slog.String("id", id), slog.Any("err", err))
+			return
+		}
+		agent, dialed = a, a
+	}
+
+	// Stop the workloads concurrently. Each one is entitled to the full grace
+	// period, so stopping them in series would multiply it by the container
+	// count and overrun the pod's own termination grace period.
+	var wg sync.WaitGroup
+	for _, wid := range t.workloadIDs {
+		wg.Add(1)
+		go func(wid string) {
+			defer wg.Done()
+			if err := stopGuestWorkload(ctx, agent, id, wid); err != nil {
+				slog.WarnContext(ctx, "Failed to stop guest workload during shutdown", slog.String("id", id), slog.String("workload", wid), slog.Any("err", err))
+			}
+		}(wid)
+	}
+	wg.Wait()
+	if dialed != nil {
+		_ = dialed.Close()
+	}
+}
+
+// stopGuestWorkload stops one guest workload, wait out workloadGracePeriod, then
+// escalate to SIGKILL and wait a bounded time for the kill to land.
+func stopGuestWorkload(ctx context.Context, agent *kata.AgentClient, id, wid string) error {
+	// Propagate SIGTERM so the actor can save state and close connections.
+	// An actor that installed no handler terminates immediately.
+	slog.InfoContext(ctx, "Sending SIGTERM to guest workload", slog.String("id", id), slog.String("workload", wid))
+	if err := signalWorkload(ctx, agent, wid, syscall.SIGTERM); err != nil {
+		return fmt.Errorf("while propagating SIGTERM to workload %q: %w", wid, err)
+	}
+
+	// One WaitProcess (the guest's waitpid) feeds both waits below, so the SIGKILL
+	// path picks up the exit the SIGTERM path timed out on rather than issuing a
+	// second, competing wait. It runs on a context of its own so the grace-period
+	// deadline bounds only our side of the wait; canceling it on return is what
+	// unblocks the goroutine and stops the guest-side wait.
+	waitCtx, waitCancel := context.WithCancel(ctx)
+	defer waitCancel()
+	done := make(chan error, 1)
+	go func() {
+		_, err := agent.WaitProcess(waitCtx, wid, wid)
+		done <- err
+	}()
+
+	termCtx, termCancel := context.WithTimeout(ctx, workloadGracePeriod)
+	defer termCancel()
+	err := waitWorkloadStop(termCtx, done)
+	if err == nil {
+		slog.InfoContext(ctx, "Guest workload exited after SIGTERM", slog.String("id", id), slog.String("workload", wid))
+		return nil
+	}
+
+	// The wait failed at the RPC layer rather than running out of time. That says
+	// nothing about the workload: WaitProcess reports the exit code in its response
+	// with a nil error, which the branch above already handled, so an error here is
+	// a dead or wedged agent connection and not a process that exited badly.
+	//
+	// Liveness is therefore unknown, so report it rather than claiming the workload
+	// exited. Do not escalate: without a working agent there is no way to reach the
+	// process anyway, and ateom is already on its way out, so the container goes
+	// down with the pod.
+	if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
+		return fmt.Errorf("while waiting for workload %q to exit: %w", wid, err)
+	}
+
+	// The parent context, not our grace period, is what expired: stop here.
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	slog.WarnContext(ctx, "Grace period expired; killing guest workload", slog.String("id", id), slog.String("workload", wid), slog.Duration("grace", workloadGracePeriod))
+	if err := signalWorkload(ctx, agent, wid, syscall.SIGKILL); err != nil {
+		slog.WarnContext(ctx, "Failed to SIGKILL guest workload (it might have already exited)", slog.String("id", id), slog.String("workload", wid), slog.Any("err", err))
+	}
+
+	killCtx, killCancel := context.WithTimeout(ctx, workloadKillTimeout)
+	defer killCancel()
+	if err := waitWorkloadStop(killCtx, done); errors.Is(err, context.DeadlineExceeded) {
+		return fmt.Errorf("workload %q failed to exit even after SIGKILL: %w", wid, err)
+	} else if errors.Is(err, context.Canceled) {
+		return err
+	}
+
+	slog.InfoContext(ctx, "Guest workload exited after SIGKILL", slog.String("id", id), slog.String("workload", wid))
+	return nil
+}
+
+// signalWorkload delivers one signal to a guest workload's init process, bounded
+// by signalDeliveryTimeout. ateom sets ExecId equal to ContainerId, so passing wid
+// for both targets that init process.
+//
+// The bound is the point: the shutdown context has no deadline of its own, and
+// DialAgent clears the socket deadline once the vsock handshake is done, so ttrpc
+// has nothing but this ctx to give up on. A guest that is merely unresponsive
+// rather than gone — a paused VM, most plausibly, since a suspend is allowed to
+// land mid-drain — leaves the unix socket to CH perfectly healthy while the agent
+// never answers, and an unbounded call there would hang until the kubelet's
+// SIGKILL at the end of the pod's termination grace period.
+func signalWorkload(ctx context.Context, agent *kata.AgentClient, wid string, sig syscall.Signal) error {
+	sigCtx, cancel := context.WithTimeout(ctx, signalDeliveryTimeout)
+	defer cancel()
+	return agent.SignalProcess(sigCtx, wid, wid, uint32(sig))
+}
+
+// waitWorkloadStop waits for the workload's exit to land on done, or for ctx to
+// terminate first.
+func waitWorkloadStop(ctx context.Context, done <-chan error) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-done:
+		return err
+	}
+}

cmd/ateom-microvm/stats_test.go

@@ -163,9 +163,11 @@ func containerStats(usage, peak, inactiveFile, cpuNanos uint64) *agentpb.CgroupS
 }
 
 // newStatsService builds a service executing testActor with the given guest
-// containers published to GetWorkloadStats.
+// containers published to GetWorkloadStats. lock is constructed like NewService
+// does, since it is a pointer with no usable zero value and
+// TestGetWorkloadStatsDoesNotTakeLock holds it.
 func newStatsService(agent containerStatsReader, workloadIDs ...string) *AteomService {
-	s := &AteomService{}
+	s := &AteomService{lock: newCancelableMutex()}
 	s.activeActor.Store(&testActor)
 	s.guestStats.Store(&guestStatsTarget{actorUID: testActor.UID, agent: agent, workloadIDs: workloadIDs})
 	return s

internal/e2e/suites/demo/termination_test.go

@@ -37,10 +37,6 @@ import (
 // (CRASHED). We assert the control-plane state
 // machine rather than any in-actor state saving, which is the application's responsibility.
 func TestGracefulWorkerTermination(t *testing.T) {
-	if isMicroVMEnvironment() {
-		t.Skip("Skipping TestGracefulWorkerTermination for microVM environment")
-	}
-
 	nsObj := e2e.CreateNamespace(t)
 
 	ctx := context.Background()
@@ -149,13 +145,9 @@ func waitForWorkerRemoved(ctx context.Context, t *testing.T, clients *e2e.Client
 
 // TestGracefulWorkerTerminationTimeout exercises the case where the workload
 // container hangs (exceeds the 1-minute workloadGracePeriod) during SIGTERM.
-// ateom-gvisor is expected to SIGKILL the container, letting the control plane
-// mark the worker removed and the actor CRASHED.
+// The ateom is expected to SIGKILL the container, letting the control plane
+// mark the worker removed and the actor CRASHED. Runs against both runtimes.
 func TestGracefulWorkerTerminationTimeout(t *testing.T) {
-	if isMicroVMEnvironment() {
-		t.Skip("Skipping TestGracefulWorkerTerminationTimeout for microVM environment")
-	}
-
 	nsObj := e2e.CreateNamespace(t)
 
 	ctx := context.Background()
@@ -243,10 +235,6 @@ func TestGracefulWorkerTerminationTimeout(t *testing.T) {
 // deleted (evicted), and while the container is in its SIGTERM shutdown phase,
 // we initiate a suspend. Suspend should succeed.
 func TestGracefulWorkerTerminationSuspend(t *testing.T) {
-	if isMicroVMEnvironment() {
-		t.Skip("Skipping TestGracefulWorkerTerminationSuspend for microVM environment")
-	}
-
 	nsObj := e2e.CreateNamespace(t)
 
 	ctx := context.Background()