267 lines
7.6 KiB
Go
267 lines
7.6 KiB
Go
package controlplane
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
domain "proxy-pool/internal/domain/outcome"
|
|
"proxy-pool/internal/domain/workerruntime"
|
|
gatewayOutcome "proxy-pool/internal/gateway/outcome"
|
|
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/protobuf/types/known/durationpb"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidOutcomeReporter = errors.New("invalid gateway outcome reporter")
|
|
ErrOutcomeUnconfirmed = errors.New("controller did not confirm gateway outcome batch")
|
|
)
|
|
|
|
const (
|
|
outcomeRetryInitialDelay = 100 * time.Millisecond
|
|
outcomeRetryMaxDelay = 5 * time.Second
|
|
)
|
|
|
|
type OutcomeStream interface {
|
|
Send(*controlplanev1.OutcomeBatch) error
|
|
CloseAndRecv() (*controlplanev1.ReportOutcomesResponse, error)
|
|
}
|
|
|
|
type OutcomeRPCClient interface {
|
|
ReportOutcomes(context.Context) (OutcomeStream, error)
|
|
}
|
|
|
|
type OutcomeReporterOptions struct {
|
|
WorkerID string
|
|
MaxBatchesPerReportRPC int
|
|
}
|
|
|
|
type OutcomeReporter struct {
|
|
client OutcomeRPCClient
|
|
queue *gatewayOutcome.Queue
|
|
options OutcomeReporterOptions
|
|
|
|
mu sync.Mutex
|
|
sessionID string
|
|
sequence uint64
|
|
pending []domain.Batch
|
|
}
|
|
|
|
func NewOutcomeReporter(client OutcomeRPCClient, queue *gatewayOutcome.Queue, options OutcomeReporterOptions) (*OutcomeReporter, error) {
|
|
if client == nil || queue == nil || !workerruntime.ValidIdentifier(options.WorkerID) {
|
|
return nil, ErrInvalidOutcomeReporter
|
|
}
|
|
if options.MaxBatchesPerReportRPC <= 0 {
|
|
options.MaxBatchesPerReportRPC = domain.MaxBatchesPerStream
|
|
}
|
|
if options.MaxBatchesPerReportRPC > domain.MaxBatchesPerStream {
|
|
return nil, ErrInvalidOutcomeReporter
|
|
}
|
|
return &OutcomeReporter{client: client, queue: queue, options: options}, nil
|
|
}
|
|
|
|
// RunRegistered continuously batches local observations for a single
|
|
// Controller session. Transient delivery failures retain the exact pending
|
|
// batches and retry them in the same session with bounded backoff. A replaced
|
|
// Controller session reissues the retained events with the new session ID and
|
|
// its sequence space.
|
|
func (reporter *OutcomeReporter) RunRegistered(ctx context.Context, registration Registration) error {
|
|
if reporter == nil || ctx == nil || !workerruntime.ValidIdentifier(registration.SessionID) {
|
|
return ErrInvalidOutcomeReporter
|
|
}
|
|
reporter.startSession(registration.SessionID)
|
|
retryDelay := outcomeRetryInitialDelay
|
|
for {
|
|
err := reporter.Report(ctx, registration.SessionID)
|
|
if err == nil {
|
|
retryDelay = outcomeRetryInitialDelay
|
|
continue
|
|
}
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if terminalOutcomeError(err) {
|
|
return err
|
|
}
|
|
if err := waitOutcomeRetry(ctx, retryDelay); err != nil {
|
|
return err
|
|
}
|
|
retryDelay = nextOutcomeRetryDelay(retryDelay)
|
|
}
|
|
}
|
|
|
|
func (reporter *OutcomeReporter) startSession(sessionID string) {
|
|
reporter.mu.Lock()
|
|
defer reporter.mu.Unlock()
|
|
if reporter.sessionID == sessionID {
|
|
return
|
|
}
|
|
reporter.sessionID = sessionID
|
|
reporter.sequence = 0
|
|
for index := range reporter.pending {
|
|
reporter.pending[index].SessionID = sessionID
|
|
reporter.pending[index].Sequence = uint64(index + 1)
|
|
}
|
|
}
|
|
|
|
func terminalOutcomeError(err error) bool {
|
|
if errors.Is(err, ErrInvalidOutcomeReporter) || errors.Is(err, ErrNotRegistered) {
|
|
return true
|
|
}
|
|
switch status.Code(err) {
|
|
case codes.Aborted, codes.AlreadyExists, codes.FailedPrecondition,
|
|
codes.InvalidArgument, codes.PermissionDenied, codes.Unauthenticated:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func waitOutcomeRetry(ctx context.Context, delay time.Duration) error {
|
|
timer := time.NewTimer(delay)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-timer.C:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func nextOutcomeRetryDelay(delay time.Duration) time.Duration {
|
|
if delay >= outcomeRetryMaxDelay/2 {
|
|
return outcomeRetryMaxDelay
|
|
}
|
|
return delay * 2
|
|
}
|
|
|
|
func (reporter *OutcomeReporter) Report(ctx context.Context, sessionID string) error {
|
|
if reporter == nil || ctx == nil || !workerruntime.ValidIdentifier(sessionID) {
|
|
return ErrInvalidOutcomeReporter
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
reporter.mu.Lock()
|
|
defer reporter.mu.Unlock()
|
|
if reporter.sessionID != sessionID {
|
|
return ErrNotRegistered
|
|
}
|
|
if err := reporter.fillPending(ctx); err != nil {
|
|
return err
|
|
}
|
|
streamCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
stream, err := reporter.client.ReportOutcomes(streamCtx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, batch := range reporter.pending {
|
|
request, requestErr := outcomeRequest(batch)
|
|
if requestErr != nil {
|
|
return requestErr
|
|
}
|
|
if sendErr := stream.Send(request); sendErr != nil {
|
|
return sendErr
|
|
}
|
|
}
|
|
response, err := stream.CloseAndRecv()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
last := reporter.pending[len(reporter.pending)-1].Sequence
|
|
if response == nil || response.GetAcceptedThroughSequence() != last {
|
|
return ErrOutcomeUnconfirmed
|
|
}
|
|
reporter.sequence = last
|
|
reporter.pending = nil
|
|
return nil
|
|
}
|
|
|
|
func (reporter *OutcomeReporter) fillPending(ctx context.Context) error {
|
|
if len(reporter.pending) == 0 {
|
|
events, err := reporter.queue.Next(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := reporter.append(events); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for len(reporter.pending) < reporter.options.MaxBatchesPerReportRPC {
|
|
events, ok := reporter.queue.TryNext()
|
|
if !ok {
|
|
return nil
|
|
}
|
|
if err := reporter.append(events); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (reporter *OutcomeReporter) append(events []domain.Event) error {
|
|
batch, _, err := domain.NormalizeBatch(domain.Batch{
|
|
WorkerID: reporter.options.WorkerID, SessionID: reporter.sessionID,
|
|
Sequence: reporter.sequence + uint64(len(reporter.pending)) + 1, Events: events,
|
|
}, reporter.queue.MaxBatch())
|
|
if err != nil {
|
|
return errors.Join(ErrInvalidOutcomeReporter, err)
|
|
}
|
|
reporter.pending = append(reporter.pending, batch)
|
|
return nil
|
|
}
|
|
|
|
func outcomeRequest(batch domain.Batch) (*controlplanev1.OutcomeBatch, error) {
|
|
proxies := make([]*controlplanev1.ProxyOutcome, len(batch.Events))
|
|
for index, event := range batch.Events {
|
|
stage, ok := outcomeStage(event.Stage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("%w: unknown outcome stage", ErrInvalidOutcomeReporter)
|
|
}
|
|
proxies[index] = &controlplanev1.ProxyOutcome{
|
|
ProxyId: event.ProxyID, RoutingName: event.RoutingName, Stage: stage, Success: event.Success,
|
|
ErrorClass: string(event.ErrorClass), Latency: durationpb.New(event.Latency), ObservedAt: timestamppb.New(event.ObservedAt),
|
|
}
|
|
}
|
|
return &controlplanev1.OutcomeBatch{
|
|
WorkerId: batch.WorkerID, SessionId: batch.SessionID, Sequence: batch.Sequence, Outcomes: proxies,
|
|
}, nil
|
|
}
|
|
|
|
func outcomeStage(stage domain.Stage) (controlplanev1.OutcomeStage, bool) {
|
|
switch stage {
|
|
case domain.StageDial:
|
|
return controlplanev1.OutcomeStage_OUTCOME_STAGE_DIAL, true
|
|
case domain.StageProxyHandshake:
|
|
return controlplanev1.OutcomeStage_OUTCOME_STAGE_PROXY_HANDSHAKE, true
|
|
case domain.StageResponseHeaders:
|
|
return controlplanev1.OutcomeStage_OUTCOME_STAGE_RESPONSE_HEADERS, true
|
|
case domain.StageTunnel:
|
|
return controlplanev1.OutcomeStage_OUTCOME_STAGE_TUNNEL, true
|
|
default:
|
|
return controlplanev1.OutcomeStage_OUTCOME_STAGE_UNSPECIFIED, false
|
|
}
|
|
}
|
|
|
|
type generatedOutcomeRPCClient struct {
|
|
client controlplanev1.WorkerControlPlaneClient
|
|
}
|
|
|
|
func NewGeneratedOutcomeRPCClient(client controlplanev1.WorkerControlPlaneClient) OutcomeRPCClient {
|
|
return generatedOutcomeRPCClient{client: client}
|
|
}
|
|
|
|
func (client generatedOutcomeRPCClient) ReportOutcomes(ctx context.Context) (OutcomeStream, error) {
|
|
if client.client == nil {
|
|
return nil, ErrInvalidOutcomeReporter
|
|
}
|
|
return client.client.ReportOutcomes(ctx)
|
|
}
|