proxy-pool/internal/domain/activitypool/pool.go

package activitypool

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"sort"
	"sync"
	"time"

	extractionDomain "proxy-pool/internal/domain/extraction"
	healthDomain "proxy-pool/internal/domain/health"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

const defaultIdempotencyTTL = 5 * time.Minute

var (
	ErrInvalidBatch        = errors.New("invalid activity pool batch")
	ErrInvalidHealthUpdate = errors.New("invalid activity pool health update")
	ErrInvalidProxyLookup  = errors.New("invalid activity pool proxy lookup")
	ErrActivityNotFound    = errors.New("activity pool proxy not found")
	ErrStaleHealthUpdate   = errors.New("stale activity pool health update")
	ErrInvalidInventory    = errors.New("invalid activity pool inventory query")
	ErrInvalidMaintenance  = errors.New("invalid activity pool maintenance request")
)

// FetchedBatch describes one ephemeral provider response. Proxies without a
// usable expiry are dropped because this pool is intentionally rebuildable.
type FetchedBatch struct {
	ObservedAt             time.Time
	ConfiguredTTL          time.Duration
	AllocationSafetyMargin time.Duration
	MaxSize                int
	Proxies                []proxyDomain.Proxy
}

type UpsertResult struct {
	Accepted  int
	Inserted  int
	Refreshed int
	Dropped   int
}

type Upserter interface {
	// A proxy unique key has one incumbent upstream for its current lifetime.
	// Duplicates from other upstreams must not replace that source lifecycle.
	UpsertFetched(context.Context, string, FetchedBatch) (UpsertResult, error)
}

type HealthUpdate struct {
	ProxyID   string
	CheckedAt time.Time
	NextState proxyDomain.State
	Latency   time.Duration
}

type Inventory struct {
	UpstreamID string
	Managed    int
}

// StateInventory is a low-cardinality operational view of one upstream.
// Expired and removed entries are intentionally excluded.
type StateInventory struct {
	UpstreamID string
	Fetched    int64
	Checking   int64
	Available  int64
	Suspect    int64
	Draining   int64
	Unhealthy  int64
	Extracted  int64
}

type HealthStore interface {
	ApplyHealth(context.Context, HealthUpdate) (Entry, error)
}

// GlobalHealthCommand carries a BASIC or EGRESS fact to the authoritative
// activity pool. TARGET observations use a separate profile store and never
// enter this command.
type GlobalHealthCommand struct {
	Observation            healthDomain.Observation
	MaxConsecutiveFailures int
}

type GlobalHealthStore interface {
	ApplyGlobalObservation(context.Context, GlobalHealthCommand) (Entry, error)
}

// TargetHealthCommand carries a TARGET fact. Its result is isolated by Proxy
// and TargetProfile and must not modify the activity-pool Proxy state.
type TargetHealthCommand struct {
	Observation            healthDomain.Observation
	MaxConsecutiveFailures int
}

type TargetHealthStore interface {
	ApplyTargetObservation(context.Context, TargetHealthCommand) (healthDomain.TargetState, error)
}

// ProxyUpstreamReader returns only the authoritative upstream that owns a
// still-live proxy. It intentionally does not expose proxy endpoints or
// credentials to Controller policy resolution.
type ProxyUpstreamReader interface {
	UpstreamForProxy(context.Context, string, time.Time) (string, error)
}

type InventoryReader interface {
	Inventory(context.Context, string, time.Time) (Inventory, error)
}

type StateInventoryReader interface {
	ReadStateInventory(context.Context, []string, time.Time) ([]StateInventory, error)
}

type Maintainer interface {
	SweepExpired(context.Context, time.Time, int) (int, error)
}

// UnhealthySweepCommand requests a bounded Controller-side sweep of proxies
// that have remained globally unhealthy past each upstream's configured grace
// period. A policy is absent when the upstream has removal disabled.
type UnhealthySweepCommand struct {
	Now                   time.Time
	Limit                 int
	RemoveAfterByUpstream map[string]time.Duration
}

// UnhealthySweepResult separates completed removals from candidates still
// owned by a Worker. Owned entries are intentionally deferred so a stale
// health result never tears down an active Gateway proxy.
type UnhealthySweepResult struct {
	Removed       int
	DeferredOwned int
}

type UnhealthyRemover interface {
	SweepUnhealthy(context.Context, UnhealthySweepCommand) (UnhealthySweepResult, error)
}

const maximumUnhealthySweepScan = 1024

type Entry struct {
	Proxy         proxyDomain.Proxy
	UsableUntil   time.Time
	OwnerWorkerID string
	State         proxyDomain.State
	GlobalHealth  healthDomain.GlobalState
}

