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 MaxTotal int64 ExpectedPerFetch int Managed int FetchedTotal int64 } type FetchBudgetSnapshot struct { Managed int PendingExpected int FetchedTotal int64 } // FetchBudget owns both current-inventory and cumulative-fetch accounting for // one upstream. Reserving the expected response before I/O closes the race // between concurrent provider calls. type FetchBudget struct { mu sync.Mutex upstreamID string maxSize int maxTotal int64 expected int usage FetchBudgetSnapshot } func NewFetchBudget(config FetchBudgetConfig) (*FetchBudget, error) { if config.UpstreamID == "" || config.MaxSize <= 0 || config.ExpectedPerFetch <= 0 || config.ExpectedPerFetch > config.MaxSize || config.MaxTotal < 0 || config.Managed < 0 || config.FetchedTotal < 0 { return nil, ErrInvalidFetchBudget } if config.MaxTotal > 0 && int64(config.ExpectedPerFetch) > config.MaxTotal { return nil, ErrInvalidFetchBudget } return &FetchBudget{ upstreamID: config.UpstreamID, maxSize: config.MaxSize, maxTotal: config.MaxTotal, expected: config.ExpectedPerFetch, usage: FetchBudgetSnapshot{ Managed: config.Managed, FetchedTotal: config.FetchedTotal, }, }, 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 if poolRoom < b.expected { return false } if b.maxTotal > 0 { totalRoom := b.maxTotal - b.usage.FetchedTotal - int64(b.usage.PendingExpected) return totalRoom >= int64(b.expected) } return true } 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 requests and cumulative // fetch usage 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(fetched, retained int) error { if p == nil || p.budget == nil { return ErrFetchPermitFinished } if fetched < 0 || retained < 0 || retained > fetched || 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 p.budget.usage.FetchedTotal += int64(fetched) 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 }