diff --git a/docs/api/control-plane.md b/docs/api/control-plane.md index fb39fc6..bcea8ba 100644 --- a/docs/api/control-plane.md +++ b/docs/api/control-plane.md @@ -32,6 +32,10 @@ Gateway 校验后 ACK 并开始 Runtime 心跳。Controller 会从 Redis 的有 增量、Gateway 进程装配、Outcome 与 Checker 闭环尚未实现。`ReportOutcomes` 仍明确 返回 `Unimplemented`；`100,000 QPS` 仍是未验证的设计目标。 +`WatchSnapshots` 建立时校验当前 session；每次签发快照引用时也把 `session_id` +交给 Redis 原子校验。重复 Register 会同时清除旧 Runtime 和已签发引用，因此迟到的 +旧 Stream 既不能覆盖新 session 的引用，也不会向旧连接发送未获授权的快照。 + ## 2. Worker 会话 ```mermaid @@ -54,6 +58,10 @@ sequenceDiagram `worker_id` 是逻辑节点，`instance_id` 区分进程重启，`session_id` 防止旧进程 继续上报。所有权 `epoch` 小于 Controller 当前值的数据必须拒绝。 +同一 `worker_id` 重注册会替换 session，并清除旧 Runtime 与已签发 Snapshot +Reference。旧 Stream 即使在替换后仍收到上游更新，其签发操作也会以 +`FailedPrecondition` 结束，不能影响新 session 的 ACK 基线。 + `ReportRuntimeRequest.report_sequence` 在当前 `session_id` 内严格单调递增。 相同序号只允许内容完全相同的幂等重放；较小序号或相同序号的不同内容必须 拒绝。`observed_at` 只用于观测，不作为乱序判定依据，运行态 TTL 统一使用 diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index 0f8b452..ac21003 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -213,7 +213,9 @@ WorkerControlPlane 现已接入 Controller 生命周期：Register、ACK 和 Run 身份、消息/流限制和有界停机已实现。`WatchSnapshots` 会发送当前 epoch 的基础完整 Snapshot 并保持连接；Gateway 已具备 Register/Watch/ACK/Runtime 会话协调组件。 按 Worker 的可下发 ownership 索引已进入 Redis 原子脚本，并可构建无凭据引用的 -已归属 Proxy payload。Routing payload、凭据分发、Gateway 命令与 Outcome 上报仍未实现。 +已归属 Proxy payload。Snapshot 签发与 session 匹配在同一 Redis Lua 操作中完成， +重注册会清除旧引用，避免迟到 Stream 覆盖新 session。Routing payload、凭据分发、 +Gateway 命令与 Outcome 上报仍未实现。 已新增公用 `domain/activitypool` 契约及并发安全内存参考实现，Provider Reconciler 通过 `UpsertFetched` 写入带供应商 TTL 和分配安全余量的批次；已覆盖 diff --git a/docs/requirements/completion-audit.md b/docs/requirements/completion-audit.md index ee2e86c..26b74dc 100644 --- a/docs/requirements/completion-audit.md +++ b/docs/requirements/completion-audit.md @@ -52,7 +52,8 @@ plaintext fixture 与 SPIFFE mTLS 服务端；基础 Snapshot 流和 Gateway 的 Register/Watch/ACK/Runtime 会话协调已实现。Redis 以 Worker 可下发 ownership 索引构建无凭据引用的已归属 Proxy payload，并以租约收紧可用期；Routing payload、 - 凭据分发、Outcome 和 Checker 尚未闭环。 + 凭据分发、Outcome 和 Checker 尚未闭环。Snapshot 签发在 Redis 中原子匹配当前 + `session_id`，重注册会清除旧引用，迟到旧 Stream 不会覆盖新 session。 - `PostgreSQL 管理面`：已定义 `adminstate` 事务 seam、并发安全 MemoryStore、 公用契约、100 并发 Routing CAS、租约 Outbox 和只含六张管理表的 Schema；pgx Adapter 已在真实 PostgreSQL 18 上通过同一契约、迁移幂等、审计/Outbox diff --git a/internal/adapters/redisactivity/runtime.go b/internal/adapters/redisactivity/runtime.go index 375617d..c02cbc7 100644 --- a/internal/adapters/redisactivity/runtime.go +++ b/internal/adapters/redisactivity/runtime.go @@ -13,13 +13,14 @@ import ( const runtimeWireVersion = 1 const ( - runtimeReplaceSession = "replace_session" - runtimeCurrentEpoch = "current_epoch" - runtimeOpenSession = "open_session" - runtimeRecordSnapshot = "record_snapshot" - runtimeAcknowledge = "acknowledge_snapshot" - runtimeReplaceReport = "replace_report" - runtimeRead = "read" + runtimeReplaceSession = "replace_session" + runtimeCurrentEpoch = "current_epoch" + runtimeOpenSession = "open_session" + runtimeValidateSession = "validate_session" + runtimeRecordSnapshot = "record_snapshot" + runtimeAcknowledge = "acknowledge_snapshot" + runtimeReplaceReport = "replace_report" + runtimeRead = "read" ) type runtimeSessionWire struct { @@ -44,6 +45,18 @@ type runtimeSnapshotReferenceWire struct { Checksum string `json:"checksum"` } +type runtimeSnapshotIssueWire struct { + Version int `json:"version"` + SessionID string `json:"sessionId"` + Reference runtimeSnapshotReferenceWire `json:"reference"` +} + +type runtimeSessionValidationWire struct { + Version int `json:"version"` + WorkerID string `json:"workerId"` + SessionID string `json:"sessionId"` +} + type runtimeAcknowledgementWire struct { Version int `json:"version"` WorkerID string `json:"workerId"` @@ -140,15 +153,46 @@ func (a *Adapter) OpenSession(ctx context.Context, session workerruntime.Session return invalidScriptReply("unexpected worker open session reply") } -func (a *Adapter) RecordIssuedSnapshot(ctx context.Context, reference workerruntime.SnapshotReference, ttl time.Duration) error { +func (a *Adapter) ValidateSession(ctx context.Context, workerID, sessionID string) error { + if err := validateRuntimeCall(ctx, a); err != nil { + return err + } + if !workerruntime.ValidIdentifier(workerID) || !workerruntime.ValidIdentifier(sessionID) { + return workerruntime.ErrInvalidSession + } + payload, err := json.Marshal(runtimeSessionValidationWire{ + Version: runtimeWireVersion, WorkerID: workerID, SessionID: sessionID, + }) + if err != nil { + return workerruntime.ErrInvalidSession + } + reply, err := a.runRuntime(ctx, runtimeValidateSession, 0, payload, "") + if err != nil { + return err + } + switch reply.Status { + case scriptOK: + return nil + case scriptInvalid: + return workerruntime.ErrInvalidSession + case scriptUnavailable: + return workerruntime.ErrStaleSession + default: + return invalidScriptReply("unexpected worker session validation reply") + } +} + +func (a *Adapter) RecordIssuedSnapshot(ctx context.Context, sessionID string, reference workerruntime.SnapshotReference, ttl time.Duration) error { if err := validateRuntimeCall(ctx, a); err != nil { return err } normalized, err := workerruntime.NormalizeSnapshotReference(reference) - if err != nil || ttl <= 0 { + if err != nil || !workerruntime.ValidIdentifier(sessionID) || ttl <= 0 { return workerruntime.ErrInvalidSnapshotReference } - payload, err := json.Marshal(referenceWire(normalized)) + payload, err := json.Marshal(runtimeSnapshotIssueWire{ + Version: runtimeWireVersion, SessionID: sessionID, Reference: referenceWire(normalized), + }) if err != nil { return workerruntime.ErrInvalidSnapshotReference } @@ -167,6 +211,8 @@ func (a *Adapter) RecordIssuedSnapshot(ctx context.Context, reference workerrunt return workerruntime.ErrConflictingSnapshotReference case scriptSnapshotMismatch: return workerruntime.ErrSnapshotMismatch + case scriptUnavailable: + return workerruntime.ErrStaleSession default: return invalidScriptReply("unexpected worker snapshot reference reply") } diff --git a/internal/adapters/redisactivity/scripts/runtime.lua b/internal/adapters/redisactivity/scripts/runtime.lua index cc7f3a8..c46fb5f 100644 --- a/internal/adapters/redisactivity/scripts/runtime.lua +++ b/internal/adapters/redisactivity/scripts/runtime.lua @@ -148,20 +148,44 @@ if operation == 'open_session' then end redis.call('HDEL', runtime_key, session.workerId) redis.call('ZREM', runtime_expiry_key, session.workerId) + redis.call('HDEL', snapshots_key, session.workerId) + redis.call('ZREM', snapshot_expiry_key, session.workerId) session.expiresAtMs = now + ttl_ms redis.call('HSET', sessions_key, session.workerId, cjson.encode(session)) redis.call('ZADD', session_expiry_key, session.expiresAtMs, session.workerId) return reply('ok') end +if operation == 'validate_session' then + local validation = decode_table(payload) + if not validation or validation.version ~= 1 or type(validation.workerId) ~= 'string' or + validation.workerId == '' or type(validation.sessionId) ~= 'string' or validation.sessionId == '' then + return reply('invalid') + end + local session = decode_table(redis.call('HGET', sessions_key, validation.workerId)) + if not valid_control_session(session) or session.workerId ~= validation.workerId or + session.sessionId ~= validation.sessionId or type(session.expiresAtMs) ~= 'number' or + session.expiresAtMs <= now then + return reply('unavailable') + end + return reply('ok') +end + if operation == 'record_snapshot' then if not ttl_ms or ttl_ms <= 0 then return reply('invalid') end - local reference = decode_table(payload) - if not valid_reference(reference) then + local issue = decode_table(payload) + if not issue or issue.version ~= 1 or type(issue.sessionId) ~= 'string' or issue.sessionId == '' or + not valid_reference(issue.reference) then return reply('invalid') end + local reference = issue.reference + local session = decode_table(redis.call('HGET', sessions_key, reference.workerId)) + if not valid_control_session(session) or session.workerId ~= reference.workerId or session.sessionId ~= issue.sessionId or + type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then + return reply('unavailable') + end local epoch = redis.call('GET', epoch_key) if not epoch then epoch = '1' diff --git a/internal/controller/worker/grpc_handler.go b/internal/controller/worker/grpc_handler.go index c7c0b1a..c29d547 100644 --- a/internal/controller/worker/grpc_handler.go +++ b/internal/controller/worker/grpc_handler.go @@ -68,6 +68,9 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot (len(request.GetLastChecksum()) != 0 && len(request.GetLastChecksum()) != sha256.Size) { return grpcError(ErrInvalidCommand) } + if err := handler.service.ValidateSession(stream.Context(), request.GetWorkerId(), request.GetSessionId()); err != nil { + return grpcError(err) + } updates, err := handler.snapshots.Watch(stream.Context(), SnapshotWatchRequest{ WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), LastAppliedVersion: request.GetLastAppliedVersion(), LastChecksum: append([]byte(nil), request.GetLastChecksum()...), @@ -83,7 +86,7 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot if !ok { return nil } - if err := handler.issueSnapshot(stream.Context(), request.GetWorkerId(), snapshot); err != nil { + if err := handler.issueSnapshot(stream.Context(), request.GetWorkerId(), request.GetSessionId(), snapshot); err != nil { return grpcError(err) } if err := stream.Send(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: snapshot}}); err != nil { @@ -93,7 +96,7 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot } } -func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID string, snapshot *controlplanev1.WorkerSnapshot) error { +func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID, sessionID string, snapshot *controlplanev1.WorkerSnapshot) error { if snapshot == nil || snapshot.GetVersion() == 0 || snapshot.GetOwnershipEpoch() == 0 || len(snapshot.GetChecksum()) != sha256.Size || snapshot.GetGeneratedAt() == nil || snapshot.GetGeneratedAt().CheckValid() != nil || snapshot.GetValidUntil() == nil || snapshot.GetValidUntil().CheckValid() != nil { @@ -101,7 +104,7 @@ func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID string, } var checksum [sha256.Size]byte copy(checksum[:], snapshot.GetChecksum()) - return handler.service.IssueSnapshot(ctx, workerruntime.SnapshotReference{ + return handler.service.IssueSnapshot(ctx, sessionID, workerruntime.SnapshotReference{ WorkerID: workerID, Version: snapshot.GetVersion(), OwnershipEpoch: snapshot.GetOwnershipEpoch(), Checksum: checksum, }) } diff --git a/internal/controller/worker/grpc_handler_test.go b/internal/controller/worker/grpc_handler_test.go index 37e87d8..0758c10 100644 --- a/internal/controller/worker/grpc_handler_test.go +++ b/internal/controller/worker/grpc_handler_test.go @@ -85,6 +85,31 @@ func TestGRPCHandlerStreamsIssuedFullSnapshots(t *testing.T) { if service.issued.WorkerID != "worker-a" || service.issued.Version != 3 || service.issued.Checksum[0] != 1 { t.Fatalf("issued snapshot = %+v", service.issued) } + if service.issuedSessionID != "session-a" { + t.Fatalf("issued session = %q, want session-a", service.issuedSessionID) + } +} + +func TestGRPCHandlerDoesNotDeliverSnapshotWhenSessionBecomesStale(t *testing.T) { + checksum := make([]byte, 32) + checksum[0] = 1 + service := &grpcServiceStub{issueErr: workerruntime.ErrStaleSession} + client, cleanup := grpcWorkerClient(t, service, allowIdentity{}, snapshotSourceStub{snapshots: []*controlplanev1.WorkerSnapshot{{ + Version: 3, OwnershipEpoch: 9, Checksum: checksum, + GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)), + }}}) + defer cleanup() + stream, err := client.WatchSnapshots(context.Background(), &controlplanev1.WatchSnapshotsRequest{WorkerId: "worker-a", SessionId: "session-a"}) + if err != nil { + t.Fatalf("WatchSnapshots(): %v", err) + } + _, err = stream.Recv() + if status.Code(err) != codes.FailedPrecondition { + t.Fatalf("Recv() error = %v, want FailedPrecondition", err) + } + if service.issued.Version != 0 || service.issuedSessionID != "" { + t.Fatalf("stale session issued snapshot = %+v for %q", service.issued, service.issuedSessionID) + } } type grpcServiceStub struct { @@ -95,14 +120,24 @@ type grpcServiceStub struct { report workerruntime.Report reportErr error issued workerruntime.SnapshotReference + issuedSessionID string + validateErr error + issueErr error } func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) { return stub.registration, stub.registerErr } func (stub *grpcServiceStub) CurrentOwnershipEpoch(context.Context) (uint64, error) { return 9, nil } -func (stub *grpcServiceStub) IssueSnapshot(_ context.Context, reference workerruntime.SnapshotReference) error { +func (stub *grpcServiceStub) ValidateSession(context.Context, string, string) error { + return stub.validateErr +} +func (stub *grpcServiceStub) IssueSnapshot(_ context.Context, sessionID string, reference workerruntime.SnapshotReference) error { + if stub.issueErr != nil { + return stub.issueErr + } stub.issued = reference + stub.issuedSessionID = sessionID return nil } func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error { diff --git a/internal/controller/worker/service.go b/internal/controller/worker/service.go index db63ea6..f2efdd3 100644 --- a/internal/controller/worker/service.go +++ b/internal/controller/worker/service.go @@ -64,7 +64,8 @@ type Options struct { type Service interface { Register(context.Context, RegisterCommand) (Registration, error) CurrentOwnershipEpoch(context.Context) (uint64, error) - IssueSnapshot(context.Context, workerruntime.SnapshotReference) error + ValidateSession(context.Context, string, string) error + IssueSnapshot(context.Context, string, workerruntime.SnapshotReference) error Acknowledge(context.Context, SnapshotAcknowledgement) error ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error) } @@ -83,7 +84,20 @@ func (service *service) CurrentOwnershipEpoch(ctx context.Context) (uint64, erro return epoch, nil } -func (service *service) IssueSnapshot(ctx context.Context, reference workerruntime.SnapshotReference) error { +func (service *service) ValidateSession(ctx context.Context, workerID, sessionID string) error { + if ctx == nil || !workerruntime.ValidIdentifier(workerID) || !workerruntime.ValidIdentifier(sessionID) { + return ErrInvalidCommand + } + if err := ctx.Err(); err != nil { + return err + } + if err := service.store.ValidateSession(ctx, workerID, sessionID); err != nil { + return classifyStoreError(err) + } + return nil +} + +func (service *service) IssueSnapshot(ctx context.Context, sessionID string, reference workerruntime.SnapshotReference) error { if ctx == nil { return ErrInvalidCommand } @@ -91,10 +105,10 @@ func (service *service) IssueSnapshot(ctx context.Context, reference workerrunti return err } normalized, err := workerruntime.NormalizeSnapshotReference(reference) - if err != nil { + if err != nil || !workerruntime.ValidIdentifier(sessionID) { return errors.Join(ErrInvalidCommand, err) } - if err := service.store.RecordIssuedSnapshot(ctx, normalized, service.options.SessionTTL); err != nil { + if err := service.store.RecordIssuedSnapshot(ctx, sessionID, normalized, service.options.SessionTTL); err != nil { return classifyStoreError(err) } return nil diff --git a/internal/controller/worker/service_test.go b/internal/controller/worker/service_test.go index d130143..afdef66 100644 --- a/internal/controller/worker/service_test.go +++ b/internal/controller/worker/service_test.go @@ -35,7 +35,7 @@ func TestServiceRegistersAcknowledgesAndReportsRuntime(t *testing.T) { WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch, Checksum: sha256.Sum256([]byte("snapshot-7")), } - if err := service.IssueSnapshot(context.Background(), reference); err != nil { + if err := service.IssueSnapshot(context.Background(), registered.SessionID, reference); err != nil { t.Fatalf("IssueSnapshot(): %v", err) } if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{ @@ -90,7 +90,10 @@ func (unavailableStore) CurrentOwnershipEpoch(context.Context) (uint64, error) { func (unavailableStore) OpenSession(context.Context, workerruntime.Session, time.Duration) error { return errors.New("redis unavailable") } -func (unavailableStore) RecordIssuedSnapshot(context.Context, workerruntime.SnapshotReference, time.Duration) error { +func (unavailableStore) ValidateSession(context.Context, string, string) error { + return errors.New("redis unavailable") +} +func (unavailableStore) RecordIssuedSnapshot(context.Context, string, workerruntime.SnapshotReference, time.Duration) error { return errors.New("redis unavailable") } func (unavailableStore) AcknowledgeSnapshot(context.Context, workerruntime.SnapshotAcknowledgement, time.Duration) error { diff --git a/internal/domain/workerruntime/contracttest/contract.go b/internal/domain/workerruntime/contracttest/contract.go index a546fbb..988ea14 100644 --- a/internal/domain/workerruntime/contracttest/contract.go +++ b/internal/domain/workerruntime/contracttest/contract.go @@ -24,6 +24,7 @@ func Run(t *testing.T, factory Factory) { t.Helper() t.Run("acknowledged runtime lifecycle", func(t *testing.T) { runLifecycle(t, newFixture(t, factory)) }) t.Run("negative acknowledgement fences runtime", func(t *testing.T) { runNegativeAck(t, newFixture(t, factory)) }) + t.Run("session replacement fences issued snapshots", func(t *testing.T) { runSessionReplacementFence(t, newFixture(t, factory)) }) } func runLifecycle(t *testing.T, fixture Fixture) { @@ -39,7 +40,7 @@ func runLifecycle(t *testing.T, fixture Fixture) { if err := fixture.Store.ReplaceRuntime(ctx, report, fixture.TTL); !errors.Is(err, workerruntime.ErrSnapshotMismatch) { t.Fatalf("ReplaceRuntime(before ACK) error = %v", err) } - if err := fixture.Store.RecordIssuedSnapshot(ctx, reference, fixture.TTL); err != nil { + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-a", reference, fixture.TTL); err != nil { t.Fatalf("RecordIssuedSnapshot(): %v", err) } ack := workerruntime.SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true} @@ -69,14 +70,14 @@ func runNegativeAck(t *testing.T, fixture Fixture) { t.Fatalf("CurrentOwnershipEpoch(): %v", err) } first := snapshot(7, epoch, "snapshot-7") - if err := fixture.Store.RecordIssuedSnapshot(ctx, first, fixture.TTL); err != nil { + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-a", first, fixture.TTL); err != nil { t.Fatalf("RecordIssuedSnapshot(first): %v", err) } if err := fixture.Store.AcknowledgeSnapshot(ctx, workerruntime.SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: first, Applied: true}, fixture.TTL); err != nil { t.Fatalf("AcknowledgeSnapshot(first): %v", err) } second := snapshot(8, epoch, "snapshot-8") - if err := fixture.Store.RecordIssuedSnapshot(ctx, second, fixture.TTL); err != nil { + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-a", second, fixture.TTL); err != nil { t.Fatalf("RecordIssuedSnapshot(second): %v", err) } if err := fixture.Store.AcknowledgeSnapshot(ctx, workerruntime.SnapshotAcknowledgement{ @@ -89,6 +90,33 @@ func runNegativeAck(t *testing.T, fixture Fixture) { } } +func runSessionReplacementFence(t *testing.T, fixture Fixture) { + t.Helper() + ctx := context.Background() + open(t, fixture.Store, fixture.TTL) + epoch, err := fixture.Store.CurrentOwnershipEpoch(ctx) + if err != nil { + t.Fatalf("CurrentOwnershipEpoch(): %v", err) + } + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-a", snapshot(3, epoch, "old"), fixture.TTL); err != nil { + t.Fatalf("RecordIssuedSnapshot(old): %v", err) + } + if err := fixture.Store.OpenSession(ctx, workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-b", SessionID: "session-b", Zone: "zone-a", ProtocolVersion: 1, + }, fixture.TTL); err != nil { + t.Fatalf("OpenSession(replacement): %v", err) + } + if err := fixture.Store.ValidateSession(ctx, "worker-a", "session-a"); !errors.Is(err, workerruntime.ErrStaleSession) { + t.Fatalf("ValidateSession(old): %v, want ErrStaleSession", err) + } + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-a", snapshot(4, epoch, "stale"), fixture.TTL); !errors.Is(err, workerruntime.ErrStaleSession) { + t.Fatalf("RecordIssuedSnapshot(stale): %v, want ErrStaleSession", err) + } + if err := fixture.Store.RecordIssuedSnapshot(ctx, "session-b", snapshot(1, epoch, "new"), fixture.TTL); err != nil { + t.Fatalf("RecordIssuedSnapshot(new): %v", err) + } +} + func newFixture(t *testing.T, factory Factory) Fixture { t.Helper() fixture := factory(t) diff --git a/internal/domain/workerruntime/memory.go b/internal/domain/workerruntime/memory.go index c72e4de..6bb09bd 100644 --- a/internal/domain/workerruntime/memory.go +++ b/internal/domain/workerruntime/memory.go @@ -84,27 +84,52 @@ func (store *MemoryStore) OpenSession(ctx context.Context, session Session, ttl 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) RecordIssuedSnapshot(ctx context.Context, reference SnapshotReference, ttl time.Duration) error { - if ctx == nil || store == nil || ttl <= 0 { - return ErrInvalidSnapshotReference +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 } - normalized, err := NormalizeSnapshotReference(reference) - 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[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 } diff --git a/internal/domain/workerruntime/memory_test.go b/internal/domain/workerruntime/memory_test.go index 2acbfab..08a424b 100644 --- a/internal/domain/workerruntime/memory_test.go +++ b/internal/domain/workerruntime/memory_test.go @@ -22,7 +22,7 @@ func TestMemoryStoreRequiresAcknowledgedSnapshotForRuntime(t *testing.T) { t.Fatalf("ReplaceRuntime(before ACK) error = %v, want ErrSnapshotMismatch", err) } reference := controlReference(7, epoch, "snapshot-7") - if err := store.RecordIssuedSnapshot(ctx, reference, time.Minute); err != nil { + if err := store.RecordIssuedSnapshot(ctx, "session-a", reference, time.Minute); err != nil { t.Fatalf("RecordIssuedSnapshot(): %v", err) } if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{ @@ -45,7 +45,7 @@ func TestMemoryStoreNegativeAcknowledgementFencesDelayedRuntime(t *testing.T) { t.Fatalf("CurrentOwnershipEpoch(): %v", err) } first := controlReference(7, epoch, "snapshot-7") - if err := store.RecordIssuedSnapshot(ctx, first, time.Minute); err != nil { + 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 { @@ -56,7 +56,7 @@ func TestMemoryStoreNegativeAcknowledgementFencesDelayedRuntime(t *testing.T) { } second := controlReference(8, epoch, "snapshot-8") - if err := store.RecordIssuedSnapshot(ctx, second, time.Minute); err != nil { + 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 { @@ -83,7 +83,7 @@ func TestMemoryStoreAcknowledgementReplayPreservesRuntimeFence(t *testing.T) { t.Fatalf("CurrentOwnershipEpoch(): %v", err) } reference := controlReference(7, epoch, "snapshot-7") - if err := store.RecordIssuedSnapshot(ctx, reference, time.Minute); err != nil { + 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} @@ -109,18 +109,43 @@ 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, controlReference(7, epoch, "first"), time.Minute); err != nil { + 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, controlReference(7, epoch, "second"), time.Minute); !errors.Is(err, ErrConflictingSnapshotReference) { + 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) diff --git a/internal/domain/workerruntime/runtime.go b/internal/domain/workerruntime/runtime.go index a7e423b..15cfe01 100644 --- a/internal/domain/workerruntime/runtime.go +++ b/internal/domain/workerruntime/runtime.go @@ -91,7 +91,8 @@ type SessionWriter interface { type ControlStore interface { CurrentOwnershipEpoch(context.Context) (uint64, error) OpenSession(context.Context, Session, time.Duration) error - RecordIssuedSnapshot(context.Context, SnapshotReference, time.Duration) error + ValidateSession(context.Context, string, string) error + RecordIssuedSnapshot(context.Context, string, SnapshotReference, time.Duration) error AcknowledgeSnapshot(context.Context, SnapshotAcknowledgement, time.Duration) error ReplaceRuntime(context.Context, Report, time.Duration) error }