llmgw/test/e2e/phase3_test.go

package e2e

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router"
	"gitea.mynodes.xyz/ilker/llmgw/internal/sim"
)

// simDispatcher is a stand-in for the agent fleet: instead of sending commands
// over a stream it starts real sim-engines on this host, on exactly the ports
// the reconciler allocated.
//
// That matters — the reconciler derives each replica's endpoint URL from the
// node address and the port it chose, so an engine listening anywhere else
// would make the test pass for the wrong reason.
type simDispatcher struct {
	t *testing.T

	mu      sync.Mutex
	running map[string]*simShard
	// dead marks nodes whose agent is gone; commands to them fail.
	dead map[string]bool

	onStatus func(replicaID, shardID string, st supervisor.State, errMsg string)
}

type simShard struct {
	eng  *sim.Engine
	srv  *http.Server
	port int
}

func newSimDispatcher(t *testing.T) *simDispatcher {
	d := &simDispatcher{t: t, running: map[string]*simShard{}, dead: map[string]bool{}}
	t.Cleanup(d.shutdown)
	return d
}

func (d *simDispatcher) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
	d.mu.Lock()
	if d.dead[cmd.NodeID] {
		d.mu.Unlock()
		return fmt.Errorf("agent on %s is unreachable", cmd.NodeID)
	}
	d.mu.Unlock()

	cfg := simConfig("", int64(cmd.Spec.HTTPPort))
	cfg.Model = cmd.Spec.Profile.ServedModelName
	eng := sim.New(cfg)

	ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.Spec.HTTPPort))
	if err != nil {
		eng.Close()
		return fmt.Errorf("binding port %d: %w", cmd.Spec.HTTPPort, err)
	}
	srv := &http.Server{Handler: sim.NewServer(eng, slog.New(slog.DiscardHandler)), ReadHeaderTimeout: time.Second}
	go func() { _ = srv.Serve(ln) }()

	d.mu.Lock()
	d.running[cmd.ShardID] = &simShard{eng: eng, srv: srv, port: cmd.Spec.HTTPPort}
	d.mu.Unlock()

	// A real agent reports Ready only once the engine serves the model; the
	// sim-engine is ready as soon as it is listening.
	go d.onStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Ready, "")
	return nil
}

func (d *simDispatcher) StopShard(_ context.Context, _, shardID string) error {
	d.mu.Lock()
	sh, ok := d.running[shardID]
	delete(d.running, shardID)
	d.mu.Unlock()
	if ok {
		sh.close()
	}
	return nil
}

func (d *simDispatcher) DrainShard(context.Context, string, string) error { return nil }

// killNode stops everything on a node and makes further commands to it fail,
// which is what losing a machine looks like from the control plane.
func (d *simDispatcher) killNode(nodeID string, shardIDs []string) {
	d.mu.Lock()
	d.dead[nodeID] = true
	var doomed []*simShard
	for _, id := range shardIDs {
		if sh, ok := d.running[id]; ok {
			doomed = append(doomed, sh)
			delete(d.running, id)
		}
	}
	d.mu.Unlock()

	for _, sh := range doomed {
		sh.close()
	}
}

func (d *simDispatcher) liveShards() int {
	d.mu.Lock()
	defer d.mu.Unlock()
	return len(d.running)
}

func (d *simDispatcher) shutdown() {
	d.mu.Lock()
	all := make([]*simShard, 0, len(d.running))
	for _, sh := range d.running {
		all = append(all, sh)
	}
	d.running = map[string]*simShard{}
	d.mu.Unlock()
	for _, sh := range all {
		sh.close()
	}
}

func (s *simShard) close() {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_ = s.srv.Shutdown(ctx)
	s.eng.Close()
}

// controlHarness is a router fed by a real control plane.
type controlHarness struct {
	router *httptest.Server
	rt     *router.Router
	ctl    *control.Controller
	disp   *simDispatcher
	inv    *inventory.Store
}

