325 lines
12 KiB
Go
325 lines
12 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"errors"
|
|
"io"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/domain/outcome"
|
|
"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"
|
|
"google.golang.org/protobuf/types/known/emptypb"
|
|
)
|
|
|
|
type IdentityAuthorizer interface {
|
|
Authorize(context.Context, string) error
|
|
}
|
|
|
|
type GRPCHandler struct {
|
|
controlplanev1.UnimplementedWorkerControlPlaneServer
|
|
service Service
|
|
identity IdentityAuthorizer
|
|
snapshots SnapshotSource
|
|
}
|
|
|
|
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) {
|
|
if request == nil || handler == nil || handler.service == nil || handler.identity == nil {
|
|
return nil, grpcError(ErrInvalidCommand)
|
|
}
|
|
if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
|
|
return nil, err
|
|
}
|
|
registration, err := handler.service.Register(ctx, RegisterCommand{
|
|
WorkerID: request.GetWorkerId(), InstanceID: request.GetInstanceId(), Zone: request.GetZone(),
|
|
ProtocolVersion: request.GetSupportedProtocolVersion(), Labels: cloneLabels(request.GetLabels()),
|
|
})
|
|
if err != nil {
|
|
return nil, grpcError(err)
|
|
}
|
|
return &controlplanev1.RegisterWorkerResponse{
|
|
WorkerId: registration.WorkerID, SessionId: registration.SessionID, OwnershipEpoch: registration.OwnershipEpoch,
|
|
HeartbeatInterval: durationpb.New(registration.HeartbeatInterval), MaxStaleAge: durationpb.New(registration.MaxStaleAge),
|
|
}, 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)
|
|
}
|
|
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()...),
|
|
})
|
|
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
|
|
}
|
|
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 {
|
|
return err
|
|
}
|
|
expiryTimer, expiry = resetSnapshotExpiryTimer(expiryTimer, snapshot.GetValidUntil().AsTime())
|
|
}
|
|
}
|
|
}
|
|
|
|
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().AsTime().After(time.Now().UTC()) {
|
|
return ErrInvalidCommand
|
|
}
|
|
var checksum [sha256.Size]byte
|
|
copy(checksum[:], snapshot.GetChecksum())
|
|
return handler.service.IssueSnapshot(ctx, sessionID, workerruntime.SnapshotReference{
|
|
WorkerID: workerID, Version: snapshot.GetVersion(), OwnershipEpoch: snapshot.GetOwnershipEpoch(), Checksum: checksum,
|
|
})
|
|
}
|
|
|
|
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)
|
|
}
|
|
if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
|
|
return nil, err
|
|
}
|
|
err := handler.service.Acknowledge(ctx, SnapshotAcknowledgement{
|
|
WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), Version: request.GetVersion(),
|
|
OwnershipEpoch: request.GetOwnershipEpoch(), Checksum: append([]byte(nil), request.GetChecksum()...),
|
|
Applied: request.GetApplied(), ErrorCode: request.GetErrorCode(), ErrorMessage: request.GetErrorMessage(),
|
|
})
|
|
if err != nil {
|
|
return nil, grpcError(err)
|
|
}
|
|
return &emptypb.Empty{}, nil
|
|
}
|
|
|
|
func (handler *GRPCHandler) ReportRuntime(ctx context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
|
|
if request == nil || handler == nil || handler.service == nil || handler.identity == nil || request.GetObservedAt() == nil || request.GetObservedAt().CheckValid() != nil {
|
|
return nil, grpcError(ErrInvalidCommand)
|
|
}
|
|
if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
|
|
return nil, err
|
|
}
|
|
counters := make([]workerruntime.Counter, len(request.GetCounters()))
|
|
for index, counter := range request.GetCounters() {
|
|
if counter == nil {
|
|
return nil, grpcError(ErrInvalidCommand)
|
|
}
|
|
counters[index] = workerruntime.Counter{
|
|
ProxyID: counter.GetProxyId(), Active: int64(counter.GetActive()), Reserved: int64(counter.GetReserved()), Draining: counter.GetDraining(),
|
|
}
|
|
}
|
|
decision, err := handler.service.ReportRuntime(ctx, workerruntime.Report{
|
|
WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), Sequence: request.GetReportSequence(),
|
|
SnapshotVersion: request.GetSnapshotVersion(), OwnershipEpoch: request.GetOwnershipEpoch(),
|
|
ObservedAt: request.GetObservedAt().AsTime(), Counters: counters,
|
|
})
|
|
if err != nil {
|
|
return nil, grpcError(err)
|
|
}
|
|
return &controlplanev1.ReportRuntimeResponse{
|
|
AcceptedOwnershipEpoch: decision.AcceptedOwnershipEpoch, RequireFullSnapshot: decision.RequireFullSnapshot,
|
|
}, nil
|
|
}
|
|
|
|
func (handler *GRPCHandler) ReportOutcomes(stream grpc.ClientStreamingServer[controlplanev1.OutcomeBatch, controlplanev1.ReportOutcomesResponse]) error {
|
|
if stream == nil || handler == nil || handler.service == nil || handler.identity == nil {
|
|
return grpcError(ErrInvalidCommand)
|
|
}
|
|
var workerID, sessionID string
|
|
var accepted uint64
|
|
batches := 0
|
|
for {
|
|
request, err := stream.Recv()
|
|
if errors.Is(err, io.EOF) {
|
|
if workerID == "" {
|
|
return grpcError(ErrInvalidCommand)
|
|
}
|
|
return stream.SendAndClose(&controlplanev1.ReportOutcomesResponse{AcceptedThroughSequence: accepted})
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if batches == outcome.MaxBatchesPerStream {
|
|
return grpcError(ErrInvalidCommand)
|
|
}
|
|
batch, err := decodeOutcomeBatch(request)
|
|
if err != nil {
|
|
return grpcError(err)
|
|
}
|
|
if workerID == "" {
|
|
if err := handler.authorize(stream.Context(), batch.WorkerID); err != nil {
|
|
return err
|
|
}
|
|
workerID, sessionID = batch.WorkerID, batch.SessionID
|
|
} else if batch.WorkerID != workerID || batch.SessionID != sessionID {
|
|
return grpcError(ErrInvalidCommand)
|
|
}
|
|
accepted, err = handler.service.ReportOutcomes(stream.Context(), batch)
|
|
if err != nil {
|
|
return grpcError(err)
|
|
}
|
|
batches++
|
|
}
|
|
}
|
|
|
|
func decodeOutcomeBatch(request *controlplanev1.OutcomeBatch) (outcome.Batch, error) {
|
|
if request == nil {
|
|
return outcome.Batch{}, ErrInvalidCommand
|
|
}
|
|
events := make([]outcome.Event, len(request.GetOutcomes()))
|
|
for index, item := range request.GetOutcomes() {
|
|
if item == nil || item.GetLatency() == nil || item.GetLatency().CheckValid() != nil ||
|
|
item.GetObservedAt() == nil || item.GetObservedAt().CheckValid() != nil {
|
|
return outcome.Batch{}, ErrInvalidCommand
|
|
}
|
|
stage, ok := decodeOutcomeStage(item.GetStage())
|
|
if !ok {
|
|
return outcome.Batch{}, ErrInvalidCommand
|
|
}
|
|
events[index] = outcome.Event{
|
|
ProxyID: item.GetProxyId(), RoutingName: item.GetRoutingName(), Stage: stage, Success: item.GetSuccess(),
|
|
ErrorClass: outcome.ErrorClass(item.GetErrorClass()), Latency: item.GetLatency().AsDuration(), ObservedAt: item.GetObservedAt().AsTime(),
|
|
}
|
|
}
|
|
return outcome.Batch{WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), Sequence: request.GetSequence(), Events: events}, nil
|
|
}
|
|
|
|
func decodeOutcomeStage(stage controlplanev1.OutcomeStage) (outcome.Stage, bool) {
|
|
switch stage {
|
|
case controlplanev1.OutcomeStage_OUTCOME_STAGE_DIAL:
|
|
return outcome.StageDial, true
|
|
case controlplanev1.OutcomeStage_OUTCOME_STAGE_PROXY_HANDSHAKE:
|
|
return outcome.StageProxyHandshake, true
|
|
case controlplanev1.OutcomeStage_OUTCOME_STAGE_RESPONSE_HEADERS:
|
|
return outcome.StageResponseHeaders, true
|
|
case controlplanev1.OutcomeStage_OUTCOME_STAGE_TUNNEL:
|
|
return outcome.StageTunnel, true
|
|
default:
|
|
return outcome.StageUnspecified, false
|
|
}
|
|
}
|
|
|
|
func (handler *GRPCHandler) authorize(ctx context.Context, workerID string) error {
|
|
if err := handler.identity.Authorize(ctx, workerID); err != nil {
|
|
return status.Error(codes.PermissionDenied, "worker identity is not authorized")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func grpcError(err error) error {
|
|
switch {
|
|
case errors.Is(err, context.Canceled):
|
|
return status.Error(codes.Canceled, "worker control request canceled")
|
|
case errors.Is(err, context.DeadlineExceeded):
|
|
return status.Error(codes.DeadlineExceeded, "worker control request deadline exceeded")
|
|
case errors.Is(err, ErrInvalidCommand), errors.Is(err, outcome.ErrInvalidBatch), errors.Is(err, workerruntime.ErrInvalidReport),
|
|
errors.Is(err, workerruntime.ErrInvalidOutcome),
|
|
errors.Is(err, workerruntime.ErrInvalidAcknowledgement), errors.Is(err, workerruntime.ErrInvalidSnapshotReference):
|
|
return status.Error(codes.InvalidArgument, "invalid worker control request")
|
|
case errors.Is(err, ErrProtocolVersion):
|
|
return status.Error(codes.FailedPrecondition, "unsupported worker protocol version")
|
|
case errors.Is(err, workerruntime.ErrStaleSession):
|
|
return status.Error(codes.FailedPrecondition, "worker session is stale")
|
|
case errors.Is(err, workerruntime.ErrSnapshotMismatch):
|
|
return status.Error(codes.FailedPrecondition, "worker snapshot does not match issued snapshot")
|
|
case errors.Is(err, workerruntime.ErrStaleAcknowledgement):
|
|
return status.Error(codes.Aborted, "worker snapshot acknowledgement is stale")
|
|
case errors.Is(err, workerruntime.ErrStaleReport):
|
|
return status.Error(codes.Aborted, "worker runtime sequence is stale")
|
|
case errors.Is(err, workerruntime.ErrConflictingReport):
|
|
return status.Error(codes.AlreadyExists, "worker runtime sequence conflicts")
|
|
case errors.Is(err, workerruntime.ErrStaleOutcome):
|
|
return status.Error(codes.Aborted, "worker outcome sequence is stale")
|
|
case errors.Is(err, workerruntime.ErrConflictingOutcome):
|
|
return status.Error(codes.AlreadyExists, "worker outcome sequence conflicts")
|
|
default:
|
|
return status.Error(codes.Unavailable, "worker control plane unavailable")
|
|
}
|
|
}
|
|
|
|
func cloneLabels(labels map[string]string) map[string]string {
|
|
result := make(map[string]string, len(labels))
|
|
for key, value := range labels {
|
|
result[key] = value
|
|
}
|
|
return result
|
|
}
|