llmgw/internal/agent/supervisor/supervisor_test.go

package supervisor

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

// scriptedHealth reports ready once the flag is set.
type scriptedHealth struct{ ready atomic.Bool }

func (h *scriptedHealth) Ready(context.Context, string, string, string) error {
	if h.ready.Load() {
		return nil
	}
	return errors.New("not ready")
}

func newSupervisor(t *testing.T) (*Supervisor, *scriptedHealth) {
	t.Helper()
	h := &scriptedHealth{}
	s, err := New(Config{
		WorkDir:        t.TempDir(),
		LogLines:       64,
		StopTimeout:    3 * time.Second,
		HealthInterval: 20 * time.Millisecond,
	}, h, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.StopAll(context.Background()) })
	return s, h
}

// shellCommand builds the single-command sequence for a shell script.
func shellCommand(script string) []engines.Command {
	return []engines.Command{{Path: "/bin/sh", Args: []string{"-c", script}, Env: os.Environ()}}
}

func waitState(t *testing.T, sh *Shard, want State, limit time.Duration) {
	t.Helper()
	deadline := time.Now().Add(limit)
	for time.Now().Before(deadline) {
		if sh.Status().State == want {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("shard stayed in %s, expected %s", sh.Status().State, want)
}

func TestStartAndStop(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	if st := sh.Status(); st.State != Starting || st.PID == 0 {
		t.Fatalf("status after start: %+v", st)
	}

	if err := s.Stop(t.Context(), "a"); err != nil {
		t.Fatal(err)
	}
	if st := sh.Status().State; st != Stopped {
		t.Errorf("state after stop: %s", st)
	}
}

func TestUnexpectedExitIsFailure(t *testing.T) {
	// An engine exiting on its own is a failure even with status 0: it is meant
	// to serve until told otherwise.
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
	if err != nil {
		t.Fatal(err)
	}
	waitState(t, sh, Failed, 3*time.Second)
	if !strings.Contains(sh.Status().Err, "unexpectedly") {
		t.Errorf("err=%q", sh.Status().Err)
	}
}

func TestNonZeroExitRecordsCode(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 42"))
	if err != nil {
		t.Fatal(err)
	}
	waitState(t, sh, Failed, 3*time.Second)
	if got := sh.Status().ExitCode; got != 42 {
		t.Errorf("exit code=%d, expected 42", got)
	}
}

func TestStopKillsChildProcesses(t *testing.T) {
	// The reason for the process group: engines fork workers that hold GPU
	// memory, and signalling only the leader leaves them behind.
	s, _ := newSupervisor(t)
	dir := t.TempDir()
	marker := filepath.Join(dir, "child.pid")

	// The shell forks a child that writes its own pid, then both sleep.
	script := fmt.Sprintf(`sh -c 'echo $$ > %s; sleep 30' & sleep 30`, marker)
	if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand(script)); err != nil {
		t.Fatal(err)
	}

	// Wait for the child to announce itself.
	var childPID int
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if b, err := os.ReadFile(marker); err == nil {
			if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &childPID); err == nil && childPID > 0 {
				break
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
	if childPID == 0 {
		t.Fatal("child never started")
	}

	if err := s.Stop(t.Context(), "a"); err != nil {
		t.Fatal(err)
	}

	deadline = time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if !processAlive(childPID) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Errorf("child process %d survived the stop; it would still hold GPU memory", childPID)
}

func TestStopKillsAfterTimeout(t *testing.T) {
	// A process ignoring SIGTERM must still die.
	s, _ := newSupervisor(t)
	s.cfg.StopTimeout = 300 * time.Millisecond

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
		shellCommand("trap '' TERM; sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	start := time.Now()
	if err := s.Stop(t.Context(), "a"); err != nil {
		t.Fatal(err)
	}
	if d := time.Since(start); d > 5*time.Second {
		t.Errorf("stop took %v", d)
	}
	if processAlive(pid) {
		t.Error("the process survived SIGKILL")
	}
}

func TestStartIsIdempotent(t *testing.T) {
	// The control plane retries commands, so a repeated start must not produce
	// a second process (plan §5.6).
	s, _ := newSupervisor(t)

	first, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	second, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if !errors.Is(err, ErrAlreadyRunning) {
		t.Fatalf("expected ErrAlreadyRunning, got %v", err)
	}
	if first.Status().PID != second.Status().PID {
		t.Error("a second process was started for the same shard")
	}
}

func TestRestartAfterTerminalState(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
	waitState(t, sh, Failed, 3*time.Second)

	// Once it has failed, the same id may be started again.
	again, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if err != nil {
		t.Fatalf("restarting a failed shard: %v", err)
	}
	if again.Status().State != Starting {
		t.Errorf("state=%s", again.Status().State)
	}
}

func TestWaitReady(t *testing.T) {
	s, h := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}

	go func() {
		time.Sleep(100 * time.Millisecond)
		h.ready.Store(true)
	}()

	if err := s.WaitReady(t.Context(), sh, "http://x", "m", 3*time.Second); err != nil {
		t.Fatal(err)
	}
	st := sh.Status()
	if st.State != Ready {
		t.Errorf("state=%s", st.State)
	}
	if st.ReadyAt.IsZero() {
		t.Error("ReadyAt was not recorded")
	}
}

func TestWaitReadyTimeoutIncludesLogs(t *testing.T) {
	// When a model fails to load, the engine output is the only diagnosis
	// available, so the error has to carry it.
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
		shellCommand("echo 'CUDA out of memory'; sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	eventually(t, 3*time.Second, "the line was never captured", func() bool {
		return strings.Contains(strings.Join(sh.Logs(10), "\n"), "CUDA out of memory")
	})

	err = s.WaitReady(t.Context(), sh, "http://x", "m", 200*time.Millisecond)
	if err == nil {
		t.Fatal("expected a timeout")
	}
	if !strings.Contains(err.Error(), "CUDA out of memory") {
		t.Errorf("the error does not carry engine output: %v", err)
	}
}

func TestWaitReadyDetectsEarlyExit(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
	if err != nil {
		t.Fatal(err)
	}
	// Should return as soon as the process dies, not after the full timeout.
	start := time.Now()
	if err := s.WaitReady(t.Context(), sh, "http://x", "m", 10*time.Second); err == nil {
		t.Fatal("expected an error")
	}
	if d := time.Since(start); d > 3*time.Second {
		t.Errorf("waited %v for a process that had already exited", d)
	}
}

func TestDrain(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
	if err := s.Drain("a"); err != nil {
		t.Fatal(err)
	}
	if got := sh.Status().State; got != Draining {
		t.Errorf("state=%s", got)
	}
	// A drained shard's process is still alive; only routing stops.
	if !processAlive(sh.Status().PID) {
		t.Error("draining killed the process; it should keep serving in-flight work")
	}
}

func TestLogsAreCaptured(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
		shellCommand("echo line1; echo line2 >&2; echo line3; sleep 0.2"))
	if err != nil {
		t.Fatal(err)
	}
	eventually(t, 3*time.Second, "not every line arrived", func() bool {
		l := strings.Join(sh.Logs(10), "\n")
		return strings.Contains(l, "line1") && strings.Contains(l, "line2") && strings.Contains(l, "line3")
	})

	logs := strings.Join(sh.Logs(10), "\n")
	for _, want := range []string{"line1", "line2", "line3"} {
		if !strings.Contains(logs, want) {
			t.Errorf("%q missing from logs:\n%s", want, logs)
		}
	}
}

