proxy-pool/internal/gateway/snapshot/store.go

package snapshot

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	proxyDomain "proxy-pool/internal/domain/proxy"
	"proxy-pool/internal/domain/routing"
	"proxy-pool/internal/domain/workerruntime"
)

var (
	ErrWrongTarget          = errors.New("snapshot targets another cluster or worker")
	ErrResyncRequired       = errors.New("snapshot sequence requires a full resync")
	ErrChecksumMismatch     = errors.New("snapshot checksum mismatch")
	ErrInvalidRuntimeReport = errors.New("invalid worker runtime report")
	ErrInvalidRuntimeLimit  = errors.New("invalid snapshot runtime limit")
	ErrRuntimeLimitExceeded = errors.New("snapshot runtime limit exceeded")
)

const defaultRuntimeLimit = 1_000_000

const activeRuntimeShardCount = 64

type activeRuntimeShard struct {
	mu      sync.Mutex
	entries map[string]*proxyDomain.Capacity
}

type activeRuntimeIndex [activeRuntimeShardCount]activeRuntimeShard

func (index *activeRuntimeIndex) track(proxyID string, runtime *proxyDomain.Capacity, nonzero bool) {
	shard := &index[activeRuntimeShardIndex(proxyID)]
	shard.mu.Lock()
	defer shard.mu.Unlock()
	if nonzero {
		if shard.entries == nil {
			shard.entries = make(map[string]*proxyDomain.Capacity)
		}
		shard.entries[proxyID] = runtime
		return
	}
	active, reserved, _ := runtime.Counters()
	if active == 0 && reserved == 0 && shard.entries[proxyID] == runtime {
		delete(shard.entries, proxyID)
	}
}

func (index *activeRuntimeIndex) remove(proxyID string, runtime *proxyDomain.Capacity) {
	shard := &index[activeRuntimeShardIndex(proxyID)]
	shard.mu.Lock()
	defer shard.mu.Unlock()
	if shard.entries[proxyID] == runtime {
		delete(shard.entries, proxyID)
	}
}

func (index *activeRuntimeIndex) rangeEntries(visit func(string, *proxyDomain.Capacity)) {
	for shardIndex := range index {
		shard := &index[shardIndex]
		shard.mu.Lock()
		for proxyID, runtime := range shard.entries {
			visit(proxyID, runtime)
		}
		shard.mu.Unlock()
	}
}

func activeRuntimeShardIndex(proxyID string) uint64 {
	const (
		offset = uint64(14695981039346656037)
		prime  = uint64(1099511628211)
	)
	hash := offset
	for index := 0; index < len(proxyID); index++ {
		hash ^= uint64(proxyID[index])
		hash *= prime
	}
	return hash % activeRuntimeShardCount
}

type runtimeRegistration struct {
	capacity *proxyDomain.Capacity
	current  atomic.Bool
}

type Envelope struct {
	ClusterID  string
	WorkerID   string
	Epoch      uint64
	Version    uint64
	Full       bool
	Checksum   string
	ValidUntil time.Time
	Proxies    []proxyDomain.Proxy
	Routing    []routing.Rule
}

type Entry struct {
	Proxy   proxyDomain.Proxy
	Runtime *proxyDomain.Capacity
}

type View struct {
	ClusterID  string
	WorkerID   string
	Epoch      uint64
	Version    uint64
	Checksum   string
	ValidUntil time.Time
	Entries    []Entry
	routing    *routing.RuleSet

	all        []int
	byScheme   map[proxyDomain.Scheme][]int
	byUpstream map[string][]int
	byTag      map[string][]int
}

type Query struct {
	Now          time.Time
	Scheme       proxyDomain.Scheme
	Upstreams    []string
	RequiredTags map[string]string
	Exclude      map[string]struct{}
	SafetyMargin time.Duration
}

type Selection struct {
	view  *View
	base  []int
	query Query
}

type Store struct {
	clusterID string
	workerID  string
	current   atomic.Pointer[View]

	mu       sync.Mutex
	runtimes map[string]*runtimeRegistration
	active   activeRuntimeIndex
	limit    int
}

func NewStore(clusterID, workerID string) *Store {
	return &Store{
		clusterID: clusterID,
		workerID:  workerID,
		runtimes:  make(map[string]*runtimeRegistration),
		limit:     defaultRuntimeLimit,
	}
}

func NewStoreWithRuntimeLimit(clusterID, workerID string, limit int) (*Store, error) {
	if limit <= 0 {
		return nil, ErrInvalidRuntimeLimit
	}
	store := NewStore(clusterID, workerID)
	store.limit = limit
	return store, nil
}

