diff --git a/internal/gateway/controlplane/reporter.go b/internal/gateway/controlplane/reporter.go new file mode 100644 index 0000000..7690ba9 --- /dev/null +++ b/internal/gateway/controlplane/reporter.go @@ -0,0 +1,177 @@ +package controlplane + +import ( + "context" + "errors" + "fmt" + "math" + "sync" + "time" + + controlplanev1 "proxy-pool/gen/controlplane/v1" + "proxy-pool/internal/domain/workerruntime" + "proxy-pool/internal/gateway/snapshot" + + "google.golang.org/protobuf/types/known/timestamppb" +) + +var ( + ErrInvalidOptions = errors.New("invalid gateway control plane reporter options") + ErrNotRegistered = errors.New("gateway worker session is not registered") + ErrFullSnapshotRequired = errors.New("controller requires a full snapshot") + ErrOwnershipChanged = errors.New("controller accepted a different ownership epoch") +) + +type Client interface { + RegisterWorker(context.Context, *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) + ReportRuntime(context.Context, *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) +} + +type RuntimeSource interface { + RuntimeReport(string, uint64, time.Time) (workerruntime.Report, error) +} + +type Options struct { + WorkerID string + InstanceID string + Zone string + ProtocolVersion uint32 + Labels map[string]string + Now func() time.Time +} + +type Registration struct { + SessionID string + OwnershipEpoch uint64 + HeartbeatInterval time.Duration + MaxStaleAge time.Duration +} + +type RuntimeReporter struct { + client Client + snapshots RuntimeSource + options Options + + mu sync.Mutex + session Registration + sequence uint64 +} + +func NewRuntimeReporter(client Client, snapshots RuntimeSource, options Options) (*RuntimeReporter, error) { + if client == nil || snapshots == nil || options.ProtocolVersion == 0 || options.Now == nil || + !workerruntime.ValidIdentifier(options.WorkerID) || !workerruntime.ValidIdentifier(options.InstanceID) || + !workerruntime.ValidIdentifier(options.Zone) { + return nil, ErrInvalidOptions + } + if _, err := workerruntime.NormalizeLabels(options.Labels); err != nil { + return nil, errors.Join(ErrInvalidOptions, err) + } + options.Labels = cloneLabels(options.Labels) + return &RuntimeReporter{client: client, snapshots: snapshots, options: options}, nil +} + +func (reporter *RuntimeReporter) Register(ctx context.Context) (Registration, error) { + if reporter == nil || ctx == nil { + return Registration{}, ErrInvalidOptions + } + if err := ctx.Err(); err != nil { + return Registration{}, err + } + response, err := reporter.client.RegisterWorker(ctx, &controlplanev1.RegisterWorkerRequest{ + WorkerId: reporter.options.WorkerID, InstanceId: reporter.options.InstanceID, Zone: reporter.options.Zone, + SupportedProtocolVersion: reporter.options.ProtocolVersion, Labels: cloneLabels(reporter.options.Labels), + }) + if err != nil { + return Registration{}, err + } + registration, err := validateRegistration(reporter.options.WorkerID, response) + if err != nil { + return Registration{}, err + } + reporter.mu.Lock() + reporter.session = registration + reporter.sequence = 0 + reporter.mu.Unlock() + return registration, nil +} + +func (reporter *RuntimeReporter) Report(ctx context.Context) error { + if reporter == nil || ctx == nil { + return ErrInvalidOptions + } + if err := ctx.Err(); err != nil { + return err + } + reporter.mu.Lock() + defer reporter.mu.Unlock() + if reporter.session.SessionID == "" { + return ErrNotRegistered + } + sequence := reporter.sequence + 1 + report, err := reporter.snapshots.RuntimeReport(reporter.session.SessionID, sequence, reporter.options.Now().UTC()) + if err != nil { + return fmt.Errorf("build gateway runtime report: %w", err) + } + request, err := runtimeRequest(report) + if err != nil { + return err + } + response, err := reporter.client.ReportRuntime(ctx, request) + if err != nil { + return err + } + if response == nil { + return ErrOwnershipChanged + } + if response.GetRequireFullSnapshot() { + reporter.sequence = sequence + return ErrFullSnapshotRequired + } + if response.GetAcceptedOwnershipEpoch() != reporter.session.OwnershipEpoch { + return ErrOwnershipChanged + } + reporter.sequence = sequence + return nil +} + +func validateRegistration(workerID string, response *controlplanev1.RegisterWorkerResponse) (Registration, error) { + if response == nil || response.GetWorkerId() != workerID || !workerruntime.ValidIdentifier(response.GetSessionId()) || + response.GetOwnershipEpoch() == 0 || response.GetHeartbeatInterval() == nil || response.GetMaxStaleAge() == nil || + response.GetHeartbeatInterval().AsDuration() <= 0 || response.GetMaxStaleAge().AsDuration() <= 0 { + return Registration{}, ErrInvalidOptions + } + return Registration{ + SessionID: response.GetSessionId(), OwnershipEpoch: response.GetOwnershipEpoch(), + HeartbeatInterval: response.GetHeartbeatInterval().AsDuration(), MaxStaleAge: response.GetMaxStaleAge().AsDuration(), + }, nil +} + +func runtimeRequest(report workerruntime.Report) (*controlplanev1.ReportRuntimeRequest, error) { + counters := make([]*controlplanev1.ProxyRuntime, len(report.Counters)) + for index, counter := range report.Counters { + if counter.Active < 0 || counter.Reserved < 0 || counter.Active > math.MaxUint32 || counter.Reserved > math.MaxUint32 { + return nil, ErrInvalidOptions + } + counters[index] = &controlplanev1.ProxyRuntime{ + ProxyId: counter.ProxyID, Active: uint32(counter.Active), Reserved: uint32(counter.Reserved), Draining: counter.Draining, + } + } + return &controlplanev1.ReportRuntimeRequest{ + WorkerId: report.WorkerID, SessionId: report.SessionID, SnapshotVersion: report.SnapshotVersion, + OwnershipEpoch: report.OwnershipEpoch, ReportSequence: report.Sequence, Counters: counters, + ObservedAt: timestamppb.New(report.ObservedAt), + }, nil +} + +func cloneLabels(source map[string]string) map[string]string { + if source == nil { + return nil + } + result := make(map[string]string, len(source)) + for key, value := range source { + result[key] = value + } + return result +} + +var _ RuntimeSource = (*snapshot.Store)(nil) diff --git a/internal/gateway/controlplane/reporter_test.go b/internal/gateway/controlplane/reporter_test.go new file mode 100644 index 0000000..1917e0a --- /dev/null +++ b/internal/gateway/controlplane/reporter_test.go @@ -0,0 +1,120 @@ +package controlplane + +import ( + "context" + "errors" + "testing" + "time" + + controlplanev1 "proxy-pool/gen/controlplane/v1" + "proxy-pool/internal/gateway/snapshot" + + "google.golang.org/protobuf/types/known/durationpb" +) + +func TestRuntimeReporterRegistersAndReportsSnapshotCounters(t *testing.T) { + store := snapshot.NewStore("cluster-a", "worker-a") + applySnapshot(t, store) + client := &clientStub{registration: &controlplanev1.RegisterWorkerResponse{ + WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7, + HeartbeatInterval: durationpb.New(10 * time.Second), MaxStaleAge: durationpb.New(30 * time.Second), + }} + now := time.Date(2026, 7, 31, 10, 0, 0, 0, time.UTC) + reporter, err := NewRuntimeReporter(client, store, Options{ + WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, + Now: func() time.Time { return now }, + }) + if err != nil { + t.Fatalf("NewRuntimeReporter(): %v", err) + } + registration, err := reporter.Register(context.Background()) + if err != nil || registration.SessionID != "session-a" || client.register.GetZone() != "zone-a" { + t.Fatalf("Register() = %+v, %v; request=%+v", registration, err, client.register) + } + if err := reporter.Report(context.Background()); err != nil { + t.Fatalf("Report(): %v", err) + } + if client.runtime.GetSessionId() != "session-a" || client.runtime.GetReportSequence() != 1 || + client.runtime.GetSnapshotVersion() != 1 || !client.runtime.GetObservedAt().AsTime().Equal(now) { + t.Fatalf("ReportRuntime request = %+v", client.runtime) + } + if err := reporter.Report(context.Background()); err != nil { + t.Fatalf("second Report(): %v", err) + } + if client.runtime.GetReportSequence() != 2 { + t.Fatalf("second report sequence = %d, want 2", client.runtime.GetReportSequence()) + } +} + +func TestRuntimeReporterRetainsSequenceOnTransportErrorAndSurfacesResync(t *testing.T) { + store := snapshot.NewStore("cluster-a", "worker-a") + applySnapshot(t, store) + transportErr := errors.New("temporary transport failure") + client := &clientStub{registration: &controlplanev1.RegisterWorkerResponse{ + WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7, + HeartbeatInterval: durationpb.New(time.Second), MaxStaleAge: durationpb.New(3 * time.Second), + }, runtimeErr: transportErr} + reporter, err := NewRuntimeReporter(client, store, Options{WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now}) + if err != nil { + t.Fatalf("NewRuntimeReporter(): %v", err) + } + if _, err := reporter.Register(context.Background()); err != nil { + t.Fatalf("Register(): %v", err) + } + if err := reporter.Report(context.Background()); !errors.Is(err, transportErr) { + t.Fatalf("Report() error = %v, want transport error", err) + } + client.runtimeErr = nil + client.runtimeResponse = &controlplanev1.ReportRuntimeResponse{RequireFullSnapshot: true} + if err := reporter.Report(context.Background()); !errors.Is(err, ErrFullSnapshotRequired) { + t.Fatalf("Report() error = %v, want ErrFullSnapshotRequired", err) + } + if client.runtime.GetReportSequence() != 1 { + t.Fatalf("retried report sequence = %d, want 1", client.runtime.GetReportSequence()) + } +} + +func TestRuntimeReporterRejectsInvalidState(t *testing.T) { + store := snapshot.NewStore("cluster-a", "worker-a") + if _, err := NewRuntimeReporter(&clientStub{}, store, Options{}); !errors.Is(err, ErrInvalidOptions) { + t.Fatalf("NewRuntimeReporter() error = %v, want ErrInvalidOptions", err) + } + reporter, err := NewRuntimeReporter(&clientStub{}, store, Options{WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now}) + if err != nil { + t.Fatalf("NewRuntimeReporter(): %v", err) + } + if err := reporter.Report(context.Background()); !errors.Is(err, ErrNotRegistered) { + t.Fatalf("Report() error = %v, want ErrNotRegistered", err) + } +} + +type clientStub struct { + register *controlplanev1.RegisterWorkerRequest + registration *controlplanev1.RegisterWorkerResponse + registerErr error + runtime *controlplanev1.ReportRuntimeRequest + runtimeResponse *controlplanev1.ReportRuntimeResponse + runtimeErr error +} + +func (client *clientStub) RegisterWorker(_ context.Context, request *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) { + client.register = request + return client.registration, client.registerErr +} + +func (client *clientStub) ReportRuntime(_ context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) { + client.runtime = request + if client.runtimeResponse == nil { + client.runtimeResponse = &controlplanev1.ReportRuntimeResponse{AcceptedOwnershipEpoch: 7} + } + return client.runtimeResponse, client.runtimeErr +} + +func applySnapshot(t *testing.T, store *snapshot.Store) { + t.Helper() + envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 7, Version: 1, Full: true} + envelope.Checksum = snapshot.Checksum(nil) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(): %v", err) + } +}