// --- orphan adoption ---

func TestFindOrphans(t *testing.T) {
	s, _ := newSupervisor(t)

	spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
	sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	// Simulate an agent restart: a new supervisor over the same work dir.
	s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
		&scriptedHealth{}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	orphans, err := s2.FindOrphans()
	if err != nil {
		t.Fatal(err)
	}
	if len(orphans) != 1 {
		t.Fatalf("expected 1 orphan, got %d", len(orphans))
	}
	o := orphans[0]
	if o.ShardID != "a" || o.PID != pid || !o.Alive {
		t.Errorf("orphan=%+v (want shard a, pid %d, alive)", o, pid)
	}
	// The command line is stored as a fingerprint, because pids get reused and
	// adopting an unrelated process is worse than leaking the original.
	if len(o.Cmdline) == 0 || !strings.Contains(strings.Join(o.Cmdline, " "), "sleep 30") {
		t.Errorf("cmdline fingerprint missing: %v", o.Cmdline)
	}
	// And the spec, without which the control plane could not tell which
	// replica the process belongs to and would place a duplicate.
	if o.ReplicaID != "rep-1" {
		t.Errorf("replica id not recovered: %q", o.ReplicaID)
	}
	if o.Spec.HTTPPort != 20100 {
		t.Errorf("spec not recovered: %+v", o.Spec)
	}

	if err := s2.KillOrphan(o); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if !processAlive(pid) {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if processAlive(pid) {
		t.Error("the orphan survived KillOrphan")
	}
	if _, err := os.Stat(s2.statePath("a")); !os.IsNotExist(err) {
		t.Error("the state file was left behind")
	}
}

func TestStateFileRemovedOnExit(t *testing.T) {
	s, _ := newSupervisor(t)

	sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
	waitState(t, sh, Failed, 3*time.Second)
	eventually(t, 3*time.Second,
		"the state file outlived the process; a later scan would find a dead orphan", func() bool {
			_, err := os.Stat(s.statePath("a"))
			return os.IsNotExist(err)
		})
}

// --- state machine and ring buffer ---

func TestStateStrings(t *testing.T) {
	for st, want := range map[State]string{
		Pending: "pending", Starting: "starting", Loading: "loading",
		Ready: "ready", Draining: "draining", Stopped: "stopped", Failed: "failed",
	} {
		if got := st.String(); got != want {
			t.Errorf("%d -> %q, want %q", int(st), got, want)
		}
	}
	if !Stopped.Terminal() || !Failed.Terminal() || Ready.Terminal() {
		t.Error("Terminal is wrong")
	}
	if !Ready.Live() || !Draining.Live() || Stopped.Live() {
		t.Error("Live is wrong")
	}
}

func TestRingBuffer(t *testing.T) {
	r := newRingBuffer(3)
	if got := r.tail(10); len(got) != 0 {
		t.Errorf("empty buffer returned %v", got)
	}

	r.add("a")
	r.add("b")
	if got := strings.Join(r.tail(10), ","); got != "a,b" {
		t.Errorf("tail=%q", got)
	}

	r.add("c")
	r.add("d") // evicts "a"
	if got := strings.Join(r.tail(10), ","); got != "b,c,d" {
		t.Errorf("after wrap tail=%q, want b,c,d", got)
	}
	if got := strings.Join(r.tail(2), ","); got != "c,d" {
		t.Errorf("tail(2)=%q", got)
	}
}

func TestAdoptTakesOverRunningShard(t *testing.T) {
	// The control plane restart fix depends on this: a shard that survived the
	// agent must come back under supervision, not be restarted or leaked.
	s, _ := newSupervisor(t)

	spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20200}
	sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	// A new supervisor over the same work dir, as after an agent restart.
	s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 3 * time.Second},
		&scriptedHealth{}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	orphans, err := s2.FindOrphans()
	if err != nil {
		t.Fatal(err)
	}
	if len(orphans) != 1 || !orphans[0].Alive {
		t.Fatalf("orphan not found: %+v", orphans)
	}

	adopted, err := s2.Adopt(orphans[0])
	if err != nil {
		t.Fatal(err)
	}
	if got := adopted.Status().PID; got != pid {
		t.Errorf("adopted pid %d, expected %d", got, pid)
	}
	if adopted.Spec.ReplicaID != "rep-1" {
		t.Errorf("spec lost in adoption: %+v", adopted.Spec)
	}
	// The earlier output is gone; saying so beats presenting an empty log as if
	// the engine had been silent.
	if logs := strings.Join(adopted.Logs(5), " "); !strings.Contains(logs, "adopted") {
		t.Errorf("adoption was not noted in the log: %q", logs)
	}

	// Stopping an adopted shard must actually kill it: it has no exec.Cmd, so
	// the signal path is the only one available.
	if err := s2.Stop(t.Context(), "a"); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if !processAlive(pid) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Error("the adopted process survived Stop")
}

