orca/internal/cli/drain.go

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.
//
// F6: allocID is parsed from remote `systemctl list-units` output and is
// therefore attacker-controlled (a malicious peer could emit a crafted
// unit name). Validate against ^[A-Za-z0-9_-]+$ before interpolation into
// the shell command to prevent stored command injection.
func stopAlloc(ctx context.Context, ex drainExecer, peer, allocID string) error {
	if !validSafeName(allocID) {
		return fmt.Errorf("stopAlloc: invalid alloc id %q (allowed: A-Z a-z 0-9 _ -)", allocID)
	}
	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, actorFromCtx(ctx), "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, actorFromCtx(ctx), "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, actorFromCtx(ctx), "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)

}