diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index 752acf2d3..7e4d40c74 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -24,9 +24,12 @@ import ( "net" "net/url" "os" + "os/signal" "sort" "strings" "sync" + "syscall" + "time" "cloud.google.com/go/compute/metadata" "github.com/agent-substrate/substrate/internal/actorlog" @@ -47,7 +50,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 ( @@ -71,6 +76,8 @@ var ( // ingress to. const actorHTTPUpstream = "http://" + ateomnet.ActorVethIP + ":80" +const workloadGracePeriod = 1 * time.Minute + func main() { pflag.Parse() if *showVersion { @@ -173,6 +180,21 @@ func do(ctx context.Context) error { 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 sandbox so the actor can save its + // state and exit cleanly before the grace period expires. + 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 sandbox has stopped. + ateomService.gracefulShutdown(context.Background()) + // Stop the server gracefully. This blocks until all in-flight RPCs have completed. + svr.GracefulStop() + }() + if err := svr.Serve(lis); err != nil { slog.ErrorContext(ctx, "Failed to serve", slog.Any("err", err)) os.Exit(1) @@ -225,6 +247,14 @@ func runAtunnel(ctx context.Context, upstream *url.URL) (*atunnel.Server, *atunn return atunnelServer, atunnelEgress, atunnelEgressPort, nil } +// workloadSession captures the in-memory metadata for the workload currently running +// in the sandbox, so the SIGTERM handler knows which containers to signal and +// wait on during graceful shutdown. The sandbox runs one workload at a time. +type workloadSession struct { + rcmd *runsc + containers []string +} + // AteomService is a service for shepherding single microvm. type AteomService struct { ateompb.UnimplementedAteomServer @@ -233,6 +263,14 @@ type AteomService struct { // subcommands are probably not safe to call concurrently. lock sync.Mutex + // shuttingDown is set once SIGTERM has been received. While true, new + // workload RPCs are rejected with codes.Unavailable. Guarded by lock. + shuttingDown bool + + // activeSession tracks the currently running workload (nil when idle). Set by + // RunWorkload/RestoreWorkload, cleared by CheckpointWorkload. Guarded by lock. + activeSession *workloadSession + interiorNetNS netns.NsHandle actorLogger *actorlog.ActorLogger atunnel *atunnel.Server @@ -259,6 +297,123 @@ func NewService(interiorNetNS netns.NsHandle, actorLogger *actorlog.ActorLogger, } } +// 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 { + return status.Error(codes.Unavailable, "worker draining: not accepting new workloads") + } + return nil +} + +// gracefulShutdown propagates SIGTERM into the sandbox and waits for the application's +// containers to exit. +func (s *AteomService) gracefulShutdown(ctx context.Context) { + s.lock.Lock() + s.shuttingDown = true + session := s.activeSession + // Release the lock so that AteomService and respond to new RPCs. + s.lock.Unlock() + + if session == nil { + slog.InfoContext(ctx, "No active workload at shutdown; exiting cleanly") + return + } + + var wg sync.WaitGroup + for _, name := range session.containers { + wg.Add(1) + go func(containerName string) { + defer wg.Done() + if err := s.killContainer(ctx, session, containerName); err != nil { + slog.WarnContext(ctx, "Failed to kill container during shutdown", slog.String("container", containerName), slog.Any("err", err)) + } + }(name) + } + wg.Wait() + + slog.InfoContext(ctx, "All application containers stopped; shutting down cleanly") +} + +// killContainer stops a container by sending SIGTERM, waiting for the grace period, +// and escalating to SIGKILL if necessary. +func (s *AteomService) killContainer(ctx context.Context, session *workloadSession, name string) error { + // Propagate SIGTERM to the application container so it can save state and close connections. + // If the actor installed no SIGTERM handler it terminates immediately. + slog.InfoContext(ctx, "Sending SIGTERM to container", slog.String("container", name)) + if err := session.rcmd.cmdKill(ctx, name, "SIGTERM"); err != nil { + slog.ErrorContext(ctx, "Failed to propagate SIGTERM to container", slog.String("container", name), slog.Any("err", err)) + return fmt.Errorf("failed to propagate SIGTERM to container %q: %w", name, err) + } + + done := make(chan error, 1) + go func() { + done <- session.rcmd.cmdWait(ctx, name) + }() + + sigTermCtx, sigTermCtxCancel := context.WithTimeout(ctx, workloadGracePeriod) + defer sigTermCtxCancel() + + err := waitContainerStop(sigTermCtx, done) + if err == nil { + return nil + } + + // If the wait completed because the container actually exited (but with an error), + // we shouldn't send SIGKILL. + if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) { + slog.InfoContext(ctx, "Container exited with error status", slog.String("container", name), slog.Any("err", err)) + return nil + } + + // If the parent context was cancelled or exceeded return immediately + if ctx.Err() != nil { + return ctx.Err() + } + + // sigTermCtx timed out. Send SIGKILL. + slog.WarnContext(ctx, "Grace period expired; killing container", slog.String("container", name)) + if err := session.rcmd.cmdKill(ctx, name, "SIGKILL"); err != nil { + slog.WarnContext(ctx, "Failed to send SIGKILL to container (it might have already exited)", slog.String("container", name), slog.Any("err", err)) + } + + // Block until the killed container actually exits, but set a short timeout (e.g. 5 seconds) + // to avoid blocking indefinitely if gVisor is completely broken. + killCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + err = waitContainerStop(killCtx, done) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + return fmt.Errorf("container %q failed to exit even after SIGKILL: %w", name, err) + } + if errors.Is(err, context.Canceled) { + return err + } + } + + slog.InfoContext(ctx, "Container exited after SIGKILL", slog.String("container", name)) + return nil +} + +// waitContainerStop waits for container exit or context termination. +func waitContainerStop(ctx context.Context, done <-chan error) error { + select { + case <-ctx.Done(): + return ctx.Err() + case err := <-done: + return err + } +} + +func containerNames(containers []*ateompb.Container) []string { + names := make([]string, 0, len(containers)) + for _, c := range containers { + names = append(names, c.GetName()) + } + return names +} + func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkloadRequest) (resp *ateompb.RunWorkloadResponse, retErr error) { s.lock.Lock() defer s.lock.Unlock() @@ -266,6 +421,10 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload return nil, err } + if err := s.rejectIfDraining(); err != nil { + return nil, err + } + actorRef := resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()} s.actorLogger.EmitLifecycleLog("Actor starting", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) @@ -345,10 +504,13 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload } s.actorLogger.EmitLifecycleLog("Actor started", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) + s.activeSession = &workloadSession{rcmd: rcmd, containers: containerNames(req.GetSpec().GetContainers())} return &ateompb.RunWorkloadResponse{}, nil } +// 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() @@ -432,6 +594,7 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec } s.actorLogger.EmitLifecycleLog("Actor checkpointed", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) + s.activeSession = nil return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles}, nil } @@ -486,6 +649,10 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore return nil, err } + if err := s.rejectIfDraining(); err != nil { + return nil, err + } + actorRef := resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()} s.actorLogger.EmitLifecycleLog("Actor restoring", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) @@ -590,6 +757,7 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore } s.actorLogger.EmitLifecycleLog("Actor restored", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName()) + s.activeSession = &workloadSession{rcmd: rcmd, containers: containerNames(req.GetSpec().GetContainers())} return &ateompb.RestoreWorkloadResponse{}, nil } diff --git a/cmd/ateom-gvisor/runsc.go b/cmd/ateom-gvisor/runsc.go index 003bd861a..5b87d7072 100644 --- a/cmd/ateom-gvisor/runsc.go +++ b/cmd/ateom-gvisor/runsc.go @@ -19,12 +19,14 @@ package main import ( "context" "encoding/json" + "errors" "fmt" "io" "log/slog" "os" "os/exec" "path/filepath" + "syscall" specs "github.com/opencontainers/runtime-spec/specs-go" @@ -303,3 +305,70 @@ func (r *runsc) cmdState(ctx context.Context, containerName string) error { } return nil } + +// killArgs builds the argv for `runsc kill `. Factored out +// so the argument construction can be unit-tested without executing runsc. +func (r *runsc) killArgs(containerName, signal string) []string { + return []string{ + "-log-format", "json", + "--alsologtostderr", + "-root", ateompath.RunSCStateDir(r.actorUID), + "kill", + containerName, + signal, + } +} + +// cmdKill sends signal to the given container's process(es) inside the gVisor +// sandbox. Used during graceful shutdown to propagate SIGTERM to the actor. +func (r *runsc) cmdKill(ctx context.Context, containerName, signal string) error { + reapLock.RLock() + defer reapLock.RUnlock() + + slog.InfoContext(ctx, "About to run runsc kill", slog.String("container", containerName), slog.String("signal", signal)) + + cmd := exec.CommandContext(ctx, r.path, r.killArgs(containerName, signal)...) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + return fmt.Errorf("while running `runsc kill`: %w", err) + } + return nil +} + +// waitArgs builds the argv for `runsc wait `. Factored out so the +// argument construction can be unit-tested without executing runsc. +func (r *runsc) waitArgs(containerName string) []string { + return []string{ + "-log-format", "json", + "--alsologtostderr", + "-root", ateompath.RunSCStateDir(r.actorUID), + "wait", + containerName, + } +} + +// cmdWait blocks until the given container's process exits. Used during +// graceful shutdown to confirm the actor has stopped before ateom exits. +// +// We deliberately DO NOT acquire reapLock here. If we held reapLock.RLock() +// during this long wait, a pending background reaper write lock (reapLock.Lock()) +// would block, starving any subsequent read lock attempts (like CheckpointWorkload +// which needs to run runsc checkpoint). +func (r *runsc) cmdWait(ctx context.Context, containerName string) error { + slog.InfoContext(ctx, "About to run runsc wait", slog.String("container", containerName)) + + cmd := exec.CommandContext(ctx, r.path, r.waitArgs(containerName)...) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + // TODO: If the background child reaper collects the runsc wait process before + // cmd.Run's own wait finishes, it returns ECHILD. We can fix this by forking + // the reap.ReapChildren() call in main. + if errors.Is(err, syscall.ECHILD) { + return nil + } + return fmt.Errorf("while running `runsc wait`: %w", err) + } + return nil +} diff --git a/demos/counter/counter.go b/demos/counter/counter.go index 107bef68d..bfebcd555 100644 --- a/demos/counter/counter.go +++ b/demos/counter/counter.go @@ -27,19 +27,22 @@ import ( "net" "net/http" "os" + "os/signal" "path/filepath" "strconv" "sync" "sync/atomic" + "syscall" "time" "github.com/spf13/pflag" ) var ( - requestCount uint64 - ready atomic.Bool - fileMutex sync.Mutex + requestCount uint64 + ready atomic.Bool + fileMutex sync.Mutex + sigtermSleepDurationSecs = 15 ) func incrementFileCounter(filePath string) int { @@ -67,6 +70,16 @@ func main() { pflag.Parse() ctx := context.Background() + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGTERM) + go func() { + sig := <-sigCh + slog.InfoContext(ctx, "Received signal, waiting before exiting", slog.String("signal", sig.String()), slog.Int("sleep_secs", sigtermSleepDurationSecs)) + time.Sleep(time.Duration(sigtermSleepDurationSecs) * time.Second) + slog.InfoContext(ctx, "Exiting now") + os.Exit(0) + }() + slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil))) defaultMux := http.NewServeMux() @@ -115,6 +128,24 @@ func main() { w.Write([]byte("ok\n")) }) + defaultMux.HandleFunc("/set-sigterm-sleep", func(w http.ResponseWriter, r *http.Request) { + durationStr := r.URL.Query().Get("duration") + if durationStr == "" { + http.Error(w, "missing duration parameter", http.StatusBadRequest) + return + } + d, err := strconv.Atoi(durationStr) + if err != nil || d < 0 { + http.Error(w, "invalid duration parameter", http.StatusBadRequest) + return + } + sigtermSleepDurationSecs = d + response := fmt.Sprintf("SIGTERM sleep duration set to %d seconds\n", d) + slog.InfoContext(r.Context(), "Updated SIGTERM sleep duration", slog.Int("duration_secs", d)) + w.WriteHeader(http.StatusOK) + w.Write([]byte(response)) + }) + go func() { slog.InfoContext(ctx, "Starting counter server on port 80") if err := http.ListenAndServe(":80", defaultMux); err != nil { diff --git a/internal/e2e/suites/demo/demo_test.go b/internal/e2e/suites/demo/demo_test.go index dcb910688..8fd3ee72a 100644 --- a/internal/e2e/suites/demo/demo_test.go +++ b/internal/e2e/suites/demo/demo_test.go @@ -996,9 +996,12 @@ func createActorTemplateWithTwoDurableDirs(ctx context.Context, t *testing.T, cl } func waitForActorStatus(ctx context.Context, t *testing.T, clients *e2e.Clients, actorName string, expectedStatus ateapipb.Actor_Status) { + waitForActorStatusWithTimeout(ctx, t, clients, actorName, expectedStatus, 60*time.Second) +} + +func waitForActorStatusWithTimeout(ctx context.Context, t *testing.T, clients *e2e.Clients, actorName string, expectedStatus ateapipb.Actor_Status, timeout time.Duration) { t.Helper() t.Logf("Waiting for Actor %q to be %v...", actorName, expectedStatus) - timeout := 60 * time.Second deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { resp, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ @@ -1016,12 +1019,16 @@ func waitForActorStatus(ctx context.Context, t *testing.T, clients *e2e.Clients, } func callActor(t *testing.T, actorRef resources.ActorRef) (string, error) { + return callActorPath(t, actorRef, "POST", "/") +} + +func callActorPath(t *testing.T, actorRef resources.ActorRef, method, path string) (string, error) { t.Helper() deadline := time.Now().Add(30 * time.Second) var lastErr error for time.Now().Before(deadline) { - resp, err := callActorOnce(t, actorRef) + resp, err := callActorPathOnce(t, actorRef, method, path) if err == nil { return resp, nil } @@ -1032,7 +1039,7 @@ func callActor(t *testing.T, actorRef resources.ActorRef) (string, error) { return "", fmt.Errorf("timed out waiting for actor response: %w", lastErr) } -func callActorOnce(t *testing.T, actorRef resources.ActorRef) (string, error) { +func callActorPathOnce(t *testing.T, actorRef resources.ActorRef, method, path string) (string, error) { t.Helper() clients := e2e.GetClients() @@ -1099,7 +1106,7 @@ func callActorOnce(t *testing.T, actorRef resources.ActorRef) (string, error) { } localPort := forwardedPorts[0].Local - reqHttp, err := http.NewRequest("POST", fmt.Sprintf("http://127.0.0.1:%d", localPort), nil) + reqHttp, err := http.NewRequest(method, fmt.Sprintf("http://127.0.0.1:%d%s", localPort, path), nil) if err != nil { return "", fmt.Errorf("failed to create request: %w", err) } diff --git a/internal/e2e/suites/demo/termination_test.go b/internal/e2e/suites/demo/termination_test.go new file mode 100644 index 000000000..9c0cb74e5 --- /dev/null +++ b/internal/e2e/suites/demo/termination_test.go @@ -0,0 +1,339 @@ +// 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. + +package demo + +import ( + "context" + "testing" + "time" + + "github.com/agent-substrate/substrate/internal/e2e" + "github.com/agent-substrate/substrate/internal/resources" + "github.com/agent-substrate/substrate/pkg/api/v1alpha1" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +// TestGracefulWorkerTermination exercises the propagated-SIGTERM eviction flow +// end to end: an actor is scheduled onto a worker pod, that pod is deleted +// (simulating a Kubernetes eviction), and the control plane is expected to mark +// the worker DRAINING, then remove it and detach the actor once the pod is gone. +// +// The demo counter actor installs a SIGTERM handler that sleeps for 5 seconds +// before exiting, simulating a real-world workload that waits for graceful +// termination. The actor eventually lands in a terminal, non-RUNNING state +// (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() + clients := e2e.GetClients() + + _, _ = clients.SubstrateAPI.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: demoAtespace}}}) + + at, err := createActorTemplate(ctx, t, clients, nsObj, v1alpha1.SnapshotScopeFull, v1alpha1.SnapshotScopeFull, v1alpha1.ResumeSourceColdBoot) + if err != nil { + t.Fatalf("failed to initialize ActorTemplate: %v", err) + } + + actorID := "graceful-term-" + nsObj.Name + if _, err := clients.SubstrateAPI.CreateActor(ctx, &ateapipb.CreateActorRequest{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: demoAtespace, Name: actorID}, + ActorTemplateNamespace: nsObj.Name, + ActorTemplateName: at.Name, + }, + }); err != nil { + t.Fatalf("failed to create Actor: %v", err) + } + defer func() { + _, _ = clients.SubstrateAPI.DeleteActor(ctx, &ateapipb.DeleteActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + }() + + // Bring the actor up on a worker so it is bound to a pod. + if _, err := clients.SubstrateAPI.ResumeActor(ctx, &ateapipb.ResumeActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }); err != nil { + t.Fatalf("failed to resume Actor: %v", err) + } + waitForActorStatus(ctx, t, clients, actorID, ateapipb.Actor_STATUS_RUNNING) + + // Set the sigterm sleep interval to 15 seconds. + if _, err := callActorPath(t, resources.ActorRef{Atespace: demoAtespace, Name: actorID}, "GET", "/set-sigterm-sleep?duration=15"); err != nil { + t.Fatalf("failed to set sigterm sleep: %v", err) + } + + running, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get running Actor: %v", err) + } + podNS := running.GetWorkerAssignment().GetWorkerNamespace() + podName := running.GetWorkerAssignment().GetWorkerPod() + if podNS == "" || podName == "" { + t.Fatalf("running actor has no bound worker pod: ns=%q name=%q", podNS, podName) + } + t.Logf("Actor %q bound to worker pod %s/%s", actorID, podNS, podName) + + // Evict the worker pod. The kubelet sends SIGTERM to ateom, which propagates + // it into the sandbox; the control plane marks the worker DRAINING on the + // DeletionTimestamp watch event and cleans up when the pod is finally gone. + if err := clients.K8s.CoreV1().Pods(podNS).Delete(ctx, podName, metav1.DeleteOptions{}); err != nil { + t.Fatalf("failed to delete worker pod %s/%s: %v", podNS, podName, err) + } + + // The worker record must eventually be removed once the pod is gone. + if err := waitForWorkerRemoved(ctx, t, clients, podName, 60*time.Second); err != nil { + t.Fatalf("worker %s not removed after pod deletion: %v", podName, err) + } + + // Verify the actor lands in STATUS_CRASHED. + waitForActorStatus(ctx, t, clients, actorID, ateapipb.Actor_STATUS_CRASHED) + + // Verify the pod assignment was cleared. + actor, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get actor: %v", err) + } + if pod := actor.GetWorkerAssignment().GetWorkerPod(); pod != "" { + t.Errorf("actor still bound to worker pod %q, expected empty", pod) + } +} + +// waitForWorkerRemoved polls ListWorkers until the named worker is absent. +func waitForWorkerRemoved(ctx context.Context, t *testing.T, clients *e2e.Clients, podName string, timeout time.Duration) error { + t.Helper() + deadline := time.Now().Add(timeout) + for { + resp, err := clients.SubstrateAPI.ListWorkers(ctx, &ateapipb.ListWorkersRequest{}) + if err == nil { + found := false + for _, w := range resp.GetWorkers() { + if w.GetWorkerPod() == podName { + found = true + break + } + } + if !found { + return nil + } + } + if time.Now().After(deadline) { + return context.DeadlineExceeded + } + time.Sleep(time.Second) + } +} + +// 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. +func TestGracefulWorkerTerminationTimeout(t *testing.T) { + if isMicroVMEnvironment() { + t.Skip("Skipping TestGracefulWorkerTerminationTimeout for microVM environment") + } + + nsObj := e2e.CreateNamespace(t) + + ctx := context.Background() + clients := e2e.GetClients() + + _, _ = clients.SubstrateAPI.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: demoAtespace}}}) + + at, err := createActorTemplate(ctx, t, clients, nsObj, v1alpha1.SnapshotScopeFull, v1alpha1.SnapshotScopeFull, v1alpha1.ResumeSourceColdBoot) + if err != nil { + t.Fatalf("failed to initialize ActorTemplate: %v", err) + } + + actorID := "graceful-term-timeout-" + nsObj.Name + if _, err := clients.SubstrateAPI.CreateActor(ctx, &ateapipb.CreateActorRequest{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: demoAtespace, Name: actorID}, + ActorTemplateNamespace: nsObj.Name, + ActorTemplateName: at.Name, + }, + }); err != nil { + t.Fatalf("failed to create Actor: %v", err) + } + defer func() { + _, _ = clients.SubstrateAPI.DeleteActor(ctx, &ateapipb.DeleteActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + }() + + // Bring the actor up on a worker so it is bound to a pod. + if _, err := clients.SubstrateAPI.ResumeActor(ctx, &ateapipb.ResumeActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }); err != nil { + t.Fatalf("failed to resume Actor: %v", err) + } + waitForActorStatus(ctx, t, clients, actorID, ateapipb.Actor_STATUS_RUNNING) + + // Set the sigterm sleep interval to 90 seconds (longer than the 1-minute grace period). + if _, err := callActorPath(t, resources.ActorRef{Atespace: demoAtespace, Name: actorID}, "GET", "/set-sigterm-sleep?duration=90"); err != nil { + t.Fatalf("failed to set sigterm sleep: %v", err) + } + + running, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get running Actor: %v", err) + } + podNS := running.GetWorkerAssignment().GetWorkerNamespace() + podName := running.GetWorkerAssignment().GetWorkerPod() + if podNS == "" || podName == "" { + t.Fatalf("running actor has no bound worker pod: ns=%q name=%q", podNS, podName) + } + t.Logf("Actor %q bound to worker pod %s/%s", actorID, podNS, podName) + + // Evict the worker pod. The kubelet sends SIGTERM to ateom, which propagates + // it into the sandbox; the container hangs, triggering the 1-minute timeout, + // followed by SIGKILL by ateom. + if err := clients.K8s.CoreV1().Pods(podNS).Delete(ctx, podName, metav1.DeleteOptions{}); err != nil { + t.Fatalf("failed to delete worker pod %s/%s: %v", podNS, podName, err) + } + + // The worker record must eventually be removed once the pod is gone. + // Since there is a 1-minute timeout + up to 5s SIGKILL wait, we need a + // larger timeout (e.g. 120 seconds). + if err := waitForWorkerRemoved(ctx, t, clients, podName, 120*time.Second); err != nil { + t.Fatalf("worker %s not removed after pod deletion: %v", podName, err) + } + + // Verify the actor lands in STATUS_CRASHED. + waitForActorStatusWithTimeout(ctx, t, clients, actorID, ateapipb.Actor_STATUS_CRASHED, 120*time.Second) + + // Verify the pod assignment was cleared. + actor, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get actor: %v", err) + } + if pod := actor.GetWorkerAssignment().GetWorkerPod(); pod != "" { + t.Errorf("actor still bound to worker pod %q, expected empty", pod) + } +} + +// TestGracefulWorkerTerminationSuspend exercises the case where a worker pod is +// 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() + clients := e2e.GetClients() + + _, _ = clients.SubstrateAPI.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: demoAtespace}}}) + + at, err := createActorTemplate(ctx, t, clients, nsObj, v1alpha1.SnapshotScopeFull, v1alpha1.SnapshotScopeFull, v1alpha1.ResumeSourceColdBoot) + if err != nil { + t.Fatalf("failed to initialize ActorTemplate: %v", err) + } + + actorID := "graceful-term-suspend-" + nsObj.Name + if _, err := clients.SubstrateAPI.CreateActor(ctx, &ateapipb.CreateActorRequest{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: demoAtespace, Name: actorID}, + ActorTemplateNamespace: nsObj.Name, + ActorTemplateName: at.Name, + }, + }); err != nil { + t.Fatalf("failed to create Actor: %v", err) + } + defer func() { + _, _ = clients.SubstrateAPI.DeleteActor(ctx, &ateapipb.DeleteActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + }() + + // Bring the actor up on a worker so it is bound to a pod. + if _, err := clients.SubstrateAPI.ResumeActor(ctx, &ateapipb.ResumeActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }); err != nil { + t.Fatalf("failed to resume Actor: %v", err) + } + waitForActorStatus(ctx, t, clients, actorID, ateapipb.Actor_STATUS_RUNNING) + + // Set the sigterm sleep interval to 30 seconds. + if _, err := callActorPath(t, resources.ActorRef{Atespace: demoAtespace, Name: actorID}, "GET", "/set-sigterm-sleep?duration=30"); err != nil { + t.Fatalf("failed to set sigterm sleep: %v", err) + } + + running, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get running Actor: %v", err) + } + podNS := running.GetWorkerAssignment().GetWorkerNamespace() + podName := running.GetWorkerAssignment().GetWorkerPod() + if podNS == "" || podName == "" { + t.Fatalf("running actor has no bound worker pod: ns=%q name=%q", podNS, podName) + } + t.Logf("Actor %q bound to worker pod %s/%s", actorID, podNS, podName) + + // Evict the worker pod. The kubelet sends SIGTERM to ateom, which propagates + // it into the sandbox; the container hangs for 30s. + if err := clients.K8s.CoreV1().Pods(podNS).Delete(ctx, podName, metav1.DeleteOptions{}); err != nil { + t.Fatalf("failed to delete worker pod %s/%s: %v", podNS, podName, err) + } + + // Wait 2 seconds to make sure the SIGTERM was sent and the container is in its sleep phase. + time.Sleep(2 * time.Second) + + // Suspend the actor. + t.Logf("Suspending actor %q during termination", actorID) + if _, err := clients.SubstrateAPI.SuspendActor(ctx, &ateapipb.SuspendActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }); err != nil { + t.Fatalf("failed to suspend Actor: %v", err) + } + + // The worker record must eventually be removed once the pod is gone. + if err := waitForWorkerRemoved(ctx, t, clients, podName, 60*time.Second); err != nil { + t.Fatalf("worker %s not removed after pod deletion: %v", podName, err) + } + + // Verify the actor lands in STATUS_SUSPENDED + waitForActorStatus(ctx, t, clients, actorID, ateapipb.Actor_STATUS_SUSPENDED) + + // Verify the pod assignment was cleared. + actor, err := clients.SubstrateAPI.GetActor(ctx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: demoAtespace, Name: actorID}, + }) + if err != nil { + t.Fatalf("failed to get actor: %v", err) + } + if pod := actor.GetWorkerAssignment().GetWorkerPod(); pod != "" { + t.Errorf("actor still bound to worker pod %q, expected empty", pod) + } +}