Jon Chery
41bcf0a6bf
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---
527 lines
15 KiB
Go
527 lines
15 KiB
Go
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))
|
|
}
|