proxy-pool/internal/controller/health/scheduler_runner.go

package health

import (
	"context"
	"errors"
	"time"
)

var (
	ErrInvalidSchedulerRunner = errors.New("invalid health scheduler runner")
	ErrInvalidDueSource       = errors.New("invalid health due source response")
	ErrInvalidTaskSink        = errors.New("invalid health task sink response")
)

// DueSource is implemented by the Controller's future Redis due-index. Both
// calls are bounded and stay outside the Gateway request path.
type DueSource interface {
	InFlight(context.Context, time.Time) (int, error)
	DueCandidates(context.Context, time.Time, int) ([]Candidate, error)
}

// UpstreamTaskSource scopes due reads and shared capacity to one configured
// upstream while all leases remain in the same Redis namespace.
type UpstreamTaskSource interface {
	InFlightForUpstream(context.Context, string, time.Time) (int, error)
	DueCandidatesForUpstream(context.Context, string, time.Time, int) ([]Candidate, error)
}

type upstreamDueSource struct {
	source     UpstreamTaskSource
	upstreamID string
}

func NewUpstreamDueSource(source UpstreamTaskSource, upstreamID string) (DueSource, error) {
	if nilInterface(source) || upstreamID == "" {
		return nil, ErrInvalidSchedulerRunner
	}
	return upstreamDueSource{source: source, upstreamID: upstreamID}, nil
}

func (source upstreamDueSource) InFlight(ctx context.Context, now time.Time) (int, error) {
	return source.source.InFlightForUpstream(ctx, source.upstreamID, now)
}

func (source upstreamDueSource) DueCandidates(ctx context.Context, now time.Time, limit int) ([]Candidate, error) {
	return source.source.DueCandidatesForUpstream(ctx, source.upstreamID, now, limit)
}

// TaskSink atomically offers an already bounded batch to the shared leased
// task store. It must leave unaccepted candidates eligible for a later tick.
type TaskSink interface {
	Offer(context.Context, []PlannedTask) (int, error)
}

type SchedulerRunnerOptions struct {
	PollInterval time.Duration
	BatchSize    int
	Now          func() time.Time
}

// SchedulerRunner performs one bounded tick at a time. Run has one ticker
// loop, and Tick performs no concurrent fan-out.
type SchedulerRunner struct {
	planner *Planner
	source  DueSource
	sink    TaskSink
	options SchedulerRunnerOptions
}

func NewSchedulerRunner(
	planner *Planner,
	source DueSource,
	sink TaskSink,
	options SchedulerRunnerOptions,
) (*SchedulerRunner, error) {
	if planner == nil || nilInterface(source) || nilInterface(sink) || options.PollInterval <= 0 ||
		options.BatchSize <= 0 || options.Now == nil {
		return nil, ErrInvalidSchedulerRunner
	}
	return &SchedulerRunner{planner: planner, source: source, sink: sink, options: options}, nil
}

type TickResult struct {
	Planned int
	Offered int
}

func (runner *SchedulerRunner) Tick(ctx context.Context) (TickResult, error) {
	if ctx == nil || runner == nil || runner.planner == nil || nilInterface(runner.source) ||
		nilInterface(runner.sink) || runner.options.Now == nil {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	if err := ctx.Err(); err != nil {
		return TickResult{}, err
	}
	now := runner.options.Now()
	if now.IsZero() {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	inFlight, err := runner.source.InFlight(ctx, now.UTC())
	if err != nil {
		return TickResult{}, err
	}
	if inFlight < 0 {
		return TickResult{}, ErrInvalidDueSource
	}
	if inFlight >= runner.planner.policy.MaxInFlight {
		return TickResult{}, nil
	}
	candidates, err := runner.source.DueCandidates(ctx, now.UTC(), runner.options.BatchSize)
	if err != nil {
		return TickResult{}, err
	}
	if len(candidates) > runner.options.BatchSize {
		return TickResult{}, ErrInvalidDueSource
	}
	plans, err := runner.planner.Plan(now.UTC(), inFlight, runner.options.BatchSize, candidates)
	if err != nil {
		return TickResult{}, err
	}
	result := TickResult{Planned: len(plans)}
	if len(plans) == 0 {
		return result, nil
	}
	offered, err := runner.sink.Offer(ctx, plans)
	if err != nil {
		return TickResult{}, err
	}
	if offered < 0 || offered > len(plans) {
		return TickResult{}, ErrInvalidTaskSink
	}
	result.Offered = offered
	return result, nil
}

func (runner *SchedulerRunner) Run(ctx context.Context) error {
	if ctx == nil || runner == nil {
		return ErrInvalidSchedulerRunner
	}
	if _, err := runner.Tick(ctx); err != nil {
		return err
	}
	ticker := time.NewTicker(runner.options.PollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			if _, err := runner.Tick(ctx); err != nil {
				return err
			}
		}
	}
}