func TestAdoptRefusesDeadProcess(t *testing.T) {
	s, _ := newSupervisor(t)
	if _, err := s.Adopt(Orphan{ShardID: "gone", PID: 999999, Alive: false}); err == nil {
		t.Error("a dead process was adopted")
	}
}

func TestAdoptedExitIsReported(t *testing.T) {
	// An adopted engine dying has to be noticed; nothing waits on it.
	s, _ := newSupervisor(t)

	spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a"}
	sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
		&scriptedHealth{}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	orphans, _ := s2.FindOrphans()
	adopted, err := s2.Adopt(orphans[0])
	if err != nil {
		t.Fatal(err)
	}

	// Kill it behind the supervisor's back.
	_ = killGroup(pid, syscall.SIGKILL)
	waitState(t, adopted, Failed, 5*time.Second)
}

func TestCorruptStateFileIsCleanedUp(t *testing.T) {
	// A half-written state file describes a process that cannot be identified,
	// so it cannot be adopted safely.
	s, _ := newSupervisor(t)
	if err := os.WriteFile(s.statePath("broken"), []byte("{not json"), 0o600); err != nil {
		t.Fatal(err)
	}

	orphans, err := s.FindOrphans()
	if err != nil {
		t.Fatal(err)
	}
	if len(orphans) != 1 || orphans[0].ShardID != "broken" || orphans[0].Alive {
		t.Fatalf("corrupt state file not reported for cleanup: %+v", orphans)
	}
	if err := s.KillOrphan(orphans[0]); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(s.statePath("broken")); !os.IsNotExist(err) {
		t.Error("the corrupt state file was left behind")
	}
}

