311 lines
10 KiB
Go
311 lines
10 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/domain/activitypool"
|
|
extractionDomain "proxy-pool/internal/domain/extraction"
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
func TestOwnershipManagerPreventsDualAssignment(t *testing.T) {
|
|
manager := newTestOwnershipManager(t, "proxy-1")
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
var succeeded atomic.Int64
|
|
var wg sync.WaitGroup
|
|
|
|
for index := range 100 {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
_, err := manager.Assign(context.Background(), now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute)
|
|
if err == nil {
|
|
succeeded.Add(1)
|
|
return
|
|
}
|
|
if !errors.Is(err, ErrAlreadyOwned) {
|
|
t.Errorf("Assign(): %v", err)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
if got := succeeded.Load(); got != 1 {
|
|
t.Fatalf("successful assignments = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerRenewsOnlyCurrentAssignment(t *testing.T) {
|
|
manager := newTestOwnershipManager(t, "proxy-1")
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
ctx := context.Background()
|
|
assigned, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Assign(): %v", err)
|
|
}
|
|
renewed, err := manager.Renew(ctx, now.Add(30*time.Second), "proxy-1", "worker-1", assigned.Epoch, time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Renew(): %v", err)
|
|
}
|
|
if renewed.Version != assigned.Version+1 || !renewed.ExpiresAt.Equal(now.Add(90*time.Second)) {
|
|
t.Fatalf("renewed assignment = %+v", renewed)
|
|
}
|
|
if _, err := manager.Renew(ctx, now, "proxy-1", "worker-2", assigned.Epoch, time.Minute); !errors.Is(err, ErrStaleAssignment) {
|
|
t.Fatalf("Renew(stale) error = %v, want ErrStaleAssignment", err)
|
|
}
|
|
if expired, err := manager.Expire(ctx, now.Add(time.Minute), 32); err != nil || len(expired) != 0 {
|
|
t.Fatalf("renewed assignment expired at old deadline: %+v", expired)
|
|
}
|
|
}
|
|
|
|
func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) {
|
|
for iteration := range 100 {
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
store := newTestActivityPool(t, now, "proxy-1")
|
|
manager, err := NewOwnershipManager(store)
|
|
if err != nil {
|
|
t.Fatalf("iteration %d NewOwnershipManager(): %v", iteration, err)
|
|
}
|
|
|
|
start := make(chan struct{})
|
|
assigned := make(chan bool, 1)
|
|
extracted := make(chan bool, 1)
|
|
go func() {
|
|
<-start
|
|
_, assignErr := manager.Assign(context.Background(), now, "proxy-1", "worker-1", time.Minute)
|
|
if assignErr != nil && !errors.Is(assignErr, ErrOwnershipUnavailable) {
|
|
t.Errorf("iteration %d Assign(): %v", iteration, assignErr)
|
|
}
|
|
assigned <- assignErr == nil
|
|
}()
|
|
go func() {
|
|
<-start
|
|
result, extractErr := store.Extract(context.Background(), extractionDomain.Command{
|
|
ClientID: "client-1", Requested: 1, Fulfillment: extractionDomain.Partial, Now: now,
|
|
})
|
|
if extractErr != nil {
|
|
t.Errorf("iteration %d Extract(): %v", iteration, extractErr)
|
|
}
|
|
extracted <- result.Returned == 1
|
|
}()
|
|
close(start)
|
|
|
|
wins := 0
|
|
if <-assigned {
|
|
wins++
|
|
}
|
|
if <-extracted {
|
|
wins++
|
|
}
|
|
if wins != 1 {
|
|
t.Fatalf("iteration %d successful ownership/extraction operations = %d, want 1", iteration, wins)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) {
|
|
manager := newTestOwnershipManager(t, "proxy-1")
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
ctx := context.Background()
|
|
assignment, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Assign(): %v", err)
|
|
}
|
|
draining, err := manager.BeginDrain(ctx, "proxy-1", "worker-1", assignment.Epoch)
|
|
if err != nil {
|
|
t.Fatalf("BeginDrain(): %v", err)
|
|
}
|
|
if !draining.Draining || draining.Version != assignment.Version+1 {
|
|
t.Fatalf("draining assignment = %+v", draining)
|
|
}
|
|
|
|
if err := manager.AcknowledgeDrain(ctx, "proxy-1", "worker-1", assignment.Epoch, 1, 0); !errors.Is(err, ErrDrainNotReady) {
|
|
t.Fatalf("AcknowledgeDrain(active) error = %v, want ErrDrainNotReady", err)
|
|
}
|
|
if err := manager.AcknowledgeDrain(ctx, "proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil {
|
|
t.Fatalf("AcknowledgeDrain(zero): %v", err)
|
|
}
|
|
if _, ok, err := manager.Get(ctx, "proxy-1"); err != nil || ok {
|
|
t.Fatal("assignment still exists after drain acknowledgement")
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) {
|
|
manager := newTestOwnershipManager(t, "proxy-1")
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
ctx := context.Background()
|
|
first, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Assign(first): %v", err)
|
|
}
|
|
if expired, err := manager.Expire(ctx, now.Add(59*time.Second), 32); err != nil || len(expired) != 0 {
|
|
t.Fatalf("expired early: %+v", expired)
|
|
}
|
|
if expired, err := manager.Expire(ctx, now.Add(time.Minute), 32); err != nil || len(expired) != 1 || expired[0].ProxyID != "proxy-1" {
|
|
t.Fatalf("Expire() = %+v, want proxy-1", expired)
|
|
}
|
|
|
|
second, err := manager.Assign(ctx, now.Add(time.Minute), "proxy-1", "worker-2", time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Assign(second): %v", err)
|
|
}
|
|
if second.Epoch <= first.Epoch {
|
|
t.Fatalf("second epoch = %d, want greater than %d", second.Epoch, first.Epoch)
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerPropagatesCanceledContextWithoutCallingRepository(t *testing.T) {
|
|
repository := &ownershipRepositoryStub{}
|
|
manager, err := NewOwnershipManager(repository)
|
|
if err != nil {
|
|
t.Fatalf("NewOwnershipManager(): %v", err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
_, err = manager.Assign(ctx, time.Now(), "proxy-1", "worker-1", time.Minute)
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Assign() error = %v, want context.Canceled", err)
|
|
}
|
|
if repository.assignCalled {
|
|
t.Fatal("repository Assign called with canceled context")
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerPassesContextAndRepositoryError(t *testing.T) {
|
|
storageErr := errors.New("redis unavailable")
|
|
ctx := context.WithValue(context.Background(), ownershipContextKey{}, "request-1")
|
|
repository := &ownershipRepositoryStub{
|
|
assign: func(got context.Context, _ time.Time, _, _ string, _ time.Duration) (Assignment, error) {
|
|
if got != ctx {
|
|
t.Fatal("Assign() did not pass the original context")
|
|
}
|
|
return Assignment{}, storageErr
|
|
},
|
|
}
|
|
manager, err := NewOwnershipManager(repository)
|
|
if err != nil {
|
|
t.Fatalf("NewOwnershipManager(): %v", err)
|
|
}
|
|
|
|
_, err = manager.Assign(ctx, time.Now(), "proxy-1", "worker-1", time.Minute)
|
|
if !errors.Is(err, storageErr) {
|
|
t.Fatalf("Assign() error = %v, want repository error", err)
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerPassesExpireLimit(t *testing.T) {
|
|
ctx := context.Background()
|
|
now := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)
|
|
want := []Assignment{{ProxyID: "proxy-1"}}
|
|
repository := &ownershipRepositoryStub{
|
|
expire: func(got context.Context, gotNow time.Time, gotLimit int) ([]Assignment, error) {
|
|
if got != ctx || !gotNow.Equal(now) || gotLimit != 32 {
|
|
t.Fatalf("Expire() arguments = (%v, %v, %d)", got, gotNow, gotLimit)
|
|
}
|
|
return want, nil
|
|
},
|
|
}
|
|
manager, err := NewOwnershipManager(repository)
|
|
if err != nil {
|
|
t.Fatalf("NewOwnershipManager(): %v", err)
|
|
}
|
|
|
|
got, err := manager.Expire(ctx, now, 32)
|
|
if err != nil || len(got) != 1 || got[0].ProxyID != want[0].ProxyID {
|
|
t.Fatalf("Expire() = %+v, %v", got, err)
|
|
}
|
|
}
|
|
|
|
func TestOwnershipManagerRejectsInvalidContextAndExpireLimit(t *testing.T) {
|
|
manager, err := NewOwnershipManager(&ownershipRepositoryStub{})
|
|
if err != nil {
|
|
t.Fatalf("NewOwnershipManager(): %v", err)
|
|
}
|
|
now := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)
|
|
|
|
if _, err := manager.Assign(nil, now, "proxy-1", "worker-1", time.Minute); !errors.Is(err, ErrInvalidOwnership) {
|
|
t.Fatalf("Assign(nil context) error = %v, want ErrInvalidOwnership", err)
|
|
}
|
|
if _, err := manager.Expire(context.Background(), now, 0); !errors.Is(err, ErrInvalidOwnership) {
|
|
t.Fatalf("Expire(zero limit) error = %v, want ErrInvalidOwnership", err)
|
|
}
|
|
}
|
|
|
|
type ownershipContextKey struct{}
|
|
|
|
type ownershipRepositoryStub struct {
|
|
assign func(context.Context, time.Time, string, string, time.Duration) (Assignment, error)
|
|
expire func(context.Context, time.Time, int) ([]Assignment, error)
|
|
assignCalled bool
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) {
|
|
r.assignCalled = true
|
|
if r.assign != nil {
|
|
return r.assign(ctx, now, proxyID, workerID, ttl)
|
|
}
|
|
return Assignment{}, nil
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) Renew(context.Context, time.Time, string, string, uint64, time.Duration) (Assignment, error) {
|
|
return Assignment{}, nil
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) BeginDrain(context.Context, string, string, uint64) (Assignment, error) {
|
|
return Assignment{}, nil
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) AcknowledgeDrain(context.Context, string, string, uint64, int64, int64) error {
|
|
return nil
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) Get(context.Context, string) (Assignment, bool, error) {
|
|
return Assignment{}, false, nil
|
|
}
|
|
|
|
func (r *ownershipRepositoryStub) Expire(ctx context.Context, now time.Time, limit int) ([]Assignment, error) {
|
|
if r.expire != nil {
|
|
return r.expire(ctx, now, limit)
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
var _ ownershipDomain.Repository = (*ownershipRepositoryStub)(nil)
|
|
|
|
func newTestOwnershipManager(t *testing.T, proxyIDs ...string) *OwnershipManager {
|
|
t.Helper()
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
manager, err := NewOwnershipManager(newTestActivityPool(t, now, proxyIDs...))
|
|
if err != nil {
|
|
t.Fatalf("NewOwnershipManager(): %v", err)
|
|
}
|
|
return manager
|
|
}
|
|
|
|
func newTestActivityPool(t *testing.T, now time.Time, proxyIDs ...string) *activitypool.MemoryPool {
|
|
t.Helper()
|
|
proxies := make([]proxyDomain.Proxy, 0, len(proxyIDs))
|
|
for index, proxyID := range proxyIDs {
|
|
proxies = append(proxies, proxyDomain.Proxy{
|
|
ID: proxyID, Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.1",
|
|
Port: uint16(8000 + index), State: proxyDomain.StateAvailable,
|
|
})
|
|
}
|
|
store := activitypool.NewMemoryPool()
|
|
result, err := store.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
|
ObservedAt: now, ConfiguredTTL: 10 * time.Minute, MaxSize: 100, Proxies: proxies,
|
|
})
|
|
if err != nil || result.Inserted != len(proxyIDs) {
|
|
t.Fatalf("UpsertFetched() = %+v, %v", result, err)
|
|
}
|
|
return store
|
|
}
|