350 lines
13 KiB
Go
350 lines
13 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/controlplane/snapshotwire"
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
platformCredentials "proxy-pool/internal/platform/credentials"
|
|
)
|
|
|
|
func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
source, err := NewInitialSnapshotSource(epochReaderStub{epoch: 9}, time.Minute, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewInitialSnapshotSource(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{
|
|
WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if full.GetVersion() != 5 || full.GetOwnershipEpoch() != 9 || !full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("snapshot = %+v", full)
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
func TestRefreshingSnapshotSourceIssuesSequentialFullSnapshots(t *testing.T) {
|
|
base := &refreshingSnapshotSourceStub{}
|
|
source, err := NewRefreshingSnapshotSource(base, 10*time.Millisecond)
|
|
if err != nil {
|
|
t.Fatalf("NewRefreshingSnapshotSource() = %v", err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
updates, err := source.Watch(ctx, SnapshotWatchRequest{
|
|
WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4, LastChecksum: []byte{4},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Watch() = %v", err)
|
|
}
|
|
first := receiveSnapshot(t, updates)
|
|
second := receiveSnapshot(t, updates)
|
|
if first.GetVersion() != 5 || second.GetVersion() != 6 || second.GetVersion() <= first.GetVersion() {
|
|
t.Fatalf("refreshed versions = (%d, %d)", first.GetVersion(), second.GetVersion())
|
|
}
|
|
if calls := base.Calls(); len(calls) < 2 || calls[0] != 4 || calls[1] != 5 {
|
|
t.Fatalf("base requested versions = %v", calls)
|
|
}
|
|
if checksums := base.Checksums(); len(checksums) < 2 || string(checksums[0]) != string([]byte{4}) ||
|
|
string(checksums[1]) != string([]byte{5}) {
|
|
t.Fatalf("base requested checksums = %v", checksums)
|
|
}
|
|
}
|
|
|
|
func TestRefreshingSnapshotSourceRefreshesImmediatelyAfterManagementNotification(t *testing.T) {
|
|
base := &refreshingSnapshotSourceStub{}
|
|
broker := NewSnapshotRefreshBroker()
|
|
source, err := NewRefreshingSnapshotSourceWithRefreshEvents(base, time.Hour, broker)
|
|
if err != nil {
|
|
t.Fatalf("NewRefreshingSnapshotSourceWithRefreshEvents() = %v", err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
updates, err := source.Watch(ctx, SnapshotWatchRequest{
|
|
WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4, LastChecksum: []byte{4},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Watch() = %v", err)
|
|
}
|
|
if first := receiveSnapshot(t, updates); first.GetVersion() != 5 {
|
|
t.Fatalf("initial version = %d, want 5", first.GetVersion())
|
|
}
|
|
broker.NotifySnapshotRefresh()
|
|
if refreshed := receiveSnapshot(t, updates); refreshed.GetVersion() != 6 {
|
|
t.Fatalf("refreshed version = %d, want 6", refreshed.GetVersion())
|
|
}
|
|
if calls := base.Calls(); len(calls) != 2 || calls[0] != 4 || calls[1] != 5 {
|
|
t.Fatalf("base requested versions = %v", calls)
|
|
}
|
|
}
|
|
|
|
func TestSnapshotRefreshBrokerBroadcastsCoalescedNotifications(t *testing.T) {
|
|
broker := NewSnapshotRefreshBroker()
|
|
first, unsubscribeFirst := broker.SubscribeSnapshotRefresh()
|
|
second, unsubscribeSecond := broker.SubscribeSnapshotRefresh()
|
|
defer unsubscribeSecond()
|
|
|
|
broker.NotifySnapshotRefresh()
|
|
broker.NotifySnapshotRefresh()
|
|
select {
|
|
case <-first:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("first subscriber did not receive notification")
|
|
}
|
|
select {
|
|
case <-second:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("second subscriber did not receive notification")
|
|
}
|
|
select {
|
|
case <-first:
|
|
t.Fatal("first subscriber received an uncoalesced notification")
|
|
default:
|
|
}
|
|
|
|
unsubscribeFirst()
|
|
broker.NotifySnapshotRefresh()
|
|
select {
|
|
case <-first:
|
|
t.Fatal("unsubscribed receiver received a notification")
|
|
default:
|
|
}
|
|
select {
|
|
case <-second:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("remaining subscriber did not receive notification")
|
|
}
|
|
}
|
|
|
|
func TestRefreshingSnapshotSourceRefreshesBeforeShortSnapshotExpiry(t *testing.T) {
|
|
source, err := NewRefreshingSnapshotSource(&refreshingSnapshotSourceStub{}, 2*time.Second)
|
|
if err != nil {
|
|
t.Fatalf("NewRefreshingSnapshotSource() = %v", err)
|
|
}
|
|
now := time.Now().UTC()
|
|
if delay := source.delayFor(&controlplanev1.WorkerSnapshot{
|
|
GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(100 * time.Millisecond)),
|
|
}); delay != 50*time.Millisecond {
|
|
t.Fatalf("delayFor(short-lived snapshot) = %s, want 50ms", delay)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceBuildsBoundedProxySnapshot(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(10 * time.Minute)
|
|
usableUntil := now.Add(5 * time.Minute)
|
|
source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
|
|
MaxConcurrency: 7, State: proxyDomain.StateAvailable, Tags: map[string]string{"region": "cn"},
|
|
},
|
|
OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute*2, 10, 4096, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSource(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if full.GetVersion() != 1 || full.GetOwnershipEpoch() != 9 || len(full.GetProxies()) != 1 {
|
|
t.Fatalf("snapshot = %+v", full)
|
|
}
|
|
proxy := full.GetProxies()[0]
|
|
if proxy.GetOwnershipEpoch() != 4 || !proxy.GetUsableUntil().AsTime().Equal(now.Add(time.Minute)) ||
|
|
!full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("wire proxy = %+v, valid until = %s", proxy, full.GetValidUntil().AsTime())
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceIncludesRoutingInFullSnapshotChecksum(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(10 * time.Minute)
|
|
usableUntil := now.Add(5 * time.Minute)
|
|
routing := &controlplanev1.RoutingRule{
|
|
Name: "gateway", Enabled: true, HostRegex: ".*", Upstreams: []string{"upstream-a"},
|
|
Strategy: &controlplanev1.RoutingStrategy{Type: controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM},
|
|
OnUnavailable: controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT,
|
|
}
|
|
source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
|
|
MaxConcurrency: 7, State: proxyDomain.StateAvailable,
|
|
},
|
|
OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute*2, 10, 4096, func() time.Time { return now }, routingSourceStub{rules: []*controlplanev1.RoutingRule{routing}})
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSource(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if len(full.GetRouting()) != 1 || full.GetRouting()[0].GetName() != "gateway" || full.GetRouting()[0] == routing {
|
|
t.Fatalf("snapshot routing = %+v", full.GetRouting())
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceRejectsCredentialReferenceUntilMaterialIsAvailable(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(time.Minute)
|
|
usableUntil := now.Add(time.Minute)
|
|
source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
|
|
MaxConcurrency: 1, State: proxyDomain.StateAvailable, CredentialVersion: "v1",
|
|
},
|
|
OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute, 10, 4096, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSource(): %v", err)
|
|
}
|
|
_, err = source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if !errors.Is(err, ErrSnapshotCredentialsUnavailable) {
|
|
t.Fatalf("Watch() error = %v, want ErrSnapshotCredentialsUnavailable", err)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceIncludesReferencedCredentialMaterial(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(time.Minute)
|
|
usableUntil := now.Add(30 * time.Second)
|
|
store, err := platformCredentials.NewMemoryStore(2)
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryStore(): %v", err)
|
|
}
|
|
reference, err := store.Put(context.Background(), "provider-a", platformCredentials.Value{Username: "upstream", Password: "secret"})
|
|
if err != nil {
|
|
t.Fatalf("Put(): %v", err)
|
|
}
|
|
source, err := NewOwnedSnapshotSourceWithCredentials(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, SourceUpstream: "upstream-a",
|
|
Username: "upstream", SecretRef: reference.SecretRef, CredentialVersion: reference.CredentialVersion,
|
|
ExpiresAt: &expiresAt, UsableUntil: &usableUntil, MaxConcurrency: 1, State: proxyDomain.StateAvailable,
|
|
}, OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute, 10, 4096, func() time.Time { return now }, store)
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSourceWithCredentials(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if len(full.GetCredentials()) != 1 || full.GetCredentials()[0].GetSecretRef() != reference.SecretRef ||
|
|
full.GetCredentials()[0].GetUsername() != "upstream" || full.GetCredentials()[0].GetPassword() != "secret" ||
|
|
full.GetProxies()[0].GetCredentialVersion() != reference.CredentialVersion {
|
|
t.Fatal("snapshot credential material or proxy credential reference is invalid")
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
type epochReaderStub struct {
|
|
epoch uint64
|
|
err error
|
|
}
|
|
|
|
type refreshingSnapshotSourceStub struct {
|
|
mu sync.Mutex
|
|
calls []uint64
|
|
checksums [][]byte
|
|
}
|
|
|
|
func (source *refreshingSnapshotSourceStub) Watch(
|
|
_ context.Context,
|
|
request SnapshotWatchRequest,
|
|
) (<-chan *controlplanev1.WorkerSnapshot, error) {
|
|
source.mu.Lock()
|
|
source.calls = append(source.calls, request.LastAppliedVersion)
|
|
source.checksums = append(source.checksums, append([]byte(nil), request.LastChecksum...))
|
|
source.mu.Unlock()
|
|
updates := make(chan *controlplanev1.WorkerSnapshot, 1)
|
|
updates <- &controlplanev1.WorkerSnapshot{
|
|
Version: request.LastAppliedVersion + 1, Checksum: []byte{byte(request.LastAppliedVersion + 1)},
|
|
}
|
|
return updates, nil
|
|
}
|
|
|
|
func (source *refreshingSnapshotSourceStub) Calls() []uint64 {
|
|
source.mu.Lock()
|
|
defer source.mu.Unlock()
|
|
return append([]uint64(nil), source.calls...)
|
|
}
|
|
|
|
func (source *refreshingSnapshotSourceStub) Checksums() [][]byte {
|
|
source.mu.Lock()
|
|
defer source.mu.Unlock()
|
|
checksums := make([][]byte, len(source.checksums))
|
|
for index := range source.checksums {
|
|
checksums[index] = append([]byte(nil), source.checksums[index]...)
|
|
}
|
|
return checksums
|
|
}
|
|
|
|
func receiveSnapshot(t *testing.T, updates <-chan *controlplanev1.WorkerSnapshot) *controlplanev1.WorkerSnapshot {
|
|
t.Helper()
|
|
select {
|
|
case snapshot, ok := <-updates:
|
|
if !ok || snapshot == nil {
|
|
t.Fatal("snapshot updates closed before next full snapshot")
|
|
}
|
|
return snapshot
|
|
case <-time.After(time.Second):
|
|
t.Fatal("timed out waiting for snapshot")
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (reader epochReaderStub) CurrentOwnershipEpoch(context.Context) (uint64, error) {
|
|
return reader.epoch, reader.err
|
|
}
|
|
|
|
type snapshotReaderStub struct {
|
|
proxies []ownershipDomain.SnapshotProxy
|
|
err error
|
|
}
|
|
|
|
type routingSourceStub struct {
|
|
rules []*controlplanev1.RoutingRule
|
|
err error
|
|
}
|
|
|
|
func (source routingSourceStub) Read(context.Context) ([]*controlplanev1.RoutingRule, error) {
|
|
return source.rules, source.err
|
|
}
|
|
|
|
func (reader snapshotReaderStub) ReadWorkerSnapshot(context.Context, string, int) ([]ownershipDomain.SnapshotProxy, error) {
|
|
return reader.proxies, reader.err
|
|
}
|