type MemoryPool struct {
	mu sync.Mutex

	entries    map[string]Entry
	keyByID    map[string]string
	idempotent map[string]idempotencyEntry
	ownership  map[string]ownershipDomain.Assignment
	drains     map[string]ownershipDomain.DrainTicket
	targets    map[targetHealthKey]healthDomain.TargetState
	unhealthy  map[string]time.Time
	nextEpoch  uint64
}

type targetHealthKey struct {
	proxyID    string
	profileKey string
}

type idempotencyEntry struct {
	command   extractionDomain.Command
	result    extractionDomain.Result
	expiresAt time.Time
}

var (
	_ Upserter                         = (*MemoryPool)(nil)
	_ HealthStore                      = (*MemoryPool)(nil)
	_ GlobalHealthStore                = (*MemoryPool)(nil)
	_ TargetHealthStore                = (*MemoryPool)(nil)
	_ ProxyUpstreamReader              = (*MemoryPool)(nil)
	_ InventoryReader                  = (*MemoryPool)(nil)
	_ StateInventoryReader             = (*MemoryPool)(nil)
	_ Maintainer                       = (*MemoryPool)(nil)
	_ UnhealthyRemover                 = (*MemoryPool)(nil)
	_ extractionDomain.Store           = (*MemoryPool)(nil)
	_ ownershipDomain.Repository       = (*MemoryPool)(nil)
	_ ownershipDomain.DrainTicketStore = (*MemoryPool)(nil)
)

func NewMemoryPool() *MemoryPool {
	return &MemoryPool{
		entries:    make(map[string]Entry),
		keyByID:    make(map[string]string),
		idempotent: make(map[string]idempotencyEntry),
		ownership:  make(map[string]ownershipDomain.Assignment),
		drains:     make(map[string]ownershipDomain.DrainTicket),
		targets:    make(map[targetHealthKey]healthDomain.TargetState),
		unhealthy:  make(map[string]time.Time),
	}
}

func (p *MemoryPool) UpsertFetched(ctx context.Context, upstreamID string, batch FetchedBatch) (UpsertResult, error) {
	var result UpsertResult
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if p == nil || upstreamID == "" || batch.ObservedAt.IsZero() || batch.ConfiguredTTL < 0 || batch.MaxSize <= 0 ||
		batch.AllocationSafetyMargin < 0 ||
		(batch.ConfiguredTTL > 0 && batch.AllocationSafetyMargin >= batch.ConfiguredTTL) {
		return result, ErrInvalidBatch
	}
	for _, candidate := range batch.Proxies {
		if candidate.SourceUpstream != "" && candidate.SourceUpstream != upstreamID {
			return result, ErrInvalidBatch
		}
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return result, err
	}
	p.purgeExpiredLocked(batch.ObservedAt)
	managedByUpstream := make(map[string]int)
	for _, entry := range p.entries {
		if managedActivityState(entry.State) {
			managedByUpstream[entry.Proxy.SourceUpstream]++
		}
	}
	seenIDs := make(map[string]string, len(batch.Proxies))
	for _, candidate := range batch.Proxies {
		if !validProxyIdentity(candidate) {
			continue
		}
		key := candidate.UniqueKey()
		proxyID := candidate.ID
		if proxyID == "" {
			proxyID = stableProxyID(key)
		}
		if existingKey, exists := p.keyByID[proxyID]; exists && existingKey != key {
			return UpsertResult{}, ErrInvalidBatch
		}
		if existingKey, exists := seenIDs[proxyID]; exists && existingKey != key {
			return UpsertResult{}, ErrInvalidBatch
		}
		seenIDs[proxyID] = key
	}

	for _, candidate := range batch.Proxies {
		if !validProxyIdentity(candidate) {
			result.Dropped++
			continue
		}
		candidate.SourceUpstream = upstreamID
		expiresAt := proxyDomain.EffectiveExpiry(batch.ObservedAt, candidate.ExpiresAt, 0, batch.ConfiguredTTL)
		if expiresAt == nil {
			result.Dropped++
			continue
		}
		usableUntil := expiresAt.Add(-batch.AllocationSafetyMargin)
		if !usableUntil.After(batch.ObservedAt) {
			result.Dropped++
			continue
		}
		candidate.ExpiresAt = expiresAt
		candidate.UsableUntil = &usableUntil
		if candidate.CreatedAt.IsZero() {
			candidate.CreatedAt = batch.ObservedAt.UTC()
		}
		if candidate.State == "" {
			candidate.State = proxyDomain.StateFetched
		}
		key := candidate.UniqueKey()
		result.Accepted++

		if current, exists := p.entries[key]; exists {
			result.Refreshed++
			if current.State == proxyDomain.StateExtracted || current.Proxy.SourceUpstream != upstreamID {
				continue
			}
			candidate.ID = current.Proxy.ID
			candidate.CreatedAt = current.Proxy.CreatedAt
			candidate.State = current.State
			candidate.LastCheckedAt = current.Proxy.LastCheckedAt
			candidate.LastSuccessAt = current.Proxy.LastSuccessAt
			candidate.Latency = current.Proxy.Latency
			p.entries[key] = Entry{
				Proxy:         cloneProxy(candidate),
				UsableUntil:   usableUntil,
				OwnerWorkerID: current.OwnerWorkerID,
				State:         current.State,
				GlobalHealth:  current.GlobalHealth,
			}
			continue
		}
		if managedByUpstream[upstreamID] >= batch.MaxSize {
			result.Dropped++
			continue
		}

		if candidate.ID == "" {
			candidate.ID = stableProxyID(key)
		}
		p.entries[key] = Entry{
			Proxy:        cloneProxy(candidate),
			UsableUntil:  usableUntil,
			State:        candidate.State,
			GlobalHealth: healthDomain.GlobalState{State: candidate.State},
		}
		p.keyByID[candidate.ID] = key
		if managedActivityState(candidate.State) {
			managedByUpstream[upstreamID]++
		}
		result.Inserted++
	}
	return result, nil
}

