proxy-pool/internal/gateway/dispatch/dispatcher.go

package dispatch

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	proxyDomain "proxy-pool/internal/domain/proxy"
	"proxy-pool/internal/domain/routing"
	"proxy-pool/internal/gateway/snapshot"
)

var (
	ErrNoCandidate          = errors.New("no local proxy candidate is available")
	ErrPreferredUnavailable = errors.New("preferred local proxy is unavailable")
)

type Request struct {
	RoutingName      string
	Action           routing.Action
	Strategy         routing.Strategy
	OnUnavailable    routing.OnUnavailableAction
	WaitTimeout      time.Duration
	Now              time.Time
	Scheme           proxyDomain.Scheme
	Upstreams        []string
	RequiredTags     map[string]string
	Exclude          map[string]struct{}
	SafetyMargin     time.Duration
	PreferredProxyID string
}

type Lease struct {
	Proxy    proxyDomain.Proxy
	Epoch    uint64
	Version  uint64
	reserved *proxyDomain.Reservation
}

func (l *Lease) Commit() error  { return l.reserved.Commit() }
func (l *Lease) Cancel() error  { return l.reserved.Cancel() }
func (l *Lease) Release() error { return l.reserved.Release() }

type Dispatcher struct {
	store        *snapshot.Store
	cursor       atomic.Uint64
	random       routing.RandomSource
	routingState atomic.Pointer[routingSelectorState]
}

type routingSelectorState struct {
	epoch     uint64
	version   uint64
	selectors sync.Map
}

func New(store *snapshot.Store, randomSources ...routing.RandomSource) *Dispatcher {
	dispatcher := &Dispatcher{store: store}
	if len(randomSources) > 0 {
		dispatcher.random = randomSources[0]
	}
	return dispatcher
}

func (d *Dispatcher) Acquire(request Request) (*Lease, error) {
	if d == nil || d.store == nil {
		return nil, ErrNoCandidate
	}
	view := d.store.Current()
	if view == nil || len(view.Entries) == 0 {
		return nil, ErrNoCandidate
	}
	if request.Now.IsZero() {
		request.Now = time.Now().UTC()
	}
	if request.PreferredProxyID != "" {
		return d.acquirePreferred(view, request)
	}
	if request.Strategy.Type != "" {
		return d.acquireRouted(view, request)
	}
	return d.acquireFromUpstreams(view, request, request.Upstreams)
}

func (d *Dispatcher) acquirePreferred(view *snapshot.View, request Request) (*Lease, error) {
	entry, found := view.EntryByID(request.PreferredProxyID, d.query(request, request.Upstreams))
	if !found {
		return nil, ErrPreferredUnavailable
	}
	reservation, reserved := entry.Runtime.Reserve()
	if !reserved {
		return nil, ErrNoCandidate
	}
	return &Lease{
		Proxy:    entry.Proxy,
		Epoch:    view.Epoch,
		Version:  view.Version,
		reserved: reservation,
	}, nil
}

// AcquireWait retries local snapshot dispatch at a bounded interval until a
// capacity slot appears, the route timeout expires, or the caller cancels.
func (d *Dispatcher) AcquireWait(ctx context.Context, request Request, timeout time.Duration) (*Lease, error) {
	if d == nil || ctx == nil || timeout <= 0 {
		return nil, ErrNoCandidate
	}
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	retry := time.NewTicker(5 * time.Millisecond)
	defer retry.Stop()
	for {
		request.Now = time.Now().UTC()
		lease, err := d.Acquire(request)
		if !errors.Is(err, ErrNoCandidate) {
			return lease, err
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-deadline.C:
			return nil, ErrNoCandidate
		case <-retry.C:
		}
	}
}

func (d *Dispatcher) acquireRouted(view *snapshot.View, request Request) (*Lease, error) {
	if request.Strategy.Type == routing.StrategySequential {
		if request.Strategy.CurrentUpstream == "" {
			return nil, ErrNoCandidate
		}
		return d.acquireFromUpstreams(view, request, []string{request.Strategy.CurrentUpstream})
	}
	candidates, err := d.routingCandidates(view, request)
	if err != nil {
		return nil, ErrNoCandidate
	}
	selector, err := d.routingSelector(view, request)
	if err != nil {
		return nil, ErrNoCandidate
	}
	for range candidates {
		candidate, selectErr := selector.Select(candidates)
		if selectErr != nil {
			return nil, ErrNoCandidate
		}
		lease, acquireErr := d.acquireFromUpstreams(view, request, []string{candidate.Name})
		if acquireErr == nil {
			return lease, nil
		}
		for index := range candidates {
			if candidates[index].Name == candidate.Name {
				candidates[index].Eligible = false
				break
			}
		}
	}
	return nil, ErrNoCandidate
}