func startControlPlane(t *testing.T, nodes int, minReplicas int) *controlHarness {
	t.Helper()

	cat, err := catalog.New([]catalog.Model{{
		Name: testModel, Source: "/models/test",
		Scaling:  catalog.Scaling{Min: minReplicas, Max: minReplicas + 2},
		Profiles: []catalog.Profile{{Name: "single", GPUClass: "h100-80g", TP: 1, Weight: 1}},
	}})
	if err != nil {
		t.Fatal(err)
	}

	cfg := config.Default()
	cfg.Router.ScrapeInterval = 200 * time.Millisecond
	cfg.Router.StaleAfter = 3 * time.Second
	cfg.Router.Admission.Enabled = false // this test is about endpoint churn

	rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	disp := newSimDispatcher(t)

	ccfg := control.DefaultConfig()
	ccfg.Reconciler.Interval = 100 * time.Millisecond
	ccfg.Reconciler.BaseBackoff = 200 * time.Millisecond
	ccfg.Reconciler.MaxBackoff = 500 * time.Millisecond
	ccfg.Reconciler.UnreachableAfter = 300 * time.Millisecond
	ccfg.Reconciler.LostAfter = 600 * time.Millisecond
	// Scaled like every other interval here rather than disabled, so the
	// settling window is part of what these tests exercise: the whole fleet is
	// registered up front, and placement should resume once it holds steady.
	ccfg.Reconciler.SettleFor = 400 * time.Millisecond
	ccfg.PublishInterval = 100 * time.Millisecond

	ctl := control.New(ccfg, catalog.NewLive(cat), disp, rt, slog.New(slog.DiscardHandler))
	disp.onStatus = ctl.Reconciler().OnShardStatus

	// Every node is this host, so the ports the reconciler allocates are the
	// ports the sim-engines actually bind.
	for i := range nodes {
		id := fmt.Sprintf("n%d", i)
		ctl.Inventory().Upsert(inventory.Node{
			ID: id, MgmtAddr: "127.0.0.1", State: inventory.Ready,
			GPUs: []discovery.GPU{{
				Index: 0, Name: "NVIDIA H100 80GB HBM3",
				MemoryTotalMB: 81559, ComputeCapability: "9.0",
			}},
		})
	}

	ctx, cancel := context.WithCancel(context.Background())
	rt.Start(ctx)
	ctl.Start(ctx)

	h := &controlHarness{
		router: httptest.NewServer(rt.Handler()),
		rt:     rt, ctl: ctl, disp: disp, inv: ctl.Inventory(),
	}
	t.Cleanup(func() {
		h.router.Close()
		ctl.Close()
		cancel()
		rt.Stop()
	})
	return h
}

// heartbeat keeps the named nodes alive.
func (h *controlHarness) heartbeat(ids ...string) {
	for _, id := range ids {
		h.inv.Touch(id)
	}
}

