proxy-pool/internal/gateway/controlplane/reporter.go

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)