// --- multi-process shards (plan §11, phase 4) ---

func TestASidecarAndAnEngineBothRunUnderOneShard(t *testing.T) {
	// The Ray shape: a supporting process runs alongside the engine for the
	// shard's lifetime, but the engine is the one whose exit ends the shard.
	s, _ := newSupervisor(t)
	dir := t.TempDir()
	order := filepath.Join(dir, "order")

	// Each writes to the order file and to stdout: the file records the
	// sequence, stdout exercises per-process log attribution.
	cmds := []engines.Command{
		{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
			fmt.Sprintf("echo sidecar >> %s; echo hello-from-sidecar; sleep 30", order)}, Env: os.Environ()},
		{Name: "engine", Path: "/bin/sh", Args: []string{"-c",
			fmt.Sprintf("echo engine >> %s; echo hello-from-engine; sleep 30", order)}, Env: os.Environ()},
	}

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
	if err != nil {
		t.Fatal(err)
	}

	// Generous, because this waits on two processes starting under the race
	// detector while every other package is competing for the same cores. Three
	// seconds was enough alone and not in a full run.
	var started []string
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) && len(started) < 2 {
		if b, err := os.ReadFile(order); err == nil {
			started = strings.Fields(string(b))
		}
		if len(started) < 2 {
			time.Sleep(20 * time.Millisecond)
		}
	}
	// Asserted rather than broken out of: the original fell through on timeout,
	// so a slow machine turned this into a no-op and the failure surfaced later
	// as something else.
	if len(started) != 2 {
		t.Fatalf("only %d of 2 processes started: %v", len(started), started)
	}
	// Which of the two appended first is deliberately not asserted. The
	// supervisor spawns the sidecar first, but spawning is not sequencing:
	// nothing here makes the engine wait, so two shells racing to append give
	// either order and a test that demanded one would fail for the scheduler's
	// reasons rather than the supervisor's. The wait that is a real guarantee
	// is WaitForLeaderPort, and TestWaitForLeaderPortBlocksUntilListening is
	// where it is checked.
	if !slices.Contains(started, "sidecar") || !slices.Contains(started, "engine") {
		t.Fatalf("expected both processes to have run, got %v", started)
	}

	// The shard's pid is the engine's, not the sidecar's: the engine's exit is
	// what ends the shard.
	if sh.Status().PID == 0 {
		t.Fatal("no engine pid recorded")
	}
	// Both processes appear in one log tail, distinguishable by prefix.
	deadline = time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		logs := strings.Join(sh.Logs(20), "\n")
		if strings.Contains(logs, "[sidecar] hello-from-sidecar") &&
			strings.Contains(logs, "[engine] hello-from-engine") {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Errorf("output is not attributed per process:\n%s", strings.Join(sh.Logs(20), "\n"))
}

