500 lines
13 KiB
Go
500 lines
13 KiB
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/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
|
|
}
|
|
|
|
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
|
|
|
|
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 != 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)
|
|
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,
|
|
}
|
|
next.buildIndexes()
|
|
s.current.Store(next)
|
|
return nil
|
|
}
|
|
|
|
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 {
|
|
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.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
|
|
}
|