proxy-pool/internal/domain/routing/sequential.go

package routing

import (
	"fmt"
	"sync"
)

type EndBehavior string

const (
	EndStayLast EndBehavior = "stayLast"
	EndStop     EndBehavior = "stop"
	EndLoop     EndBehavior = "loop"
)

type EmptyObservation struct {
	Count      int
	Generation uint64
}

type upstreamEmpty struct {
	count      int
	generation uint64
}

// UpstreamEmptyState is shared by every Routing that references an Upstream.
// A generation represents one uninterrupted empty-result episode.
type UpstreamEmptyState struct {
	mu             sync.RWMutex
	nextGeneration uint64
	states         map[string]upstreamEmpty
}

func NewUpstreamEmptyState() *UpstreamEmptyState {
	return &UpstreamEmptyState{states: make(map[string]upstreamEmpty)}
}

func (s *UpstreamEmptyState) ObserveEmpty(upstream string) EmptyObservation {
	s.mu.Lock()
	defer s.mu.Unlock()
	state := s.states[upstream]
	if state.count == 0 {
		s.nextGeneration++
		state.generation = s.nextGeneration
	}
	state.count++
	s.states[upstream] = state
	return EmptyObservation{Count: state.count, Generation: state.generation}
}

func (s *UpstreamEmptyState) ObserveValid(upstream string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	state := s.states[upstream]
	state.count = 0
	s.states[upstream] = state
}

func (s *UpstreamEmptyState) Count(upstream string) int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.states[upstream].count
}

type CursorSnapshot struct {
	Index   int
	Version uint64
	Stopped bool
}

// RoutingCursor owns only per-Routing selection state. Advance uses an
// expected version so simultaneous threshold observers can change it once.
type RoutingCursor struct {
	mu        sync.RWMutex
	state     CursorSnapshot
	processed map[string]uint64
}

func NewRoutingCursor() *RoutingCursor {
	return &RoutingCursor{
		state:     CursorSnapshot{Version: 1},
		processed: make(map[string]uint64),
	}
}

func (c *RoutingCursor) Snapshot() CursorSnapshot {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.state
}

func (c *RoutingCursor) Advance(
	expectedVersion uint64,
	currentUpstream string,
	emptyGeneration uint64,
	upstreamCount int,
	end EndBehavior,
) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.state.Stopped || c.state.Version != expectedVersion || emptyGeneration == 0 ||
		c.processed[currentUpstream] >= emptyGeneration {
		return false
	}
	c.processed[currentUpstream] = emptyGeneration
	if c.state.Index+1 < upstreamCount {
		c.state.Index++
		c.state.Version++
		return true
	}
	switch end {
	case EndStop:
		c.state.Stopped = true
		c.state.Version++
		return true
	case EndLoop:
		if upstreamCount > 1 {
			c.state.Index = 0
			c.state.Version++
			return true
		}
	}
	return false
}

type Sequential struct {
	upstreams []string
	threshold int
	end       EndBehavior
	empty     *UpstreamEmptyState
	cursor    *RoutingCursor
}

func NewSequential(upstreams []string, threshold int) (*Sequential, error) {
	return NewSequentialWithState(upstreams, threshold, EndStayLast, NewUpstreamEmptyState())
}

func NewSequentialWithState(
	upstreams []string,
	threshold int,
	end EndBehavior,
	empty *UpstreamEmptyState,
) (*Sequential, error) {
	if len(upstreams) == 0 {
		return nil, fmt.Errorf("sequential strategy requires at least one upstream")
	}
	if threshold <= 0 {
		return nil, fmt.Errorf("sequential threshold must be greater than zero")
	}
	if end == "" {
		end = EndStayLast
	}
	if end != EndStayLast && end != EndStop && end != EndLoop {
		return nil, fmt.Errorf("sequential end behavior %q is invalid", end)
	}
	if empty == nil {
		return nil, fmt.Errorf("sequential empty state is required")
	}
	for _, upstream := range upstreams {
		if upstream == "" {
			return nil, fmt.Errorf("sequential upstream is required")
		}
	}
	return &Sequential{
		upstreams: append([]string(nil), upstreams...),
		threshold: threshold,
		end:       end,
		empty:     empty,
		cursor:    NewRoutingCursor(),
	}, nil
}

func (s *Sequential) Current() string {
	current, ok, _ := s.CurrentSelection()
	if !ok {
		return ""
	}
	return current
}

func (s *Sequential) CurrentSelection() (upstream string, available bool, version uint64) {
	state := s.cursor.Snapshot()
	if state.Stopped || state.Index < 0 || state.Index >= len(s.upstreams) {
		return "", false, state.Version
	}
	return s.upstreams[state.Index], true, state.Version
}

func (s *Sequential) ObserveEmpty(upstream string) bool {
	observation := s.empty.ObserveEmpty(upstream)
	if observation.Count < s.threshold {
		return false
	}
	current, available, version := s.CurrentSelection()
	if !available || current != upstream {
		return false
	}
	return s.cursor.Advance(version, upstream, observation.Generation, len(s.upstreams), s.end)
}

func (s *Sequential) ObserveValid(upstream string) {
	s.empty.ObserveValid(upstream)
}

func (s *Sequential) EmptyCount(upstream string) int {
	return s.empty.Count(upstream)
}