319 lines
13 KiB
Go
319 lines
13 KiB
Go
package workerruntime
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestMemoryStoreRequiresAcknowledgedSnapshotForRuntime(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
openControlSession(t, store, "session-a", time.Minute)
|
|
epoch, err := store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
t.Fatalf("CurrentOwnershipEpoch(): %v", err)
|
|
}
|
|
report := controlReport(now, "session-a", 1, 7, epoch)
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
|
|
t.Fatalf("ReplaceRuntime(before ACK) error = %v, want ErrSnapshotMismatch", err)
|
|
}
|
|
reference := controlReference(7, epoch, "snapshot-7")
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", reference, time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(): %v", err)
|
|
}
|
|
if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{
|
|
WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true,
|
|
}, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(after ACK): %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreNegativeAcknowledgementFencesDelayedRuntime(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
openControlSession(t, store, "session-a", time.Minute)
|
|
epoch, err := store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
t.Fatalf("CurrentOwnershipEpoch(): %v", err)
|
|
}
|
|
first := controlReference(7, epoch, "snapshot-7")
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", first, time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(first): %v", err)
|
|
}
|
|
if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: first, Applied: true}, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(first): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 1, 7, epoch), time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(first): %v", err)
|
|
}
|
|
|
|
second := controlReference(8, epoch, "snapshot-8")
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", second, time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(second): %v", err)
|
|
}
|
|
if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: second, Applied: false, ErrorCode: "apply_failed"}, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(negative): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 2, 7, epoch), time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
|
|
t.Fatalf("ReplaceRuntime(delayed): %v, want ErrSnapshotMismatch", err)
|
|
}
|
|
if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: second, Applied: true}, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(recover): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 2, 8, epoch), time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(recovered): %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreAcknowledgementReplayPreservesRuntimeFence(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
openControlSession(t, store, "session-a", time.Minute)
|
|
epoch, err := store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
t.Fatalf("CurrentOwnershipEpoch(): %v", err)
|
|
}
|
|
reference := controlReference(7, epoch, "snapshot-7")
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", reference, time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(): %v", err)
|
|
}
|
|
ack := SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true}
|
|
if err := store.AcknowledgeSnapshot(ctx, ack, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(first): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 100, 7, epoch), time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(): %v", err)
|
|
}
|
|
if err := store.AcknowledgeSnapshot(ctx, ack, time.Minute); err != nil {
|
|
t.Fatalf("AcknowledgeSnapshot(replay): %v", err)
|
|
}
|
|
got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: epoch}})
|
|
if err != nil || len(got) != 1 || !got[0].Fresh {
|
|
t.Fatalf("ReadRuntime(after ACK replay) = %+v, %v", got, err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 99, 7, epoch), time.Minute); !errors.Is(err, ErrStaleReport) {
|
|
t.Fatalf("ReplaceRuntime(stale): %v, want ErrStaleReport", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreRejectsConflictingSnapshotReference(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
openControlSession(t, store, "session-a", time.Minute)
|
|
epoch, err := store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
t.Fatalf("CurrentOwnershipEpoch(): %v", err)
|
|
}
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", controlReference(7, epoch, "first"), time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(first): %v", err)
|
|
}
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", controlReference(7, epoch, "second"), time.Minute); !errors.Is(err, ErrConflictingSnapshotReference) {
|
|
t.Fatalf("RecordIssuedSnapshot(conflict): %v, want ErrConflictingSnapshotReference", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreFencesIssuedSnapshotsAfterSessionReplacement(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
openControlSession(t, store, "session-a", time.Minute)
|
|
epoch, err := store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
t.Fatalf("CurrentOwnershipEpoch(): %v", err)
|
|
}
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", controlReference(3, epoch, "old"), time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(old): %v", err)
|
|
}
|
|
openControlSession(t, store, "session-b", time.Minute)
|
|
if err := store.ValidateSession(ctx, "worker-a", "session-a"); !errors.Is(err, ErrStaleSession) {
|
|
t.Fatalf("ValidateSession(old) = %v, want ErrStaleSession", err)
|
|
}
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-a", controlReference(4, epoch, "stale"), time.Minute); !errors.Is(err, ErrStaleSession) {
|
|
t.Fatalf("RecordIssuedSnapshot(stale) = %v, want ErrStaleSession", err)
|
|
}
|
|
if err := store.RecordIssuedSnapshot(ctx, "session-b", controlReference(1, epoch, "new"), time.Minute); err != nil {
|
|
t.Fatalf("RecordIssuedSnapshot(new): %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreReplacesSparseRuntimeAndClearsMissingCounters(t *testing.T) {
|
|
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
registerRuntimeSession(t, store, "session-a", time.Minute)
|
|
report := Report{
|
|
WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
|
|
SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
|
|
Counters: []Counter{{ProxyID: "proxy-a", Active: 2, Reserved: 1}},
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(first): %v", err)
|
|
}
|
|
got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
|
|
if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Active: 2, Reserved: 1, Fresh: true}) {
|
|
t.Fatalf("ReadRuntime(first) = %+v, %v", got, err)
|
|
}
|
|
|
|
report.Sequence = 2
|
|
report.ObservedAt = now.Add(time.Second)
|
|
report.Counters = nil
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(empty): %v", err)
|
|
}
|
|
got, err = store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
|
|
if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Fresh: true}) {
|
|
t.Fatalf("ReadRuntime(empty) = %+v, %v", got, err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreFencesSessionsAndReportSequence(t *testing.T) {
|
|
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
registerRuntimeSession(t, store, "session-a", time.Minute)
|
|
report := Report{
|
|
WorkerID: "worker-a", SessionID: "session-a", Sequence: 2,
|
|
SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
|
|
Counters: []Counter{{ProxyID: "proxy-a", Active: 1}},
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(first): %v", err)
|
|
}
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(idempotent): %v", err)
|
|
}
|
|
conflict := report
|
|
conflict.Counters = []Counter{{ProxyID: "proxy-a", Active: 2}}
|
|
if err := store.ReplaceRuntime(ctx, conflict, time.Minute); !errors.Is(err, ErrConflictingReport) {
|
|
t.Fatalf("ReplaceRuntime(conflict) error = %v", err)
|
|
}
|
|
stale := report
|
|
stale.Sequence = 1
|
|
if err := store.ReplaceRuntime(ctx, stale, time.Minute); !errors.Is(err, ErrStaleReport) {
|
|
t.Fatalf("ReplaceRuntime(stale) error = %v", err)
|
|
}
|
|
registerRuntimeSession(t, store, "session-b", time.Minute)
|
|
if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrStaleSession) {
|
|
t.Fatalf("ReplaceRuntime(old session) error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreFailsClosedForExpiredOrOlderOwnershipReport(t *testing.T) {
|
|
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
ctx := context.Background()
|
|
registerRuntimeSession(t, store, "session-a", time.Minute)
|
|
if err := store.ReplaceRuntime(ctx, Report{
|
|
WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
|
|
SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
|
|
}, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(): %v", err)
|
|
}
|
|
queries := []OwnedProxy{
|
|
{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 10},
|
|
{ProxyID: "proxy-b", WorkerID: "worker-a", OwnershipEpoch: 9},
|
|
}
|
|
got, err := store.ReadRuntime(ctx, queries)
|
|
if err != nil || got[0].Fresh || !got[1].Fresh {
|
|
t.Fatalf("ReadRuntime(ownership fence) = %+v, %v", got, err)
|
|
}
|
|
now = now.Add(time.Minute)
|
|
got, err = store.ReadRuntime(ctx, queries[1:])
|
|
if err != nil || len(got) != 1 || got[0].Fresh {
|
|
t.Fatalf("ReadRuntime(expired) = %+v, %v", got, err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreRejectsRuntimeBeyondAcknowledgedSnapshot(t *testing.T) {
|
|
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
registerRuntimeSession(t, store, "session-a", time.Minute)
|
|
report := Report{
|
|
WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
|
|
SnapshotVersion: 4, OwnershipEpoch: 9, ObservedAt: now,
|
|
}
|
|
if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
|
|
t.Fatalf("ReplaceRuntime(ahead snapshot) error = %v, want ErrSnapshotMismatch", err)
|
|
}
|
|
report.SnapshotVersion = 3
|
|
report.OwnershipEpoch = 10
|
|
if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
|
|
t.Fatalf("ReplaceRuntime(ahead epoch) error = %v, want ErrSnapshotMismatch", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreExpiredSameIdentitySessionDoesNotReactivateOldReport(t *testing.T) {
|
|
now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
|
|
store := newRuntimeStore(t, &now)
|
|
registerRuntimeSession(t, store, "session-a", time.Second)
|
|
if err := store.ReplaceRuntime(context.Background(), Report{
|
|
WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
|
|
SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
|
|
}, time.Minute); err != nil {
|
|
t.Fatalf("ReplaceRuntime(): %v", err)
|
|
}
|
|
session := store.sessions["worker-a"]
|
|
session.expiresAt = now.Add(time.Second)
|
|
store.sessions["worker-a"] = session
|
|
now = now.Add(2 * time.Second)
|
|
registerRuntimeSession(t, store, "session-a", time.Minute)
|
|
got, err := store.ReadRuntime(context.Background(), []OwnedProxy{{
|
|
ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9,
|
|
}})
|
|
if err != nil || len(got) != 1 || got[0].Fresh {
|
|
t.Fatalf("ReadRuntime(after re-register) = %+v, %v; want stale", got, err)
|
|
}
|
|
}
|
|
|
|
func newRuntimeStore(t *testing.T, now *time.Time) *MemoryStore {
|
|
t.Helper()
|
|
store, err := NewMemoryStore(func() time.Time { return *now })
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryStore(): %v", err)
|
|
}
|
|
return store
|
|
}
|
|
|
|
func registerRuntimeSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) {
|
|
t.Helper()
|
|
if err := store.ReplaceSession(context.Background(), Session{
|
|
WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID,
|
|
AckedSnapshotVersion: 3, AckedOwnershipEpoch: 9,
|
|
}, ttl); err != nil {
|
|
t.Fatalf("ReplaceSession(): %v", err)
|
|
}
|
|
}
|
|
|
|
func openControlSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) {
|
|
t.Helper()
|
|
if err := store.OpenSession(context.Background(), Session{
|
|
WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID, Zone: "zone-a", ProtocolVersion: 1,
|
|
Labels: map[string]string{"region": "test"},
|
|
}, ttl); err != nil {
|
|
t.Fatalf("OpenSession(): %v", err)
|
|
}
|
|
}
|
|
|
|
func controlReference(version, epoch uint64, content string) SnapshotReference {
|
|
return SnapshotReference{
|
|
WorkerID: "worker-a", Version: version, OwnershipEpoch: epoch, Checksum: sha256.Sum256([]byte(content)),
|
|
}
|
|
}
|
|
|
|
func controlReport(now time.Time, sessionID string, sequence, version, epoch uint64) Report {
|
|
return Report{
|
|
WorkerID: "worker-a", SessionID: sessionID, Sequence: sequence, SnapshotVersion: version,
|
|
OwnershipEpoch: epoch, ObservedAt: now, Counters: []Counter{{ProxyID: "proxy-a", Active: 1}},
|
|
}
|
|
}
|