Aleksey Shakhmatov
8d322123a4
Implements balancer middleware with URL rewriting per-request: - RoundRobin, Failover, and WeightedRandom endpoint selection strategies - Background HealthChecker with configurable probe interval and path - Thread-safe health state tracking with sync.RWMutex
79 lines
1.9 KiB
Go
79 lines
1.9 KiB
Go
package balancer
|
|
|
|
import (
|
|
"errors"
|
|
"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)
|
|
}
|
|
|
|
// 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 {
|
|
o := &options{
|
|
strategy: RoundRobin(),
|
|
}
|
|
for _, opt := range opts {
|
|
opt(o)
|
|
}
|
|
|
|
if o.healthChecker != nil {
|
|
o.healthChecker.Start(endpoints)
|
|
}
|
|
|
|
return func(next http.RoundTripper) http.RoundTripper {
|
|
return middleware.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
|
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, err := url.Parse(ep.URL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Clone the request URL and replace scheme+host with the endpoint.
|
|
r := req.Clone(req.Context())
|
|
r.URL.Scheme = epURL.Scheme
|
|
r.URL.Host = epURL.Host
|
|
r.Host = epURL.Host
|
|
|
|
return next.RoundTrip(r)
|
|
})
|
|
}
|
|
}
|