llmgw/internal/agent/failover_test.go

package agent

import (
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

type standby bool

func (s standby) IsLeader() bool { return !bool(s) }

// startPlain runs an agent service over h2c, which is enough to test who
// accepts a connection and who refuses it.
//
// A nil leader means none was configured, which is every deployment that
// existed before leadership did and has to keep accepting agents.
func startPlain(t *testing.T, leader agentsvc.Leader) (*inventory.Store, string) {
	t.Helper()
	inv := inventory.NewStore()
	svc := agentsvc.New(inv, newRecorder(), false, slog.New(slog.DiscardHandler))
	if leader != nil {
		svc.SetLeadership(leader)
	}

	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)
	// h2c, because the agent speaks HTTP/2 with prior knowledge and a plain
	// server negotiates HTTP/1.1 — which turns a clear refusal into a frame
	// parsing error that says nothing about why.
	srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
	t.Cleanup(srv.Close)
	return inv, srv.URL
}

func TestAStandbyRefusesAnAgent(t *testing.T) {
	// Commands travel down the stream the agent opened, so whichever control
	// plane holds it is the only one that can start anything on that node. A
	// standby holding it means the leader's placements fail with "no agent
	// connected" for a node that is plainly connected.
	_, addr := startPlain(t, standby(true))

	ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
	err := ag.session(context.Background(), ag.cfg.Server)
	if err == nil {
		t.Fatal("a standby accepted an agent")
	}
	if !strings.Contains(err.Error(), "standby") {
		t.Errorf("the refusal does not say why: %v", err)
	}
}

func TestAnAgentReachesTheLeaderPastAStandby(t *testing.T) {
	// The point of the list. The standby is first, so the agent only registers
	// if it moved on rather than sitting on the refusal.
	_, standbyAddr := startPlain(t, standby(true))
	leaderInv, leaderAddr := startPlain(t, standby(false))

	ag := runAgent(t, Config{
		NodeID: "n1", MgmtAddr: "127.0.0.1",
		Servers: []string{standbyAddr, leaderAddr},
	})

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	go ag.Run(ctx)

	deadline := time.Now().Add(8 * time.Second)
	for time.Now().Before(deadline) {
		if _, ok := leaderInv.Get("n1"); ok {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("the agent never reached the leader past the standby")
}

func TestALeaderStillAcceptsWithNoLeadershipConfigured(t *testing.T) {
	// Every deployment before this had no leadership at all, and must keep
	// accepting agents with nothing configured.
	inv, addr := startPlain(t, nil)

	ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	go ag.Run(ctx)

	deadline := time.Now().Add(4 * time.Second)
	for time.Now().Before(deadline) {
		if _, ok := inv.Get("n1"); ok {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("a service with no leadership configured refused an agent")
}

// flipping is a leader that can stop leading partway through a test.
type flipping struct{ leading atomic.Bool }

func (f *flipping) IsLeader() bool { return f.leading.Load() }

func TestAnAgentLeavesAControlPlaneThatStopsLeading(t *testing.T) {
	// Leadership lost while the process lives on: the store was slow and the
	// lease lapsed, or another control plane took it. Agents that stayed on
	// the old leader would leave the new one with no nodes, so the old one
	// has to let them go rather than wait for them to leave.
	was := &flipping{}
	was.leading.Store(true)
	wasInv, wasAddr := startPlain(t, was)
	nowInv, nowAddr := startPlain(t, standby(false))

	ag := runAgent(t, Config{
		NodeID: "n1", MgmtAddr: "127.0.0.1",
		Servers: []string{wasAddr, nowAddr},
	})
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	go ag.Run(ctx)

	waitFor(t, 5*time.Second, "the agent never reached the first leader", func() bool {
		_, ok := wasInv.Get("n1")
		return ok
	})
	was.leading.Store(false)
	waitFor(t, 8*time.Second, "the agent stayed on a control plane that stopped leading", func() bool {
		_, ok := nowInv.Get("n1")
		return ok
	})
}

func waitFor(t *testing.T, d time.Duration, msg string, ok func() bool) {
	t.Helper()
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if ok() {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal(msg)
}