diff --git a/agent/internal/controller/controller.go b/agent/internal/controller/controller.go index ddd3dca1..a0dbef7f 100644 --- a/agent/internal/controller/controller.go +++ b/agent/internal/controller/controller.go @@ -13,6 +13,8 @@ import ( "errors" "fmt" "os" + "path/filepath" + "strconv" "strings" "sync" "syscall" @@ -65,10 +67,13 @@ type NodeController struct { log logr.Logger holderID string checkpointFn func(ctx context.Context, params CheckpointParams) error - restoreFn func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) - writeControlSentinelFn func(int, string) error + restoreFn func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) + writeControlSentinelFn func(int, string, []byte) error controlSentinelExistsFn func(int, string) (bool, error) + readControlSentinelFn func(int, string) ([]byte, error) + validateProcessStateFn func(string, int) error sendSignalFn func(logr.Logger, int, syscall.Signal, string) error + restoreMonitorInterval time.Duration restoreQueue workqueue.TypedDelayingInterface[client.ObjectKey] restorePodLister corev1listers.PodLister @@ -76,6 +81,7 @@ type NodeController struct { inFlightMu sync.Mutex handledRestores sync.Map + restoreMonitors sync.Map // contentIndexer is the PodSnapshotContent informer's indexer, indexed by source pod // (podRefIndex). The source-pod informer uses it to map a pod event back to its work order. @@ -152,6 +158,7 @@ const ( restorePodFinalizer = "snapshot/restore-protection" snapshotEventComponent = "snapshot" restoreSafetyRequeueInterval = 30 * time.Second + restoredProcessMonitorInterval = 1 * time.Second // snapshotContentResyncInterval re-drives every PodSnapshotContent work order so a // not-yet-Ready source pod is re-checked for quiesce without a busy loop. @@ -220,9 +227,12 @@ func newDefaultController( ), restoreFn: executor.Restore, - writeControlSentinelFn: snapshotruntime.WriteControlSentinel, + writeControlSentinelFn: snapshotruntime.WriteControlSentinelData, controlSentinelExistsFn: snapshotruntime.ControlSentinelExists, + readControlSentinelFn: snapshotruntime.ReadControlSentinel, + validateProcessStateFn: snapshotruntime.ValidateProcessState, sendSignalFn: snapshotruntime.SendSignalToPID, + restoreMonitorInterval: restoredProcessMonitorInterval, } w.checkpointFn = w.executorCheckpoint return w @@ -852,7 +862,7 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, artifa defer cancel() } - placeholderHostPID, err := op.executeRestore(restoreCtx) + result, err := op.executeRestore(restoreCtx) if err != nil { var cleanupErr *executor.RestoreCleanupError if !errors.As(err, &cleanupErr) { @@ -861,7 +871,7 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, artifa op.log.Error(cleanupErr, "Restore completed with cleanup errors") emitPodEvent(ctx, w.clientset, op.log, pod, snapshotEventComponent, corev1.EventTypeWarning, "RestoreCleanupFailed", cleanupErr.Error()) } - return op.completeRestore(ctx, placeholderHostPID) + return op.completeRestore(ctx, result) } // recoverCompletedRestore avoids replaying CRIU when the destination-scoped @@ -876,13 +886,23 @@ func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, if err != nil { return false, fmt.Errorf("resolve restore container before checking completion sentinel: %w", err) } - exists, err := op.controller.controlSentinelExistsFn(hostPID, snapshotv1alpha1.RestoreCompleteFile) + w := op.controller + exists, err := w.controlSentinelExistsFn(hostPID, snapshotv1alpha1.RestoreCompleteFile) if err != nil { return false, fmt.Errorf("check restore completion sentinel: %w", err) } if !exists { return false, nil } + if data, err := w.readControlSentinelFn(hostPID, snapshotv1alpha1.RestoreCompleteFile); err == nil { + if restoredPID, parseErr := strconv.Atoi(strings.TrimSpace(string(data))); parseErr == nil && restoredPID > 0 { + w.monitorRestoredProcess(ctx, op.log, op.monitorKey(hostPID, restoredPID), hostPID, restoredPID) + } else { + op.log.Info("Restore completion sentinel does not include a restored PID; lifecycle monitor cannot be recovered", "sentinel", snapshotv1alpha1.RestoreCompleteFile, "value", strings.TrimSpace(string(data))) + } + } else { + op.log.Error(err, "Failed to read restore completion sentinel; lifecycle monitor cannot be recovered", "sentinel", snapshotv1alpha1.RestoreCompleteFile) + } return true, nil } @@ -905,7 +925,7 @@ func (w *NodeController) newRestoreOperation( } } -func (op *restoreOperation) executeRestore(ctx context.Context) (int, error) { +func (op *restoreOperation) executeRestore(ctx context.Context) (*executor.RestoreResult, error) { w := op.controller req := executor.RestoreRequest{ ContentUID: op.artifact.ContentUID, @@ -936,20 +956,65 @@ func (op *restoreOperation) failRestore(ctx context.Context, restoreErr error) e return restoreErr } -func (op *restoreOperation) completeRestore(ctx context.Context, placeholderHostPID int) error { +func (op *restoreOperation) completeRestore(ctx context.Context, result *executor.RestoreResult) error { w := op.controller + if result == nil { + return fmt.Errorf("restore completed without result") + } // Any PID inside the container mount namespace reaches the control // volume through /host/proc//root. - if err := w.writeControlSentinelFn(placeholderHostPID, snapshotv1alpha1.RestoreCompleteFile); err != nil { + sentinelData := []byte(strconv.Itoa(result.RestoredPID) + "\n") + if err := w.writeControlSentinelFn(result.PlaceholderPID, snapshotv1alpha1.RestoreCompleteFile, sentinelData); err != nil { op.log.Error(err, "Failed to write restore-complete sentinel") - if killErr := w.sendSignalFn(op.log, placeholderHostPID, syscall.SIGKILL, "restore sentinel failed"); killErr != nil { + if killErr := w.sendSignalFn(op.log, result.PlaceholderPID, syscall.SIGKILL, "restore sentinel failed"); killErr != nil { return errors.Join(fmt.Errorf("failed to write restore-complete sentinel: %w", err), fmt.Errorf("placeholder could not be killed: %w", killErr)) } return fmt.Errorf("failed to write restore-complete sentinel: %w", err) } + w.monitorRestoredProcess(ctx, op.log, op.monitorKey(result.PlaceholderPID, result.RestoredPID), result.PlaceholderPID, result.RestoredPID) return nil } +func (op *restoreOperation) monitorKey(placeholderHostPID, restoredPID int) string { + return fmt.Sprintf("%s/%s/%d/%d", op.pod.UID, op.destination, placeholderHostPID, restoredPID) +} + +func (w *NodeController) monitorRestoredProcess(ctx context.Context, log logr.Logger, key string, placeholderHostPID, restoredPID int) { + if placeholderHostPID <= 0 || restoredPID <= 0 { + log.Error(fmt.Errorf("invalid restore process ids"), "Cannot monitor restored process lifecycle", "placeholder_host_pid", placeholderHostPID, "restored_pid", restoredPID) + return + } + if _, loaded := w.restoreMonitors.LoadOrStore(key, struct{}{}); loaded { + return + } + procRoot := filepath.Join(snapshotruntime.HostProcPath, strconv.Itoa(placeholderHostPID), "root", "proc") + interval := w.restoreMonitorInterval + if interval <= 0 { + interval = restoredProcessMonitorInterval + } + go func() { + defer w.restoreMonitors.Delete(key) + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + if err := w.validateProcessStateFn(procRoot, restoredPID); err == nil { + continue + } else { + log.Info("Restored process exited; terminating restore placeholder", "restored_pid", restoredPID, "placeholder_host_pid", placeholderHostPID, "error", err) + if signalErr := w.sendSignalFn(log, placeholderHostPID, syscall.SIGKILL, "restored process exited"); signalErr != nil { + log.Error(signalErr, "Failed to terminate restore placeholder after restored process exit", "restored_pid", restoredPID, "placeholder_host_pid", placeholderHostPID) + } + return + } + } + }() +} + // applyRestoredCondition uses server-side apply against the status subresource. // Pod conditions are an associative list keyed by type, so this field manager // owns only nvidia.com/Restored and does not replace kubelet-owned conditions. diff --git a/agent/internal/controller/controller_test.go b/agent/internal/controller/controller_test.go index 06a23ba7..b43f2146 100644 --- a/agent/internal/controller/controller_test.go +++ b/agent/internal/controller/controller_test.go @@ -108,7 +108,7 @@ func TestNewDefaultControllerSetsDefaultOperations(t *testing.T) { testr.New(t), ) t.Cleanup(w.restoreQueue.ShutDown) - if w.checkpointFn == nil || w.restoreFn == nil || w.writeControlSentinelFn == nil || w.controlSentinelExistsFn == nil || w.sendSignalFn == nil || w.restoreQueue == nil { + if w.checkpointFn == nil || w.restoreFn == nil || w.writeControlSentinelFn == nil || w.controlSentinelExistsFn == nil || w.validateProcessStateFn == nil || w.sendSignalFn == nil || w.restoreQueue == nil || w.restoreMonitorInterval <= 0 { t.Fatal("default controller operations must be initialized") } } @@ -145,9 +145,12 @@ func makeTestController(t *testing.T, pod *corev1.Pod, apiObjects ...runtime.Obj runtime: &fakeRuntime{}, injector: noopInjector{}, restoreFn: executor.Restore, - writeControlSentinelFn: func(int, string) error { return nil }, + writeControlSentinelFn: func(int, string, []byte) error { return nil }, controlSentinelExistsFn: func(int, string) (bool, error) { return false, nil }, + readControlSentinelFn: func(int, string) ([]byte, error) { return nil, errors.New("not implemented") }, + validateProcessStateFn: func(string, int) error { return nil }, sendSignalFn: func(logr.Logger, int, syscall.Signal, string) error { return nil }, + restoreMonitorInterval: time.Millisecond, restoreQueue: workqueue.NewTypedDelayingQueue[client.ObjectKey](), log: testr.New(t), holderID: "test-holder", @@ -423,16 +426,16 @@ func TestReconcileRestorePodRunsMappedDestinationsConcurrently(t *testing.T) { started := make(chan string, 2) release := make(chan struct{}) - w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { started <- req.DestinationContainerName <-release if req.DestinationContainerName == "engine-0" { - return 100, nil + return &executor.RestoreResult{PlaceholderPID: 100, RestoredPID: 43}, nil } - return 101, nil + return &executor.RestoreResult{PlaceholderPID: 101, RestoredPID: 44}, nil } sentinels := make(chan int, 2) - w.writeControlSentinelFn = func(pid int, name string) error { + w.writeControlSentinelFn = func(pid int, name string, _ []byte) error { assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) sentinels <- pid return nil @@ -470,11 +473,11 @@ func TestReconcileRestorePodReportsPartialSuccess(t *testing.T) { require.NoError(t, err) require.NoError(t, os.MkdirAll(path, 0o700)) - w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { if req.DestinationContainerName == "engine-1" { - return 0, errors.New("restore failed") + return nil, errors.New("restore failed") } - return 100, nil + return &executor.RestoreResult{PlaceholderPID: 100, RestoredPID: 43}, nil } requeue := w.reconcileRestorePod(context.Background(), pod) @@ -495,8 +498,8 @@ func TestReconcileRestorePodReportsAllDestinationsFailed(t *testing.T) { path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, string(content.UID), "main") require.NoError(t, err) require.NoError(t, os.MkdirAll(path, 0o700)) - w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { - return 0, fmt.Errorf("%s restore failed", req.DestinationContainerName) + w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { + return nil, fmt.Errorf("%s restore failed", req.DestinationContainerName) } requeue := w.reconcileRestorePod(context.Background(), pod) @@ -514,10 +517,10 @@ func TestRestorePodContainersKeepsAggregateInProgressWhileDestinationIsPending(t pod.Status.ContainerStatuses = pod.Status.ContainerStatuses[:1] w := makeTestController(t, pod) ctx, cancel := context.WithCancel(context.Background()) - w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { assert.Equal(t, "engine-0", req.DestinationContainerName) cancel() - return 100, nil + return &executor.RestoreResult{PlaceholderPID: 100, RestoredPID: 43}, nil } plan := &restorePlan{ artifact: &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"}, @@ -541,9 +544,9 @@ func TestPreflightRestoreRejectsInvalidMappingBeforeExecution(t *testing.T) { snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) restoreCalls := 0 - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ - return 0, nil + return nil, nil } requeue := w.reconcileRestorePod(context.Background(), pod) @@ -983,8 +986,8 @@ func TestReconcileRestorePodRunsPreflightOnce(t *testing.T) { }, }). Build() - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { - return 4242, nil + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { + return &executor.RestoreResult{PlaceholderPID: 4242, RestoredPID: 43}, nil } requeue := w.reconcileRestorePod(context.Background(), pod) @@ -1022,13 +1025,13 @@ func TestRestoreFinalizerProtectsExecutionAndIsRemovedAfterSuccess(t *testing.T) require.NoError(t, os.MkdirAll(path, 0o700)) restoreCalls := 0 var request executor.RestoreRequest - w.restoreFn = func(ctx context.Context, _ snapshotruntime.Runtime, _ logr.Logger, got executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + w.restoreFn = func(ctx context.Context, _ snapshotruntime.Runtime, _ logr.Logger, got executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ request = got live, getErr := w.clientset.CoreV1().Pods(pod.Namespace).Get(ctx, pod.Name, metav1.GetOptions{}) require.NoError(t, getErr) assert.True(t, hasFinalizer(live, restorePodFinalizer)) - return 4242, nil + return &executor.RestoreResult{PlaceholderPID: 4242, RestoredPID: 43}, nil } processQueuedRestorePod(t, w, pod) @@ -1052,14 +1055,15 @@ func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *test require.NoError(t, os.MkdirAll(path, 0o700)) restoreCalls := 0 - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ - return 4242, nil + return &executor.RestoreResult{PlaceholderPID: 4242, RestoredPID: 43}, nil } sentinelWritten := false - w.writeControlSentinelFn = func(pid int, name string) error { + w.writeControlSentinelFn = func(pid int, name string, data []byte) error { assert.Equal(t, 4242, pid) assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + assert.Equal(t, "43\n", string(data)) sentinelWritten = true return nil } @@ -1068,6 +1072,11 @@ func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *test assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) return sentinelWritten, nil } + w.readControlSentinelFn = func(pid int, name string) ([]byte, error) { + assert.Equal(t, 4242, pid) + assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + return []byte("43\n"), nil + } statusPatches := 0 w.clientset.(*fake.Clientset).PrependReactor("patch", "pods", func(action clientgotesting.Action) (bool, runtime.Object, error) { patch := action.(clientgotesting.PatchAction) @@ -1114,9 +1123,9 @@ func TestRestoreFinalizerRemovalRetriesWithoutReplayingRestore(t *testing.T) { require.NoError(t, err) require.NoError(t, os.MkdirAll(path, 0o700)) restoreCalls := 0 - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ - return 4242, nil + return &executor.RestoreResult{PlaceholderPID: 4242, RestoredPID: 43}, nil } metadataPatches := 0 w.clientset.(*fake.Clientset).PrependReactor("patch", "pods", func(action clientgotesting.Action) (bool, runtime.Object, error) { @@ -1196,12 +1205,12 @@ func TestRunRestoreCleanupFailureStillCompletesRestore(t *testing.T) { } var request executor.RestoreRequest - w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, got executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + w.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, got executor.RestoreRequest, _ executor.RestoreMounter) (*executor.RestoreResult, error) { request = got - return 4242, executor.NewRestoreCleanupError(errors.New("unmount checkpoint artifact: unmount failed")) + return &executor.RestoreResult{PlaceholderPID: 4242, RestoredPID: 43}, executor.NewRestoreCleanupError(errors.New("unmount checkpoint artifact: unmount failed")) } var sentinelPID int - w.writeControlSentinelFn = func(pid int, _ string) error { + w.writeControlSentinelFn = func(pid int, _ string, _ []byte) error { sentinelPID = pid return nil } @@ -1222,9 +1231,9 @@ func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { w.runtime = &fakeRuntime{resolveContainerPID: 4242} artifact := &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"} restoreCalls := 0 - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ - return 0, errors.New("criu restore failed") + return nil, errors.New("criu restore failed") } signalCalls := 0 w.sendSignalFn = func(logr.Logger, int, syscall.Signal, string) error { @@ -1251,9 +1260,9 @@ func TestRunRestoreFailureKillsPlaceholder(t *testing.T) { w.runtime = &fakeRuntime{resolveContainerPID: 4242} artifact := &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"} restoreCalls := 0 - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { restoreCalls++ - return 0, errors.New("criu restore failed") + return nil, errors.New("criu restore failed") } signalCalls := 0 w.sendSignalFn = func(logr.Logger, int, syscall.Signal, string) error { @@ -1280,9 +1289,9 @@ func TestRunRestoreFinalizesExistingCompletionSentinelWithoutReplay(t *testing.T assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) return true, nil } - w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (*executor.RestoreResult, error) { t.Fatal("restore executor must not be replayed after the completion sentinel exists") - return 0, nil + return nil, nil } artifact := &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"} @@ -1290,6 +1299,43 @@ func TestRunRestoreFinalizesExistingCompletionSentinelWithoutReplay(t *testing.T require.NoError(t, err) } +func TestMonitorRestoredProcessKillsPlaceholderWhenWorkloadExits(t *testing.T) { + pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + w := makeTestController(t, pod) + w.restoreMonitorInterval = time.Millisecond + + checks := 0 + w.validateProcessStateFn = func(procRoot string, pid int) error { + assert.Equal(t, "/host/proc/4242/root/proc", procRoot) + assert.Equal(t, 43, pid) + checks++ + if checks == 1 { + return nil + } + return errors.New("process 43 exited") + } + + signaled := make(chan syscall.Signal, 1) + w.sendSignalFn = func(_ logr.Logger, pid int, sig syscall.Signal, reason string) error { + assert.Equal(t, 4242, pid) + assert.Equal(t, "restored process exited", reason) + signaled <- sig + return nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + w.monitorRestoredProcess(ctx, testr.New(t), "test-monitor", 4242, 43) + + select { + case sig := <-signaled: + assert.Equal(t, syscall.SIGKILL, sig) + case <-time.After(time.Second): + t.Fatal("placeholder was not signaled after restored process exit") + } + assert.GreaterOrEqual(t, checks, 2) +} + func TestRestoreArtifactReady(t *testing.T) { w := makeTestController(t, nil) ready, err := w.restoreArtifactReady(testr.New(t), "inference/restore-worker", w.config.Storage.BasePath+"/missing") diff --git a/agent/internal/executor/restore.go b/agent/internal/executor/restore.go index 8dfde222..eca4c652 100644 --- a/agent/internal/executor/restore.go +++ b/agent/internal/executor/restore.go @@ -76,17 +76,21 @@ type RestoreRequest struct { Clientset kubernetes.Interface } +type RestoreResult struct { + PlaceholderPID int + RestoredPID int +} + // Restore performs external restore for the given request. -// Returns the namespace-relative PID of the restored process. // The DaemonSet side inspects the placeholder and launches nsrestore, // which handles rootfs application, CRIU restore, and CUDA restore inside the namespace. // // Returns the placeholder container's host PID so callers can reach into the -// container's mount namespace (e.g. to write sentinels under /snapshot-control) -// without re-resolving via the runtime. -func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, req RestoreRequest, mounts RestoreMounter) (placeholderPID int, retErr error) { +// container's mount namespace, and the restored process PID as seen from inside +// that namespace. +func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, req RestoreRequest, mounts RestoreMounter) (restoreResult *RestoreResult, retErr error) { if mounts == nil { - return 0, fmt.Errorf("restore mounter is required") + return nil, fmt.Errorf("restore mounter is required") } var cleanupErr error @@ -117,24 +121,24 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r artifactPath, err := nsmount.ResolveArtifact(req.BasePath, req.ContentUID, req.ArtifactContainerName) if err != nil { - return 0, fmt.Errorf("resolve checkpoint artifact: %w", err) + return nil, fmt.Errorf("resolve checkpoint artifact: %w", err) } manifest, err := types.ReadManifest(artifactPath) if err != nil { - return 0, fmt.Errorf("read checkpoint manifest: %w", err) + return nil, fmt.Errorf("read checkpoint manifest: %w", err) } if err := validateRestoreManifest(req, manifest); err != nil { - return 0, err + return nil, err } snap, gpuDeviceMapDuration, err := inspectRestore(ctx, rt, log, req, manifest) if err != nil { - return 0, err + return nil, err } bundleMount, err := mounts.MountBundle(ctx, snap.PlaceholderPID) if err != nil { - return 0, fmt.Errorf("mount agent bundle into placeholder: %w", err) + return nil, fmt.Errorf("mount agent bundle into placeholder: %w", err) } activeMounts = append(activeMounts, restoreMount{ action: "unmount agent bundle from placeholder", @@ -143,7 +147,7 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r artifactMount, err := mounts.MountArtifact(ctx, bundleMount, artifactPath) if err != nil { - return 0, fmt.Errorf("mount checkpoint artifact into placeholder: %w", err) + return nil, fmt.Errorf("mount checkpoint artifact into placeholder: %w", err) } activeMounts = append(activeMounts, restoreMount{ action: "unmount checkpoint artifact from placeholder", @@ -152,13 +156,13 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r result, err := execNSRestore(ctx, log, req, snap, bundleMount, nsmount.CheckpointDst) if err != nil { - return 0, fmt.Errorf("nsrestore failed: %w", err) + return nil, fmt.Errorf("nsrestore failed: %w", err) } if result.CleanupError != nil { cleanupErr = errors.Join(cleanupErr, result.CleanupError) } if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil { - return 0, err + return nil, err } cleanup() @@ -190,7 +194,10 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r "placeholder_host_pid", snap.PlaceholderPID, ) - return snap.PlaceholderPID, nil + return &RestoreResult{ + PlaceholderPID: snap.PlaceholderPID, + RestoredPID: result.RestoredPID, + }, nil } func remainingDuration(wall time.Duration, parts ...time.Duration) time.Duration { diff --git a/agent/internal/runtime/control.go b/agent/internal/runtime/control.go index 48565df0..a386a670 100644 --- a/agent/internal/runtime/control.go +++ b/agent/internal/runtime/control.go @@ -25,11 +25,27 @@ import ( // The write uses create-then-rename so the workload never observes a partial // file. func WriteControlSentinel(hostPID int, name string) error { + return WriteControlSentinelData(hostPID, name, []byte("done\n")) +} + +// WriteControlSentinelData writes a sentinel file with caller-provided content +// into the workload container's snapshot-control volume. +func WriteControlSentinelData(hostPID int, name string, data []byte) error { if hostPID <= 0 { return fmt.Errorf("invalid host PID %d for control sentinel %q", hostPID, name) } dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", snapshotv1alpha1.SnapshotControlMountPath) - return writeSentinelInDir(dir, name) + return writeSentinelInDir(dir, name, data) +} + +// ReadControlSentinel reads a sentinel file from the workload container's +// snapshot-control volume. +func ReadControlSentinel(hostPID int, name string) ([]byte, error) { + if hostPID <= 0 { + return nil, fmt.Errorf("invalid host PID %d for control sentinel %q", hostPID, name) + } + dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", snapshotv1alpha1.SnapshotControlMountPath) + return readSentinelInDir(dir, name) } // ControlSentinelExists reports whether a sentinel exists in the workload @@ -52,10 +68,10 @@ func RemoveControlSentinel(dir, name string) error { return removeSentinelInDir(dir, name) } -func writeSentinelInDir(dir, name string) error { +func writeSentinelInDir(dir, name string, data []byte) error { tmpPath := filepath.Join(dir, "."+name+".tmp") finalPath := filepath.Join(dir, name) - if err := os.WriteFile(tmpPath, []byte("done\n"), 0o644); err != nil { + if err := os.WriteFile(tmpPath, data, 0o644); err != nil { return fmt.Errorf("write temp sentinel %s: %w", tmpPath, err) } if err := os.Rename(tmpPath, finalPath); err != nil { @@ -65,6 +81,22 @@ func writeSentinelInDir(dir, name string) error { return nil } +func readSentinelInDir(dir, name string) ([]byte, error) { + info, err := os.Stat(dir) + if err != nil { + return nil, fmt.Errorf("control sentinel dir %s: %w", dir, err) + } + if !info.IsDir() { + return nil, fmt.Errorf("control sentinel dir %s: not a directory", dir) + } + path := filepath.Join(dir, name) + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read control sentinel %s: %w", path, err) + } + return data, nil +} + func removeSentinelInDir(dir, name string) error { info, err := os.Stat(dir) if err != nil { diff --git a/agent/internal/runtime/control_test.go b/agent/internal/runtime/control_test.go index c45a8381..94f58d3f 100644 --- a/agent/internal/runtime/control_test.go +++ b/agent/internal/runtime/control_test.go @@ -12,7 +12,7 @@ import ( func TestWriteSentinelInDir_CreatesFileAtomically(t *testing.T) { dir := t.TempDir() - if err := writeSentinelInDir(dir, "snapshot-complete"); err != nil { + if err := writeSentinelInDir(dir, "snapshot-complete", []byte("done\n")); err != nil { t.Fatalf("writeSentinelInDir failed: %v", err) } @@ -37,10 +37,10 @@ func TestWriteSentinelInDir_CreatesFileAtomically(t *testing.T) { func TestWriteSentinelInDir_Overwrites(t *testing.T) { dir := t.TempDir() - if err := writeSentinelInDir(dir, "restore-complete"); err != nil { + if err := writeSentinelInDir(dir, "restore-complete", []byte("done\n")); err != nil { t.Fatalf("first write failed: %v", err) } - if err := writeSentinelInDir(dir, "restore-complete"); err != nil { + if err := writeSentinelInDir(dir, "restore-complete", []byte("done\n")); err != nil { t.Fatalf("second write failed: %v", err) } data, err := os.ReadFile(filepath.Join(dir, "restore-complete")) @@ -52,9 +52,23 @@ func TestWriteSentinelInDir_Overwrites(t *testing.T) { } } +func TestReadSentinelInDir_ReadsPayload(t *testing.T) { + dir := t.TempDir() + if err := writeSentinelInDir(dir, "restore-complete", []byte("43\n")); err != nil { + t.Fatalf("writeSentinelInDir: %v", err) + } + data, err := readSentinelInDir(dir, "restore-complete") + if err != nil { + t.Fatalf("readSentinelInDir: %v", err) + } + if string(data) != "43\n" { + t.Fatalf("readSentinelInDir() = %q, want 43\\n", data) + } +} + func TestWriteSentinelInDir_DirMissing(t *testing.T) { missing := filepath.Join(t.TempDir(), "does-not-exist") - if err := writeSentinelInDir(missing, "snapshot-complete"); err == nil { + if err := writeSentinelInDir(missing, "snapshot-complete", []byte("done\n")); err == nil { t.Fatal("expected error writing into missing directory") } } @@ -78,7 +92,7 @@ func TestControlSentinelExistsInDir(t *testing.T) { t.Fatal("missing sentinel reported as present") } - if err := writeSentinelInDir(dir, "restore-complete"); err != nil { + if err := writeSentinelInDir(dir, "restore-complete", []byte("done\n")); err != nil { t.Fatalf("writeSentinelInDir: %v", err) } exists, err = controlSentinelExistsInDir(dir, "restore-complete") @@ -106,7 +120,7 @@ func TestRemoveControlSentinel_MissingFile(t *testing.T) { func TestRemoveControlSentinel_RemovesExisting(t *testing.T) { dir := t.TempDir() - if err := writeSentinelInDir(dir, "restore-complete"); err != nil { + if err := writeSentinelInDir(dir, "restore-complete", []byte("done\n")); err != nil { t.Fatalf("writeSentinelInDir: %v", err) } if err := RemoveControlSentinel(dir, "restore-complete"); err != nil {