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

package snapshot

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

	proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
)

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")
)

type Envelope struct {
	ClusterID string
	WorkerID  string
	Epoch     uint64
	Version   uint64
	Full      bool
	Checksum  string
	Proxies   []proxyDomain.Proxy
}

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

type View struct {
	ClusterID string
	WorkerID  string
	Epoch     uint64
	Version   uint64
	Checksum  string
	Entries   []Entry

	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]*proxyDomain.Capacity
}

func NewStore(clusterID, workerID string) *Store {
	return &Store{
		clusterID: clusterID,
		workerID:  workerID,
		runtimes:  make(map[string]*proxyDomain.Capacity),
	}
}

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

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 != Checksum(envelope.Proxies) {
		return ErrChecksumMismatch
	}

	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)
	entries := make([]Entry, 0, len(proxies))
	for _, descriptor := range proxies {
		runtime := s.runtimes[descriptor.ID]
		if runtime == nil {
			runtime = proxyDomain.NewCapacity(descriptor.MaxConcurrency)
		} else {
			runtime.SetMax(descriptor.MaxConcurrency)
		}
		// Keep runtimes for temporarily absent IDs. Old immutable views may still
		// hold in-flight leases, so reclaiming here could reset active capacity if
		// the same Proxy reappears in a later snapshot.
		s.runtimes[descriptor.ID] = runtime
		entries = append(entries, Entry{Proxy: descriptor, Runtime: runtime})
	}

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

func (v *View) Select(query Query) Selection {
	if query.Now.IsZero() {
		query.Now = time.Now().UTC()
	}
	if v == nil {
		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 {
	canonical := cloneAndSort(proxies)
	encoded, err := json.Marshal(canonical)
	if err != nil {
		panic(fmt.Sprintf("encode snapshot checksum: %v", err))
	}
	digest := sha256.Sum256(encoded)
	return hex.EncodeToString(digest[:])
}

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