288 lines
7.2 KiB
Go
288 lines
7.2 KiB
Go
package dispatch
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
"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 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 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 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 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"
|
|
}
|