func (p *MemoryPool) ApplyHealth(ctx context.Context, update HealthUpdate) (Entry, error) {
	if ctx == nil {
		return Entry{}, ErrInvalidHealthUpdate
	}
	if err := ctx.Err(); err != nil {
		return Entry{}, err
	}
	if p == nil || update.ProxyID == "" || update.CheckedAt.IsZero() || update.NextState == "" || update.Latency < 0 {
		return Entry{}, ErrInvalidHealthUpdate
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return Entry{}, err
	}
	p.purgeExpiredLocked(update.CheckedAt)
	entry, ok := p.entryByIDLocked(update.ProxyID)
	if !ok {
		return Entry{}, ErrActivityNotFound
	}
	if entry.Proxy.LastCheckedAt != nil {
		if update.CheckedAt.Before(*entry.Proxy.LastCheckedAt) {
			return Entry{}, ErrStaleHealthUpdate
		}
		if update.CheckedAt.Equal(*entry.Proxy.LastCheckedAt) {
			if entry.State != update.NextState {
				return Entry{}, ErrStaleHealthUpdate
			}
			entry.Proxy = cloneProxy(entry.Proxy)
			return entry, nil
		}
	}
	if entry.State != update.NextState && !proxyDomain.CanTransition(entry.State, update.NextState) {
		return Entry{}, ErrInvalidHealthUpdate
	}
	checkedAt := update.CheckedAt.UTC()
	entry.State = update.NextState
	entry.Proxy.State = update.NextState
	entry.GlobalHealth.State = update.NextState
	entry.Proxy.LastCheckedAt = &checkedAt
	entry.Proxy.Latency = update.Latency
	if update.NextState == proxyDomain.StateAvailable {
		lastSuccessAt := checkedAt
		entry.Proxy.LastSuccessAt = &lastSuccessAt
	}
	p.setEntryByIDLocked(update.ProxyID, entry)
	entry.Proxy = cloneProxy(entry.Proxy)
	return entry, nil
}

// ApplyGlobalObservation atomically reduces a BASIC or EGRESS Checker fact
// against the current Proxy health state and publishes the resulting state.
func (p *MemoryPool) ApplyGlobalObservation(
	ctx context.Context,
	command GlobalHealthCommand,
) (Entry, error) {
	if ctx == nil {
		return Entry{}, ErrInvalidHealthUpdate
	}
	if err := ctx.Err(); err != nil {
		return Entry{}, err
	}
	normalized, err := healthDomain.NormalizeObservation(command.Observation)
	if err != nil {
		return Entry{}, err
	}
	if normalized.Level == healthDomain.LevelTarget || command.MaxConsecutiveFailures <= 0 {
		return Entry{}, ErrInvalidHealthUpdate
	}
	if p == nil {
		return Entry{}, ErrInvalidHealthUpdate
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return Entry{}, err
	}
	p.purgeExpiredLocked(normalized.ObservedAt)
	entry, ok := p.entryByIDLocked(normalized.ProxyID)
	if !ok {
		return Entry{}, ErrActivityNotFound
	}
	current := entry.GlobalHealth
	if current.State == "" {
		current.State = entry.State
	}
	if current.State != entry.State {
		return Entry{}, ErrInvalidHealthUpdate
	}
	next, err := healthDomain.ReduceGlobal(current, normalized, command.MaxConsecutiveFailures)
	if err != nil {
		return Entry{}, err
	}
	if entry.State != next.State && !proxyDomain.CanTransition(entry.State, next.State) {
		return Entry{}, ErrInvalidHealthUpdate
	}
	entry.State = next.State
	entry.Proxy.State = next.State
	entry.Proxy.LastCheckedAt = timePointer(normalized.ObservedAt)
	entry.Proxy.Latency = normalized.Latency
	if normalized.Success {
		entry.Proxy.LastSuccessAt = timePointer(normalized.ObservedAt)
	}
	entry.GlobalHealth = next
	if next.State == proxyDomain.StateUnhealthy && !next.UnhealthySince.IsZero() {
		p.unhealthy[normalized.ProxyID] = next.UnhealthySince
	} else {
		delete(p.unhealthy, normalized.ProxyID)
	}
	p.setEntryByIDLocked(normalized.ProxyID, entry)
	entry.Proxy = cloneProxy(entry.Proxy)
	return entry, nil
}

