feat: add gateway runtime control plane reporter
This commit is contained in:
youfak
parent
9edf1a9bab
commit
5a1873a9f0
177
internal/gateway/controlplane/reporter.go
Normal file
177
internal/gateway/controlplane/reporter.go
Normal file
@ -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)
|
||||
120
internal/gateway/controlplane/reporter_test.go
Normal file
120
internal/gateway/controlplane/reporter_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user