youfak/proxy-pool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: make ownership repository context aware

Browse Source
This commit is contained in:
youfak 2026-07-29 14:29:43 +08:00
parent 179e4d5a0c
commit c8e64758da
5 changed files with 350 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
internal/controller/pool/ownership.go
View File

@ -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)
}

153
internal/controller/pool/ownership_test.go
View File

@ -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)

82
internal/domain/activitypool/pool.go
View File

@ -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 {

101
internal/domain/activitypool/pool_test.go
View File

@ -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()

13
internal/domain/ownership/ownership.go
View File

@ -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)
}