func (d *Dispatcher) routingCandidates(view *snapshot.View, request Request) ([]routing.Candidate, error) {
	if len(request.Upstreams) == 0 {
		return nil, ErrNoCandidate
	}
	result := make([]routing.Candidate, 0, len(request.Upstreams))
	seen := make(map[string]struct{}, len(request.Upstreams))
	for _, upstream := range request.Upstreams {
		if upstream == "" {
			return nil, ErrNoCandidate
		}
		if _, duplicate := seen[upstream]; duplicate {
			return nil, ErrNoCandidate
		}
		seen[upstream] = struct{}{}
		candidate := routing.Candidate{Name: upstream, Eligible: true}
		switch request.Strategy.Type {
		case routing.StrategyWeighted:
			weight, exists := request.Strategy.Weights[upstream]
			if !exists || weight == 0 || uint64(weight) > uint64(maxInt()) {
				return nil, ErrNoCandidate
			}
			candidate.Weight = int(weight)
		case routing.StrategyLeastConnections:
			candidate.Active, candidate.Eligible = view.UpstreamLoad(snapshot.Query{
				Now: request.Now, Scheme: request.Scheme, RequiredTags: request.RequiredTags,
				Exclude: request.Exclude, SafetyMargin: request.SafetyMargin,
			}, upstream)
		}
		result = append(result, candidate)
	}
	return result, nil
}

func (d *Dispatcher) routingSelector(view *snapshot.View, request Request) (routing.Selector, error) {
	state := d.selectorState(view)
	key := request.RoutingName + "\x00" + string(request.Strategy.Type)
	if existing, found := state.selectors.Load(key); found {
		return existing.(routing.Selector), nil
	}
	var selector routing.Selector
	switch request.Strategy.Type {
	case routing.StrategyRandom:
		selector = routing.NewRandom(d.random)
	case routing.StrategyRoundRobin:
		selector = routing.NewRoundRobin()
	case routing.StrategyWeighted:
		selector = routing.NewWeighted(d.random)
	case routing.StrategyLeastConnections:
		selector = routing.NewLeastConnections()
	default:
		return nil, ErrNoCandidate
	}
	actual, _ := state.selectors.LoadOrStore(key, selector)
	return actual.(routing.Selector), nil
}

func (d *Dispatcher) selectorState(view *snapshot.View) *routingSelectorState {
	for {
		current := d.routingState.Load()
		if current != nil && current.epoch == view.Epoch && current.version == view.Version {
			return current
		}
		next := &routingSelectorState{epoch: view.Epoch, version: view.Version}
		if d.routingState.CompareAndSwap(current, next) {
			return next
		}
	}
}

func (d *Dispatcher) acquireFromUpstreams(view *snapshot.View, request Request, upstreams []string) (*Lease, error) {
	selection := view.Select(d.query(request, upstreams))
	if selection.Len() == 0 {
		return nil, ErrNoCandidate
	}

	start := int((d.cursor.Add(1) - 1) % uint64(selection.Len()))
	for offset := 0; offset < selection.Len(); offset++ {
		entry, ok := selection.EntryAt((start + offset) % selection.Len())
		if !ok {
			continue
		}
		reservation, ok := entry.Runtime.Reserve()
		if !ok {
			continue
		}
		return &Lease{
			Proxy:    entry.Proxy,
			Epoch:    view.Epoch,
			Version:  view.Version,
			reserved: reservation,
		}, nil
	}
	return nil, ErrNoCandidate
}

func (d *Dispatcher) query(request Request, upstreams []string) snapshot.Query {
	return snapshot.Query{
		Now:          request.Now,
		Scheme:       request.Scheme,
		Upstreams:    upstreams,
		RequiredTags: request.RequiredTags,
		Exclude:      request.Exclude,
		SafetyMargin: request.SafetyMargin,
	}
}

func maxInt() int {
	return int(^uint(0) >> 1)
}