// ApplyTargetObservation atomically reduces an isolated TARGET profile. It
// intentionally leaves the global Proxy state, health timestamp and indexes
// untouched.
func (p *MemoryPool) ApplyTargetObservation(
	ctx context.Context,
	command TargetHealthCommand,
) (healthDomain.TargetState, error) {
	if ctx == nil {
		return healthDomain.TargetState{}, ErrInvalidHealthUpdate
	}
	if err := ctx.Err(); err != nil {
		return healthDomain.TargetState{}, err
	}
	normalized, err := healthDomain.NormalizeObservation(command.Observation)
	if err != nil {
		return healthDomain.TargetState{}, err
	}
	if normalized.Level != healthDomain.LevelTarget || command.MaxConsecutiveFailures <= 0 || p == nil {
		return healthDomain.TargetState{}, ErrInvalidHealthUpdate
	}
	profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
		RoutingName: normalized.RoutingName, TargetURL: normalized.TargetURL,
	})
	if err != nil {
		return healthDomain.TargetState{}, err
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return healthDomain.TargetState{}, err
	}
	p.purgeExpiredLocked(normalized.ObservedAt)
	if _, exists := p.entryByIDLocked(normalized.ProxyID); !exists {
		return healthDomain.TargetState{}, ErrActivityNotFound
	}
	key := targetHealthKey{proxyID: normalized.ProxyID, profileKey: profile.Key()}
	next, err := healthDomain.ReduceTarget(p.targets[key], normalized, command.MaxConsecutiveFailures)
	if err != nil {
		return healthDomain.TargetState{}, err
	}
	p.targets[key] = next
	return next, nil
}

