From 8b734b85f35c054a09207e67220f548e0f12b082 Mon Sep 17 00:00:00 2001 From: youfak Date: Fri, 31 Jul 2026 13:16:21 +0800 Subject: [PATCH] feat: add worker snapshot stream interface --- internal/controller/worker/grpc_handler.go | 67 +++++++++++++++++-- .../controller/worker/grpc_handler_test.go | 45 ++++++++++++- internal/controller/worker/server.go | 3 +- internal/controller/worker/service.go | 18 +++++ internal/controller/worker/service_test.go | 4 +- internal/controller/worker/snapshot_source.go | 21 ++++++ 6 files changed, 149 insertions(+), 9 deletions(-) create mode 100644 internal/controller/worker/snapshot_source.go diff --git a/internal/controller/worker/grpc_handler.go b/internal/controller/worker/grpc_handler.go index 62b2313..c7c0b1a 100644 --- a/internal/controller/worker/grpc_handler.go +++ b/internal/controller/worker/grpc_handler.go @@ -2,11 +2,13 @@ package worker import ( "context" + "crypto/sha256" "errors" controlplanev1 "proxy-pool/gen/controlplane/v1" "proxy-pool/internal/domain/workerruntime" + "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/durationpb" @@ -19,12 +21,17 @@ type IdentityAuthorizer interface { type GRPCHandler struct { controlplanev1.UnimplementedWorkerControlPlaneServer - service Service - identity IdentityAuthorizer + service Service + identity IdentityAuthorizer + snapshots SnapshotSource } -func NewGRPCHandler(service Service, identity IdentityAuthorizer) *GRPCHandler { - return &GRPCHandler{service: service, identity: identity} +func NewGRPCHandler(service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) *GRPCHandler { + handler := &GRPCHandler{service: service, identity: identity} + if len(snapshots) == 1 { + handler.snapshots = snapshots[0] + } + return handler } func (handler *GRPCHandler) RegisterWorker(ctx context.Context, request *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) { @@ -47,6 +54,58 @@ func (handler *GRPCHandler) RegisterWorker(ctx context.Context, request *control }, nil } +func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshotsRequest, stream grpc.ServerStreamingServer[controlplanev1.SnapshotEnvelope]) error { + if request == nil || stream == nil || handler == nil || handler.service == nil || handler.identity == nil { + return grpcError(ErrInvalidCommand) + } + if handler.snapshots == nil { + return status.Error(codes.Unimplemented, "worker snapshots are unavailable") + } + if err := handler.authorize(stream.Context(), request.GetWorkerId()); err != nil { + return err + } + if !workerruntime.ValidIdentifier(request.GetWorkerId()) || !workerruntime.ValidIdentifier(request.GetSessionId()) || + (len(request.GetLastChecksum()) != 0 && len(request.GetLastChecksum()) != sha256.Size) { + return grpcError(ErrInvalidCommand) + } + updates, err := handler.snapshots.Watch(stream.Context(), SnapshotWatchRequest{ + WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), LastAppliedVersion: request.GetLastAppliedVersion(), + LastChecksum: append([]byte(nil), request.GetLastChecksum()...), + }) + if err != nil { + return grpcError(err) + } + for { + select { + case <-stream.Context().Done(): + return stream.Context().Err() + case snapshot, ok := <-updates: + if !ok { + return nil + } + if err := handler.issueSnapshot(stream.Context(), request.GetWorkerId(), snapshot); err != nil { + return grpcError(err) + } + if err := stream.Send(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: snapshot}}); err != nil { + return err + } + } + } +} + +func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID 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 { + return ErrInvalidCommand + } + var checksum [sha256.Size]byte + copy(checksum[:], snapshot.GetChecksum()) + return handler.service.IssueSnapshot(ctx, workerruntime.SnapshotReference{ + WorkerID: workerID, Version: snapshot.GetVersion(), OwnershipEpoch: snapshot.GetOwnershipEpoch(), Checksum: checksum, + }) +} + 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 a6f4e69..b9070c5 100644 --- a/internal/controller/worker/grpc_handler_test.go +++ b/internal/controller/worker/grpc_handler_test.go @@ -64,6 +64,29 @@ func TestGRPCHandlerMapsErrorsAndLeavesStreamsUnimplemented(t *testing.T) { } } +func TestGRPCHandlerStreamsIssuedFullSnapshots(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(time.Minute)), + } + 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) + } + received, err := stream.Recv() + if err != nil || received.GetFull().GetVersion() != 3 { + t.Fatalf("Recv() = %+v, %v", received, err) + } + if service.issued.WorkerID != "worker-a" || service.issued.Version != 3 || service.issued.Checksum[0] != 1 { + t.Fatalf("issued snapshot = %+v", service.issued) + } +} + type grpcServiceStub struct { registration Registration registerErr error @@ -71,11 +94,16 @@ type grpcServiceStub struct { acknowledgeErr error report workerruntime.Report reportErr error + issued workerruntime.SnapshotReference } func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) { return stub.registration, stub.registerErr } +func (stub *grpcServiceStub) IssueSnapshot(_ context.Context, reference workerruntime.SnapshotReference) error { + stub.issued = reference + return nil +} func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error { stub.acknowledgement = acknowledgement return stub.acknowledgeErr @@ -89,11 +117,24 @@ type allowIdentity struct{} func (allowIdentity) Authorize(context.Context, string) error { return nil } -func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer) (controlplanev1.WorkerControlPlaneClient, func()) { +type snapshotSourceStub struct { + snapshots []*controlplanev1.WorkerSnapshot +} + +func (source snapshotSourceStub) Watch(_ context.Context, _ SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) { + updates := make(chan *controlplanev1.WorkerSnapshot, len(source.snapshots)) + for _, snapshot := range source.snapshots { + updates <- snapshot + } + close(updates) + return updates, nil +} + +func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) (controlplanev1.WorkerControlPlaneClient, func()) { t.Helper() listener := bufconn.Listen(1 << 20) server := grpc.NewServer() - controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity)) + controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity, snapshots...)) go func() { _ = server.Serve(listener) }() connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithInsecure()) if err != nil { diff --git a/internal/controller/worker/server.go b/internal/controller/worker/server.go index 08de9dd..68054ec 100644 --- a/internal/controller/worker/server.go +++ b/internal/controller/worker/server.go @@ -23,6 +23,7 @@ var ErrInvalidServer = errors.New("invalid worker control server configuration") type ServerOptions struct { ShutdownTimeout time.Duration + Snapshots SnapshotSource } func DefaultServerOptions() ServerOptions { @@ -60,7 +61,7 @@ func NewServer(controlPlane config.ControlPlane, service Service, options Server }), ) grpcServer := grpc.NewServer(serverOptions...) - controlplanev1.RegisterWorkerControlPlaneServer(grpcServer, NewGRPCHandler(service, identity)) + controlplanev1.RegisterWorkerControlPlaneServer(grpcServer, NewGRPCHandler(service, identity, options.Snapshots)) return &Server{listen: controlPlane.Listen, grpcServer: grpcServer, shutdownTimeout: options.ShutdownTimeout}, nil } diff --git a/internal/controller/worker/service.go b/internal/controller/worker/service.go index b51a396..6de9ff4 100644 --- a/internal/controller/worker/service.go +++ b/internal/controller/worker/service.go @@ -60,10 +60,28 @@ type Options struct { type Service interface { Register(context.Context, RegisterCommand) (Registration, error) + IssueSnapshot(context.Context, workerruntime.SnapshotReference) error Acknowledge(context.Context, SnapshotAcknowledgement) error ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error) } +func (service *service) IssueSnapshot(ctx context.Context, reference workerruntime.SnapshotReference) error { + if ctx == nil { + return ErrInvalidCommand + } + if err := ctx.Err(); err != nil { + return err + } + normalized, err := workerruntime.NormalizeSnapshotReference(reference) + if err != nil { + return errors.Join(ErrInvalidCommand, err) + } + if err := service.store.RecordIssuedSnapshot(ctx, normalized, service.options.SessionTTL); err != nil { + return classifyStoreError(err) + } + return nil +} + type service struct { store workerruntime.ControlStore options Options diff --git a/internal/controller/worker/service_test.go b/internal/controller/worker/service_test.go index f150d47..d130143 100644 --- a/internal/controller/worker/service_test.go +++ b/internal/controller/worker/service_test.go @@ -35,8 +35,8 @@ func TestServiceRegistersAcknowledgesAndReportsRuntime(t *testing.T) { WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch, Checksum: sha256.Sum256([]byte("snapshot-7")), } - if err := store.RecordIssuedSnapshot(context.Background(), reference, time.Minute); err != nil { - t.Fatalf("RecordIssuedSnapshot(): %v", err) + if err := service.IssueSnapshot(context.Background(), reference); err != nil { + t.Fatalf("IssueSnapshot(): %v", err) } if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{ WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7, diff --git a/internal/controller/worker/snapshot_source.go b/internal/controller/worker/snapshot_source.go new file mode 100644 index 0000000..23cb838 --- /dev/null +++ b/internal/controller/worker/snapshot_source.go @@ -0,0 +1,21 @@ +package worker + +import ( + "context" + "errors" + + controlplanev1 "proxy-pool/gen/controlplane/v1" +) + +var ErrSnapshotsUnavailable = errors.New("worker snapshots are unavailable") + +type SnapshotWatchRequest struct { + WorkerID string + SessionID string + LastAppliedVersion uint64 + LastChecksum []byte +} + +type SnapshotSource interface { + Watch(context.Context, SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) +}