259 lines
6.6 KiB
Go
259 lines
6.6 KiB
Go
// 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
|
|
}
|