From c8e64758da76ad622b782e1df96a280e3e97318b Mon Sep 17 00:00:00 2001 From: youfak Date: Wed, 29 Jul 2026 14:29:43 +0800 Subject: [PATCH] refactor: make ownership repository context aware --- internal/controller/pool/ownership.go | 59 +++++--- internal/controller/pool/ownership_test.go | 153 +++++++++++++++++++-- internal/domain/activitypool/pool.go | 82 +++++++++-- internal/domain/activitypool/pool_test.go | 101 +++++++++++++- internal/domain/ownership/ownership.go | 13 +- 5 files changed, 350 insertions(+), 58 deletions(-) diff --git a/internal/controller/pool/ownership.go b/internal/controller/pool/ownership.go index 1dbd04d..5a75fcc 100644 --- a/internal/controller/pool/ownership.go +++ b/internal/controller/pool/ownership.go @@ -1,6 +1,7 @@ package pool import ( + "context" "time" ownershipDomain "proxy-pool/internal/domain/ownership" @@ -30,44 +31,62 @@ func NewOwnershipManager(repository ownershipDomain.Repository) (*OwnershipManag return &OwnershipManager{repository: repository}, nil } -func (m *OwnershipManager) Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) { - if m == nil || m.repository == nil { +func (m *OwnershipManager) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) { + if m == nil || m.repository == nil || ctx == nil { return Assignment{}, ErrInvalidOwnership } - return m.repository.Assign(now, proxyID, workerID, ttl) + if err := ctx.Err(); err != nil { + return Assignment{}, err + } + return m.repository.Assign(ctx, now, proxyID, workerID, ttl) } -func (m *OwnershipManager) Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) { - if m == nil || m.repository == nil { +func (m *OwnershipManager) Renew(ctx context.Context, now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) { + if m == nil || m.repository == nil || ctx == nil { return Assignment{}, ErrInvalidOwnership } - return m.repository.Renew(now, proxyID, workerID, epoch, ttl) + if err := ctx.Err(); err != nil { + return Assignment{}, err + } + return m.repository.Renew(ctx, now, proxyID, workerID, epoch, ttl) } -func (m *OwnershipManager) BeginDrain(proxyID, workerID string, epoch uint64) (Assignment, error) { - if m == nil || m.repository == nil { +func (m *OwnershipManager) BeginDrain(ctx context.Context, proxyID, workerID string, epoch uint64) (Assignment, error) { + if m == nil || m.repository == nil || ctx == nil { return Assignment{}, ErrInvalidOwnership } - return m.repository.BeginDrain(proxyID, workerID, epoch) + if err := ctx.Err(); err != nil { + return Assignment{}, err + } + return m.repository.BeginDrain(ctx, proxyID, workerID, epoch) } -func (m *OwnershipManager) AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error { - if m == nil || m.repository == nil { +func (m *OwnershipManager) AcknowledgeDrain(ctx context.Context, proxyID, workerID string, epoch uint64, active, reserved int64) error { + if m == nil || m.repository == nil || ctx == nil { return ErrInvalidOwnership } - return m.repository.AcknowledgeDrain(proxyID, workerID, epoch, active, reserved) + if err := ctx.Err(); err != nil { + return err + } + return m.repository.AcknowledgeDrain(ctx, proxyID, workerID, epoch, active, reserved) } -func (m *OwnershipManager) Get(proxyID string) (Assignment, bool) { - if m == nil || m.repository == nil { - return Assignment{}, false +func (m *OwnershipManager) Get(ctx context.Context, proxyID string) (Assignment, bool, error) { + if m == nil || m.repository == nil || ctx == nil { + return Assignment{}, false, ErrInvalidOwnership } - return m.repository.Get(proxyID) + if err := ctx.Err(); err != nil { + return Assignment{}, false, err + } + return m.repository.Get(ctx, proxyID) } -func (m *OwnershipManager) Expire(now time.Time) []Assignment { - if m == nil || m.repository == nil { - return nil +func (m *OwnershipManager) Expire(ctx context.Context, now time.Time, limit int) ([]Assignment, error) { + if m == nil || m.repository == nil || ctx == nil || limit <= 0 { + return nil, ErrInvalidOwnership } - return m.repository.Expire(now) + if err := ctx.Err(); err != nil { + return nil, err + } + return m.repository.Expire(ctx, now, limit) } diff --git a/internal/controller/pool/ownership_test.go b/internal/controller/pool/ownership_test.go index 560ca9e..ffdb44f 100644 --- a/internal/controller/pool/ownership_test.go +++ b/internal/controller/pool/ownership_test.go @@ -11,6 +11,7 @@ import ( "proxy-pool/internal/domain/activitypool" extractionDomain "proxy-pool/internal/domain/extraction" + ownershipDomain "proxy-pool/internal/domain/ownership" proxyDomain "proxy-pool/internal/domain/proxy" ) @@ -24,7 +25,7 @@ func TestOwnershipManagerPreventsDualAssignment(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, err := manager.Assign(now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute) + _, err := manager.Assign(context.Background(), now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute) if err == nil { succeeded.Add(1) return @@ -44,21 +45,22 @@ func TestOwnershipManagerPreventsDualAssignment(t *testing.T) { func TestOwnershipManagerRenewsOnlyCurrentAssignment(t *testing.T) { manager := newTestOwnershipManager(t, "proxy-1") now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) - assigned, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + 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(now.Add(30*time.Second), "proxy-1", "worker-1", assigned.Epoch, time.Minute) + 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(now, "proxy-1", "worker-2", assigned.Epoch, time.Minute); !errors.Is(err, ErrStaleAssignment) { + 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 := manager.Expire(now.Add(time.Minute)); len(expired) != 0 { + 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) } } @@ -77,7 +79,7 @@ func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing extracted := make(chan bool, 1) go func() { <-start - _, assignErr := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + _, 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) } @@ -111,11 +113,12 @@ func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) { manager := newTestOwnershipManager(t, "proxy-1") now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) - assignment, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + 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("proxy-1", "worker-1", assignment.Epoch) + draining, err := manager.BeginDrain(ctx, "proxy-1", "worker-1", assignment.Epoch) if err != nil { t.Fatalf("BeginDrain(): %v", err) } @@ -123,13 +126,13 @@ func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) { t.Fatalf("draining assignment = %+v", draining) } - if err := manager.AcknowledgeDrain("proxy-1", "worker-1", assignment.Epoch, 1, 0); !errors.Is(err, ErrDrainNotReady) { + 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("proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil { + if err := manager.AcknowledgeDrain(ctx, "proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil { t.Fatalf("AcknowledgeDrain(zero): %v", err) } - if _, ok := manager.Get("proxy-1"); ok { + if _, ok, err := manager.Get(ctx, "proxy-1"); err != nil || ok { t.Fatal("assignment still exists after drain acknowledgement") } } @@ -137,18 +140,19 @@ func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) { func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) { manager := newTestOwnershipManager(t, "proxy-1") now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) - first, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + 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 := manager.Expire(now.Add(59 * time.Second)); len(expired) != 0 { + 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 := manager.Expire(now.Add(time.Minute)); len(expired) != 1 || expired[0].ProxyID != "proxy-1" { + 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(now.Add(time.Minute), "proxy-1", "worker-2", time.Minute) + second, err := manager.Assign(ctx, now.Add(time.Minute), "proxy-1", "worker-2", time.Minute) if err != nil { t.Fatalf("Assign(second): %v", err) } @@ -157,6 +161,125 @@ func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) { } } +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) diff --git a/internal/domain/activitypool/pool.go b/internal/domain/activitypool/pool.go index 9ea9cd3..a3a4a5d 100644 --- a/internal/domain/activitypool/pool.go +++ b/internal/domain/activitypool/pool.go @@ -284,12 +284,18 @@ func (p *MemoryPool) Extract(ctx context.Context, command extractionDomain.Comma return result, nil } -func (p *MemoryPool) Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (ownershipDomain.Assignment, error) { +func (p *MemoryPool) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (ownershipDomain.Assignment, error) { + if err := ownershipContextError(ctx); err != nil { + return ownershipDomain.Assignment{}, err + } if p == nil || now.IsZero() || proxyID == "" || workerID == "" || ttl <= 0 { return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership } p.mu.Lock() defer p.mu.Unlock() + if err := ctx.Err(); err != nil { + return ownershipDomain.Assignment{}, err + } p.purgeExpiredLocked(now) if current, exists := p.ownership[proxyID]; exists { if current.ExpiresAt.After(now) { @@ -316,12 +322,18 @@ func (p *MemoryPool) Assign(now time.Time, proxyID, workerID string, ttl time.Du return assignment, nil } -func (p *MemoryPool) Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (ownershipDomain.Assignment, error) { +func (p *MemoryPool) Renew(ctx context.Context, now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (ownershipDomain.Assignment, error) { + if err := ownershipContextError(ctx); err != nil { + return ownershipDomain.Assignment{}, err + } if p == nil || now.IsZero() || ttl <= 0 { return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership } p.mu.Lock() defer p.mu.Unlock() + if err := ctx.Err(); err != nil { + return ownershipDomain.Assignment{}, err + } p.purgeExpiredLocked(now) assignment, ok := p.ownership[proxyID] if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch || !assignment.ExpiresAt.After(now) { @@ -337,12 +349,18 @@ func (p *MemoryPool) Renew(now time.Time, proxyID, workerID string, epoch uint64 return assignment, nil } -func (p *MemoryPool) BeginDrain(proxyID, workerID string, epoch uint64) (ownershipDomain.Assignment, error) { +func (p *MemoryPool) BeginDrain(ctx context.Context, proxyID, workerID string, epoch uint64) (ownershipDomain.Assignment, error) { + if err := ownershipContextError(ctx); err != nil { + return ownershipDomain.Assignment{}, err + } if p == nil { return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment } p.mu.Lock() defer p.mu.Unlock() + if err := ctx.Err(); err != nil { + return ownershipDomain.Assignment{}, err + } assignment, ok := p.ownership[proxyID] if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch { return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment @@ -355,12 +373,18 @@ func (p *MemoryPool) BeginDrain(proxyID, workerID string, epoch uint64) (ownersh return assignment, nil } -func (p *MemoryPool) AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error { +func (p *MemoryPool) AcknowledgeDrain(ctx context.Context, proxyID, workerID string, epoch uint64, active, reserved int64) error { + if err := ownershipContextError(ctx); err != nil { + return err + } if p == nil || active < 0 || reserved < 0 { return ownershipDomain.ErrInvalidOwnership } p.mu.Lock() defer p.mu.Unlock() + if err := ctx.Err(); err != nil { + return err + } assignment, ok := p.ownership[proxyID] if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch { return ownershipDomain.ErrStaleAssignment @@ -379,28 +403,53 @@ func (p *MemoryPool) AcknowledgeDrain(proxyID, workerID string, epoch uint64, ac return nil } -func (p *MemoryPool) Get(proxyID string) (ownershipDomain.Assignment, bool) { +func (p *MemoryPool) Get(ctx context.Context, proxyID string) (ownershipDomain.Assignment, bool, error) { + if err := ownershipContextError(ctx); err != nil { + return ownershipDomain.Assignment{}, false, err + } if p == nil { - return ownershipDomain.Assignment{}, false + return ownershipDomain.Assignment{}, false, nil } p.mu.Lock() defer p.mu.Unlock() + if err := ctx.Err(); err != nil { + return ownershipDomain.Assignment{}, false, err + } assignment, ok := p.ownership[proxyID] - return assignment, ok + return assignment, ok, nil } -func (p *MemoryPool) Expire(now time.Time) []ownershipDomain.Assignment { +func (p *MemoryPool) Expire(ctx context.Context, now time.Time, limit int) ([]ownershipDomain.Assignment, error) { + if err := ownershipContextError(ctx); err != nil { + return nil, err + } + if limit <= 0 { + return nil, ownershipDomain.ErrInvalidOwnership + } if p == nil { - return nil + return nil, nil } p.mu.Lock() defer p.mu.Unlock() - expired := make([]ownershipDomain.Assignment, 0) + if err := ctx.Err(); err != nil { + return nil, err + } + eligibleIDs := make([]string, 0) for proxyID, assignment := range p.ownership { - entry, exists := p.entryByIDLocked(proxyID) + _, exists := p.entryByIDLocked(proxyID) if exists && assignment.ExpiresAt.After(now) { continue } + eligibleIDs = append(eligibleIDs, proxyID) + } + sort.Strings(eligibleIDs) + if len(eligibleIDs) > limit { + eligibleIDs = eligibleIDs[:limit] + } + expired := make([]ownershipDomain.Assignment, 0, len(eligibleIDs)) + for _, proxyID := range eligibleIDs { + assignment := p.ownership[proxyID] + entry, exists := p.entryByIDLocked(proxyID) if exists && entry.OwnerWorkerID == assignment.WorkerID { entry.OwnerWorkerID = "" p.setEntryByIDLocked(proxyID, entry) @@ -408,9 +457,14 @@ func (p *MemoryPool) Expire(now time.Time) []ownershipDomain.Assignment { expired = append(expired, assignment) delete(p.ownership, proxyID) } - p.purgeExpiredLocked(now) - sort.Slice(expired, func(i, j int) bool { return expired[i].ProxyID < expired[j].ProxyID }) - return expired + return expired, nil +} + +func ownershipContextError(ctx context.Context) error { + if ctx == nil { + return ownershipDomain.ErrInvalidOwnership + } + return ctx.Err() } func (p *MemoryPool) purgeExpiredLocked(now time.Time) int { diff --git a/internal/domain/activitypool/pool_test.go b/internal/domain/activitypool/pool_test.go index 2d480af..c4f785a 100644 --- a/internal/domain/activitypool/pool_test.go +++ b/internal/domain/activitypool/pool_test.go @@ -283,11 +283,11 @@ func TestMemoryPoolAssignRecoversExpiredLeaseWithoutSeparateSweep(t *testing.T) now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC) pool := poolWithOneProxy(t, now) proxyID := pool.Snapshot(now)[0].Proxy.ID - first, err := pool.Assign(now, proxyID, "worker-a", 5*time.Second) + first, err := pool.Assign(context.Background(), now, proxyID, "worker-a", 5*time.Second) if err != nil { t.Fatalf("Assign(first): %v", err) } - second, err := pool.Assign(now.Add(5*time.Second), proxyID, "worker-b", 5*time.Second) + second, err := pool.Assign(context.Background(), now.Add(5*time.Second), proxyID, "worker-b", 5*time.Second) if err != nil { t.Fatalf("Assign(after lease expiry): %v", err) } @@ -305,7 +305,7 @@ func TestMemoryPoolMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) { extracted := make(chan bool, 1) go func() { <-start - _, err := pool.Assign(now, pool.Snapshot(now)[0].Proxy.ID, "worker-a", time.Minute) + _, err := pool.Assign(context.Background(), now, pool.Snapshot(now)[0].Proxy.ID, "worker-a", time.Minute) if err != nil && !errors.Is(err, ownershipDomain.ErrOwnershipUnavailable) { t.Errorf("Assign(): %v", err) } @@ -336,6 +336,101 @@ func TestMemoryPoolMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) { } } +func TestMemoryPoolOwnershipMethodsPropagateCanceledContext(t *testing.T) { + now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC) + pool := poolWithOneProxy(t, now) + proxyID := pool.Snapshot(now)[0].Proxy.ID + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, err := pool.Assign(ctx, now, proxyID, "worker-a", time.Minute); !errors.Is(err, context.Canceled) { + t.Fatalf("Assign() error = %v, want context.Canceled", err) + } + if _, err := pool.Renew(ctx, now, proxyID, "worker-a", 1, time.Minute); !errors.Is(err, context.Canceled) { + t.Fatalf("Renew() error = %v, want context.Canceled", err) + } + if _, err := pool.BeginDrain(ctx, proxyID, "worker-a", 1); !errors.Is(err, context.Canceled) { + t.Fatalf("BeginDrain() error = %v, want context.Canceled", err) + } + if err := pool.AcknowledgeDrain(ctx, proxyID, "worker-a", 1, 0, 0); !errors.Is(err, context.Canceled) { + t.Fatalf("AcknowledgeDrain() error = %v, want context.Canceled", err) + } + if _, _, err := pool.Get(ctx, proxyID); !errors.Is(err, context.Canceled) { + t.Fatalf("Get() error = %v, want context.Canceled", err) + } + if _, err := pool.Expire(ctx, now, 1); !errors.Is(err, context.Canceled) { + t.Fatalf("Expire() error = %v, want context.Canceled", err) + } +} + +func TestMemoryPoolOwnershipMethodRechecksContextAfterLock(t *testing.T) { + now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC) + pool := poolWithOneProxy(t, now) + proxyID := pool.Snapshot(now)[0].Proxy.ID + base, cancel := context.WithCancel(context.Background()) + ctx := &firstCheckContext{Context: base, checked: make(chan struct{})} + result := make(chan error, 1) + + pool.mu.Lock() + go func() { + _, err := pool.Assign(ctx, now, proxyID, "worker-a", time.Minute) + result <- err + }() + <-ctx.checked + cancel() + pool.mu.Unlock() + + if err := <-result; !errors.Is(err, context.Canceled) { + t.Fatalf("Assign() error = %v, want context.Canceled", err) + } + if _, ok, err := pool.Get(context.Background(), proxyID); err != nil || ok { + t.Fatal("canceled Assign() mutated ownership") + } +} + +func TestMemoryPoolExpireUsesStableProxyIDOrderAndLimit(t *testing.T) { + now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC) + pool := NewMemoryPool() + proxies := []proxyDomain.Proxy{ + {ID: "proxy-c", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.3", Port: 8003, State: proxyDomain.StateAvailable}, + {ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.1", Port: 8001, State: proxyDomain.StateAvailable}, + {ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.2", Port: 8002, State: proxyDomain.StateAvailable}, + } + if _, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{ + ObservedAt: now, ConfiguredTTL: 10 * time.Minute, Proxies: proxies, + }); err != nil { + t.Fatalf("UpsertFetched(): %v", err) + } + for _, proxyID := range []string{"proxy-c", "proxy-a", "proxy-b"} { + if _, err := pool.Assign(context.Background(), now, proxyID, "worker-a", time.Minute); err != nil { + t.Fatalf("Assign(%s): %v", proxyID, err) + } + } + + expired, err := pool.Expire(context.Background(), now.Add(time.Minute), 2) + if err != nil { + t.Fatalf("Expire(): %v", err) + } + if len(expired) != 2 || expired[0].ProxyID != "proxy-a" || expired[1].ProxyID != "proxy-b" { + t.Fatalf("Expire() = %+v, want proxy-a then proxy-b", expired) + } + if _, ok, err := pool.Get(context.Background(), "proxy-c"); err != nil || !ok { + t.Fatalf("Get(proxy-c) = ok %v, error %v; want remaining assignment", ok, err) + } +} + +type firstCheckContext struct { + context.Context + checked chan struct{} + once sync.Once +} + +func (c *firstCheckContext) Err() error { + err := c.Context.Err() + c.once.Do(func() { close(c.checked) }) + return err +} + func poolWithOneProxy(t *testing.T, now time.Time) *MemoryPool { t.Helper() pool := NewMemoryPool() diff --git a/internal/domain/ownership/ownership.go b/internal/domain/ownership/ownership.go index f9f2342..c62a0c6 100644 --- a/internal/domain/ownership/ownership.go +++ b/internal/domain/ownership/ownership.go @@ -1,6 +1,7 @@ package ownership import ( + "context" "errors" "time" ) @@ -26,10 +27,10 @@ type Assignment struct { // Repository is the shared authority for ownership changes. Implementations // that also support extraction must serialize both operations transactionally. type Repository interface { - Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) - Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) - BeginDrain(proxyID, workerID string, epoch uint64) (Assignment, error) - AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error - Get(proxyID string) (Assignment, bool) - Expire(now time.Time) []Assignment + Assign(context.Context, time.Time, string, string, time.Duration) (Assignment, error) + Renew(context.Context, time.Time, string, string, uint64, time.Duration) (Assignment, error) + BeginDrain(context.Context, string, string, uint64) (Assignment, error) + AcknowledgeDrain(context.Context, string, string, uint64, int64, int64) error + Get(context.Context, string) (Assignment, bool, error) + Expire(context.Context, time.Time, int) ([]Assignment, error) }