417 lines
13 KiB
Go
417 lines
13 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sort"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"proxy-pool/internal/config"
|
|
)
|
|
|
|
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)
|
|
}
|
|
|
|
// EgressUpstreamTaskSource supplies bounded per-upstream candidates for one
|
|
// configured egress probe URL. Its optional nature keeps BASIC-only stores
|
|
// source compatible during rollout.
|
|
type EgressUpstreamTaskSource interface {
|
|
UpstreamTaskSource
|
|
DueEgressCandidatesForUpstream(context.Context, string, 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)
|
|
}
|
|
|
|
type egressUpstreamDueSource struct {
|
|
source EgressUpstreamTaskSource
|
|
upstreamID string
|
|
targetURL string
|
|
}
|
|
|
|
func newEgressUpstreamDueSource(source EgressUpstreamTaskSource, upstreamID, targetURL string) (DueSource, error) {
|
|
if nilInterface(source) || upstreamID == "" || targetURL == "" {
|
|
return nil, ErrInvalidSchedulerRunner
|
|
}
|
|
return egressUpstreamDueSource{source: source, upstreamID: upstreamID, targetURL: targetURL}, nil
|
|
}
|
|
|
|
func (source egressUpstreamDueSource) InFlight(ctx context.Context, now time.Time) (int, error) {
|
|
return source.source.InFlightForUpstream(ctx, source.upstreamID, now)
|
|
}
|
|
|
|
func (source egressUpstreamDueSource) DueCandidates(ctx context.Context, now time.Time, limit int) ([]Candidate, error) {
|
|
return source.source.DueEgressCandidatesForUpstream(ctx, source.upstreamID, source.targetURL, 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 || !validSchedulerRunnerParts(source, sink, options) {
|
|
return nil, ErrInvalidSchedulerRunner
|
|
}
|
|
return &SchedulerRunner{planner: planner, source: source, sink: sink, options: options}, nil
|
|
}
|
|
|
|
// ConfiguredSchedulerRunner resolves the current effective upstream check
|
|
// policy on every tick. It keeps a running Controller aligned with validated
|
|
// config reloads without putting configuration reads in the Gateway path.
|
|
type ConfiguredSchedulerRunner struct {
|
|
configuration ConfigurationSource
|
|
upstreamID string
|
|
source DueSource
|
|
sink TaskSink
|
|
options SchedulerRunnerOptions
|
|
}
|
|
|
|
// ConfiguredSchedulerSupervisor discovers enabled upstreams from each
|
|
// published configuration snapshot. It lets newly enabled upstreams begin
|
|
// bounded BASIC scheduling without a Controller restart.
|
|
type ConfiguredSchedulerSupervisor struct {
|
|
configuration ConfigurationSource
|
|
source UpstreamTaskSource
|
|
sink TaskSink
|
|
options SchedulerRunnerOptions
|
|
cursor atomic.Uint64
|
|
}
|
|
|
|
func NewConfiguredSchedulerRunner(
|
|
configuration ConfigurationSource,
|
|
upstreamID string,
|
|
source DueSource,
|
|
sink TaskSink,
|
|
options SchedulerRunnerOptions,
|
|
) (*ConfiguredSchedulerRunner, error) {
|
|
if nilInterface(configuration) || strings.TrimSpace(upstreamID) != upstreamID || upstreamID == "" ||
|
|
!validSchedulerRunnerParts(source, sink, options) {
|
|
return nil, ErrInvalidSchedulerRunner
|
|
}
|
|
return &ConfiguredSchedulerRunner{
|
|
configuration: configuration, upstreamID: upstreamID, source: source, sink: sink, options: options,
|
|
}, nil
|
|
}
|
|
|
|
func NewConfiguredSchedulerSupervisor(
|
|
configuration ConfigurationSource,
|
|
source UpstreamTaskSource,
|
|
sink TaskSink,
|
|
options SchedulerRunnerOptions,
|
|
) (*ConfiguredSchedulerSupervisor, error) {
|
|
if nilInterface(configuration) || nilInterface(source) || nilInterface(sink) || options.PollInterval <= 0 ||
|
|
options.BatchSize <= 0 || options.Now == nil {
|
|
return nil, ErrInvalidSchedulerRunner
|
|
}
|
|
return &ConfiguredSchedulerSupervisor{
|
|
configuration: configuration, 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) ||
|
|
!validSchedulerRunnerParts(runner.source, runner.sink, runner.options) {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
return tickWithPlanner(ctx, runner.planner, runner.source, runner.sink, runner.options)
|
|
}
|
|
|
|
func (runner *ConfiguredSchedulerRunner) Tick(ctx context.Context) (TickResult, error) {
|
|
if ctx == nil || runner == nil || nilInterface(runner.configuration) || runner.upstreamID == "" ||
|
|
!validSchedulerRunnerParts(runner.source, runner.sink, runner.options) {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
planner, configured, err := configuredPlanner(runner.configuration, runner.upstreamID)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
if !configured {
|
|
return TickResult{}, nil
|
|
}
|
|
return tickWithPlanner(ctx, planner, runner.source, runner.sink, schedulerOptionsForPlanner(runner.options, planner))
|
|
}
|
|
|
|
func (supervisor *ConfiguredSchedulerSupervisor) Tick(ctx context.Context) (TickResult, error) {
|
|
if ctx == nil || supervisor == nil || nilInterface(supervisor.configuration) || nilInterface(supervisor.source) ||
|
|
nilInterface(supervisor.sink) || supervisor.options.PollInterval <= 0 || supervisor.options.BatchSize <= 0 || supervisor.options.Now == nil {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
configuration := supervisor.configuration.Current()
|
|
if configuration == nil {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
names := enabledUpstreamNames(configuration)
|
|
result := TickResult{}
|
|
for _, name := range names {
|
|
planner, configured, err := configuredPlannerForConfig(configuration, name)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
if !configured {
|
|
continue
|
|
}
|
|
basicSource, err := NewUpstreamDueSource(supervisor.source, name)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
check := config.EffectiveCheck(configuration.Defaults.Check, configuration.Upstreams[name].Check)
|
|
groups, err := supervisor.upstreamDueGroups(basicSource, name, check.URLs)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
groupOptions := schedulerOptionsForPlanner(supervisor.options, planner)
|
|
groupOptions.BatchSize = schedulerGroupBatchSize(groupOptions.BatchSize, len(groups))
|
|
start := int(supervisor.cursor.Add(1)-1) % len(groups)
|
|
for offset := range groups {
|
|
source := groups[(start+offset)%len(groups)]
|
|
tick, err := tickWithPlanner(ctx, planner, source, supervisor.sink, groupOptions)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
result.Planned += tick.Planned
|
|
result.Offered += tick.Offered
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (supervisor *ConfiguredSchedulerSupervisor) upstreamDueGroups(basic DueSource, upstreamID string, urls []string) ([]DueSource, error) {
|
|
groups := []DueSource{basic}
|
|
egress, supported := supervisor.source.(EgressUpstreamTaskSource)
|
|
if !supported || len(urls) == 0 {
|
|
return groups, nil
|
|
}
|
|
for _, targetURL := range urls {
|
|
source, err := newEgressUpstreamDueSource(egress, upstreamID, targetURL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
groups = append(groups, source)
|
|
}
|
|
return groups, nil
|
|
}
|
|
|
|
func schedulerGroupBatchSize(batchSize, groups int) int {
|
|
if batchSize <= 0 || groups <= 0 {
|
|
return 0
|
|
}
|
|
result := (batchSize + groups - 1) / groups
|
|
if result < 1 {
|
|
return 1
|
|
}
|
|
return result
|
|
}
|
|
|
|
func tickWithPlanner(
|
|
ctx context.Context,
|
|
planner *Planner,
|
|
source DueSource,
|
|
sink TaskSink,
|
|
options SchedulerRunnerOptions,
|
|
) (TickResult, error) {
|
|
if ctx == nil || planner == nil || !validSchedulerRunnerParts(source, sink, options) {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
now := options.Now()
|
|
if now.IsZero() {
|
|
return TickResult{}, ErrInvalidSchedulerRunner
|
|
}
|
|
inFlight, err := source.InFlight(ctx, now.UTC())
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
if inFlight < 0 {
|
|
return TickResult{}, ErrInvalidDueSource
|
|
}
|
|
if inFlight >= planner.policy.MaxInFlight {
|
|
return TickResult{}, nil
|
|
}
|
|
candidates, err := source.DueCandidates(ctx, now.UTC(), options.BatchSize)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
if len(candidates) > options.BatchSize {
|
|
return TickResult{}, ErrInvalidDueSource
|
|
}
|
|
plans, err := planner.Plan(now.UTC(), inFlight, options.BatchSize, candidates)
|
|
if err != nil {
|
|
return TickResult{}, err
|
|
}
|
|
result := TickResult{Planned: len(plans)}
|
|
if len(plans) == 0 {
|
|
return result, nil
|
|
}
|
|
offered, err := 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 runner == nil {
|
|
return ErrInvalidSchedulerRunner
|
|
}
|
|
return runScheduler(ctx, runner.options.PollInterval, runner.Tick)
|
|
}
|
|
|
|
func (runner *ConfiguredSchedulerRunner) Run(ctx context.Context) error {
|
|
if runner == nil {
|
|
return ErrInvalidSchedulerRunner
|
|
}
|
|
return runScheduler(ctx, runner.options.PollInterval, runner.Tick)
|
|
}
|
|
|
|
func (supervisor *ConfiguredSchedulerSupervisor) Run(ctx context.Context) error {
|
|
if supervisor == nil {
|
|
return ErrInvalidSchedulerRunner
|
|
}
|
|
return runScheduler(ctx, supervisor.options.PollInterval, supervisor.Tick)
|
|
}
|
|
|
|
func validSchedulerRunnerParts(source DueSource, sink TaskSink, options SchedulerRunnerOptions) bool {
|
|
return !nilInterface(source) && !nilInterface(sink) && options.PollInterval > 0 && options.BatchSize > 0 && options.Now != nil
|
|
}
|
|
|
|
func configuredPlanner(configuration ConfigurationSource, upstreamID string) (*Planner, bool, error) {
|
|
current := configuration.Current()
|
|
if current == nil {
|
|
return nil, false, ErrInvalidSchedulerRunner
|
|
}
|
|
return configuredPlannerForConfig(current, upstreamID)
|
|
}
|
|
|
|
func configuredPlannerForConfig(current *config.Config, upstreamID string) (*Planner, bool, error) {
|
|
if current == nil {
|
|
return nil, false, ErrInvalidSchedulerRunner
|
|
}
|
|
upstream, exists := current.Upstreams[upstreamID]
|
|
if !exists || !upstream.Enabled {
|
|
return nil, false, nil
|
|
}
|
|
check := config.EffectiveCheck(current.Defaults.Check, upstream.Check)
|
|
planner, err := NewPlanner(SchedulePolicy{
|
|
Interval: check.Interval.Value(), Jitter: check.Jitter, MaxInFlight: check.MaxInFlight,
|
|
Timeout: check.Timeout.Value(), MaxAttempts: check.MaxAttempts,
|
|
})
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
return planner, true, nil
|
|
}
|
|
|
|
func schedulerOptionsForPlanner(options SchedulerRunnerOptions, planner *Planner) SchedulerRunnerOptions {
|
|
if planner.policy.MaxInFlight < options.BatchSize {
|
|
options.BatchSize = planner.policy.MaxInFlight
|
|
}
|
|
return options
|
|
}
|
|
|
|
func enabledUpstreamNames(configuration *config.Config) []string {
|
|
names := make([]string, 0, len(configuration.Upstreams))
|
|
for name, upstream := range configuration.Upstreams {
|
|
if upstream.Enabled {
|
|
names = append(names, name)
|
|
}
|
|
}
|
|
sort.Strings(names)
|
|
return names
|
|
}
|
|
|
|
func runScheduler(ctx context.Context, pollInterval time.Duration, tick func(context.Context) (TickResult, error)) error {
|
|
if ctx == nil || pollInterval <= 0 || tick == nil {
|
|
return ErrInvalidSchedulerRunner
|
|
}
|
|
if _, err := tick(ctx); err != nil {
|
|
return err
|
|
}
|
|
ticker := time.NewTicker(pollInterval)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-ticker.C:
|
|
if _, err := tick(ctx); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
}
|