Add the health poller

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
2026-09-25 02:16:38 -07:00
parent 73b24357dd
commit 68ad94d693
3 changed files with 413 additions and 0 deletions
@@ -0,0 +1,258 @@
+// Package health tracks, for every host, whether it answered its last poll, which models it has
+// loaded, and when it last answered. The proxy reads the table to choose a host and to record a
+// request failure against a host.
+package health
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"sort"
+	"sync"
+	"time"
+)
+
+// RecoveryPolls is the number of good polls in a row a host needs after a failure before it is
+// trusted again; a router loading a model flaps, so one good poll is not enough.
+const RecoveryPolls = 2
+
+// MaxModelsBody bounds how many bytes we read from either /health or /v1/models.
+const MaxModelsBody = 1 << 20
+
+// Status is a snapshot of one host's health, safe to copy.
+type Status struct {
+	Healthy     bool      `json:"healthy"`
+	Loaded      []string  `json:"loaded"`      // sorted, unique model ids from the last good poll
+	LastOK      time.Time `json:"last_ok"`     // zero if never
+	LastErr     string    `json:"last_err"`    // "" after a good poll
+	Consecutive int       `json:"consecutive"` // good polls in a row
+}
+
+type entry struct {
+	status     Status
+	everFailed bool
+}
+
+type pollResult struct {
+	ok        bool
+	cancelled bool
+	reason    string
+	loaded    []string
+}
+
+// Table maps a host name to its health status. All methods are safe for concurrent use.
+type Table struct {
+	mu       sync.Mutex
+	baseURL  map[string]string
+	entries  map[string]entry
+	client   *http.Client
+	interval time.Duration
+}
+
+// New builds a table from hosts (name -> base URL, no trailing slash). A nil client becomes a
+// five-second-timeout client. Every host starts untrusted with an empty Loaded slice.
+func New(hosts map[string]string, interval time.Duration, client *http.Client) *Table {
+	if client == nil {
+		client = &http.Client{Timeout: 5 * time.Second}
+	}
+	t := &Table{
+		baseURL:  hosts,
+		entries:  make(map[string]entry, len(hosts)),
+		client:   client,
+		interval: interval,
+	}
+	for name := range hosts {
+		t.entries[name] = entry{status: Status{Loaded: []string{}}}
+	}
+	return t
+}
+
+// Run polls once immediately, then every interval, and returns when ctx is done.
+func (t *Table) Run(ctx context.Context) {
+	t.PollOnce(ctx)
+	ticker := time.NewTicker(t.interval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			t.PollOnce(ctx)
+		}
+	}
+}
+
+// PollOnce polls every host concurrently and returns when all have finished.
+func (t *Table) PollOnce(ctx context.Context) {
+	t.mu.Lock()
+	names := make([]string, 0, len(t.baseURL))
+	for name := range t.baseURL {
+		names = append(names, name)
+	}
+	t.mu.Unlock()
+
+	var wg sync.WaitGroup
+	for _, name := range names {
+		wg.Add(1)
+		go func(name string) {
+			defer wg.Done()
+			t.pollHost(ctx, name)
+		}(name)
+	}
+	wg.Wait()
+}
+
+// Get returns a copy of one host's status; ok is false for an unknown name.
+func (t *Table) Get(name string) (Status, bool) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	e, ok := t.entries[name]
+	if !ok {
+		return Status{}, false
+	}
+	return copyStatus(e.status), true
+}
+
+// All returns copies of every host's status.
+func (t *Table) All() map[string]Status {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	out := make(map[string]Status, len(t.entries))
+	for name, e := range t.entries {
+		out[name] = copyStatus(e.status)
+	}
+	return out
+}
+
+// MarkDown records a failure seen by the proxy. Loaded is left as last seen.
+func (t *Table) MarkDown(name, reason string) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	e, ok := t.entries[name]
+	if !ok {
+		return
+	}
+	e.everFailed = true
+	e.status.Healthy = false
+	e.status.Consecutive = 0
+	e.status.LastErr = "marked down: " + reason
+	t.entries[name] = e
+}
+
+func (t *Table) pollHost(ctx context.Context, name string) {
+	base := t.baseURL[name]
+	r := t.poll(ctx, base)
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	e := t.entries[name]
+	if r.cancelled {
+		return
+	}
+	if r.ok {
+		e.status.Consecutive++
+		e.status.LastOK = time.Now()
+		e.status.LastErr = ""
+		e.status.Loaded = r.loaded
+		e.status.Healthy = !e.everFailed || e.status.Consecutive >= RecoveryPolls
+	} else {
+		e.everFailed = true
+		e.status.Healthy = false
+		e.status.Consecutive = 0
+		e.status.LastErr = r.reason
+	}
+	t.entries[name] = e
+}
+
+// poll runs the two requests for one host. A request that fails because ctx is done yields a
+// cancelled result so the caller records nothing; any other failure yields a reason.
+func (t *Table) poll(ctx context.Context, base string) pollResult {
+	loaded, r := t.check(ctx, base+"/health", "health")
+	if r.cancelled || r.reason != "" {
+		return r
+	}
+	loaded, r = t.check(ctx, base+"/v1/models", "models")
+	if r.cancelled {
+		return r
+	}
+	if r.reason != "" {
+		return r
+	}
+	return pollResult{ok: true, loaded: loaded}
+}
+
+// check performs one GET and, on success, returns the decoded model ids. Health checks use the
+// status code only; the models check also decodes the resident model ids.
+func (t *Table) check(ctx context.Context, url, prefix string) ([]string, pollResult) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, t.fail(ctx, prefix, err)
+	}
+	resp, err := t.client.Do(req)
+	if err != nil {
+		return nil, t.fail(ctx, prefix, err)
+	}
+
+	if prefix == "health" {
+		// We only need the status code; bound and drain the body.
+		func() {
+			defer resp.Body.Close()
+			_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, MaxModelsBody))
+		}()
+		if resp.StatusCode != http.StatusOK {
+			return nil, pollResult{reason: fmt.Sprintf("health: HTTP %d", resp.StatusCode)}
+		}
+		return nil, pollResult{}
+	}
+
+	// models: bound and decode the body.
+	if resp.StatusCode != http.StatusOK {
+		defer resp.Body.Close()
+		return nil, pollResult{reason: fmt.Sprintf("models: HTTP %d", resp.StatusCode)}
+	}
+	var m struct {
+		Data []struct {
+			ID string `json:"id"`
+		} `json:"data"`
+	}
+	if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&m); err != nil {
+		defer resp.Body.Close()
+		return nil, t.fail(ctx, "models", err)
+	}
+	defer resp.Body.Close()
+
+	loaded := make([]string, 0, len(m.Data))
+	seen := make(map[string]struct{}, len(m.Data))
+	for _, d := range m.Data {
+		if d.ID == "" {
+			continue
+		}
+		if _, ok := seen[d.ID]; ok {
+			continue
+		}
+		seen[d.ID] = struct{}{}
+		loaded = append(loaded, d.ID)
+	}
+	sort.Strings(loaded)
+	return loaded, pollResult{}
+}
+
+// fail builds a failure result, treating a done ctx as a cancellation rather than a failure.
+func (t *Table) fail(ctx context.Context, prefix string, err error) pollResult {
+	if ctx.Err() != nil {
+		return pollResult{cancelled: true}
+	}
+	return pollResult{reason: prefix + ": " + err.Error()}
+}
+
+func copyStatus(s Status) Status {
+	if s.Loaded == nil {
+		return s
+	}
+	out := s
+	out.Loaded = make([]string, len(s.Loaded))
+	copy(out.Loaded, s.Loaded)
+	return out
+}
@@ -0,0 +1,154 @@
+package health_test
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"git.wntrmute.dev/kyle/crossbar/internal/health"
+)
+
+// fake is a llama-server stand-in whose /health can be flipped and whose model list is fixed.
+type fake struct {
+	srv    *httptest.Server
+	down   atomic.Bool
+	models []string
+	hits   atomic.Int32
+}
+
+func newFake(t *testing.T, models ...string) *fake {
+	f := &fake{models: models}
+	mux := http.NewServeMux()
+	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+		f.hits.Add(1)
+		if f.down.Load() {
+			http.Error(w, "loading", http.StatusServiceUnavailable)
+			return
+		}
+		_, _ = w.Write([]byte(`{"status":"ok"}`))
+	})
+	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
+		type m struct {
+			ID string `json:"id"`
+		}
+		var data []m
+		for _, id := range f.models {
+			data = append(data, m{ID: id})
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
+	})
+	f.srv = httptest.NewServer(mux)
+	t.Cleanup(f.srv.Close)
+	return f
+}
+
+func TestFirstPollMakesHealthy(t *testing.T) {
+	a := newFake(t, "zeta", "alpha", "alpha")
+	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
+	if s, ok := tbl.Get("a"); !ok || s.Healthy || len(s.Loaded) != 0 {
+		t.Fatalf("before any poll: %+v %v", s, ok)
+	}
+	tbl.PollOnce(context.Background())
+	s, _ := tbl.Get("a")
+	if !s.Healthy || s.Consecutive != 1 || s.LastErr != "" || s.LastOK.IsZero() {
+		t.Errorf("after one good poll: %+v", s)
+	}
+	if len(s.Loaded) != 2 || s.Loaded[0] != "alpha" || s.Loaded[1] != "zeta" {
+		t.Errorf("Loaded = %v, want sorted, unique [alpha zeta]", s.Loaded)
+	}
+}
+
+func TestFailureThenRecoveryNeedsTwoPolls(t *testing.T) {
+	a := newFake(t, "m")
+	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
+	ctx := context.Background()
+	tbl.PollOnce(ctx)
+	a.down.Store(true)
+	tbl.PollOnce(ctx)
+	s, _ := tbl.Get("a")
+	if s.Healthy || s.Consecutive != 0 || s.LastErr == "" {
+		t.Fatalf("after failure: %+v", s)
+	}
+	if len(s.Loaded) != 1 {
+		t.Errorf("Loaded is left as last seen; got %v", s.Loaded)
+	}
+	a.down.Store(false)
+	tbl.PollOnce(ctx)
+	if s, _ := tbl.Get("a"); s.Healthy || s.Consecutive != 1 {
+		t.Errorf("one good poll after a failure must not be healthy yet: %+v", s)
+	}
+	tbl.PollOnce(ctx)
+	if s, _ := tbl.Get("a"); !s.Healthy || s.Consecutive != 2 || s.LastErr != "" {
+		t.Errorf("two good polls: %+v", s)
+	}
+}
+
+func TestMarkDown(t *testing.T) {
+	a := newFake(t, "m")
+	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
+	tbl.PollOnce(context.Background())
+	tbl.MarkDown("a", "connection refused")
+	s, _ := tbl.Get("a")
+	if s.Healthy || s.Consecutive != 0 || s.LastErr != "marked down: connection refused" {
+		t.Errorf("after MarkDown: %+v", s)
+	}
+	tbl.MarkDown("nobody", "x") // unknown hosts are ignored, not a panic
+	tbl.PollOnce(context.Background())
+	if s, _ := tbl.Get("a"); s.Healthy {
+		t.Errorf("one poll after MarkDown must not be healthy: %+v", s)
+	}
+}
+
+func TestUnreachableAndUnknown(t *testing.T) {
+	tbl := health.New(map[string]string{"a": "http://127.0.0.1:1"}, time.Hour, &http.Client{Timeout: time.Second})
+	tbl.PollOnce(context.Background())
+	s, ok := tbl.Get("a")
+	if !ok || s.Healthy || s.LastErr == "" {
+		t.Errorf("unreachable host: %+v %v", s, ok)
+	}
+	if _, ok := tbl.Get("zzz"); ok {
+		t.Errorf("unknown host must report ok=false")
+	}
+}
+
+func TestAllIsACopy(t *testing.T) {
+	a := newFake(t, "m")
+	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
+	tbl.PollOnce(context.Background())
+	all := tbl.All()
+	all["a"].Loaded[0] = "changed"
+	if s, _ := tbl.Get("a"); s.Loaded[0] != "m" {
+		t.Errorf("All must return copies")
+	}
+	if len(all) != 1 {
+		t.Errorf("All = %v", all)
+	}
+}
+
+func TestRunPollsOnStart(t *testing.T) {
+	a := newFake(t, "m")
+	tbl := health.New(map[string]string{"a": a.srv.URL}, 20*time.Millisecond, nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() { tbl.Run(ctx); close(done) }()
+	deadline := time.Now().Add(2 * time.Second)
+	for a.hits.Load() < 3 && time.Now().Before(deadline) {
+		time.Sleep(5 * time.Millisecond)
+	}
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("Run did not return after cancel")
+	}
+	if a.hits.Load() < 3 {
+		t.Errorf("Run polled %d times in 2s at 20ms interval", a.hits.Load())
+	}
+	if s, _ := tbl.Get("a"); !s.Healthy {
+		t.Errorf("not healthy after Run: %+v", s)
+	}
+}