feat: add worker snapshot stream interface

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"
@@ -21,10 +23,15 @@ type GRPCHandler struct {
 	controlplanev1.UnimplementedWorkerControlPlaneServer
 	service   Service
 	identity  IdentityAuthorizer
+	snapshots SnapshotSource
 }
 
-func NewGRPCHandler(service Service, identity IdentityAuthorizer) *GRPCHandler {
+func NewGRPCHandler(service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) *GRPCHandler {
-	return &GRPCHandler{service: service, identity: identity}
+	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)

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 {

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
 }
 

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

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 {
+	if err := service.IssueSnapshot(context.Background(), reference); err != nil {
-		t.Fatalf("RecordIssuedSnapshot(): %v", err)
+		t.Fatalf("IssueSnapshot(): %v", err)
 	}
 	if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{
 		WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7,

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)
+}