Handle graceful termination in microvm ateom

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"
@@ -51,7 +53,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 (
@@ -210,13 +214,34 @@ func do(ctx context.Context) error {
 	}()
 	slog.InfoContext(ctx, "atunnel egress serving", slog.String("address", *atunnelEgressListenAddress))
 
+	ateomService := NewService(*podUID, *chBinary, *kataConfig, *kataDebug, 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, 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)
@@ -252,13 +277,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
@@ -341,6 +427,7 @@ var _ ateompb.AteomServer = (*AteomService)(nil)
 // NewService creates a new AteomService.
 func NewService(podUID, chBinary, kataConfig string, kataDebug bool, 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,
@@ -435,3 +522,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

@@ -82,6 +82,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
 	}
@@ -345,6 +356,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

@@ -93,6 +93,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
@@ -132,6 +137,17 @@ const maxActorContainers = 25
 // "-ovl" suffix would let "x"'s workload id collide with the "x-ovl" carrier id.
 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 overlay lower's find-paths subdir), the
 // host OCI bundle rootfs that backs the RO lower, and its OCI spec. The writable
@@ -215,14 +231,24 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload
 		return nil, err
 	}
 
-	p := actorBootParams{
-		actorRef:     resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()},
-		actorUID:     req.GetActorUid(),
-		templateNS:   req.GetActorTemplateNamespace(),
-		templateName: req.GetActorTemplateName(),
-		containers:   req.GetSpec().GetContainers(),
-		assetPaths:   req.GetRuntimeAssetPaths(),
+	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()
+
+	p := actorBootParams{
+		actorRef:      resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()},
+		actorUID:      req.GetActorUid(),
+		templateNS:    req.GetActorTemplateNamespace(),
+		templateName:  req.GetActorTemplateName(),
+		containers:    req.GetSpec().GetContainers(),
+		assetPaths:    req.GetRuntimeAssetPaths(),
 		egressGateway: req.GetEgressGateway(),
 	}
 
@@ -499,7 +525,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
 	}