refactor: make ownership repository context aware

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) {
+func (m *OwnershipManager) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) {
-	if m == nil || m.repository == nil {
+	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) {
+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 {
+	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) {
+func (m *OwnershipManager) BeginDrain(ctx context.Context, proxyID, workerID string, epoch uint64) (Assignment, error) {
-	if m == nil || m.repository == nil {
+	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 {
+func (m *OwnershipManager) AcknowledgeDrain(ctx context.Context, proxyID, workerID string, epoch uint64, active, reserved int64) error {
-	if m == nil || m.repository == nil {
+	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) {
+func (m *OwnershipManager) Get(ctx context.Context, proxyID string) (Assignment, bool, error) {
-	if m == nil || m.repository == nil {
+	if m == nil || m.repository == nil || ctx == nil {
-		return Assignment{}, false
+		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 {
+func (m *OwnershipManager) Expire(ctx context.Context, now time.Time, limit int) ([]Assignment, error) {
-	if m == nil || m.repository == nil {
+	if m == nil || m.repository == nil || ctx == nil || limit <= 0 {
-		return nil
+		return nil, ErrInvalidOwnership
 	}
-	return m.repository.Expire(now)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	return m.repository.Expire(ctx, now, limit)
 }

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)

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)
+	return expired, nil
-	sort.Slice(expired, func(i, j int) bool { return expired[i].ProxyID < expired[j].ProxyID })
+}
-	return expired
+
+func ownershipContextError(ctx context.Context) error {
+	if ctx == nil {
+		return ownershipDomain.ErrInvalidOwnership
+	}
+	return ctx.Err()
 }
 
 func (p *MemoryPool) purgeExpiredLocked(now time.Time) int {

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

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)
+	Assign(context.Context, time.Time, string, string, time.Duration) (Assignment, error)
-	Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error)
+	Renew(context.Context, time.Time, string, string, uint64, time.Duration) (Assignment, error)
-	BeginDrain(proxyID, workerID string, epoch uint64) (Assignment, error)
+	BeginDrain(context.Context, string, string, uint64) (Assignment, error)
-	AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error
+	AcknowledgeDrain(context.Context, string, string, uint64, int64, int64) error
-	Get(proxyID string) (Assignment, bool)
+	Get(context.Context, string) (Assignment, bool, error)
-	Expire(now time.Time) []Assignment
+	Expire(context.Context, time.Time, int) ([]Assignment, error)
 }