diff --git a/internal/adapters/redisactivity/health_tasks.go b/internal/adapters/redisactivity/health_tasks.go index c7dc9e1..413066e 100644 --- a/internal/adapters/redisactivity/health_tasks.go +++ b/internal/adapters/redisactivity/health_tasks.go @@ -22,11 +22,14 @@ type healthTaskRecord struct { Version int `json:"version"` TaskID string `json:"taskId"` ProxyID string `json:"proxyId"` + UpstreamID string `json:"upstreamId"` + UpstreamTasksKey string `json:"upstreamTasksKey"` Level string `json:"level"` Priority int `json:"priority"` DeadlineMS int64 `json:"deadlineMs"` NextDueMS int64 `json:"nextDueMs"` Attempts int `json:"attempts"` + MaxInFlight int `json:"maxInFlight"` State string `json:"state"` RoutingName string `json:"routingName,omitempty"` TargetURL string `json:"targetUrl,omitempty"` @@ -38,15 +41,18 @@ type healthTaskRecord struct { } type healthTaskRequest struct { - Limit int `json:"limit"` - Tasks []healthTaskRecord `json:"tasks,omitempty"` - CheckerID string `json:"checkerId,omitempty"` - InstanceID string `json:"instanceId,omitempty"` - MaxInFlight int `json:"maxInFlight,omitempty"` - LeaseTTLMS int64 `json:"leaseTTLMS,omitempty"` - Levels []healthDomain.Level `json:"levels,omitempty"` - Tokens []string `json:"tokens,omitempty"` - Fact *healthTaskFact `json:"fact,omitempty"` + Limit int `json:"limit"` + Tasks []healthTaskRecord `json:"tasks,omitempty"` + CheckerID string `json:"checkerId,omitempty"` + InstanceID string `json:"instanceId,omitempty"` + MaxInFlight int `json:"maxInFlight,omitempty"` + LeaseTTLMS int64 `json:"leaseTTLMS,omitempty"` + Levels []healthDomain.Level `json:"levels,omitempty"` + UpstreamID string `json:"upstreamId,omitempty"` + UpstreamTasksKey string `json:"upstreamTasksKey,omitempty"` + ScanLimit int `json:"scanLimit,omitempty"` + Tokens []string `json:"tokens,omitempty"` + Fact *healthTaskFact `json:"fact,omitempty"` } type healthTaskFact struct { @@ -78,35 +84,68 @@ func (a *Adapter) InFlight(ctx context.Context, now time.Time) (int, error) { // DueCandidates returns a bounded BASIC due batch. Stale references are // discarded inside the Lua script before they can reach the Controller. func (a *Adapter) DueCandidates(ctx context.Context, now time.Time, limit int) ([]healthDomain.Candidate, error) { + return a.dueCandidates(ctx, "", now, limit) +} + +func (a *Adapter) DueCandidatesForUpstream(ctx context.Context, upstreamID string, now time.Time, limit int) ([]healthDomain.Candidate, error) { + if !validHealthTaskIdentifier(upstreamID) { + return nil, healthDomain.ErrInvalidTaskBroker + } + return a.dueCandidates(ctx, upstreamID, now, limit) +} + +func (a *Adapter) dueCandidates(ctx context.Context, upstreamID string, now time.Time, limit int) ([]healthDomain.Candidate, error) { if a == nil || ctx == nil || now.IsZero() || limit <= 0 || limit > a.options.MaxCheckTasks { return nil, healthDomain.ErrInvalidTaskBroker } - reply, err := a.runHealthTaskScript(ctx, "due", now, healthTaskRequest{Limit: limit}, "runtime") + reply, err := a.runHealthTaskScript(ctx, "due", now, healthTaskRequest{ + Limit: limit, ScanLimit: a.options.MaxCandidateScan, UpstreamID: upstreamID, + }, "runtime") if err != nil { return nil, err } - if reply.Status != scriptOK || len(reply.Candidates) > limit { + if reply.Status != scriptOK || reply.CandidatesJSON == "" { return nil, invalidScriptReply("invalid health task due reply") } - result := make([]healthDomain.Candidate, 0, len(reply.Candidates)) - for _, candidate := range reply.Candidates { + var candidates []healthTaskCandidateWire + if err := decodeJSON(reply.CandidatesJSON, &candidates); err != nil || len(candidates) > limit { + return nil, invalidScriptReply("invalid health task due candidate payload") + } + result := make([]healthDomain.Candidate, 0, len(candidates)) + for _, candidate := range candidates { state := proxyDomain.State(candidate.State) switch state { case proxyDomain.StateFetched, proxyDomain.StateChecking, proxyDomain.StateAvailable, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy: default: return nil, invalidScriptReply("health task due reply contains invalid state") } - if !validHealthTaskIdentifier(candidate.ProxyID) || candidate.DueAtMS <= 0 { + if !validHealthTaskIdentifier(candidate.ProxyID) || !validHealthTaskIdentifier(candidate.UpstreamID) || candidate.DueAtMS <= 0 { return nil, invalidScriptReply("health task due reply contains invalid candidate") } result = append(result, healthDomain.Candidate{ - ProxyID: candidate.ProxyID, State: state, Level: healthDomain.LevelBasic, + ProxyID: candidate.ProxyID, UpstreamID: candidate.UpstreamID, State: state, Level: healthDomain.LevelBasic, DueAt: time.UnixMilli(candidate.DueAtMS).UTC(), }) } return result, nil } +func (a *Adapter) InFlightForUpstream(ctx context.Context, upstreamID string, now time.Time) (int, error) { + if a == nil || ctx == nil || now.IsZero() || !validHealthTaskIdentifier(upstreamID) { + return 0, healthDomain.ErrInvalidTaskBroker + } + reply, err := a.runHealthTaskScript(ctx, "upstream_inflight", now, healthTaskRequest{ + Limit: a.options.MaxCheckTasks, UpstreamID: upstreamID, UpstreamTasksKey: a.keys.upstreamTasks(upstreamID), + }, "runtime") + if err != nil { + return 0, err + } + if reply.Status != scriptOK || reply.Count < 0 { + return 0, invalidScriptReply("invalid upstream health task inflight reply") + } + return reply.Count, nil +} + // Offer atomically turns due BASIC plans into shared queued tasks. A plan that // loses the race leaves no duplicate task and no in-memory queue behind. func (a *Adapter) Offer(ctx context.Context, plans []healthDomain.PlannedTask) (int, error) { @@ -130,8 +169,9 @@ func (a *Adapter) Offer(ctx context.Context, plans []healthDomain.PlannedTask) ( } tasks[index] = healthTaskRecord{ Version: healthTaskRecordVersion, TaskID: healthDomain.TaskIDFor(plan), ProxyID: plan.Candidate.ProxyID, + UpstreamID: plan.Candidate.UpstreamID, UpstreamTasksKey: a.keys.upstreamTasks(plan.Candidate.UpstreamID), Level: string(plan.Candidate.Level), Priority: int(plan.Priority), DeadlineMS: plan.Deadline.UnixMilli(), - NextDueMS: plan.NextDue.UnixMilli(), Attempts: plan.Attempts, State: "QUEUED", + NextDueMS: plan.NextDue.UnixMilli(), Attempts: plan.Attempts, MaxInFlight: plan.MaxInFlight, State: "QUEUED", } } reply, err := a.runHealthTaskScript(ctx, "offer", now, healthTaskRequest{Limit: a.options.MaxCheckTasks, Tasks: tasks}, "runtime") @@ -294,9 +334,10 @@ func (a *Adapter) runHealthTaskScript( } func validateRedisPlannedTask(task healthDomain.PlannedTask, now time.Time) error { - if !validHealthTaskIdentifier(task.Candidate.ProxyID) || task.Candidate.Level != healthDomain.LevelBasic || + if !validHealthTaskIdentifier(task.Candidate.ProxyID) || !validHealthTaskIdentifier(task.Candidate.UpstreamID) || task.Candidate.Level != healthDomain.LevelBasic || task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" || task.Candidate.DueAt.IsZero() || - task.Deadline.IsZero() || !task.Deadline.After(now) || task.NextDue.IsZero() || !task.NextDue.After(now) || task.Attempts <= 0 { + task.Deadline.IsZero() || !task.Deadline.After(now) || task.NextDue.IsZero() || !task.NextDue.After(now) || + task.Attempts <= 0 || task.MaxInFlight <= 0 { return healthDomain.ErrInvalidLeasedTask } switch task.Candidate.State { @@ -333,9 +374,10 @@ func decodeHealthTaskRecord(payload string) (healthTaskRecord, error) { return healthTaskRecord{}, err } if record.Version != healthTaskRecordVersion || !validHealthTaskIdentifier(record.TaskID) || - !validHealthTaskIdentifier(record.ProxyID) || record.Level != string(healthDomain.LevelBasic) || + !validHealthTaskIdentifier(record.ProxyID) || !validHealthTaskIdentifier(record.UpstreamID) || record.UpstreamTasksKey == "" || + record.Level != string(healthDomain.LevelBasic) || record.Priority < int(healthDomain.PriorityFetched) || record.Priority > int(healthDomain.PriorityAvailable) || - record.DeadlineMS <= 0 || record.NextDueMS <= 0 || record.Attempts <= 0 || + record.DeadlineMS <= 0 || record.NextDueMS <= 0 || record.Attempts <= 0 || record.MaxInFlight <= 0 || (record.State != "QUEUED" && record.State != "LEASED" && record.State != "DONE") { return healthTaskRecord{}, errors.New("invalid health task record") } diff --git a/internal/adapters/redisactivity/health_tasks_integration_test.go b/internal/adapters/redisactivity/health_tasks_integration_test.go index 6f56b5e..7310e3e 100644 --- a/internal/adapters/redisactivity/health_tasks_integration_test.go +++ b/internal/adapters/redisactivity/health_tasks_integration_test.go @@ -24,7 +24,7 @@ func TestRedisHealthTasksLeaseAndRescheduleBasicChecks(t *testing.T) { t.Fatalf("UpsertFetched(): %v", err) } - candidates, err := fixture.Adapter.DueCandidates(ctx, time.Now().UTC(), 1) + candidates, err := fixture.Adapter.DueCandidates(ctx, now, 1) if err != nil || len(candidates) != 1 || candidates[0].ProxyID != "proxy-a" || candidates[0].State != proxyDomain.StateFetched { t.Fatalf("DueCandidates() = (%+v, %v)", candidates, err) } @@ -34,7 +34,7 @@ func TestRedisHealthTasksLeaseAndRescheduleBasicChecks(t *testing.T) { if err != nil { t.Fatalf("newHealthTaskPlanner(): %v", err) } - plannedAt := time.Now().UTC() + plannedAt := now plans, err := planner.Plan(plannedAt, 0, 1, candidates) if err != nil || len(plans) != 1 { t.Fatalf("Plan() = (%+v, %v)", plans, err) diff --git a/internal/adapters/redisactivity/keys.go b/internal/adapters/redisactivity/keys.go index 861730a..19d1e99 100644 --- a/internal/adapters/redisactivity/keys.go +++ b/internal/adapters/redisactivity/keys.go @@ -69,6 +69,10 @@ func (keys keyspace) checkerLeases(checkerID string) string { return keys.prefix + ":health-checker-leases:" + digestToken(checkerID) } +func (keys keyspace) upstreamTasks(upstreamID string) string { + return keys.prefix + ":health-upstream-tasks:" + digestToken(upstreamID) +} + func stateInventoryField(upstreamID, state string) string { return strconv.Itoa(len(upstreamID)) + ":" + upstreamID + ":" + state } diff --git a/internal/adapters/redisactivity/scripts.go b/internal/adapters/redisactivity/scripts.go index 7305a6b..dc41afe 100644 --- a/internal/adapters/redisactivity/scripts.go +++ b/internal/adapters/redisactivity/scripts.go @@ -45,16 +45,17 @@ type healthScriptReply struct { } type healthTaskScriptReply struct { - Status scriptStatus `json:"status"` - Count int `json:"count"` - Candidates []healthTaskCandidateWire `json:"candidates"` - Tasks []healthTaskClaimWire `json:"tasks"` + Status scriptStatus `json:"status"` + Count int `json:"count"` + CandidatesJSON string `json:"candidatesJSON"` + Tasks []healthTaskClaimWire `json:"tasks"` } type healthTaskCandidateWire struct { - ProxyID string `json:"proxyId"` - State string `json:"state"` - DueAtMS int64 `json:"dueAtMs"` + ProxyID string `json:"proxyId"` + UpstreamID string `json:"upstreamId"` + State string `json:"state"` + DueAtMS int64 `json:"dueAtMs"` } type healthTaskClaimWire struct { diff --git a/internal/adapters/redisactivity/scripts/health_tasks.lua b/internal/adapters/redisactivity/scripts/health_tasks.lua index 6f68b92..d5b7626 100644 --- a/internal/adapters/redisactivity/scripts/health_tasks.lua +++ b/internal/adapters/redisactivity/scripts/health_tasks.lua @@ -62,6 +62,9 @@ local function remove_task(task_id, requeue) redis.call('ZREM', task_expiry_key, task_id) redis.call('HDEL', tasks_key, task_id) if task then + if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then + redis.call('ZREM', task.upstreamTasksKey, task_id) + end if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then redis.call('ZREM', task.checkerLeaseKey, task_id) end @@ -182,18 +185,41 @@ if operation == 'inflight' then return finish({status = 'ok', count = redis.call('ZCARD', queued_key) + redis.call('ZCARD', leases_key)}) end +if operation == 'upstream_inflight' then + if type(payload.upstreamTasksKey) ~= 'string' or payload.upstreamTasksKey == '' then + return finish({status = 'invalid'}) + end + redis.call('ZREMRANGEBYSCORE', payload.upstreamTasksKey, '-inf', now_ms) + return finish({status = 'ok', count = redis.call('ZCARD', payload.upstreamTasksKey)}) +end + if operation == 'due' then - local result = {} - local ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, limit) + local result = cjson.decode('[]') + local scan_limit = tonumber(payload.scanLimit or limit) + if not scan_limit or scan_limit < limit then + return finish({status = 'invalid'}) + end + local ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit) for _, proxy_id in ipairs(ids) do local _, record = live_record(proxy_id) if not record or not valid_state(record.state) then redis.call('ZREM', due_key, proxy_id) + elseif type(payload.upstreamId) == 'string' and payload.upstreamId ~= '' and record.sourceUpstream ~= payload.upstreamId then + -- The ref remains due for the scheduler that owns this upstream. else - result[#result + 1] = {proxyId = proxy_id, state = record.state, dueAtMs = now_ms} + result[#result + 1] = { + proxyId = proxy_id, upstreamId = record.sourceUpstream, state = record.state, dueAtMs = now_ms + } + if #result >= limit then + break + end end end - return finish({status = 'ok', candidates = result}) + local encoded_candidates = cjson.encode(result) + if #result == 0 then + encoded_candidates = '[]' + end + return finish({status = 'ok', candidatesJSON = encoded_candidates}) end if operation == 'offer' then @@ -204,13 +230,17 @@ if operation == 'offer' then for _, task in ipairs(payload.tasks) do if type(task) ~= 'table' or task.version ~= 1 or task.level ~= 'BASIC' or task.state ~= 'QUEUED' or type(task.taskId) ~= 'string' or type(task.proxyId) ~= 'string' or tonumber(task.deadlineMs or 0) <= now_ms or - tonumber(task.nextDueMs or 0) <= now_ms or tonumber(task.attempts or 0) <= 0 then + type(task.upstreamId) ~= 'string' or task.upstreamId == '' or type(task.upstreamTasksKey) ~= 'string' or + task.upstreamTasksKey == '' or tonumber(task.nextDueMs or 0) <= now_ms or tonumber(task.attempts or 0) <= 0 or + tonumber(task.maxInFlight or 0) <= 0 then return finish({status = 'invalid'}) end local score = redis.call('ZSCORE', due_key, task.proxyId) local current = redis.call('HGET', ref_task_key, task.proxyId) local raw, record = live_record(task.proxyId) - if score and tonumber(score) <= now_ms and not current and raw and valid_state(record.state) then + local active = redis.call('ZCARD', task.upstreamTasksKey) + if score and tonumber(score) <= now_ms and not current and raw and record.sourceUpstream == task.upstreamId and + active < tonumber(task.maxInFlight) and valid_state(record.state) then if record.state == 'FETCHED' or record.state == 'UNHEALTHY' then decrement_state(record.sourceUpstream, record.state) increment_state(record.sourceUpstream, 'CHECKING') @@ -223,11 +253,13 @@ if operation == 'offer' then redis.call('HSET', ref_task_key, task.proxyId, task.taskId) redis.call('ZADD', queued_key, tonumber(task.priority), task.taskId) redis.call('ZADD', task_expiry_key, tonumber(task.deadlineMs), task.taskId) + redis.call('ZADD', task.upstreamTasksKey, tonumber(task.deadlineMs), task.taskId) local expires_at_ms = tonumber(record.expiresAtMs) touch(queued_key, expires_at_ms) touch(tasks_key, expires_at_ms) touch(task_expiry_key, expires_at_ms) touch(ref_task_key, expires_at_ms) + touch(task.upstreamTasksKey, expires_at_ms) offered = offered + 1 end end @@ -247,14 +279,14 @@ if operation == 'claim' then redis.call('ZREMRANGEBYSCORE', checker_leases_key, '-inf', now_ms) local capacity = max_in_flight - redis.call('ZCARD', checker_leases_key) if capacity <= 0 then - return finish({status = 'ok', tasks = {}}) + return finish({status = 'ok', tasks = cjson.decode('[]')}) end local supported = {} for _, level in ipairs(payload.levels) do supported[level] = true end local candidates = redis.call('ZRANGE', queued_key, 0, limit - 1) - local result = {} + local result = cjson.decode('[]') for _, task_id in ipairs(candidates) do if #result >= capacity or #result >= #payload.tokens then break @@ -305,6 +337,9 @@ if operation == 'authorize' or operation == 'complete' then if operation == 'complete' and task.state ~= 'DONE' then redis.call('ZREM', queued_key, task.taskId) redis.call('ZREM', leases_key, task.taskId) + if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then + redis.call('ZREM', task.upstreamTasksKey, task.taskId) + end if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then redis.call('ZREM', task.checkerLeaseKey, task.taskId) end diff --git a/internal/controller/bootstrap/bootstrap.go b/internal/controller/bootstrap/bootstrap.go index 2b23fe9..a06a897 100644 --- a/internal/controller/bootstrap/bootstrap.go +++ b/internal/controller/bootstrap/bootstrap.go @@ -37,6 +37,11 @@ var ( ErrStartup = errors.New("controller startup failed") ) +const ( + checkSchedulerPollInterval = 250 * time.Millisecond + checkSchedulerBatchSize = 128 +) + type Options struct { ConfigPath string Resolver config.Resolver @@ -57,6 +62,11 @@ type activityStore interface { activitypool.StateInventoryReader } +type healthTaskRuntime interface { + controllerHealth.TaskBroker + controllerHealth.UpstreamTaskSource +} + type ports struct { state admin.StateRepository activity activityStore @@ -222,7 +232,7 @@ func runWithWorkerFactory( dependencies.MetricsHandler = handler } - runners := make([]lifecycle.Runner, 0, 3) + runners := make([]lifecycle.Runner, 0, 3+len(loaded.Value.Upstreams)) if hasHTTPRuntime(loaded.Value) { runner, err := factory.New(configurationStore.Current(), dependencies, controllerRuntime.Options{HTTP: options.HTTP}) if err != nil { @@ -278,7 +288,11 @@ func runWithWorkerFactory( if identityErr != nil { return fmt.Errorf("%w: build Checker identity authorizer: %w", ErrStartup, identityErr) } - checkerHandler, handlerErr := controllerHealth.NewGRPCHandler(reducer, checkerIdentity, controllerHealth.DefaultGRPCHandlerOptions()) + checkerOptions := controllerHealth.DefaultGRPCHandlerOptions() + if tasks, ok := opened.activity.(healthTaskRuntime); ok && !nilInterface(tasks) { + checkerOptions.TaskBroker = tasks + } + checkerHandler, handlerErr := controllerHealth.NewGRPCHandler(reducer, checkerIdentity, checkerOptions) if handlerErr != nil { return fmt.Errorf("%w: build Checker control handler: %w", ErrStartup, handlerErr) } @@ -292,6 +306,13 @@ func runWithWorkerFactory( return errors.Join(ErrStartup, ErrInvalidOptions) } runners = append(runners, runner) + if tasks, ok := opened.activity.(healthTaskRuntime); ok && !nilInterface(tasks) { + schedulers, schedulerErr := newHealthSchedulers(loaded.Value, tasks, options.Now) + if schedulerErr != nil { + return fmt.Errorf("%w: build Checker health schedulers: %w", ErrStartup, schedulerErr) + } + runners = append(runners, schedulers...) + } } runners = append(runners, supervisor) group, err := lifecycle.NewGroup(runners...) @@ -301,6 +322,53 @@ func runWithWorkerFactory( return group.Run(ctx) } +func newHealthSchedulers( + configuration *config.Config, + tasks healthTaskRuntime, + now func() time.Time, +) ([]lifecycle.Runner, error) { + if configuration == nil || nilInterface(tasks) || now == nil { + return nil, ErrInvalidOptions + } + names := make([]string, 0, len(configuration.Upstreams)) + for name, upstream := range configuration.Upstreams { + if upstream.Enabled { + names = append(names, name) + } + } + sort.Strings(names) + runners := make([]lifecycle.Runner, 0, len(names)) + for _, name := range names { + check := config.EffectiveCheck(configuration.Defaults.Check, configuration.Upstreams[name].Check) + planner, err := controllerHealth.NewPlanner(controllerHealth.SchedulePolicy{ + Interval: check.Interval.Value(), Jitter: check.Jitter, MaxInFlight: check.MaxInFlight, + Timeout: check.Timeout.Value(), MaxAttempts: check.MaxAttempts, + }) + if err != nil { + return nil, err + } + source, err := controllerHealth.NewUpstreamDueSource(tasks, name) + if err != nil { + return nil, err + } + batchSize := checkSchedulerBatchSize + if check.MaxInFlight < batchSize { + batchSize = check.MaxInFlight + } + runner, err := controllerHealth.NewSchedulerRunner(planner, source, tasks, controllerHealth.SchedulerRunnerOptions{ + PollInterval: checkSchedulerPollInterval, BatchSize: batchSize, Now: now, + }) + if err != nil { + return nil, err + } + runners = append(runners, runner) + } + if len(runners) == 0 { + return nil, ErrInvalidOptions + } + return runners, nil +} + func newCheckerIdentity(controlPlane config.ControlPlane) (controllerHealth.CheckerIdentityAuthorizer, error) { switch controlPlane.TLS.Mode { case "disabled": diff --git a/internal/controller/bootstrap/bootstrap_test.go b/internal/controller/bootstrap/bootstrap_test.go index 0ad52d7..457ebaf 100644 --- a/internal/controller/bootstrap/bootstrap_test.go +++ b/internal/controller/bootstrap/bootstrap_test.go @@ -9,6 +9,7 @@ import ( "proxy-pool/internal/config" "proxy-pool/internal/controller/admin" + controllerHealth "proxy-pool/internal/controller/health" "proxy-pool/internal/controller/pool" "proxy-pool/internal/controller/provider" controllerRuntime "proxy-pool/internal/controller/runtime" @@ -310,6 +311,20 @@ func TestRetainProviderStatsKeepsAllConfiguredProviders(t *testing.T) { } } +func TestNewHealthSchedulersCreatesOneRunnerPerEnabledUpstream(t *testing.T) { + configuration, err := config.Load(strings.NewReader(bootstrapTestConfig)) + if err != nil { + t.Fatalf("config.Load(): %v", err) + } + disabled := configuration.Upstreams["provider-b"] + disabled.Enabled = false + configuration.Upstreams["provider-b"] = disabled + runners, err := newHealthSchedulers(configuration, healthTaskRuntimeStub{}, time.Now) + if err != nil || len(runners) != 1 { + t.Fatalf("newHealthSchedulers() = (%d runners, %v)", len(runners), err) + } +} + type memoryResolver struct { files map[string][]byte reads int @@ -406,6 +421,49 @@ func (readyStub) Ready(context.Context) error { return nil } type stubActivityStore struct{} +type healthTaskRuntimeStub struct{} + +func (healthTaskRuntimeStub) Offer(context.Context, []healthDomain.PlannedTask) (int, error) { + return 0, nil +} + +func (healthTaskRuntimeStub) Claim(context.Context, healthDomain.TaskClaim) ([]healthDomain.LeasedTask, error) { + return nil, nil +} + +func (healthTaskRuntimeStub) AuthorizeObservation( + context.Context, + string, + string, + healthDomain.Observation, + time.Time, +) error { + return nil +} + +func (healthTaskRuntimeStub) CompleteObservation( + context.Context, + string, + string, + healthDomain.Observation, + time.Time, +) error { + return nil +} + +func (healthTaskRuntimeStub) InFlightForUpstream(context.Context, string, time.Time) (int, error) { + return 0, nil +} + +func (healthTaskRuntimeStub) DueCandidatesForUpstream( + context.Context, + string, + time.Time, + int, +) ([]controllerHealth.Candidate, error) { + return nil, nil +} + func (*stubActivityStore) Extract(_ context.Context, command extractionDomain.Command) (extractionDomain.Result, error) { return extractionDomain.Result{Requested: command.Requested}, nil } diff --git a/internal/controller/health/scheduler.go b/internal/controller/health/scheduler.go index e3f327e..5a33c74 100644 --- a/internal/controller/health/scheduler.go +++ b/internal/controller/health/scheduler.go @@ -126,7 +126,7 @@ func (planner *Planner) Plan(now time.Time, inFlight, maxTasks int, candidates [ } result[index] = PlannedTask{ Candidate: item.candidate, Priority: item.priority, Deadline: deadline, NextDue: nextDue, - Attempts: planner.policy.MaxAttempts, + Attempts: planner.policy.MaxAttempts, MaxInFlight: planner.policy.MaxInFlight, } } return result, nil diff --git a/internal/controller/health/scheduler_runner.go b/internal/controller/health/scheduler_runner.go index 90fb5de..8df67b3 100644 --- a/internal/controller/health/scheduler_runner.go +++ b/internal/controller/health/scheduler_runner.go @@ -19,6 +19,33 @@ type DueSource interface { 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 { diff --git a/internal/domain/health/task.go b/internal/domain/health/task.go index f3d7739..cce6ff8 100644 --- a/internal/domain/health/task.go +++ b/internal/domain/health/task.go @@ -15,6 +15,7 @@ import ( // carrying endpoint credentials or any persistence-specific representation. type Candidate struct { ProxyID string + UpstreamID string State proxyDomain.State Level Level RoutingName string @@ -39,11 +40,12 @@ const ( // PlannedTask is transport-neutral work ready for a shared leased task store. type PlannedTask struct { - Candidate Candidate - Priority Priority - Deadline time.Time - NextDue time.Time - Attempts int + Candidate Candidate + Priority Priority + Deadline time.Time + NextDue time.Time + Attempts int + MaxInFlight int } // TaskIDFor derives an idempotent task identity from all execution-relevant