func (p *MemoryPool) UpstreamForProxy(ctx context.Context, proxyID string, now time.Time) (string, error) {
	if ctx == nil {
		return "", ErrInvalidProxyLookup
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	if p == nil || proxyID == "" || now.IsZero() {
		return "", ErrInvalidProxyLookup
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.purgeExpiredLocked(now)
	entry, exists := p.entryByIDLocked(proxyID)
	if !exists || entry.Proxy.SourceUpstream == "" {
		return "", ErrActivityNotFound
	}
	return entry.Proxy.SourceUpstream, nil
}

func (p *MemoryPool) Inventory(ctx context.Context, upstreamID string, now time.Time) (Inventory, error) {
	result := Inventory{UpstreamID: upstreamID}
	if ctx == nil {
		return result, ErrInvalidInventory
	}
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if p == nil || upstreamID == "" || now.IsZero() {
		return result, ErrInvalidInventory
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return result, err
	}
	for _, entry := range p.entries {
		if entry.Proxy.SourceUpstream == upstreamID && managedActivityState(entry.State) &&
			entry.Proxy.ExpiresAt != nil && entry.Proxy.ExpiresAt.After(now) {
			result.Managed++
		}
	}
	return result, nil
}

func (p *MemoryPool) ReadStateInventory(
	ctx context.Context,
	upstreamIDs []string,
	now time.Time,
) ([]StateInventory, error) {
	if ctx == nil {
		return nil, ErrInvalidInventory
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if p == nil || now.IsZero() {
		return nil, ErrInvalidInventory
	}
	result := make([]StateInventory, len(upstreamIDs))
	positions := make(map[string][]int, len(upstreamIDs))
	for index, upstreamID := range upstreamIDs {
		if upstreamID == "" {
			return nil, ErrInvalidInventory
		}
		result[index].UpstreamID = upstreamID
		positions[upstreamID] = append(positions[upstreamID], index)
	}
	if len(result) == 0 {
		return result, nil
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	for _, entry := range p.entries {
		indexes := positions[entry.Proxy.SourceUpstream]
		if len(indexes) == 0 || entry.Proxy.ExpiresAt == nil || !entry.Proxy.ExpiresAt.After(now) {
			continue
		}
		for _, index := range indexes {
			incrementStateInventory(&result[index], entry.State)
		}
	}
	return result, nil
}

func (p *MemoryPool) SweepExpired(ctx context.Context, now time.Time, limit int) (int, error) {
	if ctx == nil {
		return 0, ErrInvalidMaintenance
	}
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	if p == nil || now.IsZero() || limit <= 0 {
		return 0, ErrInvalidMaintenance
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	expiredIDs := make([]string, 0)
	for _, entry := range p.entries {
		if entry.Proxy.ExpiresAt != nil && !entry.Proxy.ExpiresAt.After(now) {
			expiredIDs = append(expiredIDs, entry.Proxy.ID)
		}
	}
	sort.Strings(expiredIDs)
	if len(expiredIDs) > limit {
		expiredIDs = expiredIDs[:limit]
	}
	for _, proxyID := range expiredIDs {
		p.removeEntryByIDLocked(proxyID)
	}
	return len(expiredIDs), nil
}

// SweepUnhealthy removes only unowned proxies which remain unhealthy beyond
// the configured upstream grace period. It is bounded by Limit and never
// reaches the Gateway request path.
func (p *MemoryPool) SweepUnhealthy(
	ctx context.Context,
	command UnhealthySweepCommand,
) (UnhealthySweepResult, error) {
	if ctx == nil {
		return UnhealthySweepResult{}, ErrInvalidMaintenance
	}
	if err := ctx.Err(); err != nil {
		return UnhealthySweepResult{}, err
	}
	if p == nil || command.Now.IsZero() || command.Limit <= 0 || !validUnhealthySweepPolicies(command.RemoveAfterByUpstream) {
		return UnhealthySweepResult{}, ErrInvalidMaintenance
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return UnhealthySweepResult{}, err
	}
	p.purgeExpiredLocked(command.Now)
	type candidate struct {
		proxyID string
		dueAt   time.Time
	}
	scanLimit := command.Limit
	if scanLimit > maximumUnhealthySweepScan/4 {
		scanLimit = maximumUnhealthySweepScan
	} else {
		scanLimit *= 4
	}
	candidates := make([]candidate, 0, min(len(p.unhealthy), scanLimit))
	scanned := 0
	for proxyID, dueAt := range p.unhealthy {
		if scanned >= scanLimit {
			break
		}
		scanned++
		entry, exists := p.entryByIDLocked(proxyID)
		if !exists || entry.State != proxyDomain.StateUnhealthy {
			delete(p.unhealthy, proxyID)
			continue
		}
		removeAfter, configured := command.RemoveAfterByUpstream[entry.Proxy.SourceUpstream]
		if !configured || entry.GlobalHealth.UnhealthySince.IsZero() {
			delete(p.unhealthy, proxyID)
			continue
		}
		if dueAt.After(command.Now) || entry.GlobalHealth.UnhealthySince.Add(removeAfter).After(command.Now) {
			continue
		}
		candidates = append(candidates, candidate{proxyID: proxyID, dueAt: dueAt})
	}
	sort.Slice(candidates, func(left, right int) bool {
		if candidates[left].dueAt.Equal(candidates[right].dueAt) {
			return candidates[left].proxyID < candidates[right].proxyID
		}
		return candidates[left].dueAt.Before(candidates[right].dueAt)
	})
	if len(candidates) > command.Limit {
		candidates = candidates[:command.Limit]
	}
	result := UnhealthySweepResult{}
	for _, candidate := range candidates {
		entry, exists := p.entryByIDLocked(candidate.proxyID)
		if !exists || entry.State != proxyDomain.StateUnhealthy {
			delete(p.unhealthy, candidate.proxyID)
			continue
		}
		if entry.OwnerWorkerID != "" {
			// Retry later without losing the authoritative first-unhealthy timestamp.
			p.unhealthy[candidate.proxyID] = command.Now.UTC().Add(time.Second)
			result.DeferredOwned++
			continue
		}
		p.removeEntryByIDLocked(candidate.proxyID)
		result.Removed++
	}
	return result, nil
}

func (p *MemoryPool) Snapshot(now time.Time) []Entry {
	if p == nil {
		return nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.purgeExpiredLocked(now)
	entries := make([]Entry, 0, len(p.entries))
	for _, entry := range p.entries {
		entry.Proxy = cloneProxy(entry.Proxy)
		entries = append(entries, entry)
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].Proxy.ID < entries[j].Proxy.ID })
	return entries
}

func (p *MemoryPool) PurgeExpired(now time.Time) int {
	if p == nil {
		return 0
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.purgeExpiredLocked(now)
}

func (p *MemoryPool) Extract(ctx context.Context, command extractionDomain.Command) (extractionDomain.Result, error) {
	result := extractionDomain.Result{Requested: command.Requested}
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if p == nil || command.Now.IsZero() || command.Requested < 0 || command.ReserveForGateway < 0 ||
		command.MinRemainingTTL < 0 || command.MaxHealthCheckAge < 0 || command.IdempotencyTTL < 0 ||
		(command.IdempotencyKey != "" && command.ClientID == "") ||
		(command.Fulfillment != extractionDomain.Partial && command.Fulfillment != extractionDomain.AllOrNothing) {
		return result, extractionDomain.ErrInvalidCommand
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return result, err
	}
	p.purgeExpiredLocked(command.Now)

	idempotencyKey := command.ClientID + "\x00" + command.IdempotencyKey
	if command.IdempotencyKey != "" {
		if committed, ok := p.idempotent[idempotencyKey]; ok {
			if !sameIdempotentRequest(committed.command, command) {
				return result, extractionDomain.ErrIdempotencyConflict
			}
			return cloneResult(committed.result), nil
		}
	}
	if command.Requested == 0 {
		p.rememberExtractionLocked(idempotencyKey, command, result)
		return result, nil
	}

	eligible := make([]string, 0, len(p.entries))
	for key, entry := range p.entries {
		if eligibleForExtraction(entry, command) {
			eligible = append(eligible, key)
		}
	}
	sort.Slice(eligible, func(i, j int) bool {
		return p.entries[eligible[i]].UsableUntil.After(p.entries[eligible[j]].UsableUntil)
	})
	available := len(eligible) - command.ReserveForGateway
	if available < 0 {
		available = 0
	}
	if command.Fulfillment == extractionDomain.AllOrNothing && available < command.Requested {
		return result, extractionDomain.ErrInsufficientProxies
	}
	count := command.Requested
	if count > available {
		count = available
	}
	for _, key := range eligible[:count] {
		entry := p.entries[key]
		entry.State = proxyDomain.StateExtracted
		entry.Proxy.State = proxyDomain.StateExtracted
		entry.GlobalHealth.State = proxyDomain.StateExtracted
		p.entries[key] = entry
		result.Items = append(result.Items, extractionCandidate(entry))
	}
	result.Returned = len(result.Items)
	if result.Returned > 0 {
		result.ExtractedAt = command.Now.UTC()
	}
	p.rememberExtractionLocked(idempotencyKey, command, result)
	return result, nil
}

func timePointer(value time.Time) *time.Time {
	canonical := value.UTC()
	return &canonical
}

func (p *MemoryPool) rememberExtractionLocked(
	idempotencyKey string,
	command extractionDomain.Command,
	result extractionDomain.Result,
) {
	if command.IdempotencyKey == "" {
		return
	}
	expiresAt := command.Now.Add(idempotencyTTL(command.IdempotencyTTL))
	for _, item := range result.Items {
		if !item.ExpiresAt.IsZero() && item.ExpiresAt.Before(expiresAt) {
			expiresAt = item.ExpiresAt
		}
	}
	if expiresAt.After(command.Now) {
		p.idempotent[idempotencyKey] = idempotencyEntry{
			command: cloneCommand(command), result: cloneResult(result), expiresAt: expiresAt,
		}
	}
}

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) {
			return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned
		}
		if entry, ok := p.entryByIDLocked(proxyID); ok && entry.OwnerWorkerID == current.WorkerID {
			entry.OwnerWorkerID = ""
			p.setEntryByIDLocked(proxyID, entry)
		}
		delete(p.ownership, proxyID)
		delete(p.drains, proxyID)
	}
	entry, ok := p.entryByIDLocked(proxyID)
	if !ok || entry.State != proxyDomain.StateAvailable || entry.OwnerWorkerID != "" || !entry.UsableUntil.After(now) {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable
	}
	p.nextEpoch++
	expiresAt := minTime(now.UTC().Add(ttl), entry.UsableUntil)
	assignment := ownershipDomain.Assignment{
		ProxyID: proxyID, WorkerID: workerID, Epoch: p.nextEpoch, Version: 1, ExpiresAt: expiresAt,
	}
	entry.OwnerWorkerID = workerID
	p.setEntryByIDLocked(proxyID, entry)
	p.ownership[proxyID] = assignment
	return assignment, nil
}

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) {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	entry, ok := p.entryByIDLocked(proxyID)
	if !ok {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	assignment.ExpiresAt = minTime(now.UTC().Add(ttl), entry.UsableUntil)
	assignment.Version++
	p.ownership[proxyID] = assignment
	return assignment, nil
}

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
	}
	if !assignment.Draining {
		assignment.Draining = true
		assignment.Version++
		p.ownership[proxyID] = assignment
		p.nextEpoch++
		p.drains[proxyID] = ownershipDomain.DrainTicket{
			ProxyID: proxyID, WorkerID: workerID, AssignmentEpoch: assignment.Epoch,
			RequiredSnapshotEpoch: p.nextEpoch,
		}
	}
	return assignment, nil
}

