108 lines
3.1 KiB
Go
108 lines
3.1 KiB
Go
package proxy
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"git.wntrmute.dev/kyle/crossbar/internal/choose"
|
|
"git.wntrmute.dev/kyle/crossbar/internal/config"
|
|
"git.wntrmute.dev/kyle/crossbar/internal/health"
|
|
"git.wntrmute.dev/kyle/crossbar/internal/lease"
|
|
"git.wntrmute.dev/kyle/crossbar/internal/limiter"
|
|
)
|
|
|
|
// Hosts adapts the health table and config for the lease table, and holds the operator's drain set.
|
|
// The lease table filters candidates by Healthy && !Draining: Healthy is the health table's view of
|
|
// a host, Draining is an operator flag that keeps new leases off a host being taken out of service.
|
|
type Hosts struct {
|
|
health *health.Table
|
|
drain map[string]bool
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// HostView builds the lease-table view of the health table and config.
|
|
func HostView(h *health.Table, cfg *config.Config) *Hosts {
|
|
return &Hosts{health: h, drain: make(map[string]bool)}
|
|
}
|
|
|
|
// Healthy reports whether the health table says the host answered its last poll; an unknown host is
|
|
// not healthy.
|
|
func (h *Hosts) Healthy(name string) bool {
|
|
s, ok := h.health.Get(name)
|
|
if !ok {
|
|
return false
|
|
}
|
|
return s.Healthy
|
|
}
|
|
|
|
// Draining reports whether an operator is draining the host.
|
|
func (h *Hosts) Draining(name string) bool {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
return h.drain[name]
|
|
}
|
|
|
|
// SetDraining turns draining on or off for a host.
|
|
func (h *Hosts) SetDraining(name string, on bool) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if on {
|
|
h.drain[name] = true
|
|
return
|
|
}
|
|
delete(h.drain, name)
|
|
}
|
|
|
|
// hostChooser adapts config, health and limiter to lease.Chooser via choose.Best.
|
|
type hostChooser struct {
|
|
cfg *config.Config
|
|
health *health.Table
|
|
lim *limiter.Limiter
|
|
}
|
|
|
|
// Chooser adapts config, health and limiter to lease.Chooser using choose.Best: among the candidates
|
|
// it prefers the hosts that have the model loaded over those that can only serve it, then the one
|
|
// with the most free slots times weight, breaking ties by shortest queue. Draining is reported false
|
|
// because the lease table already filtered draining hosts out before calling Choose.
|
|
func Chooser(cfg *config.Config, h *health.Table, l *limiter.Limiter) lease.Chooser {
|
|
return &hostChooser{cfg: cfg, health: h, lim: l}
|
|
}
|
|
|
|
func (c *hostChooser) Choose(candidates []string, model string) (string, bool) {
|
|
return choose.Best(candidates, func(host string) (choose.Info, bool) {
|
|
s, ok := c.health.Get(host)
|
|
info := choose.Info{
|
|
Healthy: ok && s.Healthy,
|
|
Draining: false,
|
|
Loaded: contains(s.Loaded, model),
|
|
CanServe: c.cfg.Serves(host, model),
|
|
Free: freeForModel(c.cfg.Hosts[host], model, c.lim.InFlight(host, model)),
|
|
Queued: c.lim.Queued(host, model),
|
|
Weight: c.cfg.Hosts[host].Weight,
|
|
}
|
|
return info, ok
|
|
})
|
|
}
|
|
|
|
func contains(list []string, v string) bool {
|
|
for _, s := range list {
|
|
if s == v {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// freeForModel returns the free slots for one model on one host: its parallel minus the in-flight
|
|
// count, floored at zero, and zero when the host does not list that model.
|
|
func freeForModel(h config.Host, model string, inflight int) int {
|
|
m, ok := h.Models[model]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
free := m.Parallel - inflight
|
|
if free < 0 {
|
|
return 0
|
|
}
|
|
return free
|
|
}
|