proxy-pool/internal/gateway/dispatch/dispatcher_test.go
youfak 265026a907
Some checks are pending
ci / proto (push) Waiting to run
ci / test (ubuntu-latest) (push) Waiting to run
ci / test (windows-latest) (push) Waiting to run
ci / race (push) Waiting to run
ci / integration (push) Waiting to run
feat: handle unavailable gateway routes
2026-07-31 15:55:50 +08:00

581 lines
17 KiB
Go

package dispatch
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
proxyDomain "proxy-pool/internal/domain/proxy"
"proxy-pool/internal/domain/routing"
"proxy-pool/internal/gateway/snapshot"
)
func TestAcquireFiltersAndReservesLocalCapacity(t *testing.T) {
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
expiresSoon := now.Add(5 * time.Second)
expiresLater := now.Add(time.Minute)
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := []proxyDomain.Proxy{
{ID: "wrong-upstream", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "b", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater},
{ID: "expiring", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresSoon},
{ID: "selected", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}},
}
envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
dispatcher := New(store)
lease, err := dispatcher.Acquire(Request{
Now: now,
Scheme: proxyDomain.SchemeHTTP,
Upstreams: []string{"a"},
RequiredTags: map[string]string{"region": "cn-east"},
SafetyMargin: 10 * time.Second,
})
if err != nil {
t.Fatalf("Acquire(): %v", err)
}
if lease.Proxy.ID != "selected" {
t.Fatalf("selected proxy = %q, want selected", lease.Proxy.ID)
}
if err := lease.Commit(); err != nil {
t.Fatalf("Commit(): %v", err)
}
if err := lease.Release(); err != nil {
t.Fatalf("Release(): %v", err)
}
}
func TestAcquireRejectsSnapshotAfterOverallValidityDeadline(t *testing.T) {
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
expiresAt := now.Add(time.Hour)
store := snapshot.NewStore("cluster-a", "worker-a")
envelope := snapshot.Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
ValidUntil: now.Add(time.Second),
Proxies: []proxyDomain.Proxy{{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable,
MaxConcurrency: 1, ExpiresAt: &expiresAt,
}},
}
envelope.Checksum = snapshot.Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
if _, err := New(store).Acquire(Request{Now: now.Add(time.Second), Scheme: proxyDomain.SchemeHTTP}); !errors.Is(err, ErrNoCandidate) {
t.Fatalf("Acquire(at validity deadline) error = %v, want ErrNoCandidate", err)
}
}
func TestAcquireNeverOversubscribesSnapshotProxy(t *testing.T) {
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := []proxyDomain.Proxy{{ID: "p1", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable, MaxConcurrency: 8}}
envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
dispatcher := New(store)
var wg sync.WaitGroup
leases := make(chan *Lease, 1000)
for range 1000 {
wg.Add(1)
go func() {
defer wg.Done()
lease, err := dispatcher.Acquire(Request{Now: time.Now(), Scheme: proxyDomain.SchemeHTTP})
if err == nil {
leases <- lease
return
}
if !errors.Is(err, ErrNoCandidate) {
t.Errorf("Acquire(): %v", err)
}
}()
}
wg.Wait()
close(leases)
count := 0
for lease := range leases {
count++
if err := lease.Cancel(); err != nil {
t.Errorf("Cancel(): %v", err)
}
}
if count != 8 {
t.Fatalf("reserved = %d, want 8", count)
}
}
func TestAcquireUsesSequentialCurrentUpstream(t *testing.T) {
store := routingStrategyStore(t)
dispatcher := New(store)
lease, err := dispatcher.Acquire(Request{
Now: time.Now(), RoutingName: "sequential", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategySequential, CurrentUpstream: "upstream-b"},
})
if err != nil {
t.Fatalf("Acquire(): %v", err)
}
if lease.Proxy.SourceUpstream != "upstream-b" {
t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
}
func TestAcquireRoundRobinCyclesRoutingUpstreams(t *testing.T) {
store := routingStrategyStore(t)
dispatcher := New(store)
request := Request{
Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
}
for index, want := range []string{"upstream-a", "upstream-b", "upstream-a"} {
lease, err := dispatcher.Acquire(request)
if err != nil {
t.Fatalf("Acquire(%d): %v", index, err)
}
if lease.Proxy.SourceUpstream != want {
t.Fatalf("Acquire(%d) upstream = %q, want %q", index, lease.Proxy.SourceUpstream, want)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(%d): %v", index, err)
}
}
}
func TestAcquireResetsRoundRobinStateWhenSnapshotVersionChanges(t *testing.T) {
store := routingStrategyStore(t)
dispatcher := New(store)
request := Request{
Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
}
first, err := dispatcher.Acquire(request)
if err != nil {
t.Fatalf("Acquire(first): %v", err)
}
if first.Proxy.SourceUpstream != "upstream-a" {
t.Fatalf("first upstream = %q, want upstream-a", first.Proxy.SourceUpstream)
}
if err := first.Cancel(); err != nil {
t.Fatalf("Cancel(first): %v", err)
}
next := snapshot.Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true,
Proxies: []proxyDomain.Proxy{
{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-b", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
},
}
next.Checksum = snapshot.Checksum(next.Proxies)
if err := store.Apply(next); err != nil {
t.Fatalf("Apply(next): %v", err)
}
second, err := dispatcher.Acquire(request)
if err != nil {
t.Fatalf("Acquire(second): %v", err)
}
if second.Proxy.SourceUpstream != "upstream-a" {
t.Fatalf("second upstream = %q, want reset upstream-a", second.Proxy.SourceUpstream)
}
if err := second.Cancel(); err != nil {
t.Fatalf("Cancel(second): %v", err)
}
}
func TestAcquireFallsBackWhenSelectedRoutingUpstreamIsFull(t *testing.T) {
store := routingStrategyStore(t)
var held *proxyDomain.Reservation
for _, entry := range store.Current().Entries {
if entry.Proxy.SourceUpstream != "upstream-a" {
continue
}
reservation, ok := entry.Runtime.Reserve()
if !ok {
t.Fatal("Reserve(upstream-a) = false, want true")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(upstream-a): %v", err)
}
held = reservation
break
}
if held == nil {
t.Fatal("upstream-a reservation was not created")
}
defer func() {
if err := held.Release(); err != nil {
t.Errorf("Release(upstream-a): %v", err)
}
}()
lease, err := New(store).Acquire(Request{
Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
})
if err != nil {
t.Fatalf("Acquire(): %v", err)
}
if lease.Proxy.SourceUpstream != "upstream-b" {
t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
}
func TestAcquireUsesWeightedRoutingStrategy(t *testing.T) {
store := routingStrategyStore(t)
dispatcher := New(store, &fixedRoutingRandom{values: []int{0, 2, 3}})
request := Request{
Now: time.Now(), RoutingName: "weighted", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategyWeighted, Weights: map[string]uint32{"upstream-a": 3, "upstream-b": 1}},
}
for index, want := range []string{"upstream-a", "upstream-a", "upstream-b"} {
lease, err := dispatcher.Acquire(request)
if err != nil {
t.Fatalf("Acquire(%d): %v", index, err)
}
if lease.Proxy.SourceUpstream != want {
t.Fatalf("Acquire(%d) upstream = %q, want %q", index, lease.Proxy.SourceUpstream, want)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(%d): %v", index, err)
}
}
}
func TestAcquireUsesLeastConnectionsAndSkipsFullUpstream(t *testing.T) {
store := routingStrategyStore(t)
view := store.Current()
for _, entry := range view.Entries {
if entry.Proxy.SourceUpstream != "upstream-a" {
continue
}
reservation, ok := entry.Runtime.Reserve()
if !ok {
t.Fatal("Reserve(upstream-a) = false, want true")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(upstream-a): %v", err)
}
defer func() {
if err := reservation.Release(); err != nil {
t.Errorf("Release(upstream-a): %v", err)
}
}()
break
}
lease, err := New(store).Acquire(Request{
Now: time.Now(), RoutingName: "least", Upstreams: []string{"upstream-a", "upstream-b"},
Strategy: routing.Strategy{Type: routing.StrategyLeastConnections},
})
if err != nil {
t.Fatalf("Acquire(): %v", err)
}
if lease.Proxy.SourceUpstream != "upstream-b" {
t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
}
func TestAcquireWaitsForReleasedLocalCapacity(t *testing.T) {
store := routingStrategyStore(t)
dispatcher := New(store)
held, err := dispatcher.Acquire(Request{Now: time.Now(), Upstreams: []string{"upstream-a"}})
if err != nil {
t.Fatalf("Acquire(held): %v", err)
}
if err := held.Commit(); err != nil {
t.Fatalf("Commit(held): %v", err)
}
released := make(chan struct{})
go func() {
time.Sleep(10 * time.Millisecond)
_ = held.Release()
close(released)
}()
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
lease, err := dispatcher.AcquireWait(ctx, Request{Upstreams: []string{"upstream-a"}}, 100*time.Millisecond)
if err != nil {
t.Fatalf("AcquireWait(): %v", err)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
<-released
}
func TestAcquireSurvivesConcurrentSnapshotApply(t *testing.T) {
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
store := snapshot.NewStore("cluster-a", "worker-a")
base := makeTestProxies(256)
envelope := snapshot.Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: 1,
Full: true,
Proxies: base,
}
envelope.Checksum = snapshot.Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(base): %v", err)
}
dispatcher := New(store)
var successes atomic.Int64
var applyErr atomic.Value
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
steps := []struct {
epoch uint64
version uint64
}{
{epoch: 1, version: 2},
{epoch: 1, version: 3},
{epoch: 1, version: 4},
{epoch: 2, version: 1},
{epoch: 2, version: 2},
{epoch: 2, version: 3},
{epoch: 3, version: 1},
{epoch: 3, version: 2},
}
for _, step := range steps {
next := snapshot.Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: step.epoch,
Version: step.version,
Full: true,
Proxies: makeTestProxies(256),
}
next.Checksum = snapshot.Checksum(next.Proxies)
if err := store.Apply(next); err != nil {
applyErr.Store(err)
return
}
}
}()
for range 8 {
wg.Add(1)
go func() {
defer wg.Done()
for range 512 {
lease, err := dispatcher.Acquire(Request{
Now: now,
Scheme: proxyDomain.SchemeHTTP,
Upstreams: []string{"upstream-a"},
RequiredTags: map[string]string{
"region": "cn-east",
},
})
if err != nil {
if !errors.Is(err, ErrNoCandidate) {
t.Errorf("Acquire(): %v", err)
return
}
continue
}
successes.Add(1)
if err := lease.Cancel(); err != nil {
t.Errorf("Cancel(): %v", err)
return
}
}
}()
}
wg.Wait()
if err, _ := applyErr.Load().(error); err != nil {
t.Fatalf("Apply(): %v", err)
}
if successes.Load() == 0 {
t.Fatal("Acquire() never succeeded during concurrent Apply")
}
}
func BenchmarkAcquire100kIndexed(b *testing.B) {
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := makeTestProxies(100_000)
envelope := snapshot.Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: 1,
Full: true,
Proxies: proxies,
}
envelope.Checksum = snapshot.Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
b.Fatalf("Apply(): %v", err)
}
dispatcher := New(store)
request := Request{
Now: time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC),
Scheme: proxyDomain.SchemeHTTP,
Upstreams: []string{"upstream-a"},
RequiredTags: map[string]string{
"region": "cn-east",
"tier": "gold",
},
SafetyMargin: 5 * time.Second,
}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
lease, err := dispatcher.Acquire(request)
if err != nil {
b.Fatalf("Acquire(): %v", err)
}
if err := lease.Cancel(); err != nil {
b.Fatalf("Cancel(): %v", err)
}
}
}
func BenchmarkAcquire100kRoutingRoundRobin(b *testing.B) {
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := makeTestProxies(100_000)
envelope := snapshot.Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies,
}
envelope.Checksum = snapshot.Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
b.Fatalf("Apply(): %v", err)
}
dispatcher := New(store)
request := Request{
Now: time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC), RoutingName: "benchmark-round-robin",
Upstreams: []string{"upstream-a", "upstream-b"}, Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
Scheme: proxyDomain.SchemeHTTP, RequiredTags: map[string]string{"region": "cn-east", "tier": "gold"},
SafetyMargin: 5 * time.Second,
}
b.ReportAllocs()
b.ResetTimer()
b.RunParallel(func(parallel *testing.PB) {
for parallel.Next() {
lease, err := dispatcher.Acquire(request)
if err != nil {
b.Fatalf("Acquire(): %v", err)
}
if err := lease.Cancel(); err != nil {
b.Fatalf("Cancel(): %v", err)
}
}
})
}
func makeTestProxies(count int) []proxyDomain.Proxy {
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
expiresAt := now.Add(10 * time.Minute)
proxies := make([]proxyDomain.Proxy, 0, count)
for index := range count {
proxies = append(proxies, proxyDomain.Proxy{
ID: fmt.Sprintf("proxy-%06d", index),
Scheme: schemeForIndex(index),
Host: fmt.Sprintf("10.0.%d.%d", index/256, index%256),
Port: uint16(20000 + index%1000),
SourceUpstream: upstreamForIndex(index),
State: proxyDomain.StateAvailable,
MaxConcurrency: 8,
ExpiresAt: &expiresAt,
Tags: map[string]string{
"region": regionForIndex(index),
"tier": tierForIndex(index),
},
})
}
return proxies
}
func routingStrategyStore(t *testing.T) *snapshot.Store {
t.Helper()
store := snapshot.NewStore("cluster-a", "worker-a")
envelope := snapshot.Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Proxies: []proxyDomain.Proxy{
{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-b", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
},
}
envelope.Checksum = snapshot.Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
return store
}
type fixedRoutingRandom struct {
values []int
index int
}
func (source *fixedRoutingRandom) Intn(limit int) int {
value := source.values[source.index]
source.index++
return value % limit
}
func schemeForIndex(index int) proxyDomain.Scheme {
switch index % 3 {
case 0:
return proxyDomain.SchemeHTTP
case 1:
return proxyDomain.SchemeHTTPS
default:
return proxyDomain.SchemeSOCKS5
}
}
func upstreamForIndex(index int) string {
if index%2 == 0 {
return "upstream-a"
}
return "upstream-b"
}
func regionForIndex(index int) string {
switch index % 4 {
case 0:
return "cn-east"
case 1:
return "us-west"
case 2:
return "eu-central"
default:
return "ap-south"
}
}
func tierForIndex(index int) string {
if index%5 == 0 {
return "gold"
}
return "silver"
}