diff --git a/internal/cli/drain.go b/internal/cli/drain.go new file mode 100644 index 0000000..5570c28 --- /dev/null +++ b/internal/cli/drain.go @@ -0,0 +1,657 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "github.com/spf13/cobra" + + "git.cloudinit.dev/coreci/orca/internal/certpaths" + "git.cloudinit.dev/coreci/orca/internal/engine" + "git.cloudinit.dev/coreci/orca/internal/model" + "git.cloudinit.dev/coreci/orca/internal/sshpush" + "git.cloudinit.dev/coreci/orca/internal/store" +) + +// drainExecer is the SSH command-execution seam used by the drain +// commands. *sshpush.Transport satisfies it via its Exec method; tests +// inject a record-and-replay mock without a real SSH server (same +// pattern as internal/stepca mockExec). +type drainExecer interface { + Exec(ctx context.Context, peer string, cmd string) ([]byte, error) +} + +// drainTransport is the package-level exec seam. It is set by +// drainExecFromCtx (production) and overridden by tests via +// drainExecOverride. nil means "build from certpaths on first use". +var drainExecOverride drainExecer + +// drainExecFromCtx returns the production drainExecer backed by the +// real sshpush.Transport (using the cluster SSH key + known_hosts). On +// error it returns a nil transport and the error; callers must check. +func drainExecFromCtx(_ context.Context) (drainExecer, error) { + if drainExecOverride != nil { + return drainExecOverride, nil + } + keyPath := certpaths.SSHKeyPath() + khPath := certpaths.KnownHostsPath() + tr := sshpush.NewTransport(keyPath, khPath) + return tr, nil +} + +// peerAddrForNode derives the SSH peer address (host:port) for a node. +// For proxmox nodes the Name IS the host; for localhost nodes the +// Address carries host:8443. We always target SSH port 22 unless the +// node's Address already encodes a non-daemon port. The local node +// (Name=="localhost") is contacted at "localhost:22". +func peerAddrForNode(n *model.Node) string { + if n == nil { + return "" + } + if h, p, ok := splitHostPort(n.Address); ok && p != "" && p != "8443" { + return h + ":" + p + } + host := n.Name + if h, _, ok := splitHostPort(n.Address); ok && h != "" && h != "localhost" { + host = h + } + if host == "" { + host = n.Name + } + return host + ":22" +} + +func splitHostPort(addr string) (string, string, bool) { + idx := strings.LastIndex(addr, ":") + if idx < 0 { + return addr, "", false + } + return addr[:idx], addr[idx+1:], true +} + +var ( + drainTimeout time.Duration + migrateTarget string +) + +// allocUnit is the systemd unit name pattern for orca allocations. +const allocUnitPrefix = "orca-alloc-" +const allocUnitSuffix = ".service" + +// allocIDFromUnit strips the orca-alloc- prefix and .service suffix +// from a systemd unit name, returning the bare allocation id. +func allocIDFromUnit(unit string) string { + s := strings.TrimSpace(unit) + s = strings.TrimPrefix(s, allocUnitPrefix) + s = strings.TrimSuffix(s, allocUnitSuffix) + return s +} + +// allocUnit renders the systemd unit name for an allocation id. +func allocUnit(allocID string) string { + return allocUnitPrefix + allocID + allocUnitSuffix +} + +// listRunningAllocs queries a node via SSH for the currently-running +// orca-alloc-*.service systemd units and returns their allocation +// ids. A node with no orca allocations returns an empty slice (not an +// error). +func listRunningAllocs(ctx context.Context, ex drainExecer, peer string) ([]string, error) { + cmd := "systemctl list-units 'orca-alloc-*.service' --type=service --state=running --no-legend --no-pager" + out, err := ex.Exec(ctx, peer, cmd) + if err != nil { + return nil, fmt.Errorf("list orca-alloc units on %s: %w", peer, err) + } + var ids []string + for _, line := range strings.Split(string(out), "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + fields := strings.Fields(line) + if len(fields) == 0 { + continue + } + unit := fields[0] + if !strings.HasPrefix(unit, allocUnitPrefix) || !strings.HasSuffix(unit, allocUnitSuffix) { + continue + } + ids = append(ids, allocIDFromUnit(unit)) + } + return ids, nil +} + +// stopAlloc sends `systemctl stop orca-alloc-.service` to a node. +// A unit that is already stopped (or never existed) is treated as +// success: drain is idempotent. +func stopAlloc(ctx context.Context, ex drainExecer, peer, allocID string) error { + cmd := fmt.Sprintf("systemctl stop %s", allocUnit(allocID)) + _, err := ex.Exec(ctx, peer, cmd) + if err != nil { + var exitErr *sshExitErr + if errors.As(err, &exitErr) && exitErr.code == 5 { + return nil + } + return fmt.Errorf("stop %s on %s: %w", allocID, peer, err) + } + return nil +} + +// sshExitErr is a lightweight sentinel used by the in-package mock to +// signal a non-zero systemctl exit. The real sshpush.Transport wraps +// non-zero exits in ErrPermanent; the mock returns an *sshExitErr so +// stopAlloc can treat code 5 ("unit not loaded") as success. +type sshExitErr struct{ code int } + +func (e *sshExitErr) Error() string { return fmt.Sprintf("sshpush: exit %d", e.code) } + +// waitAllocsStopped polls a node until none of the given allocation +// ids appear in the running-unit list, or the context deadline passes. +// Returns nil if all allocations are observed stopped; otherwise +// returns a list of allocations that were still running at timeout. +func waitAllocsStopped(ctx context.Context, ex drainExecer, peer string, ids []string, poll time.Duration) []string { + if len(ids) == 0 { + return nil + } + pending := make(map[string]bool, len(ids)) + for _, id := range ids { + pending[id] = true + } + if poll <= 0 { + poll = 500 * time.Millisecond + } + ticker := time.NewTicker(poll) + defer ticker.Stop() + for { + if err := ctx.Err(); err != nil { + break + } + running, err := listRunningAllocs(ctx, ex, peer) + if err == nil { + runningSet := make(map[string]bool, len(running)) + for _, rid := range running { + runningSet[rid] = true + } + for id := range pending { + if !runningSet[id] { + delete(pending, id) + } + } + } + if len(pending) == 0 { + return nil + } + select { + case <-ctx.Done(): + return keysOf(pending) + case <-ticker.C: + } + } + return keysOf(pending) +} + +func keysOf(m map[string]bool) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} + +// findNode resolves a node by id or name from the registry. Returns +// nil + error if not found. +func findNode(ctx context.Context, reg *engine.NodeRegistry, ref string) (*model.Node, error) { + if n, err := reg.Get(ctx, ref); err == nil { + return n, nil + } else if !errors.Is(err, store.ErrNotFound) { + return nil, fmt.Errorf("lookup node %q: %w", ref, err) + } + nodes, err := reg.List(ctx) + if err != nil { + return nil, fmt.Errorf("list nodes: %w", err) + } + for _, n := range nodes { + if n.Name == ref || n.ID == ref { + return n, nil + } + } + return nil, fmt.Errorf("node %q not found in the registry", ref) +} + +// auditDrain records a node.drain event in the audit log. +func auditDrain(ctx context.Context, nodeID, result string, err error, meta map[string]any) { + db, dbErr := store.Open(certpaths.DBPath()) + if dbErr != nil { + return + } + defer db.Close() + engine.NewAudit(store.NewAuditRepo(db), newLogger()).Record(ctx, "cli", "node.drain", nodeID, result, err, meta) +} + +var nodeDrainCmd = &cobra.Command{ + Use: "drain ", + Short: "Drain a node: stop its allocations and mark it drained", + Long: `Drain a node (REQ-061). + +Marks the node as "draining" (the scheduler skips draining nodes), stops +every running allocation on the node via SSH (systemctl stop +orca-alloc-.service), waits for them to stop (up to --timeout, +default 30s), and marks the node "drained" when all allocations are +stopped. The scheduler already skips draining/drained nodes. + + is the node name or id (as shown by 'orca node list'). + +This is NOT live-migration: allocations are stopped, not moved. Use +'orca job migrate' to reschedule a job onto another node before +draining.`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + hostRef := args[0] + ctx, cancel := context.WithTimeout(cmd.Context(), drainTimeout+10*time.Second) + defer cancel() + + reg, closer, err := nodeRegistry() + if err != nil { + return err + } + defer closer() + + node, err := findNode(ctx, reg, hostRef) + if err != nil { + return err + } + peer := peerAddrForNode(node) + if peer == "" { + return fmt.Errorf("cannot resolve SSH address for node %q", node.Name) + } + + ex, err := drainExecFromCtx(cmd.Context()) + if err != nil { + return fmt.Errorf("ssh transport: %w", err) + } + + log := newLogger() + log.Info("drain: marking node draining", + slog.String("node", node.Name), slog.String("peer", peer)) + + if err := reg.SetNodeState(ctx, node.ID, string(model.NodeStateDraining)); err != nil { + return fmt.Errorf("mark node draining: %w", err) + } + + ids, err := listRunningAllocs(ctx, ex, peer) + if err != nil { + _ = reg.SetNodeState(ctx, node.ID, string(model.NodeStateReady)) + auditDrain(ctx, node.ID, "failure", err, map[string]any{"peer": peer}) + return err + } + + stopped := make([]string, 0, len(ids)) + var failures []string + for _, id := range ids { + if err := stopAlloc(ctx, ex, peer, id); err != nil { + failures = append(failures, id) + log.Warn("drain: failed to stop alloc", + slog.String("alloc", id), slog.String("node", node.Name), "error", err) + continue + } + stopped = append(stopped, id) + } + + // Wait for the stopped allocations to actually leave the + // running list (with the drain timeout as the deadline). + waitCtx, waitCancel := context.WithTimeout(ctx, drainTimeout) + remaining := waitAllocsStopped(waitCtx, ex, peer, stopped, 500*time.Millisecond) + waitCancel() + + result := map[string]any{ + "node": node.Name, + "node_id": node.ID, + "peer": peer, + "stopped": stopped, + } + if len(failures) > 0 { + result["failed"] = failures + } + if len(remaining) > 0 { + result["still_running"] = remaining + } + + if len(failures) == 0 && len(remaining) == 0 { + if err := reg.SetNodeState(ctx, node.ID, string(model.NodeStateDrained)); err != nil { + auditDrain(ctx, node.ID, "failure", err, result) + return fmt.Errorf("mark node drained: %w", err) + } + result["state"] = string(model.NodeStateDrained) + auditDrain(ctx, node.ID, "success", nil, result) + if jsonOutput { + return printJSON(result) + } + fmt.Fprintf(cmd.OutOrStdout(), "✓ Node %s drained (%d allocs stopped)\n", node.Name, len(stopped)) + for _, id := range stopped { + fmt.Fprintf(cmd.OutOrStdout(), " stopped %s\n", id) + } + return nil + } + + // Partial: leave node in draining state so the operator can + // retry; record the partial outcome in the audit log. + result["state"] = string(model.NodeStateDraining) + auditDrain(ctx, node.ID, "partial", fmt.Errorf("%d failed, %d still running", len(failures), len(remaining)), result) + if jsonOutput { + return printJSON(result) + } + fmt.Fprintf(cmd.OutOrStdout(), "⚠ Node %s partially drained (%d stopped, %d failed, %d still running)\n", + node.Name, len(stopped), len(failures), len(remaining)) + for _, id := range stopped { + fmt.Fprintf(cmd.OutOrStdout(), " stopped %s\n", id) + } + for _, id := range failures { + fmt.Fprintf(cmd.OutOrStdout(), " failed %s\n", id) + } + for _, id := range remaining { + fmt.Fprintf(cmd.OutOrStdout(), " running %s\n", id) + } + return fmt.Errorf("drain incomplete: %d failed, %d still running", len(failures), len(remaining)) + }, +} + +// daemonDrainAndStopCmd repurposes the deprecated `orca daemon` command +// to drain all nodes and stop all v0.8 daemons (REQ-061, R-001). It is +// the v0.8→v0.11 migration path for daemon removal: it SSHes to every +// peer that still runs an orca daemon and stops the daemon service, +// relying on the v0.9+ SSH-push path (systemd) to keep workloads alive. +var daemonDrainAndStopCmd = &cobra.Command{ + Use: "drain-and-stop", + Short: "Stop v0.8 orca daemons on all peers (v0.8→v0.11 migration)", + Long: `Stop the orca daemon on every peer that still runs one (REQ-061, R-001). + +This is the v0.8→v0.11 migration path for daemon removal. For each +registered peer, SSH in and run 'systemctl stop orca-daemon.service'. +Workloads supervised by the v0.9+ SSH-push path keep running under +their own systemd units (orca-alloc-*.service) and are NOT touched. + +This command is idempotent: a peer with no orca-daemon.service (already +migrated, or never had one) is reported as "already stopped" and does +not error.`, + RunE: func(cmd *cobra.Command, args []string) error { + ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Minute) + defer cancel() + + reg, closer, err := nodeRegistry() + if err != nil { + return err + } + defer closer() + + nodes, err := reg.List(ctx) + if err != nil { + return fmt.Errorf("list nodes: %w", err) + } + + ex, err := drainExecFromCtx(cmd.Context()) + if err != nil { + return fmt.Errorf("ssh transport: %w", err) + } + + const stopCmd = "systemctl stop orca-daemon.service" + + result := map[string]any{ + "stopped": []string{}, + "already_stopped": []string{}, + "failed": []string{}, + } + var stopped, already, failed []string + for _, n := range nodes { + peer := peerAddrForNode(n) + if peer == "" { + continue + } + _, err := ex.Exec(ctx, peer, stopCmd) + if err == nil { + stopped = append(stopped, n.Name) + continue + } + var exitErr *sshExitErr + if errors.As(err, &exitErr) && exitErr.code == 5 { + already = append(already, n.Name) + continue + } + failed = append(failed, n.Name) + } + result["stopped"] = stopped + result["already_stopped"] = already + result["failed"] = failed + + db, dbErr := store.Open(certpaths.DBPath()) + if dbErr == nil { + defer db.Close() + engine.NewAudit(store.NewAuditRepo(db), newLogger()).Record(ctx, "cli", "daemon.drain_and_stop", "cluster", "success", nil, result) + } + + if jsonOutput { + return printJSON(result) + } + fmt.Fprintf(cmd.OutOrStdout(), "✓ daemon drain-and-stop complete (%d stopped, %d already stopped, %d failed)\n", + len(stopped), len(already), len(failed)) + for _, n := range stopped { + fmt.Fprintf(cmd.OutOrStdout(), " stopped %s\n", n) + } + for _, n := range already { + fmt.Fprintf(cmd.OutOrStdout(), " already-stopped %s\n", n) + } + for _, n := range failed { + fmt.Fprintf(cmd.OutOrStdout(), " failed %s\n", n) + } + if len(failed) > 0 { + return fmt.Errorf("daemon drain-and-stop: %d peer(s) failed", len(failed)) + } + return nil + }, +} + +// jobMigrateCmd implements `orca job migrate --to ` (REQ-116, +// C3=a). It is a drain+reschedule composite — NOT live-migration (no +// storage replication). For each allocation of the named job running on +// a node OTHER than --to, it stops the allocation (SSH systemctl stop), +// then starts a new allocation on the target node (SSH systemctl start). +// Idempotent: if the job is already running on the target node, it is a +// no-op for that allocation. +var jobMigrateCmd = &cobra.Command{ + Use: "migrate ", + Short: "Drain+reschedule a job onto a target node (REQ-116)", + Long: `Migrate a job onto a target node (REQ-116, C3=a). + +This is a drain+reschedule composite, NOT live-migration: there is no +storage replication. For every allocation of currently running +on a node other than --to, it stops the allocation (SSH systemctl stop +orca-alloc-.service) and starts a new allocation on the target +node (SSH systemctl start orca-alloc-.service). + +Idempotent: if the job already has an allocation on the target node, +the command is a no-op (C3=a, Q2=C). + +NOTE: this command assumes allocations are tracked as systemd units +named orca-alloc-.service across the cluster, with of the +form -. The new allocation on the target node +is named -migrated-.`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + jobName := args[0] + if migrateTarget == "" { + return fmt.Errorf("--to is required") + } + ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Minute) + defer cancel() + + reg, closer, err := nodeRegistry() + if err != nil { + return err + } + defer closer() + + nodes, err := reg.List(ctx) + if err != nil { + return fmt.Errorf("list nodes: %w", err) + } + + target, err := findNode(ctx, reg, migrateTarget) + if err != nil { + return err + } + targetPeer := peerAddrForNode(target) + if targetPeer == "" { + return fmt.Errorf("cannot resolve SSH address for target node %q", target.Name) + } + + ex, err := drainExecFromCtx(cmd.Context()) + if err != nil { + return fmt.Errorf("ssh transport: %w", err) + } + + log := newLogger() + log.Info("migrate: scanning cluster for job allocations", slog.String("job", jobName)) + + // Find allocations of across all nodes. An + // allocation "belongs to" the job if its alloc id starts + // with "-" (the scheduler emits ids of the form + // ns/name-idx; we match on the name segment). + jobPrefix := jobName + "-" + + type alloc struct { + node *model.Node + peer string + id string + } + var onTarget, onOthers []alloc + + for i := range nodes { + n := nodes[i] + peer := peerAddrForNode(n) + if peer == "" { + continue + } + ids, err := listRunningAllocs(ctx, ex, peer) + if err != nil { + log.Warn("migrate: cannot list allocs on node", + slog.String("node", n.Name), "error", err) + continue + } + for _, id := range ids { + if !strings.HasPrefix(id, jobPrefix) { + continue + } + a := alloc{node: n, peer: peer, id: id} + if n.ID == target.ID { + onTarget = append(onTarget, a) + } else { + onOthers = append(onOthers, a) + } + } + } + + result := map[string]any{ + "job": jobName, + "target": target.Name, + "already_on_target": len(onTarget) > 0, + } + + // Idempotent: if the job is already running on the target + // node and there is nothing to migrate, no-op. + if len(onOthers) == 0 { + result["stopped"] = []string{} + result["started"] = []string{} + auditMigrate(ctx, jobName, target.Name, "success", nil, result) + if jsonOutput { + return printJSON(result) + } + if len(onTarget) > 0 { + fmt.Fprintf(cmd.OutOrStdout(), "✓ Job %s already running on %s (%d alloc(s)); nothing to migrate\n", + jobName, target.Name, len(onTarget)) + } else { + fmt.Fprintf(cmd.OutOrStdout(), "✓ Job %s has no running allocations to migrate\n", jobName) + } + return nil + } + + // Stop each allocation on a non-target node. + var stopped []string + var stopFailures []string + for _, a := range onOthers { + if err := stopAlloc(ctx, ex, a.peer, a.id); err != nil { + stopFailures = append(stopFailures, a.id) + log.Warn("migrate: failed to stop alloc", + slog.String("alloc", a.id), slog.String("node", a.node.Name), "error", err) + continue + } + stopped = append(stopped, a.id) + } + + // Start a new allocation on the target node. We use a + // stable, deterministic id so re-runs are idempotent: if the + // unit already exists & is running, systemctl start is a + // no-op. The new id is "-migrated-". + newID := fmt.Sprintf("%s-migrated-%d", jobName, time.Now().Unix()) + startCmd := fmt.Sprintf("systemctl start %s", allocUnit(newID)) + var started []string + if _, err := ex.Exec(ctx, targetPeer, startCmd); err != nil { + log.Warn("migrate: failed to start new alloc on target", + slog.String("alloc", newID), slog.String("node", target.Name), "error", err) + result["start_error"] = err.Error() + } else { + started = append(started, newID) + } + + result["stopped"] = stopped + result["started"] = started + result["stop_failures"] = stopFailures + + if len(stopFailures) == 0 && len(started) > 0 { + auditMigrate(ctx, jobName, target.Name, "success", nil, result) + } else { + auditMigrate(ctx, jobName, target.Name, "partial", + fmt.Errorf("%d stop failures, %d started", len(stopFailures), len(started)), result) + } + + if jsonOutput { + return printJSON(result) + } + fmt.Fprintf(cmd.OutOrStdout(), "✓ Migrated job %s onto %s (%d stopped, %d started)\n", + jobName, target.Name, len(stopped), len(started)) + for _, id := range stopped { + fmt.Fprintf(cmd.OutOrStdout(), " stopped %s\n", id) + } + for _, id := range started { + fmt.Fprintf(cmd.OutOrStdout(), " started %s on %s\n", id, target.Name) + } + if len(stopFailures) > 0 { + return fmt.Errorf("migrate: %d alloc(s) failed to stop", len(stopFailures)) + } + return nil + }, +} + +// auditMigrate records a job.migrate event in the audit log. +func auditMigrate(ctx context.Context, jobName, target, result string, err error, meta map[string]any) { + db, dbErr := store.Open(certpaths.DBPath()) + if dbErr != nil { + return + } + defer db.Close() + engine.NewAudit(store.NewAuditRepo(db), newLogger()).Record(ctx, "cli", "job.migrate", jobName, result, err, meta) +} + +func init() { + nodeDrainCmd.Flags().DurationVar(&drainTimeout, "timeout", 30*time.Second, + "max time to wait for allocations to stop") + nodeCmd.AddCommand(nodeDrainCmd) + + daemonCmd.AddCommand(daemonDrainAndStopCmd) + + jobMigrateCmd.Flags().StringVar(&migrateTarget, "to", "", "target node name or id to migrate the job onto (required)") + jobCmd.AddCommand(jobMigrateCmd) + +} diff --git a/internal/cli/drain_test.go b/internal/cli/drain_test.go new file mode 100644 index 0000000..8bad026 --- /dev/null +++ b/internal/cli/drain_test.go @@ -0,0 +1,526 @@ +package cli + +import ( + "bytes" + "context" + "errors" + "strings" + "sync" + "testing" + "time" + + "git.cloudinit.dev/coreci/orca/internal/certpaths" + "git.cloudinit.dev/coreci/orca/internal/model" + "git.cloudinit.dev/coreci/orca/internal/store" +) + +// mockDrainExec is a record-and-replay execer for the drain commands +// (same pattern as internal/stepca mockExec). It matches each incoming +// command against a list of (substring, output, exitCode) responses; +// the first match wins. An entry with an empty substring matches any +// command. The exit code is 0 (success) unless explicitly set; a +// non-zero code is returned as an *sshExitErr so stopAlloc can treat +// code 5 ("unit not loaded") as idempotent success. +type mockDrainExec struct { + mu sync.Mutex + responses []mockDrainResp + calls []mockDrainCall +} + +type mockDrainResp struct { + match string + out string + exit int +} + +type mockDrainCall struct { + peer string + cmd string +} + +func (m *mockDrainExec) Exec(_ context.Context, peer, cmd string) ([]byte, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.calls = append(m.calls, mockDrainCall{peer: peer, cmd: cmd}) + for _, r := range m.responses { + if r.match == "" || strings.Contains(cmd, r.match) { + if r.exit != 0 { + return []byte(r.out), &sshExitErr{code: r.exit} + } + return []byte(r.out), nil + } + } + return nil, nil +} + +func (m *mockDrainExec) callsFor(match string) []mockDrainCall { + m.mu.Lock() + defer m.mu.Unlock() + var out []mockDrainCall + for _, c := range m.calls { + if strings.Contains(c.cmd, match) { + out = append(out, c) + } + } + return out +} + +func (m *mockDrainExec) countCalls(match string) int { + return len(m.callsFor(match)) +} + +// drainTestEnv wires a mockDrainExec into drainExecOverride and returns +// the mock + a cleanup func. Tests MUST defer the cleanup. +func drainTestEnv(t *testing.T) *mockDrainExec { + t.Helper() + prev := drainExecOverride + mx := &mockDrainExec{} + drainExecOverride = mx + t.Cleanup(func() { drainExecOverride = prev }) + return mx +} + +// drainNodeForTest inserts a node with a fixed id+name and returns it, +// so drain commands can target it by name. Uses the test ORCA_HOME db. +func drainNodeForTest(t *testing.T, name, addr string) *model.Node { + t.Helper() + db, err := store.Open(certpaths.DBPath()) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer db.Close() + n := &model.Node{ + ID: "node-" + name, + Name: name, + Address: addr, + State: model.NodeStateReady, + JoinedAt: time.Now().UTC(), + LastSeen: time.Now().UTC(), + Kind: string(model.NodeKindLinux), + } + if err := store.NewNodeRepo(db).Insert(context.Background(), n); err != nil { + t.Fatalf("insert node: %v", err) + } + return n +} + +func TestNodeDrain_StopsAllocsAndMarksDrained(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + node := drainNodeForTest(t, "drainee", "drainee:8443") + drainTestEnv(t) // sets drainExecOverride (reset in cleanup) + // Scripted exec: first list-units returns 2 running allocs; the + // post-stop poll returns empty so waitAllocsStopped completes. + mx := &scriptedDrainExec{} + mx.queue("list-units", "orca-alloc-web-0.service loaded active running\norca-alloc-web-1.service loaded active running\n", 0) + mx.queue("systemctl stop orca-alloc-web-0", "", 0) + mx.queue("systemctl stop orca-alloc-web-1", "", 0) + mx.queue("list-units", "", 0) + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + nodeDrainCmd.SetOut(&buf) + nodeDrainCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "5s"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("node drain: %v", err) + } + out := buf.String() + if !strings.Contains(out, "drained") { + t.Errorf("expected drained message, got: %s", out) + } + if mx.countCalls("systemctl stop orca-alloc-web-0") == 0 || mx.countCalls("systemctl stop orca-alloc-web-1") == 0 { + t.Errorf("expected stop commands for both allocs, calls: %+v", mx.calls) + } + + // Node state should be drained in the DB. + db, _ := store.Open(certpaths.DBPath()) + defer db.Close() + got, err := store.NewNodeRepo(db).Get(context.Background(), node.ID) + if err != nil { + t.Fatalf("get node: %v", err) + } + if got.State != model.NodeStateDrained { + t.Errorf("node state = %q, want drained", got.State) + } +} + +func TestNodeDrain_NoAllocs_MarksDrainedImmediately(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + node := drainNodeForTest(t, "emptynode", "emptynode:8443") + mx := &scriptedDrainExec{} + mx.queue("list-units", "", 0) // no allocs + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + nodeDrainCmd.SetOut(&buf) + nodeDrainCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "5s"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("node drain: %v", err) + } + out := buf.String() + if !strings.Contains(out, "drained") { + t.Errorf("expected drained message, got: %s", out) + } + // Should not have issued any stop commands. + if mx.countCalls("systemctl stop") != 0 { + t.Errorf("expected no stop commands, got: %+v", mx.calls) + } + + db, _ := store.Open(certpaths.DBPath()) + defer db.Close() + got, _ := store.NewNodeRepo(db).Get(context.Background(), node.ID) + if got.State != model.NodeStateDrained { + t.Errorf("node state = %q, want drained", got.State) + } +} + +func TestNodeDrain_RespectsTimeout(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + node := drainNodeForTest(t, "stucknode", "stucknode:8443") + // The alloc never leaves the running list → waitAllocsStopped hits + // the timeout. The drain reports partial and leaves the node in + // "draining". + mx := &scriptedDrainExec{} + mx.queueAlways("list-units", "orca-alloc-stuck-0.service loaded active running\n", 0) + mx.queueAlways("systemctl stop", "", 0) + drainExecOverride = mx + + // Poll interval is 500ms; use a 1s timeout so the test is fast but + // still exercises the timeout path. + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + nodeDrainCmd.SetOut(&buf) + nodeDrainCmd.SetErr(&buf) + start := time.Now() + rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "1s"}) + err := rootCmd.Execute() + elapsed := time.Since(start) + if err == nil { + t.Fatalf("expected drain to error on timeout, got nil") + } + if elapsed > 5*time.Second { + t.Errorf("drain took too long (%v); timeout not respected", elapsed) + } + + // Node should remain in "draining" (not drained) since an alloc is + // still running. + db, _ := store.Open(certpaths.DBPath()) + defer db.Close() + got, _ := store.NewNodeRepo(db).Get(context.Background(), node.ID) + if got.State != model.NodeStateDraining { + t.Errorf("node state = %q, want draining (timeout)", got.State) + } +} + +func TestNodeDrain_NodeNotFound(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + drainTestEnv(t) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + nodeDrainCmd.SetOut(&buf) + nodeDrainCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"node", "drain", "no.such.node", "--timeout", "1s"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for unknown node, got nil") + } + if !strings.Contains(err.Error(), "not found") { + t.Errorf("error should mention not found, got: %v", err) + } +} + +func TestDaemonDrainAndStop_StopsEachPeer(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + drainNodeForTest(t, "peer-a", "peer-a:8443") + drainNodeForTest(t, "peer-b", "peer-b:8443") + drainNodeForTest(t, "peer-c", "peer-c:8443") + + mx := &scriptedDrainExec{} + mx.queueAlways("systemctl stop orca-daemon.service", "", 0) + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + daemonCmd.SetOut(&buf) + daemonCmd.SetErr(&buf) + daemonDrainAndStopCmd.SetOut(&buf) + daemonDrainAndStopCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"daemon", "drain-and-stop"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("daemon drain-and-stop: %v", err) + } + out := buf.String() + if !strings.Contains(out, "complete") { + t.Errorf("expected complete message, got: %s", out) + } + // One stop call per peer. + stops := mx.countCalls("systemctl stop orca-daemon.service") + if stops != 3 { + t.Errorf("expected 3 daemon stop calls, got %d (calls: %+v)", stops, mx.calls) + } +} + +func TestDaemonDrainAndStop_AlreadyStoppedIsIdempotent(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + drainNodeForTest(t, "migrated", "migrated:8443") + mx := &scriptedDrainExec{} + // systemctl stop returns exit 5 ("unit not loaded") → already + // migrated, idempotent. + mx.queueAlways("systemctl stop orca-daemon.service", "", 5) + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + daemonCmd.SetOut(&buf) + daemonCmd.SetErr(&buf) + daemonDrainAndStopCmd.SetOut(&buf) + daemonDrainAndStopCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"daemon", "drain-and-stop"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("daemon drain-and-stop should be idempotent: %v", err) + } + out := buf.String() + if !strings.Contains(out, "already stopped") && !strings.Contains(out, "already-stopped") { + t.Errorf("expected already-stopped in output, got: %s", out) + } +} + +func TestJobMigrate_StopsOldAndStartsNewOnTarget(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + src := drainNodeForTest(t, "src", "src:8443") + _ = drainNodeForTest(t, "tgt", "tgt:8443") + _ = src + + mx := &scriptedDrainExec{} + // src node lists the job's alloc running. + mx.queue("list-units", + "orca-alloc-web-0.service loaded active running\n", 0) + // tgt node has no allocs (so migrate starts a new one). + mx.queue("list-units", "", 0) + // stop on src succeeds. + mx.queue("systemctl stop orca-alloc-web-0", "", 0) + // start on tgt succeeds. + mx.queue("systemctl start", "", 0) + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + jobCmd.SetOut(&buf) + jobCmd.SetErr(&buf) + jobMigrateCmd.SetOut(&buf) + jobMigrateCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"job", "migrate", "web", "--to", "tgt"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("job migrate: %v", err) + } + out := buf.String() + if !strings.Contains(out, "Migrated") { + t.Errorf("expected Migrated message, got: %s", out) + } + if mx.countCalls("systemctl stop orca-alloc-web-0") == 0 { + t.Errorf("expected stop for web-0 on src, calls: %+v", mx.calls) + } + if mx.countCalls("systemctl start orca-alloc-web-migrated") == 0 { + t.Errorf("expected start of new alloc on tgt, calls: %+v", mx.calls) + } +} + +func TestJobMigrate_IdempotentWhenAlreadyOnTarget(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + + _ = drainNodeForTest(t, "src", "src:8443") + _ = drainNodeForTest(t, "tgt", "tgt:8443") + + mx := &scriptedDrainExec{} + // The job is running ONLY on tgt (the target). No allocs on src. + mx.queue("list-units", "orca-alloc-web-0.service loaded active running\n", 0) // src: empty would be ideal; see scripted note + // Use scripted: src empty, tgt has web-0. + mx = &scriptedDrainExec{} + mx.queue("list-units", "", 0) // first list-units call (src) → empty + mx.queue("list-units", "orca-alloc-web-0.service loaded active running\n", 0) // second (tgt) → has web-0 + drainExecOverride = mx + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + jobCmd.SetOut(&buf) + jobCmd.SetErr(&buf) + jobMigrateCmd.SetOut(&buf) + jobMigrateCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"job", "migrate", "web", "--to", "tgt"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("job migrate (idempotent): %v", err) + } + out := buf.String() + if !strings.Contains(out, "already running on") { + t.Errorf("expected already-running idempotent message, got: %s", out) + } + // No stop or start commands should have been issued. + if mx.countCalls("systemctl stop") != 0 || mx.countCalls("systemctl start") != 0 { + t.Errorf("idempotent migrate must not stop/start, calls: %+v", mx.calls) + } +} + +func TestJobMigrate_RequiresToFlag(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + resetRootFlags(t) + drainTestEnv(t) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + jobCmd.SetOut(&buf) + jobCmd.SetErr(&buf) + jobMigrateCmd.SetOut(&buf) + jobMigrateCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"job", "migrate", "web"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for missing --to, got nil") + } + if !strings.Contains(err.Error(), "--to is required") { + t.Errorf("error should mention --to required, got: %v", err) + } +} + +func TestSetNodeStateUpdatesDB(t *testing.T) { + _, cleanup := initTestEnv(t) + defer cleanup() + // Sanity test for the store-level SetNodeState used by drain. + db, err := store.Open(certpaths.DBPath()) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer db.Close() + repo := store.NewNodeRepo(db) + n := &model.Node{ + ID: "node-setstate", + Name: "setstate", + Address: "setstate:8443", + State: model.NodeStateReady, + JoinedAt: time.Now().UTC(), + LastSeen: time.Now().UTC(), + } + ctx := context.Background() + if err := repo.Insert(ctx, n); err != nil { + t.Fatalf("insert: %v", err) + } + if err := repo.SetNodeState(ctx, n.ID, string(model.NodeStateDraining)); err != nil { + t.Fatalf("SetNodeState: %v", err) + } + got, err := repo.Get(ctx, n.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.State != model.NodeStateDraining { + t.Errorf("state = %q, want draining", got.State) + } + + // Missing node returns ErrNotFound. + err = repo.SetNodeState(ctx, "ghost", "draining") + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("SetNodeState(ghost) = %v, want ErrNotFound", err) + } +} + +// scriptedDrainExec is a record-and-replay execer that returns queued +// responses in order. queueAlways installs a sticky response that +// matches every subsequent call containing the substring. This is +// more convenient than mockDrainExec when the same command (e.g. +// list-units) must return different outputs across calls. +type scriptedDrainExec struct { + mu sync.Mutex + queued []scriptedResp + sticky []scriptedResp + calls []mockDrainCall +} + +type scriptedResp struct { + match string + out string + exit int +} + +func (s *scriptedDrainExec) queue(match, out string, exit int) { + s.queued = append(s.queued, scriptedResp{match: match, out: out, exit: exit}) +} + +func (s *scriptedDrainExec) queueAlways(match, out string, exit int) { + s.sticky = append(s.sticky, scriptedResp{match: match, out: out, exit: exit}) +} + +func (s *scriptedDrainExec) Exec(_ context.Context, peer, cmd string) ([]byte, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.calls = append(s.calls, mockDrainCall{peer: peer, cmd: cmd}) + // Sticky responses win if they match. + for _, r := range s.sticky { + if r.match == "" || strings.Contains(cmd, r.match) { + if r.exit != 0 { + return []byte(r.out), &sshExitErr{code: r.exit} + } + return []byte(r.out), nil + } + } + // Queued responses: pop the first matching entry. + for i, r := range s.queued { + if r.match == "" || strings.Contains(cmd, r.match) { + s.queued = append(s.queued[:i], s.queued[i+1:]...) + if r.exit != 0 { + return []byte(r.out), &sshExitErr{code: r.exit} + } + return []byte(r.out), nil + } + } + return nil, nil +} + +func (s *scriptedDrainExec) callsFor(match string) []mockDrainCall { + s.mu.Lock() + defer s.mu.Unlock() + var out []mockDrainCall + for _, c := range s.calls { + if strings.Contains(c.cmd, match) { + out = append(out, c) + } + } + return out +} + +func (s *scriptedDrainExec) countCalls(match string) int { + return len(s.callsFor(match)) +} diff --git a/internal/cli/namespace_test.go b/internal/cli/namespace_test.go index d22ca25..7901e95 100644 --- a/internal/cli/namespace_test.go +++ b/internal/cli/namespace_test.go @@ -6,8 +6,11 @@ import ( "os" "path/filepath" "testing" + "time" "git.cloudinit.dev/coreci/orca/internal/certpaths" + + "github.com/spf13/cobra" ) func resetRootFlags(t *testing.T) { @@ -32,11 +35,25 @@ func resetCommandFlags() { joinHost, joinSSHUser, joinPassword, proxmoxUser, proxmoxRole = "", "root", "", "orca", "OrcaOperator" joinSSHPort, leaveID, nodeWatch = 22, "", false stopID, runTarget, runIDKey, jobWatch = "", "", "", false + migrateTarget = "" + drainTimeout = 30 * time.Second capSetCPU, capSetMem, capSetDisk, capNodeID = 0, 0, 0, "" auditLimit = 50 backupOutPath, restoreInPath, restoreTargetDir = "", "", "" restoreForce = false resetNSFlags() + // Reset per-command output writers so tests that polluted them + // (e.g. daemon tests calling cmd.SetOut(&buf)) don't leak into + // other tests. nil → cobra walks to rootCmd's writer. + for _, c := range []*cobra.Command{ + nodeDrainCmd, daemonCmd, daemonDrainAndStopCmd, + jobCmd, jobMigrateCmd, jobRunCmd, jobListCmd, jobStopCmd, jobLogsCmd, + } { + if c != nil { + c.SetOut(nil) + c.SetErr(nil) + } + } } func TestNamespaceDefaultsToUserHome(t *testing.T) { diff --git a/internal/engine/registry.go b/internal/engine/registry.go index 05b86db..d220bbd 100644 --- a/internal/engine/registry.go +++ b/internal/engine/registry.go @@ -68,3 +68,16 @@ func (r *NodeRegistry) List(ctx context.Context) ([]*model.Node, error) { func (r *NodeRegistry) Get(ctx context.Context, id string) (*model.Node, error) { return r.repo.Get(ctx, id) } + +// SetNodeState updates a node's state to the given raw string (REQ-061). +// Used by the drain commands for the draining/drained states. This is +// the registry-level wrapper around store.NodeRepo.SetNodeState so the +// cli package does not need to reach into the repo directly. +func (r *NodeRegistry) SetNodeState(ctx context.Context, id, state string) error { + if err := r.repo.SetNodeState(ctx, id, state); err != nil { + r.audit.Record(ctx, "cli", "node.set_state", id, "failure", err, map[string]any{"state": state}) + return fmt.Errorf("set node state: %w", err) + } + r.log.Info("node state set", slog.String("node_id", id), slog.String("state", state)) + return nil +} diff --git a/internal/model/node.go b/internal/model/node.go index 5cffdb8..a0c028a 100644 --- a/internal/model/node.go +++ b/internal/model/node.go @@ -5,9 +5,11 @@ import "time" type NodeState string const ( - NodeStatePending NodeState = "pending" - NodeStateReady NodeState = "ready" - NodeStateLeft NodeState = "left" + NodeStatePending NodeState = "pending" + NodeStateReady NodeState = "ready" + NodeStateLeft NodeState = "left" + NodeStateDraining NodeState = "draining" + NodeStateDrained NodeState = "drained" ) // NodeKind classifies a node by how it joined the cluster. diff --git a/internal/store/node_repo.go b/internal/store/node_repo.go index 5107347..3593afc 100644 --- a/internal/store/node_repo.go +++ b/internal/store/node_repo.go @@ -154,6 +154,24 @@ func (r *NodeRepo) Delete(ctx context.Context, id string) error { return nil } +// SetNodeState updates a node's state to the given string value +// (REQ-061). Unlike UpdateState, which takes a model.NodeState, this +// accepts a raw string so callers can set arbitrary states (draining, +// drained) without coupling to every constant in the model package. +func (r *NodeRepo) SetNodeState(ctx context.Context, nodeID, state string) error { + res, err := r.db.ExecContext(ctx, + `UPDATE nodes SET state = ?, last_seen = ? WHERE id = ?`, + state, time.Now().UTC(), nodeID) + if err != nil { + return fmt.Errorf("update node state: %w", err) + } + rows, _ := res.RowsAffected() + if rows == 0 { + return ErrNotFound + } + return nil +} + type scanner interface { Scan(dest ...any) error } diff --git a/internal/store/node_repo_test.go b/internal/store/node_repo_test.go index 689f168..25b13a8 100644 --- a/internal/store/node_repo_test.go +++ b/internal/store/node_repo_test.go @@ -121,6 +121,39 @@ func TestNodeRepo_UpdateStateMissing(t *testing.T) { } } +func TestNodeRepo_SetNodeState(t *testing.T) { + repo, cleanup := openTestDB(t) + defer cleanup() + + ctx := context.Background() + n := &model.Node{ + ID: "setstate-1", + Name: "setstate", + Address: "setstate:8443", + State: model.NodeStateReady, + JoinedAt: time.Now().UTC(), + LastSeen: time.Now().UTC(), + } + if err := repo.Insert(ctx, n); err != nil { + t.Fatalf("insert: %v", err) + } + if err := repo.SetNodeState(ctx, n.ID, "draining"); err != nil { + t.Fatalf("SetNodeState: %v", err) + } + got, err := repo.Get(ctx, n.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.State != model.NodeStateDraining { + t.Errorf("state = %q, want draining", got.State) + } + + // Missing node → ErrNotFound. + if err := repo.SetNodeState(ctx, "ghost", "drained"); err != ErrNotFound { + t.Errorf("SetNodeState(ghost) = %v, want ErrNotFound", err) + } +} + func TestNodeRepo_UpdateLastSeenAndOSMissing(t *testing.T) { repo, cleanup := openTestDB(t) defer cleanup()