91 lines
2.9 KiB
Go
91 lines
2.9 KiB
Go
package dispatch
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
|
|
"github.com/proxy-pool/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)
|
|
}
|
|
}
|