feat(P05): drain + daemon drain-and-stop (REQ-061) + job migrate (REQ-116)

orca node drain <host>: marks draining, stops allocs via SSH, marks drained. orca daemon drain-and-stop: stops v0.8 daemons on peers. orca job migrate <name> --to <node>: drain+reschedule composite (C3=a, not live-migrate). Node states: draining, drained. ---ci--- project: orca phase: 05 milestone: v0.11 status: execute ---/ci---
2026-08-07 05:17:01 +00:00
parent f61ef2aa9e
commit 41bcf0a6bf
7 changed files with 1269 additions and 3 deletions
@@ -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-<id>.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 <host>",
+	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-<id>.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.
+
+<host> 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 <name> --to <node>` (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 <name>",
+	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 <name> currently running
+on a node other than --to, it stops the allocation (SSH systemctl stop
+orca-alloc-<id>.service) and starts a new allocation on the target
+node (SSH systemctl start orca-alloc-<new-id>.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-<id>.service across the cluster, with <id> of the
+form <job-name>-<replica-index>. The new allocation on the target node
+is named <name>-migrated-<timestamp>.`,
+	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 <jobName> across all nodes. An
+		// allocation "belongs to" the job if its alloc id starts
+		// with "<jobName>-" (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 "<jobName>-migrated-<unix-seconds>".
+		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)
+
+}
@@ -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))
+}
@@ -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) {