381 lines
11 KiB
Go
381 lines
11 KiB
Go
package workerruntime
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type MemoryStore struct {
|
|
mu sync.Mutex
|
|
now func() time.Time
|
|
epoch uint64
|
|
sessions map[string]memorySession
|
|
references map[string]memoryReference
|
|
reports map[string]memoryReport
|
|
}
|
|
|
|
type memorySession struct {
|
|
value Session
|
|
expiresAt time.Time
|
|
}
|
|
|
|
type memoryReference struct {
|
|
value SnapshotReference
|
|
expiresAt time.Time
|
|
}
|
|
|
|
type memoryReport struct {
|
|
value Report
|
|
digest [sha256.Size]byte
|
|
expiresAt time.Time
|
|
counters map[string]Counter
|
|
}
|
|
|
|
var (
|
|
_ ControlStore = (*MemoryStore)(nil)
|
|
_ SessionWriter = (*MemoryStore)(nil)
|
|
_ ReportWriter = (*MemoryStore)(nil)
|
|
_ RuntimeReader = (*MemoryStore)(nil)
|
|
)
|
|
|
|
func NewMemoryStore(now func() time.Time) (*MemoryStore, error) {
|
|
if now == nil {
|
|
return nil, ErrInvalidStore
|
|
}
|
|
return &MemoryStore{
|
|
now: now, epoch: 1, sessions: make(map[string]memorySession),
|
|
references: make(map[string]memoryReference), reports: make(map[string]memoryReport),
|
|
}, nil
|
|
}
|
|
|
|
func (store *MemoryStore) CurrentOwnershipEpoch(ctx context.Context) (uint64, error) {
|
|
if ctx == nil || store == nil {
|
|
return 0, ErrInvalidStore
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
if store.epoch == 0 {
|
|
return 0, ErrInvalidStore
|
|
}
|
|
return store.epoch, nil
|
|
}
|
|
|
|
// OpenSession always replaces the previous Worker session and clears Runtime.
|
|
func (store *MemoryStore) OpenSession(ctx context.Context, session Session, ttl time.Duration) error {
|
|
if ctx == nil || store == nil || ttl <= 0 {
|
|
return ErrInvalidSession
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := NormalizeSession(session)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
delete(store.reports, normalized.WorkerID)
|
|
delete(store.references, normalized.WorkerID)
|
|
store.sessions[normalized.WorkerID] = memorySession{value: normalized, expiresAt: now.Add(ttl)}
|
|
return nil
|
|
}
|
|
|
|
func (store *MemoryStore) ValidateSession(ctx context.Context, workerID, sessionID string) error {
|
|
if ctx == nil || store == nil || !ValidIdentifier(workerID) || !ValidIdentifier(sessionID) {
|
|
return ErrInvalidSession
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
session, exists := store.sessions[workerID]
|
|
if !exists || !session.expiresAt.After(now) || session.value.SessionID != sessionID {
|
|
return ErrStaleSession
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (store *MemoryStore) RecordIssuedSnapshot(ctx context.Context, sessionID string, reference SnapshotReference, ttl time.Duration) error {
|
|
if ctx == nil || store == nil || ttl <= 0 {
|
|
return ErrInvalidSnapshotReference
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := NormalizeSnapshotReference(reference)
|
|
if err != nil || !ValidIdentifier(sessionID) {
|
|
return ErrInvalidSnapshotReference
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
session, exists := store.sessions[normalized.WorkerID]
|
|
if !exists || !session.expiresAt.After(now) || session.value.SessionID != sessionID {
|
|
return ErrStaleSession
|
|
}
|
|
if normalized.OwnershipEpoch != store.epoch {
|
|
return ErrSnapshotMismatch
|
|
}
|
|
current, exists := store.references[normalized.WorkerID]
|
|
if exists && !current.expiresAt.After(now) {
|
|
delete(store.references, normalized.WorkerID)
|
|
exists = false
|
|
}
|
|
if exists {
|
|
switch compareSnapshotTuple(normalized, current.value) {
|
|
case -1:
|
|
return ErrStaleSnapshotReference
|
|
case 0:
|
|
if normalized.Checksum != current.value.Checksum {
|
|
return ErrConflictingSnapshotReference
|
|
}
|
|
}
|
|
}
|
|
store.references[normalized.WorkerID] = memoryReference{value: normalized, expiresAt: now.Add(ttl)}
|
|
return nil
|
|
}
|
|
|
|
func (store *MemoryStore) AcknowledgeSnapshot(ctx context.Context, acknowledgement SnapshotAcknowledgement, ttl time.Duration) error {
|
|
if ctx == nil || store == nil || ttl <= 0 {
|
|
return ErrInvalidAcknowledgement
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := NormalizeAcknowledgement(acknowledgement)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
session, exists := store.sessions[normalized.WorkerID]
|
|
if !exists || !session.expiresAt.After(now) || session.value.SessionID != normalized.SessionID {
|
|
return ErrStaleSession
|
|
}
|
|
if session.value.AckedSnapshotVersion != 0 {
|
|
acknowledged := sessionReference(session.value)
|
|
switch compareSnapshotTuple(normalized.Reference, acknowledged) {
|
|
case -1:
|
|
return ErrStaleAcknowledgement
|
|
case 0:
|
|
if normalized.Reference.Checksum != acknowledged.Checksum {
|
|
return ErrSnapshotMismatch
|
|
}
|
|
}
|
|
}
|
|
reference, exists := store.references[normalized.WorkerID]
|
|
if !exists || !reference.expiresAt.After(now) {
|
|
return ErrSnapshotMismatch
|
|
}
|
|
switch compareSnapshotTuple(normalized.Reference, reference.value) {
|
|
case -1:
|
|
return ErrStaleAcknowledgement
|
|
case 1:
|
|
return ErrSnapshotMismatch
|
|
}
|
|
if normalized.Reference.Checksum != reference.value.Checksum {
|
|
return ErrSnapshotMismatch
|
|
}
|
|
if !normalized.Applied {
|
|
delete(store.reports, normalized.WorkerID)
|
|
session.value.RuntimeEnabled = false
|
|
session.expiresAt = now.Add(ttl)
|
|
store.sessions[normalized.WorkerID] = session
|
|
return nil
|
|
}
|
|
if session.value.AckedSnapshotVersion != 0 && compareSnapshotTuple(normalized.Reference, sessionReference(session.value)) == 0 {
|
|
if !session.value.RuntimeEnabled {
|
|
delete(store.reports, normalized.WorkerID)
|
|
session.value.RuntimeEnabled = true
|
|
}
|
|
session.expiresAt = now.Add(ttl)
|
|
store.sessions[normalized.WorkerID] = session
|
|
return nil
|
|
}
|
|
delete(store.reports, normalized.WorkerID)
|
|
session.value.AckedSnapshotVersion = normalized.Reference.Version
|
|
session.value.AckedOwnershipEpoch = normalized.Reference.OwnershipEpoch
|
|
session.value.AckedChecksum = normalized.Reference.Checksum
|
|
session.value.RuntimeEnabled = true
|
|
session.expiresAt = now.Add(ttl)
|
|
store.sessions[normalized.WorkerID] = session
|
|
return nil
|
|
}
|
|
|
|
// ReplaceSession is retained temporarily for the pre-control-plane adapters.
|
|
func (store *MemoryStore) ReplaceSession(ctx context.Context, session Session, ttl time.Duration) error {
|
|
if ctx == nil || store == nil || !validLegacySession(session) || ttl <= 0 {
|
|
return ErrInvalidSession
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
current, exists := store.sessions[session.WorkerID]
|
|
identityChanged := exists && (current.value.SessionID != session.SessionID || current.value.InstanceID != session.InstanceID)
|
|
expired := exists && !current.expiresAt.After(now)
|
|
if exists && !identityChanged && !expired && legacySessionBefore(session, current.value) {
|
|
return ErrStaleSession
|
|
}
|
|
ackAdvanced := exists && !identityChanged && !expired && legacySessionAfter(session, current.value)
|
|
if identityChanged || expired || ackAdvanced {
|
|
delete(store.reports, session.WorkerID)
|
|
}
|
|
session.RuntimeEnabled = true
|
|
store.sessions[session.WorkerID] = memorySession{value: session, expiresAt: now.Add(ttl)}
|
|
return nil
|
|
}
|
|
|
|
func (store *MemoryStore) ReplaceRuntime(ctx context.Context, report Report, ttl time.Duration) error {
|
|
if ctx == nil || store == nil || ttl <= 0 {
|
|
return ErrInvalidReport
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, digest, err := NormalizeReport(report)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
session, exists := store.sessions[normalized.WorkerID]
|
|
if !exists || !session.expiresAt.After(now) || session.value.SessionID != normalized.SessionID {
|
|
return ErrStaleSession
|
|
}
|
|
if !session.value.RuntimeEnabled || normalized.SnapshotVersion != session.value.AckedSnapshotVersion ||
|
|
normalized.OwnershipEpoch != session.value.AckedOwnershipEpoch {
|
|
return ErrSnapshotMismatch
|
|
}
|
|
current, exists := store.reports[normalized.WorkerID]
|
|
if exists && !current.expiresAt.After(now) {
|
|
delete(store.reports, normalized.WorkerID)
|
|
exists = false
|
|
}
|
|
if exists && current.value.SessionID == normalized.SessionID {
|
|
switch {
|
|
case normalized.Sequence < current.value.Sequence:
|
|
return ErrStaleReport
|
|
case normalized.Sequence == current.value.Sequence && digest != current.digest:
|
|
return ErrConflictingReport
|
|
case normalized.Sequence == current.value.Sequence:
|
|
current.expiresAt = now.Add(ttl)
|
|
store.reports[normalized.WorkerID] = current
|
|
session.expiresAt = now.Add(ttl)
|
|
store.sessions[normalized.WorkerID] = session
|
|
return nil
|
|
}
|
|
}
|
|
store.reports[normalized.WorkerID] = memoryReport{
|
|
value: normalized, digest: digest, expiresAt: now.Add(ttl), counters: countersByProxy(normalized.Counters),
|
|
}
|
|
session.expiresAt = now.Add(ttl)
|
|
store.sessions[normalized.WorkerID] = session
|
|
return nil
|
|
}
|
|
|
|
func (store *MemoryStore) ReadRuntime(ctx context.Context, proxies []OwnedProxy) ([]Snapshot, error) {
|
|
if ctx == nil || store == nil {
|
|
return nil, ErrInvalidQuery
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
seen := make(map[string]struct{}, len(proxies))
|
|
for _, proxy := range proxies {
|
|
if !ValidIdentifier(proxy.ProxyID) || !ValidIdentifier(proxy.WorkerID) || proxy.OwnershipEpoch == 0 {
|
|
return nil, ErrInvalidQuery
|
|
}
|
|
key := proxy.WorkerID + "\x00" + proxy.ProxyID
|
|
if _, exists := seen[key]; exists {
|
|
return nil, ErrInvalidQuery
|
|
}
|
|
seen[key] = struct{}{}
|
|
}
|
|
now, err := store.currentTime()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
result := make([]Snapshot, len(proxies))
|
|
for index, proxy := range proxies {
|
|
result[index].ProxyID = proxy.ProxyID
|
|
session, sessionExists := store.sessions[proxy.WorkerID]
|
|
report, reportExists := store.reports[proxy.WorkerID]
|
|
if !sessionExists || !reportExists || !session.expiresAt.After(now) || !report.expiresAt.After(now) ||
|
|
!session.value.RuntimeEnabled || report.value.SessionID != session.value.SessionID ||
|
|
report.value.SnapshotVersion != session.value.AckedSnapshotVersion ||
|
|
report.value.OwnershipEpoch != session.value.AckedOwnershipEpoch ||
|
|
report.value.OwnershipEpoch < proxy.OwnershipEpoch {
|
|
continue
|
|
}
|
|
result[index].Fresh = true
|
|
if counter, exists := report.counters[proxy.ProxyID]; exists {
|
|
result[index].Active = counter.Active
|
|
result[index].Reserved = counter.Reserved
|
|
result[index].Draining = counter.Draining
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (store *MemoryStore) currentTime() (time.Time, error) {
|
|
now := store.now().UTC()
|
|
if now.IsZero() {
|
|
return time.Time{}, ErrInvalidStore
|
|
}
|
|
return now, nil
|
|
}
|
|
|
|
func countersByProxy(counters []Counter) map[string]Counter {
|
|
indexed := make(map[string]Counter, len(counters))
|
|
for _, counter := range counters {
|
|
indexed[counter.ProxyID] = counter
|
|
}
|
|
return indexed
|
|
}
|
|
|
|
func validLegacySession(session Session) bool {
|
|
return ValidIdentifier(session.WorkerID) && ValidIdentifier(session.InstanceID) && ValidIdentifier(session.SessionID) &&
|
|
session.AckedSnapshotVersion > 0 && session.AckedOwnershipEpoch > 0
|
|
}
|
|
|
|
func legacySessionBefore(left, right Session) bool {
|
|
return compareSnapshotTuple(sessionReference(left), sessionReference(right)) < 0
|
|
}
|
|
|
|
func legacySessionAfter(left, right Session) bool {
|
|
return compareSnapshotTuple(sessionReference(left), sessionReference(right)) > 0
|
|
}
|