proxy-pool/internal/controller/pool/fetch_budget.go

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
}