Add the health poller
Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user