proxy-pool/internal/checker/controlplane/runner.go

// Package controlplane runs bounded Checker task pulls and fact reporting.
// It never accesses Controller storage; task lease validation stays server-side.
package controlplane

import (
	"context"
	"errors"
	"io"
	"sync"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/checker/probe"
	"proxy-pool/internal/domain/workerruntime"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var (
	ErrInvalidRunner        = errors.New("invalid checker runner")
	ErrTaskBatchTooLarge    = errors.New("checker task batch exceeds configured concurrency")
	ErrObservationsRejected = errors.New("checker observations were not accepted")
)

const (
	defaultPollInterval = 500 * time.Millisecond
	defaultRetryDelay   = 20 * time.Millisecond
	maximumInFlight     = 4_096
)

type TaskStream interface {
	Recv() (*controlplanev1.CheckTask, error)
}

type Client interface {
	StreamCheckTasks(context.Context, *controlplanev1.StreamCheckTasksRequest) (TaskStream, error)
	ReportObservations(context.Context, *controlplanev1.ObservationBatch) (*controlplanev1.ReportObservationsResponse, error)
}

type Executor interface {
	Execute(context.Context, *controlplanev1.CheckTask) probe.Result
}

type Options struct {
	CheckerID       string
	InstanceID      string
	MaxInFlight     int
	SupportedLevels []controlplanev1.CheckLevel
	ReportBatchSize int
	PollInterval    time.Duration
	RetryDelay      time.Duration
	Now             func() time.Time
}

// Runner consumes at most one server-bounded pull per RunOnce. Run adds a
// bounded polling delay, avoiding a hot loop while no tasks are available.
type Runner struct {
	client   Client
	executor Executor
	options  Options
}

func NewRunner(client Client, executor Executor, options Options) (*Runner, error) {
	if client == nil || executor == nil || !workerruntime.ValidIdentifier(options.CheckerID) ||
		!workerruntime.ValidIdentifier(options.InstanceID) || options.MaxInFlight <= 0 || options.MaxInFlight > maximumInFlight ||
		len(options.SupportedLevels) == 0 || options.Now == nil {
		return nil, ErrInvalidRunner
	}
	if options.ReportBatchSize == 0 {
		options.ReportBatchSize = options.MaxInFlight
	}
	if options.PollInterval == 0 {
		options.PollInterval = defaultPollInterval
	}
	if options.RetryDelay == 0 {
		options.RetryDelay = defaultRetryDelay
	}
	if options.ReportBatchSize <= 0 || options.ReportBatchSize > options.MaxInFlight || options.PollInterval <= 0 || options.RetryDelay < 0 ||
		!validLevels(options.SupportedLevels) {
		return nil, ErrInvalidRunner
	}
	return &Runner{client: client, executor: executor, options: options}, nil
}

func (runner *Runner) Run(ctx context.Context) error {
	if ctx == nil || runner == nil {
		return ErrInvalidRunner
	}
	for {
		err := runner.RunOnce(ctx)
		if err != nil && ctx.Err() != nil {
			return ctx.Err()
		}
		wait := runner.options.PollInterval
		if err != nil {
			wait = runner.options.RetryDelay
		}
		if err := waitFor(ctx, wait); err != nil {
			return err
		}
	}
}

func (runner *Runner) RunOnce(ctx context.Context) error {
	if ctx == nil || runner == nil || runner.client == nil || runner.executor == nil || runner.options.Now == nil {
		return ErrInvalidRunner
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	stream, err := runner.client.StreamCheckTasks(ctx, &controlplanev1.StreamCheckTasksRequest{
		CheckerId: runner.options.CheckerID, InstanceId: runner.options.InstanceID,
		MaxInFlight: uint32(runner.options.MaxInFlight), SupportedLevels: append([]controlplanev1.CheckLevel(nil), runner.options.SupportedLevels...),
	})
	if err != nil {
		return err
	}
	tasks, err := collectTasks(stream, runner.options.MaxInFlight)
	if err != nil || len(tasks) == 0 {
		return err
	}
	observations := runner.executeTasks(ctx, tasks)
	for start := 0; start < len(observations); start += runner.options.ReportBatchSize {
		end := start + runner.options.ReportBatchSize
		if end > len(observations) {
			end = len(observations)
		}
		response, reportErr := runner.client.ReportObservations(ctx, &controlplanev1.ObservationBatch{
			CheckerId: runner.options.CheckerID, Observations: observations[start:end],
		})
		if reportErr != nil {
			return reportErr
		}
		if response == nil || int(response.GetAccepted()) != end-start || response.GetRejected() != 0 {
			return ErrObservationsRejected
		}
	}
	return nil
}

func collectTasks(stream TaskStream, maximum int) ([]*controlplanev1.CheckTask, error) {
	if stream == nil || maximum <= 0 {
		return nil, ErrInvalidRunner
	}
	tasks := make([]*controlplanev1.CheckTask, 0, maximum)
	for {
		task, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return tasks, nil
		}
		if err != nil {
			return nil, err
		}
		if task == nil {
			return nil, ErrInvalidRunner
		}
		tasks = append(tasks, proto.Clone(task).(*controlplanev1.CheckTask))
		if len(tasks) > maximum {
			return nil, ErrTaskBatchTooLarge
		}
	}
}

func (runner *Runner) executeTasks(ctx context.Context, tasks []*controlplanev1.CheckTask) []*controlplanev1.HealthObservation {
	results := make([]*controlplanev1.HealthObservation, len(tasks))
	jobs := make(chan int)
	var group sync.WaitGroup
	workers := runner.options.MaxInFlight
	if workers > len(tasks) {
		workers = len(tasks)
	}
	for worker := 0; worker < workers; worker++ {
		group.Add(1)
		go func() {
			defer group.Done()
			for index := range jobs {
				results[index] = runner.executeTask(ctx, tasks[index])
			}
		}()
	}
	for index := range tasks {
		jobs <- index
	}
	close(jobs)
	group.Wait()
	return results
}

func (runner *Runner) executeTask(ctx context.Context, task *controlplanev1.CheckTask) *controlplanev1.HealthObservation {
	result := probe.Result{FailureClass: probe.FailureInvalidTask}
	if task != nil && task.GetAttempt() > 0 && task.GetMaxAttempts() >= task.GetAttempt() {
		for attempt := task.GetAttempt(); attempt <= task.GetMaxAttempts(); attempt++ {
			attemptTask := proto.Clone(task).(*controlplanev1.CheckTask)
			attemptTask.Attempt = attempt
			result = runner.executor.Execute(ctx, attemptTask)
			if result.Success || ctx.Err() != nil || attempt == task.GetMaxAttempts() {
				break
			}
			if waitFor(ctx, runner.options.RetryDelay) != nil {
				break
			}
		}
	}
	if result.Latency < 0 {
		result = probe.Result{FailureClass: probe.FailureInvalidTask}
	}
	if result.Success {
		result.FailureClass = ""
	} else if result.FailureClass == "" {
		result.FailureClass = probe.FailureProxyRequest
	}
	observation := &controlplanev1.HealthObservation{
		TaskId: task.GetTaskId(), LeaseToken: task.GetLeaseToken(), ProxyId: task.GetProxyId(), Level: task.GetLevel(), Success: result.Success,
		FailureClass: result.FailureClass, Latency: durationpb.New(result.Latency),
		ObservedAt: timestamppb.New(runner.options.Now().UTC()),
	}
	if task.GetLevel() == controlplanev1.CheckLevel_CHECK_LEVEL_TARGET {
		observation.RoutingName = task.GetRoutingName()
		observation.TargetUrl = task.GetTargetUrl()
	}
	return observation
}

func validLevels(levels []controlplanev1.CheckLevel) bool {
	seen := make(map[controlplanev1.CheckLevel]struct{}, len(levels))
	for _, level := range levels {
		switch level {
		case controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS, controlplanev1.CheckLevel_CHECK_LEVEL_TARGET:
		default:
			return false
		}
		if _, duplicate := seen[level]; duplicate {
			return false
		}
		seen[level] = struct{}{}
	}
	return true
}

func waitFor(ctx context.Context, duration time.Duration) error {
	if duration <= 0 {
		return ctx.Err()
	}
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}