func TestStopKillsSidecar(t *testing.T) {
	// A Ray head left behind holds the port the next replica needs, and would
	// let a later shard join a cluster that serves nothing.
	s, _ := newSupervisor(t)
	dir := t.TempDir()
	marker := filepath.Join(dir, "sidecar.pid")

	cmds := []engines.Command{
		{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
			fmt.Sprintf("echo $$ > %s; sleep 30", marker)}, Env: os.Environ()},
		{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
	}
	if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds); err != nil {
		t.Fatal(err)
	}

	var sidePID int
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if b, err := os.ReadFile(marker); err == nil {
			if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &sidePID); err == nil && sidePID > 0 {
				break
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
	if sidePID == 0 {
		t.Fatal("the sidecar never started")
	}

	if err := s.Stop(t.Context(), "a"); err != nil {
		t.Fatal(err)
	}
	deadline = time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if !processAlive(sidePID) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Errorf("the sidecar (%d) survived the stop and still holds its port", sidePID)
}

func TestSidecarFailureFailsTheShard(t *testing.T) {
	s, _ := newSupervisor(t)
	cmds := []engines.Command{
		{Name: "sidecar", Path: "/nonexistent/binary", Env: os.Environ()},
		{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
	}
	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
	if err == nil {
		t.Fatal("a shard whose sidecar could not start was reported healthy")
	}
	if !strings.Contains(err.Error(), "sidecar") {
		t.Errorf("the error does not identify the sidecar: %v", err)
	}
	if sh.Status().State != Failed {
		t.Errorf("state=%s", sh.Status().State)
	}
}

func TestWaitForLeaderPortBlocksUntilListening(t *testing.T) {
	// A Ray worker cannot join a head that is not up, and the leader's server
	// cannot claim GPUs from a cluster that does not exist yet.
	s, _ := newSupervisor(t)
	s.cfg.RendezvousTimeout = 5 * time.Second

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	_ = ln.Close() // nothing listening yet

	spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: port}
	cmds := []engines.Command{
		{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
			fmt.Sprintf("sleep 0.4; nc -l %d", port)}, Env: os.Environ(),
			WaitForLeaderPort: true},
		{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
	}

	start := time.Now()
	if _, err := s.Start(t.Context(), "a", spec, cmds); err != nil {
		t.Skipf("nc is unavailable for the rendezvous stub: %v", err)
	}
	if d := time.Since(start); d < 300*time.Millisecond {
		t.Errorf("the engine started after %v, before the rendezvous point was listening", d)
	}
}

func TestRendezvousTimeout(t *testing.T) {
	// Nothing ever listens, so the shard must fail rather than hang forever.
	s, _ := newSupervisor(t)
	s.cfg.RendezvousTimeout = 300 * time.Millisecond

	spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: 1}
	cmds := []engines.Command{
		{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", "sleep 30"},
			Env: os.Environ(), WaitForLeaderPort: true},
		{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
	}
	if _, err := s.Start(t.Context(), "a", spec, cmds); err == nil {
		t.Fatal("the shard did not fail when the rendezvous point never appeared")
	}
}

func TestWaitAliveForNonServingRank(t *testing.T) {
	// A worker has no endpoint to poll, so staying up is the only evidence
	// available. This is what stopped multi-node replicas coming up at all.
	s, _ := newSupervisor(t)
	s.cfg.AliveGrace = 150 * time.Millisecond

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	if err := s.WaitAlive(t.Context(), sh); err != nil {
		t.Fatal(err)
	}
	if got := sh.Status().State; got != Ready {
		t.Errorf("state=%s", got)
	}
}

func TestWaitAliveFailsOnEarlyExit(t *testing.T) {
	s, _ := newSupervisor(t)
	s.cfg.AliveGrace = 2 * time.Second

	sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("exit 1"))
	if err != nil {
		t.Fatal(err)
	}
	start := time.Now()
	if err := s.WaitAlive(t.Context(), sh); err == nil {
		t.Fatal("a worker that died immediately was reported ready")
	}
	// Below AliveGrace (5s), which is what waiting it out would take; the
	// headroom is for a loaded runner, where 1.5s was sometimes not enough.
	if d := time.Since(start); d > 4*time.Second {
		t.Errorf("waited %v for a process that had already exited", d)
	}
}

