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

package pool

import (
	"errors"
	"fmt"
	"sync"

	"github.com/proxy-pool/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")
)

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
}

// 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
}