package pool import ( "errors" "fmt" "sync" "proxy-pool/internal/domain/upstream" ) var ( ErrInvalidFetchBudget = errors.New("invalid fetch budget") ErrInvalidFetchCompletion = errors.New("invalid fetch completion") ErrFetchPermitFinished = errors.New("fetch permit is already finished") ErrInvalidManagedRelease = errors.New("invalid managed proxy release") ErrInvalidManagedSynchronization = errors.New("invalid managed proxy synchronization") ErrManagedSynchronizationInFlight = errors.New("managed proxy synchronization has pending fetches") ) type FetchBudgetConfig struct { UpstreamID string MaxSize int ExpectedPerFetch int Managed int } type FetchBudgetSnapshot struct { Managed int PendingExpected int } // FetchBudget owns current-inventory accounting for one upstream. Reserving // the expected response before I/O closes the race between concurrent calls; // distributed cumulative quota belongs to the Provider coordination permit. type FetchBudget struct { mu sync.Mutex upstreamID string maxSize int expected int usage FetchBudgetSnapshot } func NewFetchBudget(config FetchBudgetConfig) (*FetchBudget, error) { if config.UpstreamID == "" || config.MaxSize <= 0 || config.ExpectedPerFetch <= 0 || config.ExpectedPerFetch > config.MaxSize || config.Managed < 0 { return nil, ErrInvalidFetchBudget } return &FetchBudget{ upstreamID: config.UpstreamID, maxSize: config.MaxSize, expected: config.ExpectedPerFetch, usage: FetchBudgetSnapshot{ Managed: config.Managed, }, }, nil } func (b *FetchBudget) ReserveFetch(upstreamID string) (upstream.FetchPermit, bool, error) { if b == nil || upstreamID == "" || upstreamID != b.upstreamID { return nil, false, fmt.Errorf("%w: upstream %q", ErrInvalidFetchBudget, upstreamID) } b.mu.Lock() defer b.mu.Unlock() if !b.canReserveLocked() { return nil, false, nil } b.usage.PendingExpected += b.expected return &fetchPermit{budget: b, expected: b.expected}, true, nil } // FetchAllowance is an advisory snapshot for reconciliation. ReserveFetch is // still the atomic authority immediately before provider I/O. func (b *FetchBudget) FetchAllowance() int { if b == nil { return 0 } b.mu.Lock() defer b.mu.Unlock() if b.canReserveLocked() { return b.expected } return 0 } func (b *FetchBudget) canReserveLocked() bool { poolRoom := b.maxSize - b.usage.Managed - b.usage.PendingExpected return poolRoom >= b.expected } func (b *FetchBudget) Snapshot() FetchBudgetSnapshot { if b == nil { return FetchBudgetSnapshot{} } b.mu.Lock() defer b.mu.Unlock() return b.usage } // SynchronizeManaged replaces the local current-inventory count with the // authoritative activity-store observation. Pending reservations remain owned // by this budget. func (b *FetchBudget) SynchronizeManaged(managed int) error { if b == nil || managed < 0 { return ErrInvalidManagedSynchronization } b.mu.Lock() defer b.mu.Unlock() if b.usage.PendingExpected > 0 { return ErrManagedSynchronizationInFlight } b.usage.Managed = managed return nil } // ReleaseManaged returns current-inventory capacity after extraction, expiry, // or removal. It deliberately does not restore the cumulative fetch quota. func (b *FetchBudget) ReleaseManaged(count int) error { if b == nil || count < 0 { return ErrInvalidManagedRelease } b.mu.Lock() defer b.mu.Unlock() if count > b.usage.Managed { return ErrInvalidManagedRelease } b.usage.Managed -= count return nil } type fetchPermit struct { budget *FetchBudget expected int finished bool } func (p *fetchPermit) Expected() int { if p == nil { return 0 } return p.expected } func (p *fetchPermit) Complete(retained int) error { if p == nil || p.budget == nil { return ErrFetchPermitFinished } if retained < 0 || retained > p.expected { return ErrInvalidFetchCompletion } p.budget.mu.Lock() defer p.budget.mu.Unlock() if p.finished { return ErrFetchPermitFinished } p.finished = true p.budget.usage.PendingExpected -= p.expected p.budget.usage.Managed += retained return nil } func (p *fetchPermit) Cancel() error { if p == nil || p.budget == nil { return ErrFetchPermitFinished } p.budget.mu.Lock() defer p.budget.mu.Unlock() if p.finished { return ErrFetchPermitFinished } p.finished = true p.budget.usage.PendingExpected -= p.expected return nil }