func TestAnAdoptedShardCanBecomeReady(t *testing.T) {
	// Adoption stopped at "the process is alive", leaving the shard in Starting
	// with nothing that would ever move it. The control plane then waited out
	// the startup timeout and killed a perfectly healthy engine to start an
	// identical one — so an agent restart, which exists precisely so engines
	// are not interrupted, cost every model on the node a full reload.
	//
	// Seen on node05: the engine answered /health throughout and the replica
	// sat in Starting for three minutes until the timeout reaped it.
	s, _ := newSupervisor(t)
	spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20201}
	if _, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30")); err != nil {
		t.Fatal(err)
	}

	health := &scriptedHealth{}
	s2, err := New(Config{
		WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second,
		HealthInterval: 10 * time.Millisecond,
	}, health, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	orphans, err := s2.FindOrphans()
	if err != nil || len(orphans) != 1 {
		t.Fatalf("orphan not found: %+v (%v)", orphans, err)
	}
	adopted, err := s2.Adopt(orphans[0])
	if err != nil {
		t.Fatal(err)
	}
	if got := adopted.Status().State; got != Starting {
		t.Fatalf("a freshly adopted shard is %s, expected Starting", got)
	}

	// The engine answers, as node05's did. The shard must be able to reach
	// Ready from here; before this it had no path to it at all.
	health.ready.Store(true)
	if err := s2.WaitReady(t.Context(), adopted, "http://127.0.0.1:20201", "m", 2*time.Second); err != nil {
		t.Fatalf("an adopted, healthy engine never became ready: %v", err)
	}
	if got := adopted.Status().State; got != Ready {
		t.Errorf("state %s after a passing health check", got)
	}
	_ = s2.Stop(t.Context(), "a")
}

func TestAWorkersOnlyProcessWaitsForTheLeader(t *testing.T) {
	// A Ray worker's join is its only process. The wait was honoured only for
	// sidecars, so the join started at once, found no head, exited, and took
	// the replica down with it.
	s, _ := newSupervisor(t)
	s.cfg.RendezvousTimeout = 5 * time.Second

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	_ = ln.Close()
	go func() {
		time.Sleep(400 * time.Millisecond)
		l, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
		if err == nil {
			t.Cleanup(func() { _ = l.Close() })
		}
	}()

	spec := engines.ShardSpec{Rank: 1, LeaderAddr: "127.0.0.1", LeaderPort: port}
	cmds := []engines.Command{{Name: "join", Path: "/bin/sh", Args: []string{"-c", "sleep 30"},
		Env: os.Environ(), WaitForLeaderPort: true}}
	start := time.Now()
	if _, err := s.Start(t.Context(), "w", spec, cmds); err != nil {
		t.Fatal(err)
	}
	if d := time.Since(start); d < 300*time.Millisecond {
		t.Errorf("the worker started after %v, before the leader was listening", d)
	}
}

// eventually waits for cond, failing with msg if it does not hold within
// limit. Sleeping a fixed time and then checking measured the machine the
// test ran on, and a loaded CI runner under the race detector lost that race.
func eventually(t *testing.T, limit time.Duration, msg string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(limit)
	for !cond() {
		if time.Now().After(deadline) {
			t.Fatal(msg)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestWhatAShardComputesOnSurvivesTheAgent(t *testing.T) {
	s, _ := newSupervisor(t)
	spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
	sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Stop(context.Background(), "a") })
	_ = sh
	s.SetCompute("a", engines.Compute{Backend: "CUDA", Offloaded: 49, Total: 49})

	s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
		&scriptedHealth{}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	orphans, err := s2.FindOrphans()
	if err != nil || len(orphans) != 1 {
		t.Fatalf("%v %d", err, len(orphans))
	}
	if c := orphans[0].Spec.Compute; c == nil || c.Backend != "CUDA" || c.Offloaded != 49 {
		t.Errorf("not kept: %+v", c)
	}
}