func (p *MemoryPool) PendingDrains(ctx context.Context, workerID string, limit int) ([]ownershipDomain.DrainTicket, error) {
	if err := ownershipContextError(ctx); err != nil {
		return nil, err
	}
	if p == nil || workerID == "" || limit <= 0 {
		return nil, ownershipDomain.ErrInvalidDrainTicket
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	proxyIDs := make([]string, 0)
	for proxyID, ticket := range p.drains {
		assignment, exists := p.ownership[proxyID]
		if !exists || !assignment.Draining || assignment.WorkerID != ticket.WorkerID || assignment.Epoch != ticket.AssignmentEpoch {
			delete(p.drains, proxyID)
			continue
		}
		if ticket.WorkerID == workerID {
			proxyIDs = append(proxyIDs, proxyID)
		}
	}
	sort.Strings(proxyIDs)
	if len(proxyIDs) > limit {
		proxyIDs = proxyIDs[:limit]
	}
	tickets := make([]ownershipDomain.DrainTicket, 0, len(proxyIDs))
	for _, proxyID := range proxyIDs {
		tickets = append(tickets, p.drains[proxyID])
	}
	return tickets, nil
}

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
	}
	if !assignment.Draining {
		return ownershipDomain.ErrNotDraining
	}
	if active > 0 || reserved > 0 {
		return ownershipDomain.ErrDrainNotReady
	}
	if entry, exists := p.entryByIDLocked(proxyID); exists && entry.OwnerWorkerID == workerID {
		entry.OwnerWorkerID = ""
		p.setEntryByIDLocked(proxyID, entry)
	}
	delete(p.ownership, proxyID)
	delete(p.drains, proxyID)
	return nil
}

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, 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, nil
}

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, nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	eligibleIDs := make([]string, 0)
	for proxyID, assignment := range p.ownership {
		_, 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)
		}
		expired = append(expired, assignment)
		delete(p.ownership, proxyID)
		delete(p.drains, proxyID)
	}
	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 {
	removed := 0
	for key, entry := range p.entries {
		if entry.Proxy.ExpiresAt == nil || entry.Proxy.ExpiresAt.After(now) {
			continue
		}
		p.removeEntryLocked(key, entry)
		removed++
	}
	for key, entry := range p.idempotent {
		if !entry.expiresAt.After(now) {
			delete(p.idempotent, key)
		}
	}
	return removed
}

