diff --git a/README.md b/README.md index 294bdbf..597948c 100644 --- a/README.md +++ b/README.md @@ -88,10 +88,10 @@ flowchart LR 提取与限流、Controller 的 Admin/Distribution/Metrics 监听，以及 PostgreSQL 管理状态；WorkerControlPlane 的 Register、Snapshot ACK、Runtime 心跳接收和 Redis 会话栅栏，以及 Gateway Outcome 上报的有界队列、序列确认与重试； - Controller 的 Redis 共享 BASIC 检查任务、按上游的有界调度、HTTP/HTTPS/SOCKS5 + Controller 的 Redis 共享 BASIC/EGRESS 检查任务、按上游的有界轮转调度、HTTP/HTTPS/SOCKS5 Checker 探测和 Observation 状态归并。 -- **部分完成**：EGRESS 与 TARGET 的任务编排和配置建模，Docker Compose/Kubernetes +- **部分完成**：EGRESS 的出口身份响应解析、TARGET 的任务编排和配置建模，Docker Compose/Kubernetes 运行时 mTLS Overlay。 - **待完成**：CONNECT 长连接/Extract 压测场景、故障演练和代表性集群压测。 @@ -169,9 +169,10 @@ Checker 的参数也可通过 `PROXY_POOL_CONTROL_PLANE_ADDRESS`、 `PROXY_POOL_CHECKER_ID`、`PROXY_POOL_CHECKER_INSTANCE_ID` 与 `PROXY_POOL_CHECKER_MAX_IN_FLIGHT` 提供。它不会访问 Redis/PostgreSQL；生产 Controller 在启用控制面时装配 Redis 共享任务 broker，并按启用的 Upstream 调度 -HTTP/HTTPS/SOCKS5 BASIC 检查。调度监督器每轮读取已发布配置，因此 reload 后的上游启停、 -检查间隔、抖动、超时、重试次数和 `maxInFlight` 会在下一轮生效；新启用的上游无需 -重启 Controller。EGRESS 与 TARGET 尚未进入生产调度。 +HTTP/HTTPS/SOCKS5 BASIC 检查，并按每个 `check.urls` 创建 EGRESS 任务。调度监督器每轮读取 +已发布配置，因此 reload 后的上游启停、检查间隔、抖动、超时、重试次数、`maxInFlight` 和 +EGRESS URL 都会在下一轮生效；BASIC 与 EGRESS 以有界轮转组共享上游并发上限。新启用的上游 +无需重启 Controller。TARGET 尚未进入生产调度，EGRESS 的出口身份响应解析也尚未配置化。 初版 HTTP 容量工具可按固定请求数或固定时长运行，并将 HTTPS 目标经 Gateway 的请求 交给标准 HTTP Transport 建立 CONNECT： @@ -227,8 +228,8 @@ go run ./cmd/proxy-loadgen ` - **P0 - Worker 控制面闭环**：Worker session、Snapshot ledger、ACK、运行态接收、 ownership 索引，以及 Gateway 快照客户端。 -- **P0 - Checker 健康链**：BASIC 的共享调度、实际探测、Observation reducer 和 - `FETCHED -> AVAILABLE / SUSPECT / UNHEALTHY` 状态链已完成；继续补齐 EGRESS、 +- **P0 - Checker 健康链**：BASIC/EGRESS 的共享调度、实际探测、Observation reducer 和 + `FETCHED -> AVAILABLE / SUSPECT / UNHEALTHY` 状态链已完成；继续补齐 EGRESS 身份解析、 TARGET。 - **P1 - Gateway 与 Routing**：Gateway 进程、快照凭据分发、五种 Routing 策略与 `onUnavailable` 已接入；动态容量调整和 Drain 闭环待完成。 diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index 5189bc4..5a32ba0 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -261,8 +261,8 @@ Controller 多副本共享同一计数。Client 身份只以 SHA-256 摘要进 Gateway 请求热路径仍只使用本地准入，不增加 Redis/PostgreSQL 调用。 WorkerControlPlane gRPC 接收端、session 签发/心跳、Snapshot ACK 账本、基础 Snapshot 流、Gateway 会话客户端与快照凭据分发已完成；权威 Proxy/Routing 发布、 -Outcome 上报已完成为有界队列、批次序列/摘要栅栏和确认重试；健康 BASIC 执行链已 -完成，EGRESS/TARGET 编排和 REMOVE 生命周期仍待完成，因此 Task 10 尚未全部完成。 +Outcome 上报已完成为有界队列、批次序列/摘要栅栏和确认重试；健康 BASIC/EGRESS 执行链已 +完成，EGRESS 身份解析、TARGET 编排和 REMOVE 生命周期仍待完成，因此 Task 10 尚未全部完成。 ## Task 11: Checker and Health Reducer @@ -284,14 +284,16 @@ Checker Observation 上报 RPC 已复用既有控制面监听接入 Controller 领取者加不可预测 lease token 的栅栏和完成后重放；任务凭据仅由认证流在执行期下发。 Redis 共享 due-index/租约持久化、每 Upstream 的跨副本 in-flight 限制和生产 broker 已完成； `proxy-checker` 独立进程、固定大小 worker-pool、任务期重试/微批上报和 HTTP/HTTPS/SOCKS5 -BASIC 探测器已完成并有测试。TARGET 探测器具备任务执行能力，但尚无生产任务调度； -EGRESS 已具备任务 URL 传输、HTTP/HTTPS 探测和全局事实回传契约，但出口身份响应解析、 -EGRESS/TARGET 多维任务索引及部署运行态仍未实现， +BASIC/EGRESS 探测器已完成并有测试。Redis 已按固定 EGRESS due-index 保存任务执行 URL， +与 BASIC 独立引用并通过上游共享并发限制；配置化监督器以有界轮转组调度每个 `check.urls`。 +TARGET 探测器具备任务执行能力，但尚无生产任务调度；EGRESS 的出口身份响应解析与部署运行态 +仍未实现， 因此本任务保持未完成。 补充进度（2026-08-02）：BASIC 调度已改为配置驱动监督器。它每轮读取已发布快照并复用 有界派发逻辑，所以 reload 后已启用上游的策略变更、停用，以及新启用上游都无需重启 -Controller 即可生效；Redis 任务存储仍仅承载 BASIC，未扩展 EGRESS/TARGET 的多维索引。 +Controller 即可生效。Redis 任务存储现已扩展 BASIC/EGRESS 的独立有界索引；TARGET 的多维 +索引仍待实现。 补充进度（2026-08-02）：已新增 `proxy-loadgen` HTTP 场景。固定请求数和固定时长两种 模式均通过固定 worker 数与有界派发通道执行，可选 QPS 限速；报告使用固定大小延迟直方图， diff --git a/docs/operations/runbook.md b/docs/operations/runbook.md index 2477fda..7aa2cee 100644 --- a/docs/operations/runbook.md +++ b/docs/operations/runbook.md @@ -23,9 +23,9 @@ `cmd/proxy-controller` 已完成配置单次加载、PostgreSQL 迁移、Redis 活动池、 Distribution/Admin/Metrics 独立监听和有界停机装配。Provider 自动补池、分布式 -配额、动态重载和 Admin 低基数统计已装配；Controller 已装配 Redis BASIC 任务 broker， -`proxy-checker` 可执行 HTTP/HTTPS/SOCKS5 BASIC 探测。`proxy-loadgen` 已提供有界 HTTP -请求场景；EGRESS/TARGET 调度、CONNECT 长连接/Extract 压测与完整 mTLS 环境 Overlay 仍属于 +配额、动态重载和 Admin 低基数统计已装配；Controller 已装配 Redis BASIC/EGRESS 任务 broker， +`proxy-checker` 可执行 HTTP/HTTPS/SOCKS5 BASIC/EGRESS 探测。`proxy-loadgen` 已提供有界 HTTP +请求场景；EGRESS 身份解析、TARGET 调度、CONNECT 长连接/Extract 压测与完整 mTLS 环境 Overlay 仍属于 `implementation-plan.md` 后续任务。 因此 Compose/Kubernetes 资产当前仍用于评审网络、资源、探针和依赖关系，不能 视为完整可运行拓扑。 diff --git a/docs/requirements/completion-audit.md b/docs/requirements/completion-audit.md index 8ab1790..cbba61f 100644 --- a/docs/requirements/completion-audit.md +++ b/docs/requirements/completion-audit.md @@ -62,8 +62,8 @@ Outcome 已实现为 Gateway 本地有界队列、微批确认重试和 Controll 记录归并，不改写 Proxy 全局状态。Controller 公用 Reducer 已作为 Observation 的唯一状态 归并边界，Checker Observation RPC 已在同一控制面监听以独立 SPIFFE 身份接入；有界任务领取、 租约归属与任务期凭据传输已由通用契约和 gRPC 往返测试覆盖。Redis 共享 due-index、 -任务/租约持久化、按上游的 in-flight 限制以及独立 Checker 的 HTTP/HTTPS/SOCKS5 BASIC 执行 -进程已经闭环；EGRESS、TARGET 的多维任务索引与生产调度尚未实现。Snapshot 签发在 Redis 中原子匹配当前 +任务/租约持久化、按上游的 in-flight 限制以及独立 Checker 的 HTTP/HTTPS/SOCKS5 BASIC/EGRESS 执行 +进程已经闭环；EGRESS 已具备多维任务索引与生产调度，出口身份响应解析和 TARGET 的多维索引仍未实现。Snapshot 签发在 Redis 中原子匹配当前 `session_id`，重注册会清除旧引用，迟到旧 Stream 不会覆盖新 session。Controller 在最近成功下发的 Snapshot `valid_until` 到达时关闭流；Gateway 的公用 `SessionSupervisor` 已实现可恢复错误的有界退避重连。Gateway 会校验并执行 Snapshot diff --git a/docs/requirements/traceability.md b/docs/requirements/traceability.md index 6845a25..b296966 100644 --- a/docs/requirements/traceability.md +++ b/docs/requirements/traceability.md @@ -7,7 +7,7 @@ | ID | 最终需求 | 来源 | 验证证据 | |---|---|---|---| -| ARCH-001 | 数据面 Worker 与控制面 Controller 分离 | 1-70 | 包、协议和部署拓扑已分离；Controller 已运行 Worker Register/Watch/ACK/Runtime/Outcome 与 Checker Observation gRPC，并发布 Proxy/Gateway Routing/按引用去重凭据完整快照；Checker 任务流已具备有界领取、租约栅栏和任务期凭据契约。Gateway 已将快照编译为同版本动态 View，并由独立进程维护控制面会话。`proxy-checker` 与 Redis BASIC 共享任务运行态、`proxy-loadgen` 有界 HTTP 场景均已实现；EGRESS/TARGET 生产编排待实现 | +| ARCH-001 | 数据面 Worker 与控制面 Controller 分离 | 1-70 | 包、协议和部署拓扑已分离；Controller 已运行 Worker Register/Watch/ACK/Runtime/Outcome 与 Checker Observation gRPC，并发布 Proxy/Gateway Routing/按引用去重凭据完整快照；Checker 任务流已具备有界领取、租约栅栏和任务期凭据契约。Gateway 已将快照编译为同版本动态 View，并由独立进程维护控制面会话。`proxy-checker` 与 Redis BASIC/EGRESS 共享任务运行态、`proxy-loadgen` 有界 HTTP 场景均已实现；TARGET 生产编排待实现 | | ARCH-002 | 热路径只做认证、本地路由和网络转发 | 1-70, 380-430 | Gateway bootstrap 集成测试验证启动期控制面会话与快照就绪，HTTP 请求只走本地 Snapshot/Dispatch；Outcome 仅写入有界非阻塞本地队列，代表性性能剖析待完成 | | ARCH-003 | Gateway、Distribution、Admin、Metrics 独立入口 | 8904-8958 | Controller 命令已装配 Distribution/Admin/Metrics 三个独立监听及联动停机；Gateway 命令已装配代理与 Metrics 监听，运行时 mTLS 部署 Overlay 待完成 | | ARCH-004 | Controller 集中 Provider 获取与切换 | 1403-1580 | Redis Leader、动态 Provider Supervisor 与 Bootstrap 生产装配已完成；Admin disable/reload 驱动取消替换，多副本按权威 HMAC 指纹和 revision 栅栏收敛并拒绝旧配置换主；Routing 切换到 Drain 的编排待完成 | @@ -80,7 +80,7 @@ | ID | 最终需求 | 来源 | 验证证据 | |---|---|---|---| | HEALTH-001 | 全局健康与 Routing/目标健康分离 | 221-270, 8679-8708 | `domain/health` 已将 BASIC/EGRESS 全局 Reducer 与 TARGET Profile Reducer 分离；TARGET 在 Memory 和 Redis 独立、随代理 TTL 归并，不改写 Proxy 全局状态；Routing 消费待实现 | -| HEALTH-002 | 健康调度有 jitter、maxInFlight 和分级频率 | 8679-8736 | 配置有效合并、URL 校验、稳定抖动/优先级 Planner、有界 Scheduler tick、Redis due-index、跨副本 in-flight 原子限制，以及 HTTP/HTTPS/SOCKS5 BASIC 生产执行器已完成；EGRESS/TARGET 多维任务索引与调度待实现 | +| HEALTH-002 | 健康调度有 jitter、maxInFlight 和分级频率 | 8679-8736 | 配置有效合并、URL 校验、稳定抖动/优先级 Planner、有界 Scheduler tick、Redis BASIC/EGRESS due-index、跨副本 in-flight 原子限制，以及 HTTP/HTTPS/SOCKS5 BASIC/EGRESS 生产执行器已完成；TARGET 多维任务索引与调度待实现 | | HEALTH-003 | 失败分级 SUSPECT -> UNHEALTHY -> REMOVE | 8679-8736 | Controller 公用 Reducer 已通过 Memory/Redis 活动池原子提交全局连续失败、精确重放和成功恢复；BASIC 任务调度已完成，REMOVE 编排待实现 | | SEC-001 | API 认证与 Proxy 认证分离，Secret 统一脱敏 | 7528-8111, 8904-8945 | Config 脱敏、Provider Store -> SecretRef -> Gateway Resolver 跨包测试与格式化泄漏回归测试 | | SEC-002 | 非回环监听无保护时严格模式启动失败 | 8112-8441 | 配置校验测试 | diff --git a/internal/adapters/redisactivity/health_tasks.go b/internal/adapters/redisactivity/health_tasks.go index 413066e..88d53ed 100644 --- a/internal/adapters/redisactivity/health_tasks.go +++ b/internal/adapters/redisactivity/health_tasks.go @@ -21,6 +21,8 @@ var _ healthDomain.TaskBroker = (*Adapter)(nil) type healthTaskRecord struct { Version int `json:"version"` TaskID string `json:"taskId"` + Reference string `json:"reference"` + DueMember string `json:"dueMember"` ProxyID string `json:"proxyId"` UpstreamID string `json:"upstreamId"` UpstreamTasksKey string `json:"upstreamTasksKey"` @@ -49,6 +51,8 @@ type healthTaskRequest struct { LeaseTTLMS int64 `json:"leaseTTLMS,omitempty"` Levels []healthDomain.Level `json:"levels,omitempty"` UpstreamID string `json:"upstreamId,omitempty"` + Level string `json:"level,omitempty"` + TargetURL string `json:"targetUrl,omitempty"` UpstreamTasksKey string `json:"upstreamTasksKey,omitempty"` ScanLimit int `json:"scanLimit,omitempty"` Tokens []string `json:"tokens,omitempty"` @@ -84,22 +88,39 @@ 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) + return a.dueCandidates(ctx, "", healthDomain.LevelBasic, "", 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) + return a.dueCandidates(ctx, upstreamID, healthDomain.LevelBasic, "", now, limit) } -func (a *Adapter) dueCandidates(ctx context.Context, upstreamID string, now time.Time, limit int) ([]healthDomain.Candidate, error) { +// DueEgressCandidatesForUpstream lazily initializes bounded EGRESS references +// from live BASIC references. The independent reference lets BASIC and EGRESS +// checks for one proxy be leased at the same time. +func (a *Adapter) DueEgressCandidatesForUpstream(ctx context.Context, upstreamID, targetURL string, now time.Time, limit int) ([]healthDomain.Candidate, error) { + if !validHealthTaskIdentifier(upstreamID) { + return nil, healthDomain.ErrInvalidTaskBroker + } + normalizedTarget, err := healthDomain.NormalizeEgressTarget(targetURL) + if err != nil || normalizedTarget != targetURL { + return nil, healthDomain.ErrInvalidTaskBroker + } + return a.dueCandidates(ctx, upstreamID, healthDomain.LevelEgress, targetURL, now, limit) +} + +func (a *Adapter) dueCandidates(ctx context.Context, upstreamID string, level healthDomain.Level, targetURL 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 } + if level != healthDomain.LevelBasic && level != healthDomain.LevelEgress { + return nil, healthDomain.ErrInvalidTaskBroker + } reply, err := a.runHealthTaskScript(ctx, "due", now, healthTaskRequest{ - Limit: limit, ScanLimit: a.options.MaxCandidateScan, UpstreamID: upstreamID, + Limit: limit, ScanLimit: a.options.MaxCandidateScan, UpstreamID: upstreamID, Level: string(level), TargetURL: targetURL, }, "runtime") if err != nil { return nil, err @@ -119,11 +140,18 @@ func (a *Adapter) dueCandidates(ctx context.Context, upstreamID string, now time default: return nil, invalidScriptReply("health task due reply contains invalid state") } - if !validHealthTaskIdentifier(candidate.ProxyID) || !validHealthTaskIdentifier(candidate.UpstreamID) || candidate.DueAtMS <= 0 { + if !validHealthTaskIdentifier(candidate.ProxyID) || !validHealthTaskIdentifier(candidate.UpstreamID) || candidate.DueAtMS <= 0 || + healthDomain.Level(candidate.Level) != level { return nil, invalidScriptReply("health task due reply contains invalid candidate") } + if level == healthDomain.LevelEgress { + normalizedTarget, err := healthDomain.NormalizeEgressTarget(candidate.TargetURL) + if err != nil || normalizedTarget != targetURL { + return nil, invalidScriptReply("health task due reply contains invalid EGRESS target") + } + } result = append(result, healthDomain.Candidate{ - ProxyID: candidate.ProxyID, UpstreamID: candidate.UpstreamID, State: state, Level: healthDomain.LevelBasic, + ProxyID: candidate.ProxyID, UpstreamID: candidate.UpstreamID, State: state, Level: level, TargetURL: candidate.TargetURL, DueAt: time.UnixMilli(candidate.DueAtMS).UTC(), }) } @@ -168,10 +196,12 @@ func (a *Adapter) Offer(ctx context.Context, plans []healthDomain.PlannedTask) ( return 0, err } tasks[index] = healthTaskRecord{ - Version: healthTaskRecordVersion, TaskID: healthDomain.TaskIDFor(plan), ProxyID: plan.Candidate.ProxyID, + Version: healthTaskRecordVersion, TaskID: healthDomain.TaskIDFor(plan), Reference: redisHealthTaskReference(plan.Candidate), + DueMember: redisHealthDueMember(plan.Candidate), 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, MaxInFlight: plan.MaxInFlight, State: "QUEUED", + RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL, } } reply, err := a.runHealthTaskScript(ctx, "offer", now, healthTaskRequest{Limit: a.options.MaxCheckTasks, Tasks: tasks}, "runtime") @@ -227,13 +257,14 @@ func (a *Adapter) Claim(ctx context.Context, claim healthDomain.TaskClaim) ([]he return nil, invalidScriptReply("health task claim reply contains invalid record") } if task.State != "LEASED" || task.LeaseCheckerID != claim.CheckerID || task.LeaseToken == "" || - task.Level != string(healthDomain.LevelBasic) || task.ProxyID != record.ID || task.DeadlineMS <= now.UnixMilli() { + task.ProxyID != record.ID || task.DeadlineMS <= now.UnixMilli() || !validRedisTaskExecution(task) { return nil, invalidScriptReply("health task claim reply violates lease contract") } result = append(result, healthDomain.LeasedTask{ TaskID: task.TaskID, LeaseToken: task.LeaseToken, ProxyID: task.ProxyID, Protocol: proxyDomain.Scheme(record.Scheme), Host: record.Host, Port: uint16(record.Port), Username: record.Username, Password: record.Password, Level: healthDomain.Level(task.Level), + RoutingName: task.RoutingName, TargetURL: task.TargetURL, Deadline: time.UnixMilli(task.DeadlineMS).UTC(), Attempts: task.Attempts, }) } @@ -321,7 +352,7 @@ func (a *Adapter) runHealthTaskScript( result, err := runScript(ctx, a.client, healthTasksScript, []string{ a.keys.records, a.keys.expiry, a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases, a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.checkerLeases(checkerID), - a.keys.stateInventory, + a.keys.stateInventory, a.keys.healthEgressDue, }, operation, now.UTC().UnixMilli(), string(payload)) if err != nil { return healthTaskScriptReply{}, err @@ -334,12 +365,24 @@ func (a *Adapter) runHealthTaskScript( } func validateRedisPlannedTask(task healthDomain.PlannedTask, now time.Time) error { - if !validHealthTaskIdentifier(task.Candidate.ProxyID) || !validHealthTaskIdentifier(task.Candidate.UpstreamID) || task.Candidate.Level != healthDomain.LevelBasic || - task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" || task.Candidate.DueAt.IsZero() || + if !validHealthTaskIdentifier(task.Candidate.ProxyID) || !validHealthTaskIdentifier(task.Candidate.UpstreamID) || task.Candidate.DueAt.IsZero() || 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.Level { + case healthDomain.LevelBasic: + if task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" { + return healthDomain.ErrInvalidLeasedTask + } + case healthDomain.LevelEgress: + normalizedTarget, err := healthDomain.NormalizeEgressTarget(task.Candidate.TargetURL) + if task.Candidate.RoutingName != "" || err != nil || normalizedTarget != task.Candidate.TargetURL { + return healthDomain.ErrInvalidLeasedTask + } + default: + return healthDomain.ErrInvalidLeasedTask + } switch task.Candidate.State { case proxyDomain.StateFetched, proxyDomain.StateAvailable, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy: return nil @@ -373,9 +416,17 @@ func decodeHealthTaskRecord(payload string) (healthTaskRecord, error) { if err := decodeJSON(payload, &record); err != nil { return healthTaskRecord{}, err } + if record.Level == string(healthDomain.LevelBasic) { + if record.Reference == "" { + record.Reference = record.ProxyID + } + if record.DueMember == "" { + record.DueMember = record.ProxyID + } + } if record.Version != healthTaskRecordVersion || !validHealthTaskIdentifier(record.TaskID) || !validHealthTaskIdentifier(record.ProxyID) || !validHealthTaskIdentifier(record.UpstreamID) || record.UpstreamTasksKey == "" || - record.Level != string(healthDomain.LevelBasic) || + !validHealthTaskReference(record.Reference) || !validHealthTaskReference(record.DueMember) || record.Priority < int(healthDomain.PriorityFetched) || record.Priority > int(healthDomain.PriorityAvailable) || record.DeadlineMS <= 0 || record.NextDueMS <= 0 || record.Attempts <= 0 || record.MaxInFlight <= 0 || (record.State != "QUEUED" && record.State != "LEASED" && record.State != "DONE") { @@ -384,6 +435,36 @@ func decodeHealthTaskRecord(payload string) (healthTaskRecord, error) { return record, nil } +func redisHealthTaskReference(candidate healthDomain.Candidate) string { + if candidate.Level == healthDomain.LevelBasic { + return candidate.ProxyID + } + return healthDomain.CandidateIdentity(candidate) +} + +func redisHealthDueMember(candidate healthDomain.Candidate) string { + if candidate.Level == healthDomain.LevelBasic { + return candidate.ProxyID + } + return candidate.TargetURL + "\x00" + candidate.ProxyID +} + +func validRedisTaskExecution(task healthTaskRecord) bool { + switch healthDomain.Level(task.Level) { + case healthDomain.LevelBasic: + return task.RoutingName == "" && task.TargetURL == "" + case healthDomain.LevelEgress: + targetURL, err := healthDomain.NormalizeEgressTarget(task.TargetURL) + return task.RoutingName == "" && err == nil && targetURL == task.TargetURL + default: + return false + } +} + +func validHealthTaskReference(value string) bool { + return value != "" && len(value) <= 2048 && strings.TrimSpace(value) == value && !strings.ContainsRune(value, '\x7f') +} + func newRedisLeaseToken() (string, error) { var entropy [24]byte if _, err := rand.Read(entropy[:]); err != nil { diff --git a/internal/adapters/redisactivity/health_tasks_integration_test.go b/internal/adapters/redisactivity/health_tasks_integration_test.go index 39007e1..0d7c185 100644 --- a/internal/adapters/redisactivity/health_tasks_integration_test.go +++ b/internal/adapters/redisactivity/health_tasks_integration_test.go @@ -123,3 +123,59 @@ func TestRedisHealthTaskSweepReleasesLeaseAndUpstreamCapacity(t *testing.T) { t.Fatalf("AuthorizeObservation(after sweep) error = %v, want ErrTaskNotFound", err) } } + +func TestRedisHealthTasksLeaseAndRescheduleEgressChecks(t *testing.T) { + fixture := newRedisTestFixture(t) + ctx := context.Background() + now := time.Now().UTC() + const targetURL = "https://egress.example/identity" + if _, err := fixture.Adapter.UpsertFetched(ctx, "provider-a", activitypool.FetchedBatch{ + ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1, + Proxies: []proxyDomain.Proxy{testProxy("proxy-a", "192.0.2.10")}, + }); err != nil { + t.Fatalf("UpsertFetched(): %v", err) + } + candidates, err := fixture.Adapter.DueEgressCandidatesForUpstream(ctx, "provider-a", targetURL, now, 1) + if err != nil || len(candidates) != 1 || candidates[0].Level != healthDomain.LevelEgress || candidates[0].TargetURL != targetURL { + t.Fatalf("DueEgressCandidatesForUpstream() = (%+v, %v)", candidates, err) + } + planner, err := controllerHealth.NewPlanner(controllerHealth.SchedulePolicy{ + Interval: 30 * time.Second, MaxInFlight: 2, Timeout: 5 * time.Second, MaxAttempts: 1, + }) + if err != nil { + t.Fatalf("NewPlanner(): %v", err) + } + plans, err := planner.Plan(now, 0, 1, candidates) + if err != nil || len(plans) != 1 { + t.Fatalf("Plan() = (%+v, %v)", plans, err) + } + if offered, err := fixture.Adapter.Offer(ctx, plans); err != nil || offered != 1 { + t.Fatalf("Offer() = (%d, %v)", offered, err) + } + claimed, err := fixture.Adapter.Claim(ctx, healthDomain.TaskClaim{ + CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1, + SupportedLevels: []healthDomain.Level{healthDomain.LevelEgress}, + }) + if err != nil || len(claimed) != 1 || claimed[0].Level != healthDomain.LevelEgress || claimed[0].TargetURL != targetURL { + t.Fatalf("Claim() = (%+v, %v)", claimed, err) + } + observation := healthDomain.Observation{ + TaskID: claimed[0].TaskID, ProxyID: claimed[0].ProxyID, Level: healthDomain.LevelEgress, + Success: true, Latency: time.Millisecond, ObservedAt: time.Now().UTC(), + } + if err := fixture.Adapter.AuthorizeObservation(ctx, "checker-a", claimed[0].LeaseToken, observation, time.Now().UTC()); err != nil { + t.Fatalf("AuthorizeObservation() = %v", err) + } + if _, err := fixture.Adapter.ApplyGlobalObservation(ctx, activitypool.GlobalHealthCommand{ + Observation: observation, MaxConsecutiveFailures: 2, + }); err != nil { + t.Fatalf("ApplyGlobalObservation() = %v", err) + } + if err := fixture.Adapter.CompleteObservation(ctx, "checker-a", claimed[0].LeaseToken, observation, time.Now().UTC()); err != nil { + t.Fatalf("CompleteObservation() = %v", err) + } + next, err := fixture.Adapter.DueEgressCandidatesForUpstream(ctx, "provider-a", targetURL, plans[0].NextDue, 1) + if err != nil || len(next) != 1 || next[0].ProxyID != "proxy-a" || next[0].TargetURL != targetURL { + t.Fatalf("DueEgressCandidatesForUpstream(next) = (%+v, %v)", next, err) + } +} diff --git a/internal/adapters/redisactivity/keys.go b/internal/adapters/redisactivity/keys.go index 19d1e99..46eacd3 100644 --- a/internal/adapters/redisactivity/keys.go +++ b/internal/adapters/redisactivity/keys.go @@ -28,6 +28,7 @@ type keyspace struct { workerRuntimeExpiry string workerOutcomes string healthDue string + healthEgressDue string healthQueued string healthLeases string healthTasks string @@ -57,6 +58,7 @@ func newKeyspace(namespace string) keyspace { workerRuntimeExpiry: prefix + ":worker-runtime-expiry", workerOutcomes: prefix + ":worker-outcomes", healthDue: prefix + ":health-due", + healthEgressDue: prefix + ":health-egress-due", healthQueued: prefix + ":health-queued", healthLeases: prefix + ":health-leases", healthTasks: prefix + ":health-tasks", diff --git a/internal/adapters/redisactivity/scripts.go b/internal/adapters/redisactivity/scripts.go index dc41afe..378b33f 100644 --- a/internal/adapters/redisactivity/scripts.go +++ b/internal/adapters/redisactivity/scripts.go @@ -55,6 +55,8 @@ type healthTaskCandidateWire struct { ProxyID string `json:"proxyId"` UpstreamID string `json:"upstreamId"` State string `json:"state"` + Level string `json:"level"` + TargetURL string `json:"targetUrl,omitempty"` DueAtMS int64 `json:"dueAtMs"` } diff --git a/internal/adapters/redisactivity/scripts/health_tasks.lua b/internal/adapters/redisactivity/scripts/health_tasks.lua index 8412cfa..0f0a454 100644 --- a/internal/adapters/redisactivity/scripts/health_tasks.lua +++ b/internal/adapters/redisactivity/scripts/health_tasks.lua @@ -8,6 +8,7 @@ local task_expiry_key = KEYS[7] local ref_task_key = KEYS[8] local checker_leases_key = KEYS[9] local state_inventory_key = KEYS[10] +local egress_due_key = KEYS[11] local operation = ARGV[1] local now_ms = tonumber(ARGV[2]) @@ -68,14 +69,20 @@ local function remove_task(task_id, requeue) if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then redis.call('ZREM', task.checkerLeaseKey, task_id) end - if type(task.proxyId) == 'string' and redis.call('HGET', ref_task_key, task.proxyId) == task_id then - redis.call('HDEL', ref_task_key, task.proxyId) + local reference = task.reference or task.proxyId + if type(reference) == 'string' and redis.call('HGET', ref_task_key, reference) == task_id then + redis.call('HDEL', ref_task_key, reference) end - if requeue and type(task.proxyId) == 'string' then + local due_member = task.dueMember or task.proxyId + if requeue and type(task.proxyId) == 'string' and type(due_member) == 'string' then local _, record = live_record(task.proxyId) if record then - redis.call('ZADD', due_key, now_ms, task.proxyId) - touch(due_key, tonumber(record.expiresAtMs)) + local task_due_key = due_key + if task.level == 'EGRESS' then + task_due_key = egress_due_key + end + redis.call('ZADD', task_due_key, now_ms, due_member) + touch(task_due_key, tonumber(record.expiresAtMs)) end end end @@ -144,8 +151,13 @@ local function increment_state(upstream, state) end local function matches_task(task, fact) - return task and task.taskId == fact.taskId and task.proxyId == fact.proxyId and task.level == fact.level and - (task.routingName or '') == (fact.routingName or '') and (task.targetUrl or '') == (fact.targetUrl or '') + if not task or task.taskId ~= fact.taskId or task.proxyId ~= fact.proxyId or task.level ~= fact.level then + return false + end + if task.level == 'EGRESS' then + return (fact.routingName or '') == '' and (fact.targetUrl or '') == '' + end + return (task.routingName or '') == (fact.routingName or '') and (task.targetUrl or '') == (fact.targetUrl or '') end local function authorize(fact) @@ -198,22 +210,72 @@ if operation == 'due' then 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, upstreamId = record.sourceUpstream, state = record.state, dueAtMs = now_ms - } - if #result >= limit then - break + local level = payload.level or 'BASIC' + if level ~= 'BASIC' and level ~= 'EGRESS' then + return finish({status = 'invalid'}) + end + if level == 'EGRESS' and (type(payload.targetUrl) ~= 'string' or payload.targetUrl == '') then + return finish({status = 'invalid'}) + end + local result_members = {} + local function append_candidate(proxy_id, record, member) + if result_members[member] then + return + end + result_members[member] = true + result[#result + 1] = { + proxyId = proxy_id, upstreamId = record.sourceUpstream, state = record.state, level = level, + targetUrl = level == 'EGRESS' and payload.targetUrl or '', dueAtMs = now_ms + } + end + if level == 'EGRESS' then + local basic_ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit) + for _, proxy_id in ipairs(basic_ids) do + local _, record = live_record(proxy_id) + if record and valid_state(record.state) and + (type(payload.upstreamId) ~= 'string' or payload.upstreamId == '' or record.sourceUpstream == payload.upstreamId) then + local member = payload.targetUrl .. '\0' .. proxy_id + local reference = proxy_id .. '\0EGRESS\0\0' .. payload.targetUrl + if not redis.call('HGET', ref_task_key, reference) then + redis.call('ZADD', egress_due_key, now_ms, member) + touch(egress_due_key, tonumber(record.expiresAtMs)) + end end - end - end + end + end + local source_due_key = due_key + if level == 'EGRESS' then + source_due_key = egress_due_key + end + local ids = redis.call('ZRANGEBYSCORE', source_due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit) + for _, member in ipairs(ids) do + local proxy_id = member + local matching_target = true + if level == 'EGRESS' then + local separator = string.find(member, '\0', 1, true) + if not separator or string.sub(member, 1, separator - 1) ~= payload.targetUrl then + matching_target = false + else + proxy_id = string.sub(member, separator + 1) + end + end + if matching_target then + local _, record = live_record(proxy_id) + local reference = proxy_id .. '\0EGRESS\0\0' .. payload.targetUrl + if not record or not valid_state(record.state) then + redis.call('ZREM', source_due_key, member) + elseif level == 'EGRESS' and redis.call('HGET', ref_task_key, reference) then + redis.call('ZREM', source_due_key, member) + 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 + append_candidate(proxy_id, record, member) + if #result >= limit then + break + end + end + end + end local encoded_candidates = cjson.encode(result) if #result == 0 then encoded_candidates = '[]' @@ -227,29 +289,40 @@ if operation == 'offer' then end local offered = 0 for _, task in ipairs(payload.tasks) do - if type(task) ~= 'table' or task.version ~= 1 or task.level ~= 'BASIC' or task.state ~= 'QUEUED' or + if type(task) ~= 'table' or task.version ~= 1 or (task.level ~= 'BASIC' and task.level ~= 'EGRESS') or task.state ~= 'QUEUED' or type(task.taskId) ~= 'string' or type(task.proxyId) ~= 'string' or tonumber(task.deadlineMs or 0) <= now_ms or 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 + task.upstreamTasksKey == '' or type(task.reference) ~= 'string' or task.reference == '' or + type(task.dueMember) ~= 'string' or task.dueMember == '' 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) + if task.level == 'BASIC' and ((task.routingName or '') ~= '' or (task.targetUrl or '') ~= '') then + return finish({status = 'invalid'}) + end + if task.level == 'EGRESS' and ((task.routingName or '') ~= '' or type(task.targetUrl) ~= 'string' or task.targetUrl == '') then + return finish({status = 'invalid'}) + end + local task_due_key = due_key + if task.level == 'EGRESS' then + task_due_key = egress_due_key + end + local score = redis.call('ZSCORE', task_due_key, task.dueMember) + local current = redis.call('HGET', ref_task_key, task.reference) local raw, record = live_record(task.proxyId) 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 + if task.level == 'BASIC' and (record.state == 'FETCHED' or record.state == 'UNHEALTHY') then decrement_state(record.sourceUpstream, record.state) increment_state(record.sourceUpstream, 'CHECKING') record.state = 'CHECKING' raw = cjson.encode(record) redis.call('HSET', records_key, task.proxyId, raw) end - redis.call('ZREM', due_key, task.proxyId) + redis.call('ZREM', task_due_key, task.dueMember) redis.call('HSET', tasks_key, task.taskId, cjson.encode(task)) - redis.call('HSET', ref_task_key, task.proxyId, task.taskId) + redis.call('HSET', ref_task_key, task.reference, 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) @@ -344,8 +417,9 @@ if operation == 'authorize' or operation == 'complete' then end task.state = 'DONE' redis.call('HSET', tasks_key, task.taskId, cjson.encode(task)) - if redis.call('HGET', ref_task_key, task.proxyId) == task.taskId then - redis.call('HDEL', ref_task_key, task.proxyId) + local reference = task.reference or task.proxyId + if redis.call('HGET', ref_task_key, reference) == task.taskId then + redis.call('HDEL', ref_task_key, reference) end local _, record = live_record(task.proxyId) if record then @@ -353,8 +427,12 @@ if operation == 'authorize' or operation == 'complete' then if next_due_ms <= now_ms then next_due_ms = now_ms end - redis.call('ZADD', due_key, next_due_ms, task.proxyId) - touch(due_key, tonumber(record.expiresAtMs)) + local task_due_key = due_key + if task.level == 'EGRESS' then + task_due_key = egress_due_key + end + redis.call('ZADD', task_due_key, next_due_ms, task.dueMember or task.proxyId) + touch(task_due_key, tonumber(record.expiresAtMs)) end end return finish({status = 'ok'}) diff --git a/internal/controller/health/scheduler_runner.go b/internal/controller/health/scheduler_runner.go index c6c13a5..07556cb 100644 --- a/internal/controller/health/scheduler_runner.go +++ b/internal/controller/health/scheduler_runner.go @@ -5,6 +5,7 @@ import ( "errors" "sort" "strings" + "sync/atomic" "time" "proxy-pool/internal/config" @@ -30,6 +31,14 @@ type UpstreamTaskSource interface { 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 @@ -50,6 +59,27 @@ func (source upstreamDueSource) DueCandidates(ctx context.Context, now time.Time 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 { @@ -102,6 +132,7 @@ type ConfiguredSchedulerSupervisor struct { source UpstreamTaskSource sink TaskSink options SchedulerRunnerOptions + cursor atomic.Uint64 } func NewConfiguredSchedulerRunner( @@ -188,20 +219,58 @@ func (supervisor *ConfiguredSchedulerSupervisor) Tick(ctx context.Context) (Tick if !configured { continue } - source, err := NewUpstreamDueSource(supervisor.source, name) + basicSource, err := NewUpstreamDueSource(supervisor.source, name) if err != nil { return TickResult{}, err } - tick, err := tickWithPlanner(ctx, planner, source, supervisor.sink, schedulerOptionsForPlanner(supervisor.options, planner)) + check := config.EffectiveCheck(configuration.Defaults.Check, configuration.Upstreams[name].Check) + groups, err := supervisor.upstreamDueGroups(basicSource, name, check.URLs) if err != nil { return TickResult{}, err } - result.Planned += tick.Planned - result.Offered += tick.Offered + 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, diff --git a/internal/controller/health/scheduler_runner_test.go b/internal/controller/health/scheduler_runner_test.go index 9f7a1f0..5a289f0 100644 --- a/internal/controller/health/scheduler_runner_test.go +++ b/internal/controller/health/scheduler_runner_test.go @@ -164,12 +164,44 @@ func TestConfiguredSchedulerSupervisorDiscoversNewEnabledUpstream(t *testing.T) } } +func TestConfiguredSchedulerSupervisorSchedulesBoundedEgressGroups(t *testing.T) { + now := time.Date(2026, 8, 2, 11, 0, 0, 0, time.UTC) + configuration := &config.Config{ + Defaults: config.Defaults{Check: config.Check{ + Interval: config.Duration(time.Minute), MaxInFlight: 6, Timeout: config.Duration(time.Second), MaxAttempts: 1, + URLs: []string{"https://egress-one.example/identity", "https://egress-two.example/identity"}, + }}, + Upstreams: map[string]config.Upstream{"provider-a": {Enabled: true}}, + } + source := &upstreamTaskSourceStub{sources: map[string]*dueSourceStub{ + "provider-a": {candidates: []Candidate{{ + ProxyID: "basic", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now, + }}, egressCandidates: map[string][]Candidate{ + "https://egress-one.example/identity": {{ProxyID: "egress-one", State: proxyDomain.StateFetched, Level: healthDomain.LevelEgress, TargetURL: "https://egress-one.example/identity", DueAt: now}}, + "https://egress-two.example/identity": {{ProxyID: "egress-two", State: proxyDomain.StateFetched, Level: healthDomain.LevelEgress, TargetURL: "https://egress-two.example/identity", DueAt: now}}, + }}, + }} + sink := &taskSinkStub{} + supervisor, err := NewConfiguredSchedulerSupervisor(&configurationSourceStub{configuration: configuration}, source, sink, + SchedulerRunnerOptions{PollInterval: time.Second, BatchSize: 6, Now: func() time.Time { return now }}) + if err != nil { + t.Fatalf("NewConfiguredSchedulerSupervisor() = %v", err) + } + result, err := supervisor.Tick(context.Background()) + if err != nil || result != (TickResult{Planned: 3, Offered: 3}) || len(sink.tasks) != 1 || + sink.tasks[0].Candidate.Level != healthDomain.LevelEgress || source.sources["provider-a"].egressCalls != 2 { + t.Fatalf("Tick() = (%+v, %v); source=%+v sink=%+v", result, err, source, sink.tasks) + } +} + type dueSourceStub struct { - inFlight int - candidates []Candidate - limit int - inFlightCalls int - dueCalls int + inFlight int + candidates []Candidate + limit int + inFlightCalls int + dueCalls int + egressCandidates map[string][]Candidate + egressCalls int } func (source *dueSourceStub) InFlight(context.Context, time.Time) (int, error) { @@ -183,6 +215,12 @@ func (source *dueSourceStub) DueCandidates(_ context.Context, _ time.Time, limit return source.candidates, nil } +func (source *dueSourceStub) DueEgressCandidates(_ context.Context, targetURL string, _ time.Time, limit int) ([]Candidate, error) { + source.egressCalls++ + source.limit = limit + return source.egressCandidates[targetURL], nil +} + type upstreamTaskSourceStub struct { sources map[string]*dueSourceStub } @@ -203,6 +241,14 @@ func (source *upstreamTaskSourceStub) DueCandidatesForUpstream(ctx context.Conte return item.DueCandidates(ctx, now, limit) } +func (source *upstreamTaskSourceStub) DueEgressCandidatesForUpstream(ctx context.Context, upstreamID, targetURL string, now time.Time, limit int) ([]Candidate, error) { + item, exists := source.sources[upstreamID] + if !exists { + return nil, errors.New("missing upstream source") + } + return item.DueEgressCandidates(ctx, targetURL, now, limit) +} + type taskSinkStub struct { tasks []PlannedTask offer int