Aleksey Shakhmatov
01478be0dc
A malformed endpoint URL panicked inside Transport, crashing the host app (often at startup from external config). Capture the parse error and surface it from the transport on first use instead. Add the previously untested HealthChecker coverage (initial probe, recovery, Stop termination, unknown endpoint), raising balancer coverage from ~41% to ~87%. Default the health probe path to /healthz to match this library's own server.
116 lines
3.0 KiB
Go
116 lines
3.0 KiB
Go
package balancer
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
|
|
"git.codelab.vc/pkg/httpx/middleware"
|
|
)
|
|
|
|
// ErrNoHealthy is returned when no healthy endpoints are available.
|
|
var ErrNoHealthy = errors.New("httpx: no healthy endpoints available")
|
|
|
|
// Endpoint represents a backend server that can handle requests.
|
|
type Endpoint struct {
|
|
URL string
|
|
Weight int
|
|
Meta map[string]string
|
|
}
|
|
|
|
// Strategy selects an endpoint from the list of healthy endpoints.
|
|
type Strategy interface {
|
|
Next(healthy []Endpoint) (Endpoint, error)
|
|
}
|
|
|
|
// Closer can be used to shut down resources associated with a balancer
|
|
// transport (e.g. background health checker goroutines).
|
|
type Closer struct {
|
|
healthChecker *HealthChecker
|
|
}
|
|
|
|
// Close stops background goroutines. Safe to call multiple times.
|
|
func (c *Closer) Close() {
|
|
if c.healthChecker != nil {
|
|
c.healthChecker.Stop()
|
|
}
|
|
}
|
|
|
|
// Transport returns a middleware that load-balances requests across the
|
|
// provided endpoints using the configured strategy.
|
|
//
|
|
// For each request the middleware picks an endpoint via the strategy,
|
|
// replaces the request URL scheme and host with the endpoint's URL,
|
|
// and forwards the request to the underlying RoundTripper.
|
|
//
|
|
// If active health checking is enabled (WithHealthCheck), a background
|
|
// goroutine periodically probes endpoints. Otherwise all endpoints are
|
|
// assumed healthy.
|
|
func Transport(endpoints []Endpoint, opts ...Option) (middleware.Middleware, *Closer) {
|
|
o := &options{
|
|
strategy: RoundRobin(),
|
|
}
|
|
for _, opt := range opts {
|
|
opt(o)
|
|
}
|
|
|
|
// Pre-parse endpoint URLs once at construction time. A malformed URL is a
|
|
// configuration error: rather than panicking (which would crash the host
|
|
// application, often at startup from external config), we capture the
|
|
// error and surface it from the transport on first use.
|
|
parsed := make(map[string]*url.URL, len(endpoints))
|
|
var parseErr error
|
|
for _, ep := range endpoints {
|
|
u, err := url.Parse(ep.URL)
|
|
if err != nil {
|
|
if parseErr == nil {
|
|
parseErr = fmt.Errorf("balancer: invalid endpoint URL %q: %w", ep.URL, err)
|
|
}
|
|
continue
|
|
}
|
|
parsed[ep.URL] = u
|
|
}
|
|
|
|
if o.healthChecker != nil {
|
|
o.healthChecker.Start(endpoints)
|
|
}
|
|
|
|
closer := &Closer{healthChecker: o.healthChecker}
|
|
|
|
return func(next http.RoundTripper) http.RoundTripper {
|
|
return middleware.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
|
if parseErr != nil {
|
|
return nil, parseErr
|
|
}
|
|
|
|
healthy := endpoints
|
|
if o.healthChecker != nil {
|
|
healthy = o.healthChecker.Healthy(endpoints)
|
|
}
|
|
|
|
if len(healthy) == 0 {
|
|
return nil, ErrNoHealthy
|
|
}
|
|
|
|
ep, err := o.strategy.Next(healthy)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
epURL := parsed[ep.URL]
|
|
|
|
// Shallow-copy request and URL to avoid mutating the original,
|
|
// without the expense of req.Clone's deep header copy.
|
|
r := *req
|
|
u := *req.URL
|
|
r.URL = &u
|
|
r.URL.Scheme = epURL.Scheme
|
|
r.URL.Host = epURL.Host
|
|
r.Host = epURL.Host
|
|
|
|
return next.RoundTrip(&r)
|
|
})
|
|
}, closer
|
|
}
|