func (p *MemoryPool) removeEntryByIDLocked(proxyID string) {
	key, ok := p.keyByID[proxyID]
	if !ok {
		return
	}
	entry, ok := p.entries[key]
	if !ok {
		delete(p.keyByID, proxyID)
		return
	}
	p.removeEntryLocked(key, entry)
}

func (p *MemoryPool) removeEntryLocked(key string, entry Entry) {
	delete(p.ownership, entry.Proxy.ID)
	delete(p.drains, entry.Proxy.ID)
	delete(p.unhealthy, entry.Proxy.ID)
	for target := range p.targets {
		if target.proxyID == entry.Proxy.ID {
			delete(p.targets, target)
		}
	}
	delete(p.keyByID, entry.Proxy.ID)
	delete(p.entries, key)
}

func validUnhealthySweepPolicies(policies map[string]time.Duration) bool {
	if len(policies) == 0 {
		return false
	}
	for upstreamID, removeAfter := range policies {
		if upstreamID == "" || removeAfter <= 0 {
			return false
		}
	}
	return true
}

func (p *MemoryPool) entryByIDLocked(proxyID string) (Entry, bool) {
	key, ok := p.keyByID[proxyID]
	if !ok {
		return Entry{}, false
	}
	entry, ok := p.entries[key]
	return entry, ok
}

func (p *MemoryPool) setEntryByIDLocked(proxyID string, entry Entry) {
	if key, ok := p.keyByID[proxyID]; ok {
		p.entries[key] = entry
	}
}

func eligibleForExtraction(entry Entry, command extractionDomain.Command) bool {
	if entry.State != proxyDomain.StateAvailable || entry.OwnerWorkerID != "" || !entry.UsableUntil.After(command.Now) {
		return false
	}
	if entry.Proxy.ExpiresAt != nil && entry.Proxy.ExpiresAt.Sub(command.Now) < command.MinRemainingTTL {
		return false
	}
	if command.MaxHealthCheckAge > 0 {
		if entry.Proxy.LastCheckedAt == nil || command.Now.Sub(*entry.Proxy.LastCheckedAt) > command.MaxHealthCheckAge {
			return false
		}
	}
	return matches(command.Protocols, string(entry.Proxy.Scheme)) &&
		matches(command.Regions, entry.Proxy.Tags["region"]) &&
		matches(command.Carriers, entry.Proxy.Tags["carrier"]) &&
		matches(command.Upstreams, entry.Proxy.SourceUpstream)
}

