From c2d0c08fcb142b35cf0528d1ccf69941980a4593 Mon Sep 17 00:00:00 2001 From: youfak Date: Fri, 31 Jul 2026 15:02:25 +0800 Subject: [PATCH] fix: close expired worker snapshot streams --- docs/api/control-plane.md | 5 ++ docs/development/implementation-plan.md | 5 +- docs/requirements/completion-audit.md | 5 +- internal/controller/worker/grpc_handler.go | 43 +++++++++++++- .../controller/worker/grpc_handler_test.go | 57 +++++++++++++++++++ 5 files changed, 110 insertions(+), 5 deletions(-) diff --git a/docs/api/control-plane.md b/docs/api/control-plane.md index c46aeec..a643d34 100644 --- a/docs/api/control-plane.md +++ b/docs/api/control-plane.md @@ -92,6 +92,11 @@ Gateway 接收完整快照时必须拒绝缺失、格式错误或已到期的 `v 保存在本地不可变视图。该整体期限到达后，调度直接按无候选处理，不再使用旧视图 发起新的上游连接，也不查询 Redis 或 PostgreSQL 补偿。 +Controller 只会下发尚未到期的完整快照，并在最近一次成功下发快照的 +`valid_until` 到达时结束 `WatchSnapshots` 流。Gateway 的会话调用方必须在流结束后 +按带 jitter 的有界退避重建 Register/Watch 会话；在新快照通过校验并原子替换前， +旧视图仍按其整体有效期 fail-closed。 + Delta 声明 `base_version`。Worker 只有在本地版本恰好等于 base 且 checksum 验证成功时才能应用；否则丢弃 Delta 并请求完整 Snapshot。构建在后台完成， 热路径只读取一次原子指针。 diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index ac21003..089361e 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -214,8 +214,9 @@ WorkerControlPlane 现已接入 Controller 生命周期：Register、ACK 和 Run Snapshot 并保持连接；Gateway 已具备 Register/Watch/ACK/Runtime 会话协调组件。 按 Worker 的可下发 ownership 索引已进入 Redis 原子脚本，并可构建无凭据引用的 已归属 Proxy payload。Snapshot 签发与 session 匹配在同一 Redis Lua 操作中完成， -重注册会清除旧引用，避免迟到 Stream 覆盖新 session。Routing payload、凭据分发、 -Gateway 命令与 Outcome 上报仍未实现。 +重注册会清除旧引用，避免迟到 Stream 覆盖新 session。Worker 服务端会在最近完整 +Snapshot 的 `valid_until` 到达时结束流，令 Gateway 调用方可按退避策略重建会话； +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 b4cd58a..41c11fa 100644 --- a/docs/requirements/completion-audit.md +++ b/docs/requirements/completion-audit.md @@ -53,8 +53,9 @@ Register/Watch/ACK/Runtime 会话协调已实现。Redis 以 Worker 可下发 ownership 索引构建无凭据引用的已归属 Proxy payload，并以租约收紧可用期；Routing payload、 凭据分发、Outcome 和 Checker 尚未闭环。Snapshot 签发在 Redis 中原子匹配当前 - `session_id`，重注册会清除旧引用，迟到旧 Stream 不会覆盖新 session。Gateway 会 - 校验并执行 Snapshot 整体 `valid_until`，过期视图不再分配新 Proxy。 + `session_id`，重注册会清除旧引用，迟到旧 Stream 不会覆盖新 session。Controller + 在最近成功下发的 Snapshot `valid_until` 到达时关闭流，Gateway 会校验并执行 + Snapshot 整体 `valid_until`，过期视图不再分配新 Proxy。 - `PostgreSQL 管理面`：已定义 `adminstate` 事务 seam、并发安全 MemoryStore、 公用契约、100 并发 Routing CAS、租约 Outbox 和只含六张管理表的 Schema；pgx Adapter 已在真实 PostgreSQL 18 上通过同一契约、迁移幂等、审计/Outbox diff --git a/internal/controller/worker/grpc_handler.go b/internal/controller/worker/grpc_handler.go index c29d547..d102de5 100644 --- a/internal/controller/worker/grpc_handler.go +++ b/internal/controller/worker/grpc_handler.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "errors" + "time" controlplanev1 "proxy-pool/gen/controlplane/v1" "proxy-pool/internal/domain/workerruntime" @@ -78,10 +79,17 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot if err != nil { return grpcError(err) } + var expiryTimer *time.Timer + var expiry <-chan time.Time + defer func() { + stopSnapshotExpiryTimer(expiryTimer) + }() for { select { case <-stream.Context().Done(): return stream.Context().Err() + case <-expiry: + return nil case snapshot, ok := <-updates: if !ok { return nil @@ -92,6 +100,7 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot if err := stream.Send(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: snapshot}}); err != nil { return err } + expiryTimer, expiry = resetSnapshotExpiryTimer(expiryTimer, snapshot.GetValidUntil().AsTime()) } } } @@ -99,7 +108,8 @@ func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshot 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 { + snapshot.GetValidUntil() == nil || snapshot.GetValidUntil().CheckValid() != nil || + !snapshot.GetValidUntil().AsTime().After(time.Now().UTC()) { return ErrInvalidCommand } var checksum [sha256.Size]byte @@ -109,6 +119,37 @@ func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID, session }) } +func resetSnapshotExpiryTimer(timer *time.Timer, deadline time.Time) (*time.Timer, <-chan time.Time) { + delay := time.Until(deadline) + if delay <= 0 { + delay = time.Nanosecond + } + if timer == nil { + timer = time.NewTimer(delay) + return timer, timer.C + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(delay) + return timer, timer.C +} + +func stopSnapshotExpiryTimer(timer *time.Timer) { + if timer == nil { + return + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } +} + func (handler *GRPCHandler) AcknowledgeSnapshot(ctx context.Context, request *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) { if request == nil || handler == nil || handler.service == nil || handler.identity == nil { return nil, grpcError(ErrInvalidCommand) diff --git a/internal/controller/worker/grpc_handler_test.go b/internal/controller/worker/grpc_handler_test.go index 0758c10..17427e8 100644 --- a/internal/controller/worker/grpc_handler_test.go +++ b/internal/controller/worker/grpc_handler_test.go @@ -2,6 +2,7 @@ package worker import ( "context" + "io" "net" "testing" "time" @@ -90,6 +91,52 @@ func TestGRPCHandlerStreamsIssuedFullSnapshots(t *testing.T) { } } +func TestGRPCHandlerClosesSnapshotStreamAtValidityDeadline(t *testing.T) { + checksum := make([]byte, 32) + checksum[0] = 1 + full := &controlplanev1.WorkerSnapshot{ + Version: 3, OwnershipEpoch: 9, Checksum: checksum, + GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(100 * time.Millisecond)), + } + service := &grpcServiceStub{} + client, cleanup := grpcWorkerClient(t, service, allowIdentity{}, holdingSnapshotSource{snapshot: full}) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + stream, err := client.WatchSnapshots(ctx, &controlplanev1.WatchSnapshotsRequest{WorkerId: "worker-a", SessionId: "session-a"}) + if err != nil { + t.Fatalf("WatchSnapshots(): %v", err) + } + if received, err := stream.Recv(); err != nil || received.GetFull().GetVersion() != 3 { + t.Fatalf("first Recv() = %+v, %v", received, err) + } + if _, err := stream.Recv(); err != io.EOF { + t.Fatalf("Recv(after validity deadline) error = %v, want EOF", err) + } +} + +func TestGRPCHandlerRejectsExpiredSnapshot(t *testing.T) { + checksum := make([]byte, 32) + checksum[0] = 1 + full := &controlplanev1.WorkerSnapshot{ + Version: 3, OwnershipEpoch: 9, Checksum: checksum, + GeneratedAt: timestamppb.New(time.Now().Add(-time.Minute)), ValidUntil: timestamppb.New(time.Now().Add(-time.Second)), + } + service := &grpcServiceStub{} + client, cleanup := grpcWorkerClient(t, service, allowIdentity{}, snapshotSourceStub{snapshots: []*controlplanev1.WorkerSnapshot{full}}) + defer cleanup() + stream, err := client.WatchSnapshots(context.Background(), &controlplanev1.WatchSnapshotsRequest{WorkerId: "worker-a", SessionId: "session-a"}) + if err != nil { + t.Fatalf("WatchSnapshots(): %v", err) + } + if _, err := stream.Recv(); status.Code(err) != codes.InvalidArgument { + t.Fatalf("Recv() code = %s, want InvalidArgument; error=%v", status.Code(err), err) + } + if service.issued.Version != 0 { + t.Fatalf("expired snapshot was issued: %+v", service.issued) + } +} + func TestGRPCHandlerDoesNotDeliverSnapshotWhenSessionBecomesStale(t *testing.T) { checksum := make([]byte, 32) checksum[0] = 1 @@ -157,6 +204,16 @@ type snapshotSourceStub struct { snapshots []*controlplanev1.WorkerSnapshot } +type holdingSnapshotSource struct { + snapshot *controlplanev1.WorkerSnapshot +} + +func (source holdingSnapshotSource) Watch(_ context.Context, _ SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) { + updates := make(chan *controlplanev1.WorkerSnapshot, 1) + updates <- source.snapshot + return updates, nil +} + func (source snapshotSourceStub) Watch(_ context.Context, _ SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) { updates := make(chan *controlplanev1.WorkerSnapshot, len(source.snapshots)) for _, snapshot := range source.snapshots {