remotes: add least-connections mirror strategy (round-robin remains default)
Add a per-remote mirror_strategy selector for the mirrorlist, supporting
round_robin (default, unchanged) and least_conn.
- models.Remote gains MirrorStrategy string + MirrorStrategy{RoundRobin,LeastConn}
constants and ValidateMirrorStrategy (enum check; least_conn requires a
non-empty mirrorlist). Empty behaves as round_robin for back-compat.
- DB: additive mirror_strategy TEXT NOT NULL DEFAULT 'round_robin' column
(CREATE TABLE + ADD COLUMN IF NOT EXISTS), wired through remoteCols/scanRemote/
CreateRemote/UpdateRemote; empty normalized to round_robin on write.
- Engine: least_conn starts each attempt with the pool URL holding the fewest
in-flight requests via a per-remote/per-URL atomic gauge (incremented around
each upstream call in head/fetch/checkUpstream), ties broken by the existing
round-robin rotation. Round-robin path and failover order unchanged;
single-URL pools are a no-op.
- Tests: unit tests for least-loaded selection, round-robin default, gauge
inc/dec, single-URL no-op, and strategy validation; DB round-trip covers the
new column; docker e2e adds a least_conn distribution test and a real dnf
install through a least_conn remote.
This commit is contained in:
@@ -10,6 +10,7 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -41,6 +42,11 @@ type Engine struct {
|
||||
// *atomic.Uint64) used to rotate the starting mirror across upstream base
|
||||
// URLs. Distribution is per-replica and approximate, which is fine.
|
||||
rrCounters sync.Map
|
||||
// inflight holds a per-remote, per-upstream-URL in-flight request gauge
|
||||
// (key "remoteName\x00baseURL" -> *atomic.Int64) used by the least_conn
|
||||
// mirror strategy to prefer the mirror currently handling the fewest
|
||||
// requests. Per-replica and approximate, which is fine.
|
||||
inflight sync.Map
|
||||
}
|
||||
|
||||
func NewEngine(db *database.DB, c *cache.Redis, s *storage.S3) *Engine {
|
||||
@@ -237,7 +243,9 @@ func (e *Engine) headUpstream(ctx context.Context, remote models.Remote, path st
|
||||
}
|
||||
var lastErr error
|
||||
for i, url := range order {
|
||||
ctr := e.beginAttempt(remote, url)
|
||||
result, err := e.headUpstreamOnce(ctx, withBaseURL(remote, url), path, prov)
|
||||
endAttempt(ctr)
|
||||
if err == nil {
|
||||
return result, nil
|
||||
}
|
||||
@@ -316,7 +324,9 @@ func (e *Engine) fetchFromUpstream(ctx context.Context, remote models.Remote, pa
|
||||
}
|
||||
var lastErr error
|
||||
for i, url := range order {
|
||||
ctr := e.beginAttempt(remote, url)
|
||||
result, err := e.fetchFromUpstreamOnce(ctx, withBaseURL(remote, url), path, prov, class, ttl, clientHeaders)
|
||||
endAttempt(ctr)
|
||||
if err == nil {
|
||||
return result, nil
|
||||
}
|
||||
@@ -517,7 +527,9 @@ func (e *Engine) checkUpstream(ctx context.Context, remote models.Remote, path,
|
||||
}
|
||||
var lastErr error
|
||||
for i, url := range order {
|
||||
ctr := e.beginAttempt(remote, url)
|
||||
notModified, err := e.checkUpstreamOnce(ctx, withBaseURL(remote, url), path, etag, prov)
|
||||
endAttempt(ctr)
|
||||
if err == nil {
|
||||
return notModified, nil
|
||||
}
|
||||
@@ -729,23 +741,61 @@ func isNetworkError(err error) bool {
|
||||
|
||||
// baseURLAttemptOrder returns the ordered upstream base URLs to try for a single
|
||||
// request, drawn from the remote's pool ([base_url] + mirrorlist). A multi-mirror
|
||||
// remote starts at the next round-robin position and advances linearly for
|
||||
// remote starts at a strategy-chosen position and advances linearly for
|
||||
// failover; a remote with no mirrorlist yields exactly [base_url], preserving the
|
||||
// original single-attempt behavior.
|
||||
// original single-attempt behavior. The default (round_robin) rotates the
|
||||
// starting mirror; least_conn starts with the mirror handling the fewest
|
||||
// in-flight requests. Failover order after the first pick is unchanged.
|
||||
func (e *Engine) baseURLAttemptOrder(remote models.Remote) []string {
|
||||
urls := remote.UpstreamPool()
|
||||
if len(urls) <= 1 {
|
||||
return urls
|
||||
}
|
||||
// Rotate by the round-robin cursor first so equal-load mirrors still spread
|
||||
// evenly; least_conn then stable-sorts this rotation by in-flight count.
|
||||
v, _ := e.rrCounters.LoadOrStore(remote.Name, new(atomic.Uint64))
|
||||
start := int(v.(*atomic.Uint64).Add(1) - 1)
|
||||
ordered := make([]string, len(urls))
|
||||
for i := range urls {
|
||||
ordered[i] = urls[(start+i)%len(urls)]
|
||||
}
|
||||
if remote.MirrorStrategy == models.MirrorStrategyLeastConn {
|
||||
sort.SliceStable(ordered, func(a, b int) bool {
|
||||
return e.inflightCounter(remote.Name, ordered[a]).Load() < e.inflightCounter(remote.Name, ordered[b]).Load()
|
||||
})
|
||||
}
|
||||
return ordered
|
||||
}
|
||||
|
||||
// inflightCounter returns the shared in-flight request gauge for a given
|
||||
// (remote, upstream URL), creating it on first use.
|
||||
func (e *Engine) inflightCounter(remoteName, url string) *atomic.Int64 {
|
||||
v, _ := e.inflight.LoadOrStore(remoteName+"\x00"+url, new(atomic.Int64))
|
||||
return v.(*atomic.Int64)
|
||||
}
|
||||
|
||||
// beginAttempt increments the in-flight gauge for a least_conn multi-mirror
|
||||
// remote before an upstream call and returns the counter to release; it is a
|
||||
// no-op (returns nil) for round-robin remotes and single-URL pools.
|
||||
func (e *Engine) beginAttempt(remote models.Remote, url string) *atomic.Int64 {
|
||||
if remote.MirrorStrategy != models.MirrorStrategyLeastConn {
|
||||
return nil
|
||||
}
|
||||
if len(remote.UpstreamPool()) <= 1 {
|
||||
return nil
|
||||
}
|
||||
ctr := e.inflightCounter(remote.Name, url)
|
||||
ctr.Add(1)
|
||||
return ctr
|
||||
}
|
||||
|
||||
// endAttempt decrements a gauge returned by beginAttempt, tolerating nil.
|
||||
func endAttempt(ctr *atomic.Int64) {
|
||||
if ctr != nil {
|
||||
ctr.Add(-1)
|
||||
}
|
||||
}
|
||||
|
||||
// withBaseURL narrows a remote's active BaseURL to a single selected mirror so
|
||||
// providers (UpstreamURL/AuthHeaders/RewriteResponse) operate on exactly that
|
||||
// upstream for this attempt.
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"git.unkin.net/unkin/artifactapi/pkg/models"
|
||||
)
|
||||
|
||||
// baseURLAttemptOrder and the in-flight gauge only touch the engine's sync.Map
|
||||
// fields, so these tests run against a zero-value Engine without a DB/S3/redis
|
||||
// stack and can drive the gauge deterministically.
|
||||
|
||||
// TestLeastConnPicksLeastLoaded pre-loads one mirror's in-flight gauge and
|
||||
// asserts a least_conn remote starts its attempt order with the idle mirror.
|
||||
func TestLeastConnPicksLeastLoaded(t *testing.T) {
|
||||
e := &Engine{}
|
||||
r := models.Remote{
|
||||
Name: "lc",
|
||||
BaseURL: "https://a.example",
|
||||
Mirrorlist: []string{"https://b.example"},
|
||||
MirrorStrategy: models.MirrorStrategyLeastConn,
|
||||
}
|
||||
|
||||
// Make A appear busy: least_conn must prefer B regardless of RR rotation.
|
||||
e.inflightCounter(r.Name, "https://a.example").Add(3)
|
||||
for i := 0; i < 5; i++ {
|
||||
order := e.baseURLAttemptOrder(r)
|
||||
if len(order) != 2 {
|
||||
t.Fatalf("attempt %d: order len = %d, want 2", i, len(order))
|
||||
}
|
||||
if order[0] != "https://b.example" {
|
||||
t.Fatalf("attempt %d: least_conn started with %q, want idle mirror https://b.example", i, order[0])
|
||||
}
|
||||
}
|
||||
|
||||
// Once B is the busier mirror, the starting pick flips to A.
|
||||
e.inflightCounter(r.Name, "https://b.example").Add(10)
|
||||
if order := e.baseURLAttemptOrder(r); order[0] != "https://a.example" {
|
||||
t.Fatalf("after loading B, least_conn started with %q, want https://a.example", order[0])
|
||||
}
|
||||
}
|
||||
|
||||
// TestRoundRobinDefaultUnchanged asserts an unset strategy still rotates the
|
||||
// starting mirror across the pool and ignores the in-flight gauge.
|
||||
func TestRoundRobinDefaultUnchanged(t *testing.T) {
|
||||
e := &Engine{}
|
||||
r := models.Remote{
|
||||
Name: "rr",
|
||||
BaseURL: "https://a.example",
|
||||
Mirrorlist: []string{"https://b.example"},
|
||||
}
|
||||
|
||||
// Even with A heavily loaded, round-robin must still rotate (not avoid A).
|
||||
e.inflightCounter(r.Name, "https://a.example").Add(100)
|
||||
starts := map[string]int{}
|
||||
for i := 0; i < 4; i++ {
|
||||
starts[e.baseURLAttemptOrder(r)[0]]++
|
||||
}
|
||||
if starts["https://a.example"] == 0 || starts["https://b.example"] == 0 {
|
||||
t.Fatalf("round-robin did not rotate starting mirror: %v", starts)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLeastConnSingleURLNoOp asserts a single-URL pool yields exactly [base_url]
|
||||
// and beginAttempt is a no-op there and for round-robin remotes.
|
||||
func TestLeastConnSingleURLNoOp(t *testing.T) {
|
||||
e := &Engine{}
|
||||
solo := models.Remote{Name: "solo", BaseURL: "https://a.example", MirrorStrategy: models.MirrorStrategyLeastConn}
|
||||
if order := e.baseURLAttemptOrder(solo); len(order) != 1 || order[0] != "https://a.example" {
|
||||
t.Fatalf("single-url order = %v, want [base_url]", order)
|
||||
}
|
||||
if ctr := e.beginAttempt(solo, "https://a.example"); ctr != nil {
|
||||
t.Fatal("beginAttempt on single-url pool should be a no-op (nil)")
|
||||
}
|
||||
|
||||
rr := models.Remote{Name: "rr2", BaseURL: "https://a.example", Mirrorlist: []string{"https://b.example"}}
|
||||
if ctr := e.beginAttempt(rr, "https://a.example"); ctr != nil {
|
||||
t.Fatal("beginAttempt on round-robin remote should be a no-op (nil)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBeginEndAttemptGauge asserts the gauge increments on begin and returns to
|
||||
// zero after endAttempt, so it tracks live in-flight requests.
|
||||
func TestBeginEndAttemptGauge(t *testing.T) {
|
||||
e := &Engine{}
|
||||
r := models.Remote{
|
||||
Name: "g",
|
||||
BaseURL: "https://a.example",
|
||||
Mirrorlist: []string{"https://b.example"},
|
||||
MirrorStrategy: models.MirrorStrategyLeastConn,
|
||||
}
|
||||
c1 := e.beginAttempt(r, "https://a.example")
|
||||
c2 := e.beginAttempt(r, "https://a.example")
|
||||
if got := e.inflightCounter(r.Name, "https://a.example").Load(); got != 2 {
|
||||
t.Fatalf("gauge after two begins = %d, want 2", got)
|
||||
}
|
||||
endAttempt(c1)
|
||||
endAttempt(c2)
|
||||
if got := e.inflightCounter(r.Name, "https://a.example").Load(); got != 0 {
|
||||
t.Fatalf("gauge after matching ends = %d, want 0", got)
|
||||
}
|
||||
endAttempt(nil) // tolerated
|
||||
}
|
||||
Reference in New Issue
Block a user