func extractionCandidate(entry Entry) extractionDomain.Candidate {
	expiresAt := time.Time{}
	if entry.Proxy.ExpiresAt != nil {
		expiresAt = *entry.Proxy.ExpiresAt
	}
	checkedAt := time.Time{}
	if entry.Proxy.LastCheckedAt != nil {
		checkedAt = *entry.Proxy.LastCheckedAt
	}
	return extractionDomain.Candidate{
		ID: entry.Proxy.ID, Protocol: string(entry.Proxy.Scheme), Host: entry.Proxy.Host,
		Port: entry.Proxy.Port, Username: entry.Proxy.Username, Region: entry.Proxy.Tags["region"],
		Carrier: entry.Proxy.Tags["carrier"], Upstream: entry.Proxy.SourceUpstream,
		OwnerWorkerID: entry.OwnerWorkerID, State: extractionDomain.Extracted,
		ExpiresAt: expiresAt, LastCheckedAt: checkedAt,
	}
}

func stableProxyID(key string) string {
	digest := sha256.Sum256([]byte(key))
	return "px_" + hex.EncodeToString(digest[:12])
}

func validProxyIdentity(candidate proxyDomain.Proxy) bool {
	if candidate.Host == "" || candidate.Port == 0 {
		return false
	}
	switch candidate.Scheme {
	case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5:
		return true
	default:
		return false
	}
}

func managedActivityState(state proxyDomain.State) bool {
	switch state {
	case proxyDomain.StateFetched, proxyDomain.StateChecking, proxyDomain.StateAvailable,
		proxyDomain.StateSuspect, proxyDomain.StateDraining:
		return true
	default:
		return false
	}
}

func incrementStateInventory(inventory *StateInventory, state proxyDomain.State) {
	if inventory == nil {
		return
	}
	switch state {
	case proxyDomain.StateFetched:
		inventory.Fetched++
	case proxyDomain.StateChecking:
		inventory.Checking++
	case proxyDomain.StateAvailable:
		inventory.Available++
	case proxyDomain.StateSuspect:
		inventory.Suspect++
	case proxyDomain.StateDraining:
		inventory.Draining++
	case proxyDomain.StateUnhealthy:
		inventory.Unhealthy++
	case proxyDomain.StateExtracted:
		inventory.Extracted++
	}
}

func cloneProxy(candidate proxyDomain.Proxy) proxyDomain.Proxy {
	if candidate.ExpiresAt != nil {
		value := *candidate.ExpiresAt
		candidate.ExpiresAt = &value
	}
	if candidate.UsableUntil != nil {
		value := *candidate.UsableUntil
		candidate.UsableUntil = &value
	}
	if candidate.LastCheckedAt != nil {
		value := *candidate.LastCheckedAt
		candidate.LastCheckedAt = &value
	}
	if candidate.LastSuccessAt != nil {
		value := *candidate.LastSuccessAt
		candidate.LastSuccessAt = &value
	}
	if candidate.Tags != nil {
		tags := make(map[string]string, len(candidate.Tags))
		for key, value := range candidate.Tags {
			tags[key] = value
		}
		candidate.Tags = tags
	}
	return candidate
}

func cloneResult(result extractionDomain.Result) extractionDomain.Result {
	result.Items = append([]extractionDomain.Candidate(nil), result.Items...)
	return result
}

func cloneCommand(command extractionDomain.Command) extractionDomain.Command {
	command.Protocols = append([]string(nil), command.Protocols...)
	command.Regions = append([]string(nil), command.Regions...)
	command.Carriers = append([]string(nil), command.Carriers...)
	command.Upstreams = append([]string(nil), command.Upstreams...)
	return command
}

func sameIdempotentRequest(left, right extractionDomain.Command) bool {
	return left.Requested == right.Requested && left.Fulfillment == right.Fulfillment &&
		equalSet(left.Protocols, right.Protocols) && equalSet(left.Regions, right.Regions) &&
		equalSet(left.Carriers, right.Carriers) && equalSet(left.Upstreams, right.Upstreams)
}

func equalSet(left, right []string) bool {
	leftSet := make(map[string]struct{}, len(left))
	for _, value := range left {
		leftSet[value] = struct{}{}
	}
	rightSet := make(map[string]struct{}, len(right))
	for _, value := range right {
		rightSet[value] = struct{}{}
	}
	if len(leftSet) != len(rightSet) {
		return false
	}
	for value := range leftSet {
		if _, ok := rightSet[value]; !ok {
			return false
		}
	}
	return true
}

func matches(allowed []string, value string) bool {
	if len(allowed) == 0 {
		return true
	}
	for _, candidate := range allowed {
		if candidate == value {
			return true
		}
	}
	return false
}

func idempotencyTTL(configured time.Duration) time.Duration {
	if configured > 0 {
		return configured
	}
	return defaultIdempotencyTTL
}

func minTime(left, right time.Time) time.Time {
	if left.Before(right) {
		return left
	}
	return right
}