150 lines
3.3 KiB
Go
150 lines
3.3 KiB
Go
package snapshot
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"sort"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
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
|
|
}
|
|
|
|
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)
|
|
s.runtimes[descriptor.ID] = runtime
|
|
} else {
|
|
runtime.SetMax(descriptor.MaxConcurrency)
|
|
}
|
|
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,
|
|
}
|
|
s.current.Store(next)
|
|
return nil
|
|
}
|
|
|
|
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
|
|
}
|