proxy-pool/internal/gateway/snapshot/store_test.go
youfak 1846c98e09
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: consume gateway routing snapshots
2026-07-31 15:36:56 +08:00

618 lines
20 KiB
Go

package snapshot
import (
"errors"
"fmt"
"reflect"
"testing"
"time"
proxyDomain "proxy-pool/internal/domain/proxy"
"proxy-pool/internal/domain/routing"
"proxy-pool/internal/domain/workerruntime"
)
func TestStoreAppliesCompleteSnapshotsInOrder(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
first := Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: 1,
Full: true,
Proxies: []proxyDomain.Proxy{{
ID: "p1",
Scheme: proxyDomain.SchemeHTTP,
Host: "127.0.0.1",
Port: 18080,
MaxConcurrency: 2,
State: proxyDomain.StateAvailable,
}},
}
first.Checksum = Checksum(first.Proxies)
if err := store.Apply(first); err != nil {
t.Fatalf("Apply(first): %v", err)
}
view := store.Current()
if view == nil || view.Version != 1 || len(view.Entries) != 1 {
t.Fatalf("Current() = %+v", view)
}
if view.Entries[0].Runtime == nil {
t.Fatal("snapshot entry has no local runtime capacity")
}
second := first
second.Version = 2
second.Proxies = append([]proxyDomain.Proxy(nil), first.Proxies...)
second.Proxies[0].Host = "localhost"
second.Checksum = Checksum(second.Proxies)
if err := store.Apply(second); err != nil {
t.Fatalf("Apply(second): %v", err)
}
if got := store.Current().Entries[0].Proxy.Host; got != "localhost" {
t.Fatalf("host = %q, want localhost", got)
}
}
func TestStoreRejectsWrongWorkerVersionGapAndChecksum(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
base := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true}
base.Checksum = Checksum(base.Proxies)
if err := store.Apply(base); err != nil {
t.Fatalf("Apply(base): %v", err)
}
wrongWorker := base
wrongWorker.Version = 2
wrongWorker.WorkerID = "worker-b"
if err := store.Apply(wrongWorker); !errors.Is(err, ErrWrongTarget) {
t.Fatalf("wrong worker error = %v, want ErrWrongTarget", err)
}
gap := base
gap.Version = 3
if err := store.Apply(gap); !errors.Is(err, ErrResyncRequired) {
t.Fatalf("version gap error = %v, want ErrResyncRequired", err)
}
badChecksum := base
badChecksum.Version = 2
badChecksum.Checksum = "bad"
if err := store.Apply(badChecksum); !errors.Is(err, ErrChecksumMismatch) {
t.Fatalf("checksum error = %v, want ErrChecksumMismatch", err)
}
}
func TestStorePublishesRoutingWithTheSameSnapshotVersion(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
routes := []routing.Rule{{
Name: "gateway-api",
Match: routing.Match{HostRegex: `^api\.example\.test$`},
Upstreams: []string{"provider-a"},
Action: routing.ActionProxy,
Strategy: routing.Strategy{Type: routing.StrategyRandom},
OnUnavailable: routing.OnUnavailableReject,
}}
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Routing: routes,
}
envelope.Checksum = ChecksumWithRouting(envelope.Proxies, envelope.Routing)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
view := store.Current()
matched, ok := view.MatchRouting(routing.Request{Host: "api.example.test", Method: "GET", Path: "/"})
if !ok || matched.Name != "gateway-api" || matched.Strategy.Type != routing.StrategyRandom ||
!reflect.DeepEqual(matched.Upstreams, []string{"provider-a"}) {
t.Fatalf("MatchRouting() = %+v, %v", matched, ok)
}
changedRouting := envelope
changedRouting.Version = 2
changedRouting.Routing = []routing.Rule{{
Name: "other", Match: routing.Match{HostRegex: `^other\.example\.test$`},
Upstreams: []string{"provider-b"}, Action: routing.ActionProxy,
Strategy: routing.Strategy{Type: routing.StrategyRoundRobin}, OnUnavailable: routing.OnUnavailableReject,
}}
changedRouting.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(changedRouting); !errors.Is(err, ErrChecksumMismatch) {
t.Fatalf("Apply(changed routing with old checksum) error = %v, want ErrChecksumMismatch", err)
}
if current := store.Current(); current.Version != 1 {
t.Fatalf("Current().Version = %d, want unchanged version 1", current.Version)
}
}
func TestViewSelectFiltersBySchemeUpstreamTagAndExclude(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 := NewStore("cluster-a", "worker-a")
proxies := []proxyDomain.Proxy{
{ID: "http-a-cn", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east", "tier": "gold"}},
{ID: "http-b-cn", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "b", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}},
{ID: "socks-a-cn", Scheme: proxyDomain.SchemeSOCKS5, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}},
{ID: "http-a-us", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "us-west"}},
{ID: "expiring", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresSoon, Tags: map[string]string{"region": "cn-east"}},
}
envelope := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
view := store.Current()
if view == nil {
t.Fatal("Current() returned nil view")
}
selection := view.Select(Query{
Now: now,
Scheme: proxyDomain.SchemeHTTP,
Upstreams: []string{"a", "c"},
RequiredTags: map[string]string{"region": "cn-east"},
Exclude: map[string]struct{}{"expiring": {}},
SafetyMargin: 10 * time.Second,
})
if got, want := collectSelectionIDs(selection), []string{"http-a-cn"}; !reflect.DeepEqual(got, want) {
t.Fatalf("selection ids = %v, want %v", got, want)
}
}
func TestViewSelectRejectsProxyAfterProviderUsableDeadline(t *testing.T) {
now := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)
expiresAt := now.Add(30 * time.Second)
usableUntil := now.Add(27 * time.Second)
proxy := proxyDomain.Proxy{
ID: "short-lived", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable,
MaxConcurrency: 1, ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
}
store := NewStore("cluster-a", "worker-a")
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Proxies: []proxyDomain.Proxy{proxy},
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
before := store.Current().Select(Query{Now: now.Add(26 * time.Second)})
if _, ok := before.EntryAt(0); !ok {
t.Fatal("EntryAt(before usable deadline) = false, want true")
}
after := store.Current().Select(Query{Now: now.Add(27 * time.Second)})
if _, ok := after.EntryAt(0); ok {
t.Fatal("EntryAt(at usable deadline) = true, want false")
}
}
func TestStoreApplyReusesCapacityAcrossVersionsAndEpochs(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
firstProxies := []proxyDomain.Proxy{{
ID: "stable",
Scheme: proxyDomain.SchemeHTTP,
Host: "127.0.0.1",
Port: 18080,
MaxConcurrency: 2,
State: proxyDomain.StateAvailable,
}}
first := Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: 1,
Full: true,
Proxies: firstProxies,
}
first.Checksum = Checksum(first.Proxies)
if err := store.Apply(first); err != nil {
t.Fatalf("Apply(first): %v", err)
}
initialView := store.Current()
initialRuntime := initialView.Entries[0].Runtime
reservation, ok := initialRuntime.Reserve()
if !ok {
t.Fatal("Reserve() = false, want true")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(): %v", err)
}
second := first
second.Version = 2
second.Proxies = []proxyDomain.Proxy{{
ID: "stable",
Scheme: proxyDomain.SchemeHTTP,
Host: "localhost",
Port: 18080,
MaxConcurrency: 5,
State: proxyDomain.StateAvailable,
}}
second.Checksum = Checksum(second.Proxies)
if err := store.Apply(second); err != nil {
t.Fatalf("Apply(second): %v", err)
}
third := second
third.Epoch = 2
third.Version = 1
third.Checksum = Checksum(third.Proxies)
if err := store.Apply(third); err != nil {
t.Fatalf("Apply(third): %v", err)
}
current := store.Current()
if got := current.Entries[0].Runtime; got != initialRuntime {
t.Fatal("runtime pointer was replaced, want capacity reuse")
}
if got := current.Entries[0].Runtime.Max(); got != 5 {
t.Fatalf("runtime max = %d, want 5", got)
}
if got := current.Entries[0].Runtime.Active(); got != 1 {
t.Fatalf("runtime active = %d, want 1", got)
}
if err := reservation.Release(); err != nil {
t.Fatalf("Release(): %v", err)
}
}
func TestStoreReusesRuntimeWhenProxyDisappearsAndReappears(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
proxy := proxyDomain.Proxy{
ID: "p1", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable, MaxConcurrency: 1,
}
first := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Proxies: []proxyDomain.Proxy{proxy},
}
first.Checksum = Checksum(first.Proxies)
if err := store.Apply(first); err != nil {
t.Fatalf("Apply(first): %v", err)
}
runtime := store.Current().Entries[0].Runtime
reservation, ok := runtime.Reserve()
if !ok {
t.Fatal("Reserve() = false, want true")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(): %v", err)
}
removed := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true}
removed.Checksum = Checksum(removed.Proxies)
if err := store.Apply(removed); err != nil {
t.Fatalf("Apply(removed): %v", err)
}
reappeared := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 3, Full: true,
Proxies: []proxyDomain.Proxy{proxy},
}
reappeared.Checksum = Checksum(reappeared.Proxies)
if err := store.Apply(reappeared); err != nil {
t.Fatalf("Apply(reappeared): %v", err)
}
reused := store.Current().Entries[0].Runtime
if reused != runtime || reused.Active() != 1 {
t.Fatalf("reappeared runtime = %p active=%d, want %p active=1", reused, reused.Active(), runtime)
}
if _, ok := reused.Reserve(); ok {
t.Fatal("Reserve() succeeded despite inherited active capacity")
}
if err := reservation.Release(); err != nil {
t.Fatalf("Release(): %v", err)
}
}
func TestStoreRuntimeReportKeepsRemovedActiveProxyUntilRelease(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
proxy := proxyDomain.Proxy{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "provider-a",
State: proxyDomain.StateAvailable, MaxConcurrency: 1,
}
initial := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Proxies: []proxyDomain.Proxy{proxy},
}
initial.Checksum = Checksum(initial.Proxies)
if err := store.Apply(initial); err != nil {
t.Fatalf("Apply(initial): %v", err)
}
runtime := store.Current().Entries[0].Runtime
reservation, ok := runtime.Reserve()
if !ok {
t.Fatal("Reserve() = false")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(): %v", err)
}
removed := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true}
removed.Checksum = Checksum(nil)
if err := store.Apply(removed); err != nil {
t.Fatalf("Apply(remove): %v", err)
}
if got := store.activeRuntimeCount(); got != 1 {
t.Fatalf("activeRuntimeCount(removed) = %d, want 1", got)
}
report, err := store.RuntimeReport("session-a", 7, now)
if err != nil {
t.Fatalf("RuntimeReport(): %v", err)
}
if report.WorkerID != "worker-a" || report.SessionID != "session-a" || report.Sequence != 7 ||
report.SnapshotVersion != 2 || report.OwnershipEpoch != 1 || !report.ObservedAt.Equal(now) ||
len(report.Counters) != 1 || report.Counters[0] != (workerruntime.Counter{
ProxyID: "proxy-a", Active: 1, Draining: true,
}) {
t.Fatalf("RuntimeReport() = %+v", report)
}
if err := reservation.Release(); err != nil {
t.Fatalf("Release(): %v", err)
}
if got := store.activeRuntimeCount(); got != 0 {
t.Fatalf("activeRuntimeCount(released) = %d, want 0", got)
}
report, err = store.RuntimeReport("session-a", 8, now.Add(time.Second))
if err != nil || len(report.Counters) != 0 {
t.Fatalf("RuntimeReport(after release) = %+v, %v", report, err)
}
}
func TestStoreRuntimeReportMarksCurrentDrainingProxy(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
proxy := proxyDomain.Proxy{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP,
State: proxyDomain.StateDraining, MaxConcurrency: 1,
}
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
Proxies: []proxyDomain.Proxy{proxy},
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
reservation, ok := store.Current().Entries[0].Runtime.Reserve()
if !ok {
t.Fatal("Reserve() = false")
}
report, err := store.RuntimeReport("session-a", 1, now)
if err != nil || len(report.Counters) != 1 || !report.Counters[0].Draining {
t.Fatalf("RuntimeReport() = %+v, %v", report, err)
}
if err := reservation.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
}
func TestStoreReclaimsRetiredZeroRuntimeBeforeHittingLimit(t *testing.T) {
store, err := NewStoreWithRuntimeLimit("cluster-a", "worker-a", 1)
if err != nil {
t.Fatalf("NewStoreWithRuntimeLimit(): %v", err)
}
first := proxyDomain.Proxy{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP,
State: proxyDomain.StateAvailable, MaxConcurrency: 1,
}
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1,
Full: true, Proxies: []proxyDomain.Proxy{first},
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(first): %v", err)
}
removed := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true,
}
removed.Checksum = Checksum(nil)
if err := store.Apply(removed); err != nil {
t.Fatalf("Apply(removed): %v", err)
}
second := first
second.ID = "proxy-b"
overLimit := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 3,
Full: true, Proxies: []proxyDomain.Proxy{second},
}
overLimit.Checksum = Checksum(overLimit.Proxies)
if err := store.Apply(overLimit); err != nil {
t.Fatalf("Apply(replacement after retired runtime) error = %v", err)
}
if current := store.Current(); current.Version != 3 || len(current.Entries) != 1 || current.Entries[0].Proxy.ID != "proxy-b" {
t.Fatalf("Current() after replacement = %+v", current)
}
if len(store.runtimes) != 1 || store.runtimes["proxy-b"] == nil || store.runtimes["proxy-a"] != nil {
t.Fatalf("runtime registry = %+v, want only proxy-b", store.runtimes)
}
}
func TestStoreRetainsActiveRuntimeAndKeepsApplyTransactionalAtLimit(t *testing.T) {
store, err := NewStoreWithRuntimeLimit("cluster-a", "worker-a", 1)
if err != nil {
t.Fatalf("NewStoreWithRuntimeLimit(): %v", err)
}
first := proxyDomain.Proxy{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP,
State: proxyDomain.StateAvailable, MaxConcurrency: 2,
}
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1,
Full: true, Proxies: []proxyDomain.Proxy{first},
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(first): %v", err)
}
reservation, ok := store.Current().Entries[0].Runtime.Reserve()
if !ok {
t.Fatal("Reserve() = false")
}
if err := reservation.Commit(); err != nil {
t.Fatalf("Commit(): %v", err)
}
second := first
second.ID = "proxy-b"
overLimit := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2,
Full: true, Proxies: []proxyDomain.Proxy{second},
}
overLimit.Checksum = Checksum(overLimit.Proxies)
if err := store.Apply(overLimit); !errors.Is(err, ErrRuntimeLimitExceeded) {
t.Fatalf("Apply(over limit) error = %v, want ErrRuntimeLimitExceeded", err)
}
current := store.Current()
if current.Version != 1 || len(current.Entries) != 1 || current.Entries[0].Proxy.ID != "proxy-a" {
t.Fatalf("Current() after rejected apply = %+v", current)
}
extra, ok := current.Entries[0].Runtime.Reserve()
if !ok {
t.Fatal("Reserve() after rejected Apply = false, want current Proxy to remain enabled")
}
if err := extra.Cancel(); err != nil {
t.Fatalf("Cancel(extra): %v", err)
}
if err := reservation.Release(); err != nil {
t.Fatalf("Release(): %v", err)
}
}
func TestStoreRuntimeActiveIndexDropsZeroCounters(t *testing.T) {
store := NewStore("cluster-a", "worker-a")
proxy := proxyDomain.Proxy{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP,
State: proxyDomain.StateAvailable, MaxConcurrency: 1,
}
envelope := Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1,
Full: true, Proxies: []proxyDomain.Proxy{proxy},
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
reservation, ok := store.Current().Entries[0].Runtime.Reserve()
if !ok {
t.Fatal("Reserve() = false")
}
if got := store.activeRuntimeCount(); got != 0 {
t.Fatalf("activeRuntimeCount(current) = %d, want 0", got)
}
if err := reservation.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
if got := store.activeRuntimeCount(); got != 0 {
t.Fatalf("activeRuntimeCount() = %d, want 0", got)
}
}
func (s *Store) activeRuntimeCount() int {
count := 0
s.active.rangeEntries(func(_ string, _ *proxyDomain.Capacity) {
count++
})
return count
}
func collectSelectionIDs(selection Selection) []string {
ids := make([]string, 0, selection.Len())
for index := 0; index < selection.Len(); index++ {
entry, ok := selection.EntryAt(index)
if !ok {
continue
}
ids = append(ids, entry.Proxy.ID)
}
return ids
}
func BenchmarkStoreApply100k(b *testing.B) {
store := NewStore("cluster-a", "worker-a")
proxies := makeBenchmarkProxies(100_000)
b.ReportAllocs()
for index := 0; index < b.N; index++ {
envelope := Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: uint64(index + 1),
Full: true,
Proxies: proxies,
}
envelope.Checksum = Checksum(envelope.Proxies)
if err := store.Apply(envelope); err != nil {
b.Fatalf("Apply(): %v", err)
}
}
}
func makeBenchmarkProxies(count int) []proxyDomain.Proxy {
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
expiresAt := now.Add(30 * time.Minute)
proxies := make([]proxyDomain.Proxy, 0, count)
for index := range count {
proxies = append(proxies, proxyDomain.Proxy{
ID: fmt.Sprintf("proxy-%06d", index),
Scheme: schemeForBenchmarkIndex(index),
Host: "127.0.0.1",
Port: uint16(20000 + index%1000),
SourceUpstream: upstreamForBenchmarkIndex(index),
State: proxyDomain.StateAvailable,
MaxConcurrency: 8,
ExpiresAt: &expiresAt,
Tags: map[string]string{
"region": regionForBenchmarkIndex(index),
"tier": tierForBenchmarkIndex(index),
},
})
}
return proxies
}
func schemeForBenchmarkIndex(index int) proxyDomain.Scheme {
switch index % 3 {
case 0:
return proxyDomain.SchemeHTTP
case 1:
return proxyDomain.SchemeHTTPS
default:
return proxyDomain.SchemeSOCKS5
}
}
func upstreamForBenchmarkIndex(index int) string {
if index%2 == 0 {
return "upstream-a"
}
return "upstream-b"
}
func regionForBenchmarkIndex(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 tierForBenchmarkIndex(index int) string {
if index%5 == 0 {
return "gold"
}
return "silver"
}