func (s *Store) Current() *View {
	if s == nil {
		return nil
	}
	return s.current.Load()
}

func (s *Store) RuntimeReport(sessionID string, sequence uint64, observedAt time.Time) (workerruntime.Report, error) {
	if s == nil || strings.TrimSpace(sessionID) != sessionID || sessionID == "" || sequence == 0 || observedAt.IsZero() {
		return workerruntime.Report{}, ErrInvalidRuntimeReport
	}
	current := s.current.Load()
	if current == nil {
		return workerruntime.Report{}, ErrInvalidRuntimeReport
	}
	visible := make(map[string]struct{}, len(current.Entries))
	counters := make([]workerruntime.Counter, 0)
	for _, entry := range current.Entries {
		visible[entry.Proxy.ID] = struct{}{}
		active, reserved, _ := entry.Runtime.Counters()
		if active == 0 && reserved == 0 {
			continue
		}
		counters = append(counters, workerruntime.Counter{
			ProxyID: entry.Proxy.ID, Active: active, Reserved: reserved,
			Draining: entry.Proxy.State == proxyDomain.StateDraining,
		})
	}
	s.active.rangeEntries(func(proxyID string, runtime *proxyDomain.Capacity) {
		if _, currentProxy := visible[proxyID]; currentProxy {
			return
		}
		active, reserved, _ := runtime.Counters()
		if active == 0 && reserved == 0 {
			return
		}
		counters = append(counters, workerruntime.Counter{
			ProxyID: proxyID, Active: active, Reserved: reserved, Draining: true,
		})
	})
	sort.Slice(counters, func(left, right int) bool {
		return counters[left].ProxyID < counters[right].ProxyID
	})
	return workerruntime.Report{
		WorkerID: s.workerID, SessionID: sessionID, Sequence: sequence,
		SnapshotVersion: current.Version, OwnershipEpoch: current.Epoch,
		ObservedAt: observedAt.UTC(), Counters: counters,
	}, nil
}

