From 11f9d05001308f8fb5f57f5b9d09dc2f97a7d0cd Mon Sep 17 00:00:00 2001 From: Riley Rice Date: Tue, 4 Aug 2026 12:38:11 -0700 Subject: [PATCH] feat(agent): complete Go agent orchestration and entrypoint Signed-off-by: Riley Rice --- agent/README.md | 57 ++- agent/go/cmd/agent/main.go | 19 +- agent/go/internal/agent/agent.go | 382 ++++++++++++++++ agent/go/internal/agent/agent_test.go | 513 ++++++++++++++++++++++ agent/go/internal/agent/interrupt.go | 283 ++++++++++++ agent/go/internal/agent/interrupt_test.go | 216 +++++++++ agent/go/internal/agent/package.go | 116 +++++ agent/go/internal/agent/package_test.go | 79 ++++ agent/go/internal/agent/steps.go | 296 +++++++++++++ agent/go/internal/agent/steps_test.go | 279 ++++++++++++ 10 files changed, 2221 insertions(+), 19 deletions(-) create mode 100644 agent/go/internal/agent/agent.go create mode 100644 agent/go/internal/agent/agent_test.go create mode 100644 agent/go/internal/agent/interrupt.go create mode 100644 agent/go/internal/agent/interrupt_test.go create mode 100644 agent/go/internal/agent/package.go create mode 100644 agent/go/internal/agent/package_test.go create mode 100644 agent/go/internal/agent/steps.go create mode 100644 agent/go/internal/agent/steps_test.go diff --git a/agent/README.md b/agent/README.md index 8b21d10b..ce79c45a 100644 --- a/agent/README.md +++ b/agent/README.md @@ -20,16 +20,39 @@ A basic example of using a container overlay The Go rewrite under `agent/go` shares execution policy between steps and interrupts through `execution.Config`. A `Config` composes the host root mount, -the child-visible step and package directories, and the stdout and stderr -writers that receive raw command output. Operations report `execution.Status`: +the package directories inside that host, and the stdout and stderr writers +that receive raw command output. Non-host steps resolve those directories +through the mounted host root before execution. Operations report +`execution.Status`: `execution.StatusSuccess` means the operation satisfied its execution policy, while `execution.StatusFailed` means it did not. Each Go interrupt owns its command construction and execution. The `Interrupt` contract exposes `Type` for the wire identity, `Run` for execution using an `execution.Config`, and `Serialize` for the operator-facing representation. -Retry state and completion flags remain orchestration concerns outside the -interrupt implementations. +The orchestration layer writes one completion marker per interrupt type and +resource ID after successful execution. Node restarts use a pending marker +containing the host boot ID: a changed boot ID promotes the marker to complete, +while an unchanged boot ID retries the restart. This keeps reboot completion +independent of the signal used to terminate the agent or its child process. + +The Go entrypoint accepts the current operator forms: + +```text +agent MODE ROOT_MOUNT COPY_DIR +agent interrupt ROOT_MOUNT COPY_DIR INTERRUPT_DATA +``` + +It also accepts the legacy forms, which default `ROOT_MOUNT` to `/root`: + +```text +agent MODE COPY_DIR +agent interrupt COPY_DIR INTERRUPT_DATA +``` + +SIGTERM cancels the active step or interrupt and prevents later steps from +starting. A failed operation or runtime error exits with status 1; malformed +arguments exit with status 2. ### Container Image Build @@ -40,14 +63,26 @@ interrupt implementations. ## Environment variables -There are a number of environment variables that can be used to control how the controller works +There are a number of environment variables that can be used to control how the agent works. 1. `COPY_RESOLV` if set to `"false"` it will NOT copy the container's `/etc/resolv.conf` to the host. -1. `OVERLAY_ALWAYS_RUN_STEP` if set to `"true"` it will ignore any step flags and always run every step. A warning will be printed to stdout if it sees a flag file. -1. `SKYHOOK_AGENT_BUFFER_LIMIT` defaults to 8KB. This is how much of the log of each step it will read before syncing the data to stdout/stderr and the log file. It is recommended to keep this somewhat low to avoid excessive delay between a step emitting some information and seeing it in the docker logs or in the log file. +1. `OVERLAY_ALWAYS_RUN_STEP` if set to `"true"` it will ignore any step flags and always run every step. A warning is logged if it sees a flag file. +1. `SKYHOOK_AGENT_WRITE_LOGS` defaults to `"true"`. Step and interrupt output is streamed directly to stdout/stderr and also written under `SKYHOOK_LOG_DIR`. Set it to `"false"` to stream without retaining host log files. + +`SKYHOOK_AGENT_BUFFER_LIMIT` applies only to the Python agent. The Go agent +streams command output directly and does not buffer it. + +The following environment variables are required and are expected to be set by either the build system or skyhook-operator. It is not recommended that they be changed manually. + +1. `OVERLAY_FRAMEWORK_VERSION` is the version of the current overlay. It is expected that this gets set by the docker build system. It is required to be able to manage the history file. It must be in the format of `{package name}-{version}`. +1. `SKYHOOK_RESOURCE_ID` is used to determine if an interrupt should be rerun. Interrupts are only run once per `SKYHOOK_RESOURCE_ID`. Skyhook operator should make this unique per configuration of the package. + +The following environment variables are optional and use the documented defaults when unset: + +1. `SKYHOOK_DATA_DIR` is the package data source used by legacy invocations when the operator has not already populated `COPY_DIR`. It defaults to `/skyhook-package`. +1. `SKYHOOK_ROOT_DIR` is the host state root for flags, interrupt markers, and history. It defaults to `/etc/skyhook`. +1. `SKYHOOK_LOG_DIR` is the host log root. It defaults to `/var/log/skyhook`. -The following are enviroment variables expected to be set by either the build system or skyhook-operator. It is not recommended they be changed manually. +The following environment variable is optional: -1. `OVERLAY_FRAMEWORK_VERSION` this the version of the current overlay. It is expected that this gets set by the docker build system. It is required to be able to manage the history file. It must be in the format of `{package name}-{version}` -1. `SKYHOOK_RESOURCE_ID` this is used to determine if an interrupt should be rerun. Interrupts are only run once per `SKYHOOK_RESOURCE_ID`. Skyhook operator should make this unique per conifguration of the package. -1. `SKYHOOK_NODE_ORDER` zero-indexed monotonic position of this node in the rollout. The first batch's nodes get `0, 1, 2, ...` and subsequent batches continue from where the previous batch left off. Useful for kubeadm upgrade workflows where the first node (`SKYHOOK_NODE_ORDER=0`) runs a different command than subsequent nodes. See [Node Order Within a Rollout](../docs/ordering_of_skyhooks.md#node-order-within-a-rollout) for details. +1. `SKYHOOK_NODE_ORDER` is a zero-indexed monotonic position of this node in the rollout. The first batch's nodes get `0, 1, 2, ...` and subsequent batches continue from where the previous batch left off. Useful for kubeadm upgrade workflows where the first node (`SKYHOOK_NODE_ORDER=0`) runs a different command than subsequent nodes. See [Node Order Within a Rollout](../docs/ordering_of_skyhooks.md#node-order-within-a-rollout) for details. diff --git a/agent/go/cmd/agent/main.go b/agent/go/cmd/agent/main.go index 028d9ff7..06ad5cd8 100644 --- a/agent/go/cmd/agent/main.go +++ b/agent/go/cmd/agent/main.go @@ -19,17 +19,20 @@ package main import ( - "fmt" - "log/slog" + "context" "os" + "os/signal" + "syscall" + + "github.com/NVIDIA/nodewright/agent/internal/agent" ) func main() { - // Establish the agent's structured-logging seam. Packages (e.g. - // config.Loader.Load) log through *slog.Logger and fall back to - // slog.Default() when passed nil, so wiring it here once keeps that - // default sane for the whole process. - slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, nil))) + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM) + defer stop() + + nodewrightAgent := agent.New() + exitCode := nodewrightAgent.Run(ctx, os.Args[1:], os.Stdout, os.Stderr) - fmt.Println("Hello, World!") + os.Exit(int(exitCode)) } diff --git a/agent/go/internal/agent/agent.go b/agent/go/internal/agent/agent.go new file mode 100644 index 00000000..e4e45552 --- /dev/null +++ b/agent/go/internal/agent/agent.go @@ -0,0 +1,382 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent coordinates package preparation, lifecycle steps, interrupts, +// completion flags, history, and logs for one agent invocation. +package agent + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "os" + "path/filepath" + "strconv" + + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/history" + "github.com/NVIDIA/nodewright/agent/internal/hostfs" + "github.com/NVIDIA/nodewright/agent/internal/stage" +) + +const ( + usage = "usage: agent MODE [ROOT_MOUNT] COPY_DIR [INTERRUPT_DATA]" + defaultRootMount = "/root" + defaultDataDir = "/skyhook-package" + + copyResolverEnv = "COPY_RESOLV" + alwaysRunEnv = "OVERLAY_ALWAYS_RUN_STEP" + resourceIDEnv = "SKYHOOK_RESOURCE_ID" + dataDirEnv = "SKYHOOK_DATA_DIR" + stateRootEnv = "SKYHOOK_ROOT_DIR" + logRootEnv = "SKYHOOK_LOG_DIR" + writeLogsEnv = "SKYHOOK_AGENT_WRITE_LOGS" +) + +// ExitCode is the process result returned by Agent.Run. +type ExitCode int + +const ( + ExitSuccess ExitCode = iota + ExitFailure + ExitUsage +) + +type request struct { + stage stage.Stage + rootMount string + copyDir string + interruptData string +} + +type runtimeConfig struct { + dataDir string + stateRoot string + logRoot string + resourceID string + alwaysRunStep bool + writeLogs bool + copyResolver bool + stdout io.Writer + stderr io.Writer + logger *slog.Logger +} + +type unsupportedStageError struct { + stage string +} + +func (err *unsupportedStageError) Error() string { + return fmt.Sprintf("unsupported stage %q", err.stage) +} + +// Agent executes one operator invocation. +type Agent interface { + Run(context.Context, []string, io.Writer, io.Writer) ExitCode +} + +type orchestrator struct{} + +var _ Agent = orchestrator{} + +// New constructs the agent orchestrator. +func New() Agent { + return orchestrator{} +} + +// Run executes one operator invocation and returns its process exit code. +func (agent orchestrator) Run( + ctx context.Context, + arguments []string, + stdout, stderr io.Writer, +) ExitCode { + if stdout == nil { + stdout = os.Stdout + } + if stderr == nil { + stderr = os.Stderr + } + logger := slog.New(slog.NewTextHandler(stderr, nil)) + + req, err := parseRequest(arguments) + if err != nil { + var unsupportedStage *unsupportedStageError + if errors.As(err, &unsupportedStage) { + logger.Warn( + "this agent version does not support the requested stage; treating it as a no-op", + "stage", unsupportedStage.stage, + ) + return ExitSuccess + } + _, _ = fmt.Fprintln(stderr, err) + _, _ = fmt.Fprintln(stderr, usage) + return ExitUsage + } + + runtime := runtimeFromEnvironment(stdout, stderr, logger) + logger.Info( + "starting agent", + "stage", req.stage, + "rootMount", req.rootMount, + "copyDir", req.copyDir, + "resourceID", runtime.resourceID, + "dataDir", runtime.dataDir, + "stateRoot", runtime.stateRoot, + "logRoot", runtime.logRoot, + "copyResolv", runtime.copyResolver, + "alwaysRunStep", runtime.alwaysRunStep, + "writeLogs", runtime.writeLogs, + ) + + status, err := agent.runRequest(ctx, req, runtime) + if err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.Info("agent stopped after receiving a termination signal", "error", err) + } else { + logger.Error("agent execution failed", "error", err) + } + return ExitFailure + } + + if status == execution.StatusFailed { + return ExitFailure + } + + return ExitSuccess +} + +func parseRequest(arguments []string) (request, error) { + var ( + mode string + rootMount string + copyDir string + interruptData string + ) + switch len(arguments) { + case 2: + mode = arguments[0] + rootMount = defaultRootMount + copyDir = arguments[1] + case 3: + mode = arguments[0] + if mode == string(stage.Interrupt) { + rootMount = defaultRootMount + copyDir = arguments[1] + interruptData = arguments[2] + } else { + rootMount = arguments[1] + copyDir = arguments[2] + } + case 4: + mode = arguments[0] + rootMount = arguments[1] + copyDir = arguments[2] + interruptData = arguments[3] + default: + return request{}, fmt.Errorf( + "expected 2, 3, or 4 arguments; received %d", + len(arguments), + ) + } + + currentStage, err := stage.ParseStage(mode) + if err != nil { + return request{}, &unsupportedStageError{stage: mode} + } + parsed := request{ + stage: currentStage, + rootMount: rootMount, + copyDir: copyDir, + interruptData: interruptData, + } + if err := validateRequest(parsed); err != nil { + return request{}, err + } + return parsed, nil +} + +func runtimeFromEnvironment(stdout, stderr io.Writer, logger *slog.Logger) runtimeConfig { + runtime := normalizeRuntime(runtimeConfig{ + dataDir: envOrDefault(dataDirEnv, defaultDataDir), + stateRoot: envOrDefault(stateRootEnv, flags.DefaultStateRoot), + logRoot: envOrDefault(logRootEnv, flags.DefaultLogRoot), + resourceID: os.Getenv(resourceIDEnv), + stdout: stdout, + stderr: stderr, + logger: logger, + }) + runtime.alwaysRunStep = envBool(alwaysRunEnv, false, runtime.logger) + runtime.writeLogs = envBool(writeLogsEnv, true, runtime.logger) + runtime.copyResolver = envBool(copyResolverEnv, true, runtime.logger) + return runtime +} + +func envOrDefault(name, fallback string) string { + value := os.Getenv(name) + if value == "" { + return fallback + } + return value +} + +func envBool(name string, fallback bool, logger *slog.Logger) bool { + value, exists := os.LookupEnv(name) + if !exists { + return fallback + } + parsed, err := strconv.ParseBool(value) + if err != nil { + logger.Warn( + "invalid boolean environment variable; using default", + "name", name, + "value", value, + "default", fallback, + ) + return fallback + } + return parsed +} + +func normalizeRuntime(runtime runtimeConfig) runtimeConfig { + if runtime.dataDir == "" { + runtime.dataDir = defaultDataDir + } + if runtime.stateRoot == "" { + runtime.stateRoot = flags.DefaultStateRoot + } + if runtime.logRoot == "" { + runtime.logRoot = flags.DefaultLogRoot + } + if runtime.stdout == nil { + runtime.stdout = os.Stdout + } + if runtime.stderr == nil { + runtime.stderr = os.Stderr + } + if runtime.logger == nil { + runtime.logger = slog.New(slog.NewTextHandler(runtime.stderr, nil)) + } + return runtime +} + +func validateRun(ctx context.Context, req request, runtime runtimeConfig) error { + if ctx == nil { + return errors.New("running agent: context is nil") + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("running agent: %w", err) + } + if err := validateRequest(req); err != nil { + return fmt.Errorf("running agent: %w", err) + } + if !filepath.IsAbs(runtime.dataDir) { + return fmt.Errorf("running agent: data directory %q is not absolute", runtime.dataDir) + } + if req.stage == stage.Interrupt { + if err := validatePathComponent("resource ID", runtime.resourceID); err != nil { + return fmt.Errorf("running agent: %w", err) + } + } + return nil +} + +func validateRequest(req request) error { + if _, err := stage.ParseStage(string(req.stage)); err != nil { + return err + } + if !filepath.IsAbs(req.rootMount) { + return fmt.Errorf("root mount %q is not absolute", req.rootMount) + } + if !filepath.IsAbs(req.copyDir) { + return fmt.Errorf("copy directory %q is not absolute", req.copyDir) + } + if req.stage == stage.Interrupt { + if req.interruptData == "" { + return errors.New("interrupt data must not be empty") + } + } else if req.interruptData != "" { + return fmt.Errorf("stage %q does not accept interrupt data", req.stage) + } + return nil +} + +func validatePathComponent(name, value string) error { + if value == "" || value == "." || !filepath.IsLocal(value) || filepath.Base(value) != value { + return fmt.Errorf("%s %q must be a single path component", name, value) + } + return nil +} + +func (orchestrator) runRequest( + ctx context.Context, + req request, + runtime runtimeConfig, +) (execution.Status, error) { + runtime = normalizeRuntime(runtime) + if err := validateRun(ctx, req, runtime); err != nil { + return execution.StatusFailed, err + } + + if runtime.copyResolver { + if err := copyResolverConfig(req.rootMount); err != nil { + return execution.StatusFailed, fmt.Errorf("copying resolver configuration: %w", err) + } + } + + copyRoot, err := hostfs.Resolve(req.rootMount, req.copyDir) + if err != nil { + return execution.StatusFailed, fmt.Errorf("resolving package copy directory: %w", err) + } + if err := ensurePackageData(req.rootMount, copyRoot, runtime.dataDir); err != nil { + return execution.StatusFailed, fmt.Errorf("preparing package data: %w", err) + } + + layout, err := flags.NewLayout(req.rootMount, runtime.stateRoot, runtime.logRoot) + if err != nil { + return execution.StatusFailed, fmt.Errorf("preparing agent filesystem layout: %w", err) + } + if req.stage == stage.Interrupt { + return runDecodedInterrupt(ctx, req, runtime, layout) + } + + cfg, err := loadConfig(copyRoot, runtime.logger) + if err != nil { + return execution.StatusFailed, err + } + + if req.stage != stage.Uninstall && req.stage != stage.UninstallCheck { + if err := prepareHost(copyRoot, req.rootMount, *cfg); err != nil { + return execution.StatusFailed, err + } + } + + flagStore, err := flags.NewStore(layout, *cfg) + if err != nil { + return execution.StatusFailed, fmt.Errorf("constructing flag store: %w", err) + } + + historyStore, err := history.NewStore(req.rootMount, layout.HistoryDir(), *cfg, runtime.logger) + if err != nil { + return execution.StatusFailed, fmt.Errorf("constructing history store: %w", err) + } + + return runSteps(ctx, req, runtime, layout, *cfg, flagStore, historyStore) +} diff --git a/agent/go/internal/agent/agent_test.go b/agent/go/internal/agent/agent_test.go new file mode 100644 index 00000000..243d79c3 --- /dev/null +++ b/agent/go/internal/agent/agent_test.go @@ -0,0 +1,513 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log/slog" + "os" + "path/filepath" + "testing" + "time" + + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/interrupts" + "github.com/NVIDIA/nodewright/agent/internal/stage" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAgent(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Agent Suite") +} + +var _ = Describe("request parsing", func() { + DescribeTable( + "supports current and legacy operator forms", + func(arguments []string, expected request) { + actual, err := parseRequest(arguments) + Expect(err).NotTo(HaveOccurred()) + Expect(actual).To(Equal(expected)) + }, + Entry( + "legacy stage", + []string{"apply", "/packages/value"}, + request{stage: stage.Apply, rootMount: "/root", copyDir: "/packages/value"}, + ), + Entry( + "legacy interrupt", + []string{"interrupt", "/packages/value", "payload"}, + request{ + stage: stage.Interrupt, rootMount: "/root", + copyDir: "/packages/value", interruptData: "payload", + }, + ), + Entry( + "current stage", + []string{"apply", "/host", "/packages/value"}, + request{stage: stage.Apply, rootMount: "/host", copyDir: "/packages/value"}, + ), + Entry( + "current interrupt", + []string{"interrupt", "/host", "/packages/value", "payload"}, + request{ + stage: stage.Interrupt, rootMount: "/host", + copyDir: "/packages/value", interruptData: "payload", + }, + ), + ) + + It("identifies stages unknown to this agent version", func() { + _, err := parseRequest([]string{"future-stage", "/host", "/package"}) + + Expect(err).To(MatchError(`unsupported stage "future-stage"`)) + var unsupported *unsupportedStageError + Expect(errors.As(err, &unsupported)).To(BeTrue()) + Expect(unsupported.stage).To(Equal("future-stage")) + }) + + DescribeTable( + "rejects malformed requests", + func(arguments []string, expected string) { + _, err := parseRequest(arguments) + Expect(err).To(MatchError(ContainSubstring(expected))) + }, + Entry("argument count", []string{"apply"}, "expected 2, 3, or 4 arguments"), + Entry("relative root", []string{"apply", "host", "/package"}, "root mount"), + Entry("relative copy directory", []string{"apply", "/host", "package"}, "copy directory"), + Entry( + "empty interrupt data", + []string{"interrupt", "/host", "/package", ""}, + "interrupt data must not be empty", + ), + Entry( + "interrupt data on a normal stage", + []string{"apply", "/host", "/package", "payload"}, + `stage "apply" does not accept interrupt data`, + ), + ) +}) + +var _ = Describe("runtime environment", func() { + It("uses documented defaults when environment values are absent", func() { + for _, name := range []string{ + copyResolverEnv, + alwaysRunEnv, + resourceIDEnv, + dataDirEnv, + stateRootEnv, + logRootEnv, + writeLogsEnv, + } { + unsetEnvironment(name) + } + + runtime := runtimeFromEnvironment(nil, nil, nil) + + Expect(runtime.dataDir).To(Equal(defaultDataDir)) + Expect(runtime.stateRoot).To(Equal(flags.DefaultStateRoot)) + Expect(runtime.logRoot).To(Equal(flags.DefaultLogRoot)) + Expect(runtime.resourceID).To(BeEmpty()) + Expect(runtime.alwaysRunStep).To(BeFalse()) + Expect(runtime.writeLogs).To(BeTrue()) + Expect(runtime.copyResolver).To(BeTrue()) + Expect(runtime.stdout).To(Equal(io.Writer(os.Stdout))) + Expect(runtime.stderr).To(Equal(io.Writer(os.Stderr))) + Expect(runtime.logger).NotTo(BeNil()) + }) + + It("loads configured values and output composition", func() { + GinkgoT().Setenv(copyResolverEnv, "false") + GinkgoT().Setenv(alwaysRunEnv, "TRUE") + GinkgoT().Setenv(resourceIDEnv, "resource") + GinkgoT().Setenv(dataDirEnv, "/data") + GinkgoT().Setenv(stateRootEnv, "/state") + GinkgoT().Setenv(logRootEnv, "/logs") + GinkgoT().Setenv(writeLogsEnv, "false") + stdout := &bytes.Buffer{} + stderr := &bytes.Buffer{} + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + + runtime := runtimeFromEnvironment(stdout, stderr, logger) + + Expect(runtime.dataDir).To(Equal("/data")) + Expect(runtime.stateRoot).To(Equal("/state")) + Expect(runtime.logRoot).To(Equal("/logs")) + Expect(runtime.resourceID).To(Equal("resource")) + Expect(runtime.alwaysRunStep).To(BeTrue()) + Expect(runtime.writeLogs).To(BeFalse()) + Expect(runtime.copyResolver).To(BeFalse()) + Expect(runtime.stdout).To(BeIdenticalTo(stdout)) + Expect(runtime.stderr).To(BeIdenticalTo(stderr)) + Expect(runtime.logger).To(BeIdenticalTo(logger)) + }) + + It("uses the fallback and warns when a boolean value is invalid", func() { + GinkgoT().Setenv(copyResolverEnv, " true ") + logOutput := &bytes.Buffer{} + logger := slog.New(slog.NewTextHandler(logOutput, nil)) + + runtime := runtimeFromEnvironment(io.Discard, io.Discard, logger) + + Expect(runtime.copyResolver).To(BeTrue()) + Expect(logOutput.String()).To(ContainSubstring( + "invalid boolean environment variable; using default", + )) + Expect(logOutput.String()).To(ContainSubstring(copyResolverEnv)) + }) + + It("normalizes empty state and log roots", func() { + runtime := normalizeRuntime(runtimeConfig{}) + + Expect(runtime.stateRoot).To(Equal(flags.DefaultStateRoot)) + Expect(runtime.logRoot).To(Equal(flags.DefaultLogRoot)) + }) +}) + +var _ = Describe("Agent.Run", func() { + BeforeEach(func() { + GinkgoT().Setenv(copyResolverEnv, "false") + GinkgoT().Setenv(alwaysRunEnv, "false") + GinkgoT().Setenv(resourceIDEnv, "") + GinkgoT().Setenv(stateRootEnv, "/state") + GinkgoT().Setenv(logRootEnv, "/logs") + GinkgoT().Setenv(writeLogsEnv, "false") + }) + + It("defines process exit codes", func() { + Expect(ExitSuccess).To(Equal(ExitCode(0))) + Expect(ExitFailure).To(Equal(ExitCode(1))) + Expect(ExitUsage).To(Equal(ExitCode(2))) + }) + + It("runs a valid operator invocation", func() { + root := GinkgoT().TempDir() + dataDir := GinkgoT().TempDir() + writePackageFixture(dataDir, "[]", true) + GinkgoT().Setenv(dataDirEnv, dataDir) + stdout := &bytes.Buffer{} + stderr := &bytes.Buffer{} + + exitCode := New().Run( + context.Background(), + []string{"config", root, "/package"}, + stdout, + stderr, + ) + + Expect(exitCode).To(Equal(ExitSuccess)) + Expect(filepath.Join(root, "package", configFileName)).To(BeAnExistingFile()) + Expect(stderr.String()).To(ContainSubstring("starting agent")) + }) + + It("treats an unsupported stage as a successful no-op", func() { + stderr := &bytes.Buffer{} + + exitCode := New().Run( + context.Background(), + []string{"future-stage", "/host", "/package"}, + io.Discard, + stderr, + ) + + Expect(exitCode).To(Equal(ExitSuccess)) + Expect(stderr.String()).To(ContainSubstring("does not support the requested stage")) + }) + + It("returns a usage exit code for an invalid invocation", func() { + stderr := &bytes.Buffer{} + + exitCode := New().Run(context.Background(), []string{"apply"}, io.Discard, stderr) + + Expect(exitCode).To(Equal(ExitUsage)) + Expect(stderr.String()).To(ContainSubstring("expected 2, 3, or 4 arguments")) + Expect(stderr.String()).To(ContainSubstring(usage)) + }) + + It("returns a failure exit code when execution fails", func() { + root := GinkgoT().TempDir() + GinkgoT().Setenv(dataDirEnv, filepath.Join(root, "missing")) + stderr := &bytes.Buffer{} + + exitCode := New().Run( + context.Background(), + []string{"apply", root, "/package"}, + io.Discard, + stderr, + ) + + Expect(exitCode).To(Equal(ExitFailure)) + Expect(stderr.String()).To(ContainSubstring("agent execution failed")) + }) + + It("returns a failure exit code when execution is canceled", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + stderr := &bytes.Buffer{} + + exitCode := New().Run( + ctx, + []string{"apply", GinkgoT().TempDir(), "/package"}, + io.Discard, + stderr, + ) + + Expect(exitCode).To(Equal(ExitFailure)) + Expect(stderr.String()).To(ContainSubstring("agent stopped after receiving a termination signal")) + }) + + It("cancels an active step and does not start the next step", func() { + root := GinkgoT().TempDir() + dataDir := GinkgoT().TempDir() + writeCancellationPackageFixture(dataDir) + GinkgoT().Setenv(dataDirEnv, dataDir) + GinkgoT().Setenv("NODEWRIGHT_AGENT_TEST_MARKER_DIR", filepath.Join(root, "package")) + stderr := &bytes.Buffer{} + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + exitCodes := make(chan ExitCode, 1) + var earlyExit error + + go func() { + exitCodes <- New().Run( + ctx, + []string{"apply", root, "/package"}, + io.Discard, + stderr, + ) + }() + + firstStarted := filepath.Join(root, "package", "first-started") + Eventually(func() error { + if _, err := os.Stat(firstStarted); err == nil { + return nil + } else if !errors.Is(err, os.ErrNotExist) { + return err + } + if earlyExit == nil { + select { + case exitCode := <-exitCodes: + earlyExit = fmt.Errorf( + "agent exited with code %d before the first step started: %s", + exitCode, + stderr, + ) + default: + } + } + if earlyExit != nil { + return earlyExit + } + return errors.New("first step has not started") + }).WithTimeout(5 * time.Second).Should(Succeed()) + cancel() + + var exitCode ExitCode + Eventually(exitCodes).WithTimeout(5 * time.Second).Should(Receive(&exitCode)) + Expect(exitCode).To(Equal(ExitFailure)) + Expect(filepath.Join(root, "package", "first-finished")).NotTo(BeAnExistingFile()) + Expect(filepath.Join(root, "package", "second-started")).NotTo(BeAnExistingFile()) + Expect(stderr.String()).To(ContainSubstring("agent stopped after receiving a termination signal")) + }) +}) + +var _ = Describe("request orchestration", func() { + It("copies legacy package data before running a stage", func() { + root := GinkgoT().TempDir() + dataDir := GinkgoT().TempDir() + writePackageFixture(dataDir, "[]", true) + _, resolverErr := os.Stat("/etc/resolv.conf") + copyResolver := resolverErr == nil + if resolverErr != nil && !errors.Is(resolverErr, os.ErrNotExist) { + Expect(resolverErr).NotTo(HaveOccurred()) + } + runtime := runtimeConfig{ + dataDir: dataDir, + stateRoot: "/state", + logRoot: "/logs", + copyResolver: copyResolver, + writeLogs: false, + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + + status, err := orchestrator{}.runRequest(context.Background(), request{ + stage: stage.Config, + rootMount: root, + copyDir: "/package", + }, runtime) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join(root, "package", configFileName)).To(BeAnExistingFile()) + if copyResolver { + Expect(filepath.Join(root, "etc", "resolv.conf")).To(BeAnExistingFile()) + } + Expect(filepath.Join(root, "state", "flags", startFlagName)).To(BeAnExistingFile()) + }) + + It("does not apply host preparation requirements during uninstall", func() { + root := GinkgoT().TempDir() + copyRoot := filepath.Join(root, "package") + writePackageFixture(copyRoot, `["missing"]`, true) + Expect(os.Mkdir(filepath.Join(copyRoot, rootOverlayDirName), 0o755)).To(Succeed()) + Expect(os.WriteFile( + filepath.Join(copyRoot, rootOverlayDirName, "should-not-copy"), + []byte("value"), + 0o600, + )).To(Succeed()) + runtime := runtimeConfig{ + stateRoot: "/state", + logRoot: "/logs", + writeLogs: false, + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + + status, err := orchestrator{}.runRequest(context.Background(), request{ + stage: stage.Uninstall, + rootMount: root, + copyDir: "/package", + }, runtime) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join(root, "should-not-copy")).NotTo(BeAnExistingFile()) + }) + + It("runs an encoded interrupt through the Agent interface", func() { + root := GinkgoT().TempDir() + Expect(os.Mkdir(filepath.Join(root, "package"), 0o755)).To(Succeed()) + encoded, err := interrupts.Encode(interrupts.NoOp{}) + Expect(err).NotTo(HaveOccurred()) + + status, err := orchestrator{}.runRequest( + context.Background(), + request{ + stage: stage.Interrupt, + rootMount: root, + copyDir: "/package", + interruptData: encoded, + }, + runtimeConfig{ + stateRoot: "/state", + logRoot: "/logs", + resourceID: "resource_package_1.0.0", + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join( + root, + "state", + interruptsDirName, + interruptFlagsDirName, + "resource_package_1.0.0", + "no_op.complete", + )).To(BeAnExistingFile()) + }) +}) + +func writePackageFixture(directory, expectedConfigFiles string, onHost bool) { + Expect(os.MkdirAll(filepath.Join(directory, "skyhook_dir"), 0o755)).To(Succeed()) + for _, name := range []string{"apply", "second", "apply-check"} { + Expect(os.WriteFile(filepath.Join(directory, "skyhook_dir", name), nil, 0o700)).To(Succeed()) + } + data := fmt.Sprintf(`{ + "schema_version": "v1", + "root_dir": "/", + "expected_config_files": %s, + "package_name": "package", + "package_version": "1.0.0", + "modes": { + "apply": [ + { + "name": "apply", + "path": "apply", + "arguments": [], + "returncodes": [0], + "on_host": %t, + "idempotence": false, + "upgrade_step": false + }, + { + "name": "second", + "path": "second", + "arguments": [], + "returncodes": [0], + "on_host": %t, + "idempotence": false, + "upgrade_step": false + } + ], + "apply-check": [ + { + "name": "apply-check", + "path": "apply-check", + "arguments": [], + "returncodes": [0], + "on_host": %t, + "idempotence": false, + "upgrade_step": false + } + ] + } + }`, expectedConfigFiles, onHost, onHost, onHost) + Expect(os.WriteFile(filepath.Join(directory, configFileName), []byte(data), 0o600)).To(Succeed()) +} + +func writeCancellationPackageFixture(directory string) { + writePackageFixture(directory, "[]", false) + stepsDir := filepath.Join(directory, "skyhook_dir") + Expect(os.WriteFile( + filepath.Join(stepsDir, "apply"), + []byte("#!/bin/sh\n: > \"$NODEWRIGHT_AGENT_TEST_MARKER_DIR/first-started\"\nwhile :; do :; done\n: > \"$NODEWRIGHT_AGENT_TEST_MARKER_DIR/first-finished\"\n"), + 0o700, + )).To(Succeed()) + Expect(os.WriteFile( + filepath.Join(stepsDir, "second"), + []byte("#!/bin/sh\n: > \"$NODEWRIGHT_AGENT_TEST_MARKER_DIR/second-started\"\n"), + 0o700, + )).To(Succeed()) +} + +func unsetEnvironment(name string) { + value, exists := os.LookupEnv(name) + Expect(os.Unsetenv(name)).To(Succeed()) + DeferCleanup(func() { + if exists { + Expect(os.Setenv(name, value)).To(Succeed()) + return + } + Expect(os.Unsetenv(name)).To(Succeed()) + }) +} diff --git a/agent/go/internal/agent/interrupt.go b/agent/go/internal/agent/interrupt.go new file mode 100644 index 00000000..4f6c2418 --- /dev/null +++ b/agent/go/internal/agent/interrupt.go @@ -0,0 +1,283 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "context" + "errors" + "fmt" + "io" + "path/filepath" + "strings" + "time" + + "github.com/NVIDIA/nodewright/agent/internal/config" + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/hostfs" + "github.com/NVIDIA/nodewright/agent/internal/interrupts" +) + +const ( + interruptsDirName = "interrupts" + interruptFlagsDirName = "flags" + markerFileMode = 0o600 + hostBootIDPath = "/proc/sys/kernel/random/boot_id" + restartMarkerPrefix = "boot-id:" +) + +func runDecodedInterrupt( + ctx context.Context, + req request, + runtime runtimeConfig, + layout flags.Layout, +) (execution.Status, error) { + value, err := interrupts.Decode(req.interruptData) + if err != nil { + return execution.StatusFailed, fmt.Errorf("decoding interrupt: %w", err) + } + return runInterrupt(ctx, req, runtime, layout, value) +} + +func runInterrupt( + ctx context.Context, + req request, + runtime runtimeConfig, + layout flags.Layout, + value interrupts.Interrupt, +) (status execution.Status, retErr error) { + interruptType := value.Type() + completeMarker := filepath.Join( + layout.StateDir(), + interruptsDirName, + interruptFlagsDirName, + runtime.resourceID, + string(interruptType)+".complete", + ) + exists, err := hostfs.RegularFileExists(req.rootMount, completeMarker) + if err != nil { + return execution.StatusFailed, fmt.Errorf("checking interrupt completion: %w", err) + } + if exists { + runtime.logger.Info( + "skipping completed interrupt", + "interrupt", interruptType, + "resourceID", runtime.resourceID, + ) + return execution.StatusSuccess, nil + } + + pendingMarker := "" + retainPendingMarker := false + if interruptType == interrupts.NodeRestartType { + var completed bool + pendingMarker, completed, err = prepareNodeRestartMarker(req, runtime, completeMarker) + if err != nil { + return execution.StatusFailed, err + } + if completed { + runtime.logger.Info( + "skipping completed interrupt after host reboot", + "interrupt", interruptType, + "resourceID", runtime.resourceID, + ) + return execution.StatusSuccess, nil + } + } + + defer func() { + if pendingMarker == "" || retainPendingMarker { + return + } + if err := hostfs.RemoveFile(req.rootMount, pendingMarker); err != nil { + status = execution.StatusFailed + retErr = errors.Join( + retErr, + fmt.Errorf("removing failed node restart marker: %w", err), + ) + } + }() + + stdout, stderr := runtime.stdout, runtime.stderr + closeLog := func() error { return nil } + closeLogErr := func() error { + if err := closeLog(); err != nil { + return fmt.Errorf("closing interrupt log: %w", err) + } + return nil + } + var logFiles flags.LogFiles + if runtime.writeLogs { + logConfig, err := configFromResourceID(runtime.resourceID) + if err != nil { + return execution.StatusFailed, fmt.Errorf("resolving interrupt log identity: %w", err) + } + logName := filepath.Join(interruptsDirName, string(interruptType)) + _, file, err := layout.CreateLogFile(logConfig, logName, time.Now(), logFileMode) + if err != nil { + return execution.StatusFailed, fmt.Errorf("preparing log for interrupt %q: %w", interruptType, err) + } + stdout = io.MultiWriter(stdout, file) + stderr = io.MultiWriter(stderr, file) + closeLog = file.Close + logFiles, err = layout.LogFilePattern(logConfig, logName) + if err != nil { + return execution.StatusFailed, errors.Join( + fmt.Errorf("resolving log retention for interrupt %q: %w", interruptType, err), + closeLogErr(), + ) + } + } + + runConfig, err := execution.NewConfig( + execution.WithRootMount(req.rootMount), + execution.WithStepRoot(flags.StepsDir(req.copyDir)), + execution.WithSkyhookDir(req.copyDir), + execution.WithRunOutput(stdout, stderr), + ) + if err != nil { + return execution.StatusFailed, errors.Join( + fmt.Errorf("configuring interrupt %q execution: %w", interruptType, err), + closeLogErr(), + ) + } + if interruptType == interrupts.NodeRestartType { + // A successful reboot can terminate the agent by several signal paths. + // Keep the pending marker once execution begins and decide completion by + // comparing host boot IDs on the next invocation. + retainPendingMarker = true + } + status, runErr := value.Run(ctx, runConfig) + closeErr := closeLogErr() + var cleanupErr error + if runtime.writeLogs { + cleanupErr = flags.CleanupOldLogs(logFiles, flags.DefaultLogRetention) + } + if runErr != nil || closeErr != nil || cleanupErr != nil { + if runErr != nil { + runErr = fmt.Errorf("running interrupt %q: %w", interruptType, runErr) + } + if cleanupErr != nil { + cleanupErr = fmt.Errorf("cleaning old interrupt logs: %w", cleanupErr) + } + return execution.StatusFailed, errors.Join( + runErr, + closeErr, + cleanupErr, + ) + } + if status != execution.StatusSuccess { + return execution.StatusFailed, nil + } + + if interruptType != interrupts.NodeRestartType { + if err := hostfs.CreateFile( + req.rootMount, + completeMarker, + []byte(time.Now().UTC().Format(time.RFC3339Nano)), + markerFileMode, + ); err != nil { + return execution.StatusFailed, fmt.Errorf("marking interrupt complete: %w", err) + } + } + return execution.StatusSuccess, nil +} + +func prepareNodeRestartMarker( + req request, + runtime runtimeConfig, + completeMarker string, +) (string, bool, error) { + bootID, err := hostBootID(req.rootMount) + if err != nil { + return "", false, fmt.Errorf("reading host boot ID: %w", err) + } + pendingMarker := strings.TrimSuffix(completeMarker, ".complete") + ".pending" + pending, err := hostfs.RegularFileExists(req.rootMount, pendingMarker) + if err != nil { + return "", false, fmt.Errorf("checking pending node restart: %w", err) + } + if pending { + data, err := hostfs.ReadFile(req.rootMount, pendingMarker) + if err != nil { + return "", false, fmt.Errorf("reading pending node restart: %w", err) + } + previousBootID, valid := strings.CutPrefix(strings.TrimSpace(string(data)), restartMarkerPrefix) + if valid && previousBootID != "" && previousBootID != bootID { + if err := hostfs.RenameFile(req.rootMount, pendingMarker, completeMarker); err != nil { + return "", false, fmt.Errorf("completing node restart marker: %w", err) + } + return pendingMarker, true, nil + } + if !valid || previousBootID == "" { + runtime.logger.Warn( + "discarding malformed pending node restart marker", + "resourceID", runtime.resourceID, + ) + } + if err := hostfs.RemoveFile(req.rootMount, pendingMarker); err != nil { + return "", false, fmt.Errorf("removing stale node restart marker: %w", err) + } + } + if err := hostfs.CreateFile( + req.rootMount, + pendingMarker, + []byte(restartMarkerPrefix+bootID+"\n"), + markerFileMode, + ); err != nil { + return "", false, fmt.Errorf("marking node restart pending: %w", err) + } + return pendingMarker, false, nil +} + +func hostBootID(rootMount string) (string, error) { + path, err := hostfs.Resolve(rootMount, hostBootIDPath) + if err != nil { + return "", err + } + data, err := hostfs.ReadFile(rootMount, path) + if err != nil { + return "", err + } + bootID := strings.TrimSpace(string(data)) + if bootID == "" { + return "", errors.New("host boot ID is empty") + } + return bootID, nil +} + +func configFromResourceID(resourceID string) (config.Config, error) { + versionSeparator := strings.LastIndex(resourceID, "_") + if versionSeparator <= 0 || versionSeparator == len(resourceID)-1 { + return config.Config{}, fmt.Errorf("resource ID %q must end in _PACKAGE_VERSION", resourceID) + } + packageSeparator := strings.Index(resourceID[:versionSeparator], "_") + if packageSeparator <= 0 || packageSeparator == versionSeparator-1 { + return config.Config{}, fmt.Errorf("resource ID %q must end in _PACKAGE_VERSION", resourceID) + } + packageName := resourceID[packageSeparator+1 : versionSeparator] + packageVersion := resourceID[versionSeparator+1:] + if err := validatePathComponent("package name", packageName); err != nil { + return config.Config{}, err + } + if err := validatePathComponent("package version", packageVersion); err != nil { + return config.Config{}, err + } + return config.Config{PackageName: packageName, PackageVersion: packageVersion}, nil +} diff --git a/agent/go/internal/agent/interrupt_test.go b/agent/go/internal/agent/interrupt_test.go new file mode 100644 index 00000000..85370801 --- /dev/null +++ b/agent/go/internal/agent/interrupt_test.go @@ -0,0 +1,216 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "context" + "errors" + "io" + "log/slog" + "os" + "path/filepath" + + "github.com/NVIDIA/nodewright/agent/internal/config" + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/interrupts" + interruptsmock "github.com/NVIDIA/nodewright/agent/internal/interrupts/mock" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/stretchr/testify/mock" +) + +func newMockInterrupt( + interruptType interrupts.InterruptType, + status execution.Status, +) *interruptsmock.MockInterrupt { + value := interruptsmock.NewMockInterrupt(GinkgoT()) + value.EXPECT().Type().Return(interruptType).Maybe() + value.EXPECT(). + Run(mock.Anything, mock.Anything). + Return(status, nil). + Once() + return value +} + +var _ = Describe("interrupt orchestration", func() { + It("preserves underscores in package names parsed from resource IDs", func() { + cfg, err := configFromResourceID("nodewright-uid-1_package_with_underscores_1.2.3") + + Expect(err).NotTo(HaveOccurred()) + Expect(cfg).To(Equal(config.Config{ + PackageName: "package_with_underscores", + PackageVersion: "1.2.3", + })) + }) + + It("marks a successful interrupt and skips it on the next invocation", func() { + root := GinkgoT().TempDir() + layout := flags.DefaultLayout(root) + req := request{rootMount: root, copyDir: "/package"} + runtime := runtimeConfig{ + resourceID: "resource_package_1.0.0", + writeLogs: true, + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + value := newMockInterrupt(interrupts.NoOpType, execution.StatusSuccess) + + status, err := runInterrupt(context.Background(), req, runtime, layout, value) + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + status, err = runInterrupt(context.Background(), req, runtime, layout, value) + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join( + layout.StateDir(), interruptsDirName, interruptFlagsDirName, + runtime.resourceID, "no_op.complete", + )).To(BeAnExistingFile()) + logEntries, err := os.ReadDir(filepath.Join( + layout.LogDir(), "package", "1.0.0", interruptsDirName, + )) + Expect(err).NotTo(HaveOccurred()) + Expect(logEntries).To(HaveLen(1)) + }) + + It("does not expose an ordinary interrupt as complete while it is running", func() { + root := GinkgoT().TempDir() + layout := flags.DefaultLayout(root) + req := request{rootMount: root, copyDir: "/package"} + runtime := runtimeConfig{ + resourceID: "resource_package_1.0.0", + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + marker := filepath.Join( + layout.StateDir(), interruptsDirName, interruptFlagsDirName, + runtime.resourceID, "service_restart.complete", + ) + value := interruptsmock.NewMockInterrupt(GinkgoT()) + value.EXPECT().Type().Return(interrupts.ServiceRestartType).Maybe() + value.EXPECT().Run(mock.Anything, mock.Anything). + Run(func(context.Context, execution.Config) { + Expect(marker).NotTo(BeAnExistingFile()) + }). + Return(execution.StatusSuccess, nil). + Once() + + status, err := runInterrupt(context.Background(), req, runtime, layout, value) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(marker).To(BeAnExistingFile()) + }) + + It("does not mark a failed ordinary interrupt complete", func() { + root := GinkgoT().TempDir() + layout := flags.DefaultLayout(root) + req := request{rootMount: root, copyDir: "/package"} + runtime := runtimeConfig{ + resourceID: "resource_package_1.0.0", + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + value := newMockInterrupt(interrupts.ServiceRestartType, execution.StatusFailed) + + status, err := runInterrupt(context.Background(), req, runtime, layout, value) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusFailed)) + Expect(filepath.Join( + layout.StateDir(), interruptsDirName, interruptFlagsDirName, + runtime.resourceID, "service_restart.complete", + )).NotTo(BeAnExistingFile()) + }) + + It("retries a canceled node restart until the host boot ID changes", func() { + root := GinkgoT().TempDir() + bootIDPath := filepath.Join(root, "proc", "sys", "kernel", "random", "boot_id") + Expect(os.MkdirAll(filepath.Dir(bootIDPath), 0o755)).To(Succeed()) + Expect(os.WriteFile(bootIDPath, []byte("boot-a\n"), 0o600)).To(Succeed()) + layout := flags.DefaultLayout(root) + req := request{rootMount: root, copyDir: "/package"} + runtime := runtimeConfig{ + resourceID: "resource_package_1.0.0", + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + markerBase := filepath.Join( + layout.StateDir(), interruptsDirName, interruptFlagsDirName, + runtime.resourceID, "node_restart", + ) + pendingMarker := markerBase + ".pending" + completeMarker := markerBase + ".complete" + canceled := interruptsmock.NewMockInterrupt(GinkgoT()) + canceled.EXPECT().Type().Return(interrupts.NodeRestartType).Maybe() + canceled.EXPECT().Run(mock.Anything, mock.Anything). + Return(execution.StatusFailed, context.Canceled). + Once() + + status, err := runInterrupt(context.Background(), req, runtime, layout, canceled) + Expect(status).To(Equal(execution.StatusFailed)) + Expect(errors.Is(err, context.Canceled)).To(BeTrue()) + Expect(pendingMarker).To(BeAnExistingFile()) + Expect(completeMarker).NotTo(BeAnExistingFile()) + + retried := newMockInterrupt(interrupts.NodeRestartType, execution.StatusSuccess) + status, err = runInterrupt(context.Background(), req, runtime, layout, retried) + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(pendingMarker).To(BeAnExistingFile()) + Expect(completeMarker).NotTo(BeAnExistingFile()) + + Expect(os.WriteFile(bootIDPath, []byte("boot-b\n"), 0o600)).To(Succeed()) + completed := interruptsmock.NewMockInterrupt(GinkgoT()) + completed.EXPECT().Type().Return(interrupts.NodeRestartType).Maybe() + status, err = runInterrupt(context.Background(), req, runtime, layout, completed) + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(pendingMarker).NotTo(BeAnExistingFile()) + Expect(completeMarker).To(BeAnExistingFile()) + }) + + It("decodes the operator wire payload before running an interrupt", func() { + root := GinkgoT().TempDir() + layout := flags.DefaultLayout(root) + encoded, err := interrupts.Encode(interrupts.NoOp{}) + Expect(err).NotTo(HaveOccurred()) + req := request{ + rootMount: root, + copyDir: "/package", + interruptData: encoded, + } + runtime := runtimeConfig{ + resourceID: "resource_package_1.0.0", + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + + status, err := runDecodedInterrupt(context.Background(), req, runtime, layout) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + }) +}) diff --git a/agent/go/internal/agent/package.go b/agent/go/internal/agent/package.go new file mode 100644 index 00000000..a481431e --- /dev/null +++ b/agent/go/internal/agent/package.go @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "errors" + "fmt" + "io/fs" + "log/slog" + "os" + "path/filepath" + + "github.com/NVIDIA/nodewright/agent/internal/config" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/hostfs" +) + +const ( + legacyNodeFilesDir = "/etc/nvidia-bootstrap/node-files" + configFileName = "config.json" + configMapsDirName = "configmaps" + rootOverlayDirName = "root_dir" +) + +func ensurePackageData(rootMount, copyRoot, dataDir string) error { + info, err := os.Stat(copyRoot) + if err == nil { + if !info.IsDir() { + return fmt.Errorf("package copy path %q is not a directory", copyRoot) + } + return nil + } + if !errors.Is(err, fs.ErrNotExist) { + return fmt.Errorf("stating package copy path %q: %w", copyRoot, err) + } + if err := hostfs.CopyTreeIfExists(rootMount, dataDir, copyRoot); err != nil { + return fmt.Errorf("copying package data from %q: %w", dataDir, err) + } + if _, err := os.Stat(copyRoot); errors.Is(err, fs.ErrNotExist) { + return fmt.Errorf("package data directory %q does not exist", dataDir) + } else if err != nil { + return fmt.Errorf("stating copied package data %q: %w", copyRoot, err) + } + if err := hostfs.CopyTreeIfExists(rootMount, legacyNodeFilesDir, copyRoot); err != nil { + return fmt.Errorf("copying legacy node files from %q: %w", legacyNodeFilesDir, err) + } + return nil +} + +func loadConfig(copyRoot string, logger *slog.Logger) (*config.Config, error) { + configPath := filepath.Join(copyRoot, configFileName) + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("reading package config %q: %w", configPath, err) + } + cfg, err := config.NewLoader().Load(data, flags.StepsDir(copyRoot), logger) + if err != nil { + return nil, fmt.Errorf("loading package config %q: %w", configPath, err) + } + return cfg, nil +} + +func prepareHost(copyRoot, rootMount string, cfg config.Config) error { + if err := hostfs.CopyTreeIfExists( + rootMount, + filepath.Join(copyRoot, rootOverlayDirName), + rootMount, + ); err != nil { + return fmt.Errorf("copying package root overlay: %w", err) + } + for _, expected := range cfg.ExpectedConfigFiles { + if !filepath.IsLocal(expected) { + return fmt.Errorf("expected config file %q must be relative to the configmaps directory", expected) + } + path := filepath.Join(copyRoot, configMapsDirName, expected) + info, err := os.Stat(path) + if errors.Is(err, fs.ErrNotExist) { + return fmt.Errorf("expected config file %q was not found in the configmaps directory", expected) + } + if err != nil { + return fmt.Errorf("stating expected config file %q: %w", expected, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("expected config file %q is not a regular file", expected) + } + } + return nil +} + +func copyResolverConfig(rootMount string) error { + source := filepath.Join(string(filepath.Separator), "etc", "resolv.conf") + destination, err := hostfs.Resolve(rootMount, source) + if err != nil { + return fmt.Errorf("resolving host resolver path: %w", err) + } + if err := hostfs.CopyFile(rootMount, source, destination); err != nil { + return fmt.Errorf("copying resolver configuration: %w", err) + } + return nil +} diff --git a/agent/go/internal/agent/package_test.go b/agent/go/internal/agent/package_test.go new file mode 100644 index 00000000..7fb48878 --- /dev/null +++ b/agent/go/internal/agent/package_test.go @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "errors" + "os" + "path/filepath" + + "github.com/NVIDIA/nodewright/agent/internal/config" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("package preparation", func() { + It("copies the container resolver into the mounted root", func() { + _, err := os.Stat("/etc/resolv.conf") + if errors.Is(err, os.ErrNotExist) { + Skip("/etc/resolv.conf is not available") + } + Expect(err).NotTo(HaveOccurred()) + + root := GinkgoT().TempDir() + + Expect(copyResolverConfig(root)).To(Succeed()) + + expected, err := os.ReadFile("/etc/resolv.conf") + Expect(err).NotTo(HaveOccurred()) + actual, err := os.ReadFile(filepath.Join(root, "etc", "resolv.conf")) + Expect(err).NotTo(HaveOccurred()) + Expect(actual).To(Equal(expected)) + }) + + It("copies root overlays and validates expected configuration files", func() { + root := GinkgoT().TempDir() + copyRoot := filepath.Join(root, "package") + Expect(os.MkdirAll(filepath.Join(copyRoot, rootOverlayDirName, "etc"), 0o755)).To(Succeed()) + Expect(os.WriteFile( + filepath.Join(copyRoot, rootOverlayDirName, "etc", "package.conf"), + []byte("configured"), + 0o600, + )).To(Succeed()) + Expect(os.MkdirAll(filepath.Join(copyRoot, configMapsDirName), 0o755)).To(Succeed()) + Expect(os.WriteFile( + filepath.Join(copyRoot, configMapsDirName, "input.conf"), + []byte("input"), + 0o600, + )).To(Succeed()) + + Expect(prepareHost(copyRoot, root, config.Config{ + ExpectedConfigFiles: []string{"input.conf"}, + })).To(Succeed()) + Expect(filepath.Join(root, "etc", "package.conf")).To(BeAnExistingFile()) + + err := prepareHost(copyRoot, root, config.Config{ + ExpectedConfigFiles: []string{"missing.conf"}, + }) + Expect(err).To(MatchError(ContainSubstring( + `expected config file "missing.conf" was not found`, + ))) + }) +}) diff --git a/agent/go/internal/agent/steps.go b/agent/go/internal/agent/steps.go new file mode 100644 index 00000000..98e2d385 --- /dev/null +++ b/agent/go/internal/agent/steps.go @@ -0,0 +1,296 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "context" + "errors" + "fmt" + "io" + "path/filepath" + "strings" + "time" + + "github.com/NVIDIA/nodewright/agent/internal/config" + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/history" + "github.com/NVIDIA/nodewright/agent/internal/stage" + "github.com/NVIDIA/nodewright/agent/internal/step" +) + +const ( + startFlagName = "START" + checkResultsFlagName = "check_results" + logFileMode = 0o600 +) + +func runSteps( + ctx context.Context, + req request, + runtime runtimeConfig, + layout flags.Layout, + cfg config.Config, + flagStore flags.Store, + historyStore history.Store, +) (execution.Status, error) { + if err := flagStore.Write(filepath.Join(layout.FlagDir(), startFlagName), nil); err != nil { + return execution.StatusFailed, fmt.Errorf("writing agent start flag: %w", err) + } + + configuredSteps := cfg.Modes[req.stage] + if len(configuredSteps) == 0 { + runtime.logger.Warn("stage has no configured steps; treating it as a no-op", "stage", req.stage) + } + + var versions history.Versions + if isUpgradeStage(req.stage) { + var err error + versions, err = historyStore.Read() + if err != nil { + return execution.StatusFailed, fmt.Errorf( + "reading package history for stage %q: %w", + req.stage, + err, + ) + } + } + + _, isCheck := stage.CheckToApply[req.stage] + results := make([]checkResult, 0, len(configuredSteps)) + for _, configuredStep := range configuredSteps { + if err := ctx.Err(); err != nil { + return execution.StatusFailed, fmt.Errorf("running stage %q: %w", req.stage, err) + } + + runnableStep := configuredStep + if isUpgradeStage(req.stage) { + runnableStep = configuredStep.WithVersions(versions.Previous, versions.Current) + } + + if !isCheck { + decision, err := flagStore.Check(configuredStep, runtime.alwaysRunStep, req.stage) + if err != nil { + return execution.StatusFailed, fmt.Errorf( + "checking completion for step %q: %w", + configuredStep.Path(), + err, + ) + } + if decision.Reason == flags.ReasonAlwaysRun { + runtime.logger.Warn( + "running completed step because always-run policy is enabled", + "stage", req.stage, + "step", configuredStep.Path(), + ) + } + if !decision.Run { + runtime.logger.Info( + "skipping completed step", + "stage", req.stage, + "step", configuredStep.Path(), + "reason", decision.Reason, + ) + continue + } + } + + status, err := runStep(ctx, req, runtime, layout, cfg, runnableStep) + if err != nil { + return execution.StatusFailed, err + } + if isCheck { + results = append(results, checkResult{ + path: configuredStep.Path(), + failed: status != execution.StatusSuccess, + }) + continue + } + if status != execution.StatusSuccess { + return execution.StatusFailed, nil + } + message := fmt.Sprintf( + "last_run: %s\nstep_always_runs: %t", + time.Now().UTC().Format(time.RFC3339Nano), + configuredStep.Idempotence() == step.Disabled, + ) + if _, err := flagStore.Mark(configuredStep, message); err != nil { + return execution.StatusFailed, fmt.Errorf( + "marking step %q complete: %w", + configuredStep.Path(), + err, + ) + } + } + + if isCheck && len(configuredSteps) > 0 { + status, err := summarizeChecks(req.stage, results, flagStore, layout) + if err != nil || status == execution.StatusFailed { + return status, err + } + } + if recordsHistory(req.stage) { + if err := historyStore.Record(req.stage, time.Now()); err != nil { + return execution.StatusFailed, fmt.Errorf( + "recording package history for stage %q: %w", + req.stage, + err, + ) + } + } + if req.stage == stage.UninstallCheck { + if err := removeStepFlags(cfg, flagStore); err != nil { + return execution.StatusFailed, err + } + } + return execution.StatusSuccess, nil +} + +func runStep( + ctx context.Context, + req request, + runtime runtimeConfig, + layout flags.Layout, + cfg config.Config, + value step.Step, +) (execution.Status, error) { + stdout, stderr := runtime.stdout, runtime.stderr + closeLog := func() error { return nil } + if runtime.writeLogs { + _, file, err := layout.CreateLogFile(cfg, value.Path(), time.Now(), logFileMode) + if err != nil { + return execution.StatusFailed, fmt.Errorf("preparing log for step %q: %w", value.Path(), err) + } + stdout = io.MultiWriter(stdout, file) + stderr = io.MultiWriter(stderr, file) + closeLog = file.Close + } + + runConfig, err := execution.NewConfig( + execution.WithRootMount(req.rootMount), + execution.WithStepRoot(flags.StepsDir(req.copyDir)), + execution.WithSkyhookDir(req.copyDir), + execution.WithRunOutput(stdout, stderr), + ) + if err != nil { + closeErr := closeLog() + if closeErr != nil { + closeErr = fmt.Errorf("closing log for step %q: %w", value.Path(), closeErr) + } + return execution.StatusFailed, errors.Join( + fmt.Errorf("configuring step %q execution: %w", value.Path(), err), + closeErr, + ) + } + status, runErr := value.Run(ctx, runConfig) + closeErr := closeLog() + if runErr != nil || closeErr != nil { + if runErr != nil { + runErr = fmt.Errorf("running step %q: %w", value.Path(), runErr) + } + if closeErr != nil { + closeErr = fmt.Errorf("closing log for step %q: %w", value.Path(), closeErr) + } + return execution.StatusFailed, errors.Join( + runErr, + closeErr, + ) + } + if runtime.writeLogs { + logFiles, err := layout.LogFilePattern(cfg, value.Path()) + if err != nil { + return execution.StatusFailed, fmt.Errorf( + "resolving log retention for step %q: %w", + value.Path(), + err, + ) + } + if err := flags.CleanupOldLogs(logFiles, flags.DefaultLogRetention); err != nil { + return execution.StatusFailed, fmt.Errorf("cleaning old logs for step %q: %w", value.Path(), err) + } + } + return status, nil +} + +type checkResult struct { + path string + failed bool +} + +func summarizeChecks( + currentStage stage.Stage, + results []checkResult, + flagStore flags.Store, + layout flags.Layout, +) (execution.Status, error) { + lines := make([]string, 0, len(results)) + failed := false + for _, result := range results { + value := "False" + if result.failed { + value = "True" + failed = true + } + lines = append(lines, result.path+" "+value) + } + if err := flagStore.Write( + filepath.Join(layout.FlagDir(), checkResultsFlagName), + []byte(strings.Join(lines, "\n")), + ); err != nil { + return execution.StatusFailed, fmt.Errorf("writing check results for stage %q: %w", currentStage, err) + } + if failed { + return execution.StatusFailed, nil + } + if err := flagStore.Write( + filepath.Join(layout.FlagDir(), string(currentStage)+"_ALL_CHECKED"), + nil, + ); err != nil { + return execution.StatusFailed, fmt.Errorf( + "writing completed-check flag for stage %q: %w", + currentStage, + err, + ) + } + return execution.StatusSuccess, nil +} + +func isUpgradeStage(currentStage stage.Stage) bool { + return currentStage == stage.Upgrade || currentStage == stage.UpgradeCheck +} + +func recordsHistory(currentStage stage.Stage) bool { + switch currentStage { + case stage.ApplyCheck, stage.UpgradeCheck, stage.UninstallCheck: + return true + default: + return false + } +} + +func removeStepFlags(cfg config.Config, flagStore flags.Store) error { + for _, currentStage := range stage.All { + for _, value := range cfg.Modes[currentStage] { + if err := flagStore.Remove(value); err != nil { + return fmt.Errorf("removing completion flag for step %q: %w", value.Path(), err) + } + } + } + return nil +} diff --git a/agent/go/internal/agent/steps_test.go b/agent/go/internal/agent/steps_test.go new file mode 100644 index 00000000..9a2cd0ce --- /dev/null +++ b/agent/go/internal/agent/steps_test.go @@ -0,0 +1,279 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * 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 agent + +import ( + "bytes" + "context" + "errors" + "io" + "log/slog" + "os" + "path/filepath" + + "github.com/NVIDIA/nodewright/agent/internal/config" + "github.com/NVIDIA/nodewright/agent/internal/execution" + "github.com/NVIDIA/nodewright/agent/internal/flags" + "github.com/NVIDIA/nodewright/agent/internal/history" + historymock "github.com/NVIDIA/nodewright/agent/internal/history/mock" + "github.com/NVIDIA/nodewright/agent/internal/stage" + "github.com/NVIDIA/nodewright/agent/internal/step" + stepmock "github.com/NVIDIA/nodewright/agent/internal/step/mock" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/stretchr/testify/mock" +) + +func newMockStep(path string) *stepmock.MockStep { + value := stepmock.NewMockStep(GinkgoT()) + value.EXPECT().Path().Return(path).Maybe() + value.EXPECT(). + Fingerprint(). + Return("fingerprint-"+filepath.Base(path), nil). + Maybe() + value.EXPECT().Idempotence().Return(step.Auto).Maybe() + return value +} + +func newRunnableMockStep(path string, status execution.Status) *stepmock.MockStep { + value := newMockStep(path) + value.EXPECT(). + Run(mock.Anything, mock.Anything). + Return(status, nil). + Once() + return value +} + +var _ = Describe("step orchestration", func() { + var ( + root string + layout flags.Layout + cfg config.Config + runtime runtimeConfig + req request + flagStore flags.Store + historyStore *historymock.MockStore + ) + + BeforeEach(func() { + root = GinkgoT().TempDir() + layout = flags.DefaultLayout(root) + cfg = config.Config{ + PackageName: "package", + PackageVersion: "1.2.3", + Modes: map[stage.Stage][]step.Step{}, + } + runtime = runtimeConfig{ + stdout: io.Discard, + stderr: io.Discard, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } + req = request{ + stage: stage.Apply, + rootMount: root, + copyDir: "/packages/package", + } + var err error + flagStore, err = flags.NewStore(layout, cfg) + Expect(err).NotTo(HaveOccurred()) + historyStore = historymock.NewMockStore(GinkgoT()) + }) + + It("writes START and completion flags around successful steps", func() { + first := newRunnableMockStep("first.sh", execution.StatusSuccess) + second := newRunnableMockStep("second.sh", execution.StatusSuccess) + cfg.Modes[stage.Apply] = []step.Step{first, second} + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join(layout.FlagDir(), startFlagName)).To(BeAnExistingFile()) + firstFlag, err := flagStore.Path(first) + Expect(err).NotTo(HaveOccurred()) + Expect(firstFlag).To(BeAnExistingFile()) + }) + + It("skips an idempotent step that is already complete", func() { + value := newMockStep("apply.sh") + cfg.Modes[stage.Apply] = []step.Step{value} + _, err := flagStore.Mark(value, "complete") + Expect(err).NotTo(HaveOccurred()) + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + }) + + It("runs a completed step when the always-run policy is enabled", func() { + value := newRunnableMockStep("apply.sh", execution.StatusSuccess) + cfg.Modes[stage.Apply] = []step.Step{value} + runtime.alwaysRunStep = true + logOutput := &bytes.Buffer{} + runtime.logger = slog.New(slog.NewTextHandler(logOutput, nil)) + _, err := flagStore.Mark(value, "complete") + Expect(err).NotTo(HaveOccurred()) + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(logOutput.String()).To(ContainSubstring( + "running completed step because always-run policy is enabled", + )) + }) + + It("stops a normal stage after the first failed step", func() { + first := newRunnableMockStep("first.sh", execution.StatusFailed) + second := newMockStep("second.sh") + cfg.Modes[stage.Apply] = []step.Step{first, second} + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusFailed)) + }) + + It("runs every check and persists mixed results", func() { + req.stage = stage.ApplyCheck + first := newRunnableMockStep("first.sh", execution.StatusSuccess) + second := newRunnableMockStep("second.sh", execution.StatusFailed) + cfg.Modes[stage.ApplyCheck] = []step.Step{first, second} + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusFailed)) + data, err := os.ReadFile(filepath.Join(layout.FlagDir(), checkResultsFlagName)) + Expect(err).NotTo(HaveOccurred()) + Expect(string(data)).To(Equal("first.sh False\nsecond.sh True")) + Expect(filepath.Join(layout.FlagDir(), "apply-check_ALL_CHECKED")).NotTo(BeAnExistingFile()) + }) + + It("marks successful checks and records package history", func() { + req.stage = stage.UpgradeCheck + value := newRunnableMockStep("check.sh", execution.StatusSuccess) + cfg.Modes[stage.UpgradeCheck] = []step.Step{value} + versions := history.Versions{Previous: "1.0.0", Current: "1.2.3"} + historyStore.EXPECT().Read().Return(versions, nil).Once() + historyStore.EXPECT(). + Record(stage.UpgradeCheck, mock.Anything). + Return(nil). + Once() + value.EXPECT(). + WithVersions(versions.Previous, versions.Current). + Return(value). + Once() + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + Expect(filepath.Join(layout.FlagDir(), "upgrade-check_ALL_CHECKED")).To(BeAnExistingFile()) + }) + + It("removes every step completion flag after uninstall-check", func() { + req.stage = stage.UninstallCheck + apply := newMockStep("apply.sh") + uninstallCheck := newRunnableMockStep("uninstall-check.sh", execution.StatusSuccess) + historyStore.EXPECT(). + Record(stage.UninstallCheck, mock.Anything). + Return(nil). + Once() + cfg.Modes[stage.Apply] = []step.Step{apply} + cfg.Modes[stage.UninstallCheck] = []step.Step{uninstallCheck} + for _, value := range []step.Step{apply, uninstallCheck} { + _, err := flagStore.Mark(value, "complete") + Expect(err).NotTo(HaveOccurred()) + } + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + for _, value := range []step.Step{apply, uninstallCheck} { + path, err := flagStore.Path(value) + Expect(err).NotTo(HaveOccurred()) + Expect(path).NotTo(BeAnExistingFile()) + } + }) + + It("propagates cancellation and does not start another step", func() { + ctx, cancel := context.WithCancel(context.Background()) + first := newMockStep("first.sh") + first.EXPECT(). + Run(mock.Anything, mock.Anything). + Run(func(context.Context, execution.Config) { + cancel() + }). + Return(execution.StatusSuccess, nil). + Once() + second := newMockStep("second.sh") + cfg.Modes[stage.Apply] = []step.Step{first, second} + + status, err := runSteps(ctx, req, runtime, layout, cfg, flagStore, historyStore) + + Expect(status).To(Equal(execution.StatusFailed)) + Expect(errors.Is(err, context.Canceled)).To(BeTrue()) + }) + + It("composes command output with a retained log file", func() { + runtime.writeLogs = true + output := "step output\n" + value := newMockStep("apply.sh") + value.EXPECT(). + Run(mock.Anything, mock.Anything). + Run(func(_ context.Context, cfg execution.Config) { + _, err := io.WriteString(cfg.Stdout(), output) + Expect(err).NotTo(HaveOccurred()) + }). + Return(execution.StatusSuccess, nil). + Once() + cfg.Modes[stage.Apply] = []step.Step{value} + + status, err := runSteps( + context.Background(), req, runtime, layout, cfg, flagStore, historyStore, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(status).To(Equal(execution.StatusSuccess)) + entries, err := os.ReadDir(filepath.Join(layout.LogDir(), cfg.PackageName, cfg.PackageVersion)) + Expect(err).NotTo(HaveOccurred()) + Expect(entries).To(HaveLen(1)) + data, err := os.ReadFile(filepath.Join(layout.LogDir(), cfg.PackageName, cfg.PackageVersion, entries[0].Name())) + Expect(err).NotTo(HaveOccurred()) + Expect(string(data)).To(Equal(output)) + }) +})