// waitForEndpoints blocks until the router sees n endpoints.
func (h *controlHarness) waitForEndpoints(t *testing.T, n int, limit time.Duration) {
	t.Helper()
	deadline := time.Now().Add(limit)
	for time.Now().Before(deadline) {
		if len(h.rt.Registry().Snapshots(testModel)) == n {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("router has %d endpoints, expected %d",
		len(h.rt.Registry().Snapshots(testModel)), n)
}

// TestControlPlaneFeedsRouter checks the wiring: replicas the reconciler brings
// up become endpoints the router can route to, without any static config.
func TestControlPlaneFeedsRouter(t *testing.T) {
	h := startControlPlane(t, 2, 2)

	// Keep the nodes alive while the reconciler works.
	stop := make(chan struct{})
	go func() {
		tk := time.NewTicker(50 * time.Millisecond)
		defer tk.Stop()
		for {
			select {
			case <-stop:
				return
			case <-tk.C:
				h.heartbeat("n0", "n1")
			}
		}
	}()
	defer close(stop)

	h.waitForEndpoints(t, 2, 30*time.Second)

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if _, _, err := streamRequest(ctx, h.router.URL, chatBody("hello", 4, true), nil); err != nil {
		t.Fatalf("the router could not serve from control-plane endpoints: %v", err)
	}
}

// TestRouterKeepsServingWhileNodeIsLost is the other half of the phase 3 exit
// criterion: a node disappears, its replica is rebuilt elsewhere, and traffic
// keeps flowing throughout.
func TestRouterKeepsServingWhileNodeIsLost(t *testing.T) {
	if testing.Short() {
		t.Skip("timing-dependent")
	}
	// Three nodes, two replicas: losing one still leaves somewhere to rebuild.
	h := startControlPlane(t, 3, 2)

	alive := map[string]bool{"n0": true, "n1": true, "n2": true}
	var aliveMu sync.Mutex

	stop := make(chan struct{})
	go func() {
		tk := time.NewTicker(50 * time.Millisecond)
		defer tk.Stop()
		for {
			select {
			case <-stop:
				return
			case <-tk.C:
				aliveMu.Lock()
				for id, ok := range alive {
					if ok {
						h.inv.Touch(id)
					}
				}
				aliveMu.Unlock()
			}
		}
	}()
	defer close(stop)

	h.waitForEndpoints(t, 2, 30*time.Second)

	// Continuous traffic for the whole experiment.
	ctx, cancel := context.WithCancel(context.Background())
	var ok, failed atomic.Int64
	var firstErr atomic.Value
	var wg sync.WaitGroup

	for range 4 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for ctx.Err() == nil {
				rctx, rcancel := context.WithTimeout(ctx, 10*time.Second)
				_, _, err := streamRequest(rctx, h.router.URL, chatBody("ping", 4, true), nil)
				rcancel()
				switch {
				case err == nil:
					ok.Add(1)
				case ctx.Err() != nil: // shutting down
				default:
					failed.Add(1)
					firstErr.CompareAndSwap(nil, err)
				}
				time.Sleep(20 * time.Millisecond)
			}
		}()
	}

	// Let traffic settle before disturbing anything.
	time.Sleep(600 * time.Millisecond)
	baseline := ok.Load()
	if baseline == 0 {
		cancel()
		wg.Wait()
		t.Fatalf("no requests succeeded before the failure: %v", firstErr.Load())
	}
	beforeFailures := failed.Load()

	// Find a node hosting a replica and kill it.
	var victim string
	var victimShards []string
	for _, rep := range h.ctl.Reconciler().Replicas() {
		if rep.State != reconciler.Ready {
			continue
		}
		victim = rep.Shards[0].NodeID
		for _, s := range rep.Shards {
			victimShards = append(victimShards, s.ID)
		}
		break
	}
	if victim == "" {
		cancel()
		wg.Wait()
		t.Fatal("no ready replica to disturb")
	}

	aliveMu.Lock()
	alive[victim] = false
	aliveMu.Unlock()
	h.disp.killNode(victim, victimShards)
	t.Logf("killed node %s (shards %v)", victim, victimShards)

	// The control plane should notice, rebuild elsewhere, and get back to two
	// endpoints. The generous window covers unreachable -> lost -> backoff.
	deadline := time.Now().Add(25 * time.Second)
	recovered := false
	for time.Now().Before(deadline) {
		if len(h.rt.Registry().Snapshots(testModel)) == 2 && h.disp.liveShards() == 2 {
			recovered = true
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	// Let traffic run a little longer on the rebuilt topology.
	time.Sleep(500 * time.Millisecond)
	cancel()
	wg.Wait()

	t.Logf("requests ok=%d failed=%d (failures before the kill: %d)",
		ok.Load(), failed.Load(), beforeFailures)

	if !recovered {
		var placed []string
		for _, rep := range h.ctl.Reconciler().Replicas() {
			placed = append(placed, fmt.Sprintf("%s=%s@%v", rep.ID, rep.State, rep.NodeIDs()))
		}
		t.Fatalf("the cluster did not return to 2 replicas: %v", placed)
	}

	// The victim must not be hosting anything any more.
	for _, rep := range h.ctl.Reconciler().Replicas() {
		if rep.OnNode(victim) {
			t.Errorf("replica %s is still placed on the lost node %s", rep.ID, victim)
		}
	}

	// Traffic kept flowing: more succeeded after the failure than before it.
	if ok.Load() <= baseline {
		t.Errorf("no requests succeeded after the node was lost (%d before, %d total)",
			baseline, ok.Load())
	}
	// Some failures are expected in the window where a replica is gone and its
	// replacement has not finished loading. What must not happen is total loss
	// of service, which is what having a second replica is for.
	if failed.Load() > ok.Load() {
		t.Errorf("more requests failed (%d) than succeeded (%d); service was not maintained: %v",
			failed.Load(), ok.Load(), firstErr.Load())
	}
}