func (s *Store) Apply(envelope Envelope) error {
	if s == nil {
		return fmt.Errorf("apply snapshot: nil store")
	}
	if envelope.ClusterID != s.clusterID || envelope.WorkerID != s.workerID {
		return ErrWrongTarget
	}
	if !envelope.Full || envelope.Epoch == 0 || envelope.Version == 0 {
		return ErrResyncRequired
	}
	if envelope.Checksum != ChecksumWithRouting(envelope.Proxies, envelope.Routing) {
		return ErrChecksumMismatch
	}
	routes, err := routing.Compile(envelope.Routing)
	if err != nil {
		return fmt.Errorf("compile snapshot routing: %w", err)
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	current := s.current.Load()
	if current != nil {
		switch {
		case envelope.Epoch < current.Epoch:
			return ErrResyncRequired
		case envelope.Epoch == current.Epoch && envelope.Version != current.Version+1:
			return ErrResyncRequired
		case envelope.Epoch > current.Epoch && envelope.Version != 1:
			return ErrResyncRequired
		}
	}

	proxies := cloneAndSort(envelope.Proxies)
	nextProxyIDs := make(map[string]struct{}, len(proxies))
	for _, descriptor := range proxies {
		nextProxyIDs[descriptor.ID] = struct{}{}
	}

	retiredCurrent := make(map[string]*runtimeRegistration)
	if current != nil {
		for _, entry := range current.Entries {
			if _, retained := nextProxyIDs[entry.Proxy.ID]; retained {
				continue
			}
			registration := s.runtimes[entry.Proxy.ID]
			registration.capacity.SetReservationEnabled(false)
			registration.current.Store(false)
			registration.capacity.SetActivityObservationEnabled(true)
			active, reserved, _ := registration.capacity.Counters()
			s.active.track(entry.Proxy.ID, registration.capacity, active+reserved > 0)
			retiredCurrent[entry.Proxy.ID] = registration
		}
	}

	reclaimable := s.reclaimableRuntimes(nextProxyIDs)
	newRuntimeCount := 0
	for _, descriptor := range proxies {
		if s.runtimes[descriptor.ID] == nil {
			newRuntimeCount++
		}
	}
	if len(s.runtimes)-len(reclaimable)+newRuntimeCount > s.limit {
		s.restoreCurrentRuntimes(retiredCurrent)
		return ErrRuntimeLimitExceeded
	}
	for proxyID, registration := range reclaimable {
		s.active.remove(proxyID, registration.capacity)
		delete(s.runtimes, proxyID)
	}
	entries := make([]Entry, 0, len(proxies))
	for _, descriptor := range proxies {
		registration := s.runtimes[descriptor.ID]
		if registration == nil {
			proxyID := descriptor.ID
			registration = &runtimeRegistration{}
			registration.capacity = proxyDomain.NewCapacityWithActivityObserver(descriptor.MaxConcurrency, func(nonzero bool) {
				if registration.current.Load() {
					return
				}
				s.active.track(proxyID, registration.capacity, nonzero)
			})
		} else {
			registration.capacity.SetMax(descriptor.MaxConcurrency)
		}
		registration.current.Store(true)
		registration.capacity.SetActivityObservationEnabled(false)
		registration.capacity.SetReservationEnabled(true)
		s.active.remove(descriptor.ID, registration.capacity)
		// A retired runtime remains until Active/Reserved reaches zero. Once it is
		// quiescent, the next Apply reclaims it before enforcing the registry limit.
		s.runtimes[descriptor.ID] = registration
		entries = append(entries, Entry{Proxy: descriptor, Runtime: registration.capacity})
	}

	next := &View{
		ClusterID:  envelope.ClusterID,
		WorkerID:   envelope.WorkerID,
		Epoch:      envelope.Epoch,
		Version:    envelope.Version,
		Checksum:   envelope.Checksum,
		ValidUntil: envelope.ValidUntil.UTC(),
		Entries:    entries,
		routing:    routes,
	}
	next.buildIndexes()
	s.current.Store(next)
	return nil
}

// MatchRouting matches a request against the immutable rules published with
// this proxy view. A View is the atomic consistency boundary for both sets.
func (v *View) MatchRouting(request routing.Request) (routing.Rule, bool) {
	if v == nil || v.routing == nil {
		return routing.Rule{}, false
	}
	return v.routing.Match(request)
}

// UpstreamLoad returns the current in-flight load and whether an eligible
// capacity slot remains for one Upstream under the supplied request filters.
func (v *View) UpstreamLoad(query Query, upstream string) (int64, bool) {
	if v == nil || upstream == "" {
		return 0, false
	}
	query.Upstreams = []string{upstream}
	selection := v.Select(query)
	var active int64
	available := false
	for index := 0; index < selection.Len(); index++ {
		entry, ok := selection.EntryAt(index)
		if !ok {
			continue
		}
		entryActive, reserved, maximum := entry.Runtime.Counters()
		active += entryActive + reserved
		if entryActive+reserved < maximum {
			available = true
		}
	}
	return active, available
}

func (s *Store) reclaimableRuntimes(nextProxyIDs map[string]struct{}) map[string]*runtimeRegistration {
	result := make(map[string]*runtimeRegistration)
	for proxyID, registration := range s.runtimes {
		if _, retained := nextProxyIDs[proxyID]; retained || registration.current.Load() || !registration.capacity.Reclaimable() {
			continue
		}
		result[proxyID] = registration
	}
	return result
}

func (s *Store) restoreCurrentRuntimes(retired map[string]*runtimeRegistration) {
	for proxyID, registration := range retired {
		registration.current.Store(true)
		registration.capacity.SetActivityObservationEnabled(false)
		registration.capacity.SetReservationEnabled(true)
		s.active.remove(proxyID, registration.capacity)
	}
}

func (v *View) Select(query Query) Selection {
	if query.Now.IsZero() {
		query.Now = time.Now().UTC()
	}
	if v == nil || (!v.ValidUntil.IsZero() && !v.ValidUntil.After(query.Now)) {
		return Selection{}
	}

	base := v.all
	if query.Scheme != "" {
		base = chooseSmaller(base, v.byScheme[query.Scheme])
	}
	if len(query.Upstreams) == 1 {
		base = chooseSmaller(base, v.byUpstream[query.Upstreams[0]])
	} else if len(query.Upstreams) > 1 {
		if merged := v.unionUpstreams(query.Upstreams); len(merged) > 0 {
			base = chooseSmaller(base, merged)
		}
	}
	for key, value := range query.RequiredTags {
		base = chooseSmaller(base, v.byTag[tagKey(key, value)])
	}
	return Selection{view: v, base: base, query: query}
}

func (s Selection) Len() int {
	return len(s.base)
}

func (s Selection) EntryAt(index int) (Entry, bool) {
	if s.view == nil || index < 0 || index >= len(s.base) {
		return Entry{}, false
	}
	entry := s.view.Entries[s.base[index]]
	if !matchesQuery(entry.Proxy, s.query) {
		return Entry{}, false
	}
	return entry, true
}

func Checksum(proxies []proxyDomain.Proxy) string {
	return ChecksumWithRouting(proxies, nil)
}

func ChecksumWithRouting(proxies []proxyDomain.Proxy, rules []routing.Rule) string {
	payload := struct {
		Proxies []proxyDomain.Proxy `json:"proxies"`
		Routing []routing.Rule      `json:"routing"`
	}{
		Proxies: cloneAndSort(proxies),
		Routing: cloneRoutingRules(rules),
	}
	encoded, err := json.Marshal(payload)
	if err != nil {
		panic(fmt.Sprintf("encode snapshot checksum: %v", err))
	}
	digest := sha256.Sum256(encoded)
	return hex.EncodeToString(digest[:])
}

func cloneRoutingRules(source []routing.Rule) []routing.Rule {
	result := make([]routing.Rule, len(source))
	for index, rule := range source {
		result[index] = rule
		result[index].WaitTimeout = rule.WaitTimeout
		result[index].Match.Methods = append([]string(nil), rule.Match.Methods...)
		result[index].Upstreams = append([]string(nil), rule.Upstreams...)
		if rule.Match.Headers != nil {
			result[index].Match.Headers = make(map[string]string, len(rule.Match.Headers))
			for name, value := range rule.Match.Headers {
				result[index].Match.Headers[name] = value
			}
		}
		if rule.Strategy.Weights != nil {
			result[index].Strategy.Weights = make(map[string]uint32, len(rule.Strategy.Weights))
			for upstream, weight := range rule.Strategy.Weights {
				result[index].Strategy.Weights[upstream] = weight
			}
		}
	}
	return result
}

func cloneAndSort(source []proxyDomain.Proxy) []proxyDomain.Proxy {
	cloned := make([]proxyDomain.Proxy, len(source))
	for index, descriptor := range source {
		cloned[index] = descriptor
		if descriptor.Tags != nil {
			cloned[index].Tags = make(map[string]string, len(descriptor.Tags))
			for key, value := range descriptor.Tags {
				cloned[index].Tags[key] = value
			}
		}
	}
	sort.Slice(cloned, func(i, j int) bool {
		return cloned[i].ID < cloned[j].ID
	})
	return cloned
}

func (v *View) buildIndexes() {
	if v == nil {
		return
	}
	v.all = make([]int, len(v.Entries))
	v.byScheme = make(map[proxyDomain.Scheme][]int)
	v.byUpstream = make(map[string][]int)
	v.byTag = make(map[string][]int)
	for index, entry := range v.Entries {
		v.all[index] = index
		v.byScheme[entry.Proxy.Scheme] = append(v.byScheme[entry.Proxy.Scheme], index)
		v.byUpstream[entry.Proxy.SourceUpstream] = append(v.byUpstream[entry.Proxy.SourceUpstream], index)
		for key, value := range entry.Proxy.Tags {
			v.byTag[tagKey(key, value)] = append(v.byTag[tagKey(key, value)], index)
		}
	}
}

func (v *View) unionUpstreams(upstreams []string) []int {
	total := 0
	for _, upstream := range upstreams {
		total += len(v.byUpstream[upstream])
	}
	if total == 0 {
		return nil
	}
	merged := make([]int, 0, total)
	for _, upstream := range upstreams {
		merged = append(merged, v.byUpstream[upstream]...)
	}
	sort.Ints(merged)
	return merged
}

func chooseSmaller(current, candidate []int) []int {
	if len(current) == 0 {
		return candidate
	}
	if len(candidate) == 0 {
		return candidate
	}
	if len(candidate) < len(current) {
		return candidate
	}
	return current
}

func matchesQuery(candidate proxyDomain.Proxy, query Query) bool {
	if candidate.State != proxyDomain.StateAvailable {
		return false
	}
	if query.Scheme != "" && candidate.Scheme != query.Scheme {
		return false
	}
	if _, excluded := query.Exclude[candidate.ID]; excluded {
		return false
	}
	if candidate.UsableUntil != nil && !candidate.UsableUntil.After(query.Now) {
		return false
	}
	if candidate.ExpiresAt != nil && !candidate.ExpiresAt.After(query.Now.Add(query.SafetyMargin)) {
		return false
	}
	if !contains(query.Upstreams, candidate.SourceUpstream) {
		return false
	}
	for key, value := range query.RequiredTags {
		if candidate.Tags[key] != value {
			return false
		}
	}
	return true
}

func contains(allowed []string, value string) bool {
	if len(allowed) == 0 {
		return true
	}
	for _, candidate := range allowed {
		if candidate == value {
			return true
		}
	}
	return false
}

func tagKey(key, value string) string {
	return key + "\x00" + value
}