proxy-pool/internal/gateway/controlplane/outcome_reporter.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.
func (reporter *OutcomeReporter) RunRegistered(ctx context.Context, registration Registration) error {
	if reporter == nil || ctx == nil || !workerruntime.ValidIdentifier(registration.SessionID) {
		return ErrInvalidOutcomeReporter
	}
	reporter.mu.Lock()
	reporter.sessionID = registration.SessionID
	reporter.sequence = 0
	reporter.pending = nil
	reporter.mu.Unlock()
	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 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)
}