From a427842954defd3cc47da68400255d9b3ac05fe1 Mon Sep 17 00:00:00 2001 From: youfak Date: Fri, 31 Jul 2026 20:58:41 +0800 Subject: [PATCH] feat: stream leased checker tasks --- README.md | 10 +- api/proto/controlplane/v1/controlplane.proto | 6 + docs/api/control-plane.md | 14 +- docs/configuration/reference.md | 4 +- docs/development/implementation-plan.md | 6 +- docs/requirements/completion-audit.md | 5 +- docs/requirements/traceability.md | 4 +- gen/controlplane/v1/controlplane.pb.go | 72 ++- internal/controller/health/grpc_handler.go | 159 +++++- .../controller/health/grpc_handler_test.go | 94 ++++ internal/controller/health/task_broker.go | 487 ++++++++++++++++++ .../controller/health/task_broker_test.go | 108 ++++ 12 files changed, 931 insertions(+), 38 deletions(-) create mode 100644 internal/controller/health/task_broker.go create mode 100644 internal/controller/health/task_broker_test.go diff --git a/README.md b/README.md index 5a3f446..0bdef02 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ flowchart LR Controller -->|Fetch| Provider[Provider API] Controller --> Redis[(Redis)] Controller --> PostgreSQL[(PostgreSQL)] - Checker[Checker
健康链待闭环] -. Observation .-> Controller + Checker[Checker
任务协议已接入] -. Observation .-> Controller Controller -. gRPC Snapshot .-> Gateway ``` @@ -88,8 +88,8 @@ flowchart LR 提取与限流、Controller 的 Admin/Distribution/Metrics 监听，以及 PostgreSQL 管理状态；WorkerControlPlane 的 Register、Snapshot ACK、Runtime 心跳接收和 Redis 会话栅栏，以及 Gateway Outcome 上报的有界队列、序列确认与重试。 -- **部分完成**：Checker 调度与健康状态链、Docker Compose/ - Kubernetes 运行时 mTLS Overlay，以及 protobuf 契约。 +- **部分完成**：Checker 的任务协议、租约契约和健康状态归并；Redis 共享调度、 + 实际探测器、Docker Compose/Kubernetes 运行时 mTLS Overlay。 - **待完成**：Checker 调度与健康状态链，以及 loadgen 和代表性集群压测。 检查项数量不等于生产就绪度。静态部署清单与 protobuf descriptor 验证也不代表 @@ -164,8 +164,8 @@ Snapshot 时 `/readyz` 返回 `503`。Checker 与 loadgen 命令尚未实现。 - [Admin API](docs/api/admin.md)：默认本地入口 `http://127.0.0.1:8082` - Controller Metrics：默认本地入口 `http://127.0.0.1:9090`，提供 `/livez`、 `/readyz` 和 `/metrics` -- [控制面协议](docs/api/control-plane.md)：Worker/Checker 的 protobuf 契约； - gRPC 运行链尚未闭环 +- [控制面协议](docs/api/control-plane.md)：Worker/Checker 的 protobuf 契约与 + 已验证的有界任务领取/Observation 上报边界 - [运维手册](docs/operations/runbook.md)：依赖、探针、发布边界与故障处置 ## 核心不变量 diff --git a/api/proto/controlplane/v1/controlplane.proto b/api/proto/controlplane/v1/controlplane.proto index 79132d0..0ad8edd 100644 --- a/api/proto/controlplane/v1/controlplane.proto +++ b/api/proto/controlplane/v1/controlplane.proto @@ -229,6 +229,12 @@ message CheckTask { google.protobuf.Duration timeout = 10; uint32 attempt = 11; google.protobuf.Timestamp deadline = 12; + // Credential material is delivered only over the authenticated Checker mTLS + // stream and must be retained only for this task's lifetime. + string username = 13; + string password = 14; + string credential_version = 15; + uint32 max_attempts = 16; } enum CheckLevel { diff --git a/docs/api/control-plane.md b/docs/api/control-plane.md index a1caebc..1207d76 100644 --- a/docs/api/control-plane.md +++ b/docs/api/control-plane.md @@ -41,7 +41,9 @@ View 的其余候选中回退。`wait_timeout` 随 `on_unavailable=WAIT` 下发 `(session_id, sequence, SHA-256)` 建立 Redis 原子栅栏，并只返回最后确认的序列。 Redis 只保留每个当前会话的一条序列和摘要；原始 Outcome、代理明细与逐请求记录均不写入 Redis 或 PostgreSQL。Gateway 只将结果写入本地有界队列，队列满时丢弃样本，不等待控制面 -或存储。Checker 闭环尚未实现；`100,000 QPS` 仍是未验证的设计目标。 +或存储。Checker 的有界任务领取、任务租约归属校验和 Observation 上报已经由 +gRPC 契约测试覆盖；Redis 共享 due-index、生产任务 broker 与独立 Checker 探测进程 +尚未接入启动拓扑。`100,000 QPS` 仍是未验证的设计目标。 `WatchSnapshots` 建立时校验当前 session；每次签发快照引用时也把 `session_id` 交给 Redis 原子校验。重复 Register 会同时清除旧 Runtime 和已签发引用，因此迟到的 @@ -162,6 +164,14 @@ Checker 只返回 `HealthObservation`。Controller reducer 按 Proxy、检查层 Routing 决定 AVAILABLE、SUSPECT 或 UNHEALTHY，并更新 Redis 活动池，避免 多个 Checker 并发写状态。 +`StreamCheckTasks` 是有界 pull：请求中的 `max_in_flight` 受服务端单次领取上限 +限制，且同一 `checker_id` 的未完成租约会占用该窗口，重连不会扩大并发。任务仅在 +被领取时通过认证的 mTLS 流携带 endpoint、`secret_ref`、版本和任务期凭据；Checker +不访问 Redis 或 PostgreSQL。`ReportObservations` 在调用 Reducer 前校验 task、Proxy、 +检查层级、目标 Profile 与领取者一致，Reducer 成功后才确认任务；相同领取者对已确认 +任务的同一事实可重放，由活动池摘要幂等处理。当前生产配置尚未装配 Redis 共享 broker， +因此无共享 broker 的服务会以 `Unavailable` 拒绝任务流，而不下发无租约任务。 + ## 7. 兼容与演进 - Proto 字段号一旦发布不得复用。 @@ -177,6 +187,8 @@ Routing 决定 AVAILABLE、SUSPECT 或 UNHEALTHY，并更新 Redis 活动池， 消息中的逻辑 ID 与证书授权一致，设置单消息大小、流持续时间、并发 Stream 和上报批次上限。`secret_ref` 是受控引用；完整 Snapshot 的 `credentials` 在 mTLS 会话中 携带引用对应材料，Controller 和 Gateway 仅在内存处理，禁止写入 Redis/PostgreSQL、日志或指标。 +Checker 任务也遵循相同边界：连接凭据只存在于 Controller 的解析过程、任务流和 Checker +的单次执行期，不进入任务日志、指标或独立任务存储。 ## 9. Gateway 启动参数 diff --git a/docs/configuration/reference.md b/docs/configuration/reference.md index e5bab38..eb1f9ad 100644 --- a/docs/configuration/reference.md +++ b/docs/configuration/reference.md @@ -22,8 +22,8 @@ go run ./cmd/proxy-controller -config CONFIG_FILE PostgreSQL 管理面迁移、Redis 活动池、Distribution/Admin 独立监听与优雅停机； Controller Metrics 独立监听、`/livez`、`/readyz` 和基础 Prometheus 运行时指标； Provider 自动补池、分布式配额、动态重载和 Admin 低基数统计已装配。Gateway 进程、 -Worker 控制面会话、Snapshot 凭据分发和 Snapshot 就绪探针已装配；Checker 与代表性 -负载验证仍在后续实施范围。 +Worker 控制面会话、Snapshot 凭据分发和 Snapshot 就绪探针已装配；Checker 的任务流与 +租约契约已接入控制面，但 Redis 共享调度、独立执行器与代表性负载验证仍在后续实施范围。 所有时间值使用 Go duration，例如 `500ms`、`30s`、`5m`。示例中的 `${TOKEN}`、`${PASSWORD}`、`${POSTGRES_URL}` 等由加载器从同名环境变量 diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index acf1375..e17e6af 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -279,8 +279,10 @@ Lua 同一原子边界归并，覆盖首次 CHECKING 失败进入 UNHEALTHY、AV 阈值降级、精确重放、冲突拒绝与成功恢复。TARGET Profile 也已在 Memory 和 Redis 中独立 归并，以哈希键保存并随代理 TTL 过期，绝不写入 Proxy 全局状态或选择索引。调度器、 Checker Observation 上报 RPC 已复用既有控制面监听接入 Controller：SPIFFE `checker` 身份、 -每批上限、配置阈值解析、代理归属查询与 Reducer 均已闭环。`StreamCheckTasks` 的租约调度、 -Checker 独立进程与探测器尚未实现，因此本任务保持未完成。 +每批上限、配置阈值解析、代理归属查询与 Reducer 均已闭环。`StreamCheckTasks` 现已实现为 +有界 pull，并通过通用任务 broker 契约完成能力协商、同 Checker 并发窗口、租约到期回收、 +领取者栅栏和完成后重放；任务凭据仅由认证流在执行期下发。Redis 共享 due-index/租约持久化、 +生产 broker、Checker 独立进程与探测器尚未实现，因此本任务保持未完成。 ## Task 12: Machine-readable Contracts diff --git a/docs/requirements/completion-audit.md b/docs/requirements/completion-audit.md index 33a8972..50c1611 100644 --- a/docs/requirements/completion-audit.md +++ b/docs/requirements/completion-audit.md @@ -60,8 +60,9 @@ Outcome 已实现为 Gateway 本地有界队列、微批确认重试和 Controll 摘要 Redis 栅栏；原始事件不落 Redis 或 PostgreSQL。Checker 已有全局健康 Reducer 与 Memory/Redis 原子状态提交基础；TARGET Profile 以独立、随代理 TTL 过期的 Redis 记录归并，不改写 Proxy 全局状态。Controller 公用 Reducer 已作为 Observation 的唯一状态 -归并边界，Checker Observation RPC 已在同一控制面监听以独立 SPIFFE 身份接入；任务调度和 -独立执行进程尚未闭环。Snapshot 签发在 Redis 中原子匹配当前 +归并边界，Checker Observation RPC 已在同一控制面监听以独立 SPIFFE 身份接入；有界任务领取、 +租约归属与任务期凭据传输已有通用契约和 gRPC 往返测试，Redis 共享任务运行态与独立执行进程 +尚未闭环。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 e1480d8..31a9c14 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/按引用去重凭据完整快照；Gateway 已将它们编译为同版本动态 View，并由独立进程维护控制面会话。Checker/Loadgen 构建产物与任务流待实现 | +| ARCH-001 | 数据面 Worker 与控制面 Controller 分离 | 1-70 | 包、协议和部署拓扑已分离；Controller 已运行 Worker Register/Watch/ACK/Runtime/Outcome 与 Checker Observation gRPC，并发布 Proxy/Gateway Routing/按引用去重凭据完整快照；Checker 任务流已具备有界领取、租约栅栏和任务期凭据契约。Gateway 已将快照编译为同版本动态 View，并由独立进程维护控制面会话。Checker/Loadgen 构建产物与 Redis 共享任务运行态待实现 | | 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 校验和 Observation 上报已完成；有界调度器、抖动和分级频率测试待实现 | +| HEALTH-002 | 健康调度有 jitter、maxInFlight 和分级频率 | 8679-8736 | 配置有效合并、URL 校验、稳定抖动/优先级 Planner、有界 Scheduler tick，以及 Checker 任务的能力协商和租约窗口测试已完成；Redis due-index、跨副本 in-flight 原子限制和生产执行器待实现 | | HEALTH-003 | 失败分级 SUSPECT -> UNHEALTHY -> REMOVE | 8679-8736 | Controller 公用 Reducer 已通过 Memory/Redis 活动池原子提交全局连续失败、精确重放和成功恢复；任务调度与 REMOVE 编排待实现 | | SEC-001 | API 认证与 Proxy 认证分离，Secret 统一脱敏 | 7528-8111, 8904-8945 | Config 脱敏、Provider Store -> SecretRef -> Gateway Resolver 跨包测试与格式化泄漏回归测试 | | SEC-002 | 非回环监听无保护时严格模式启动失败 | 8112-8441 | 配置校验测试 | diff --git a/gen/controlplane/v1/controlplane.pb.go b/gen/controlplane/v1/controlplane.pb.go index f1f0808..1077128 100644 --- a/gen/controlplane/v1/controlplane.pb.go +++ b/gen/controlplane/v1/controlplane.pb.go @@ -1782,21 +1782,27 @@ func (x *StreamCheckTasksRequest) GetSupportedLevels() []CheckLevel { } type CheckTask struct { - state protoimpl.MessageState `protogen:"open.v1"` - TaskId string `protobuf:"bytes,1,opt,name=task_id,json=taskId,proto3" json:"task_id,omitempty"` - ProxyId string `protobuf:"bytes,2,opt,name=proxy_id,json=proxyId,proto3" json:"proxy_id,omitempty"` - Protocol ProxyProtocol `protobuf:"varint,3,opt,name=protocol,proto3,enum=proxy_pool.controlplane.v1.ProxyProtocol" json:"protocol,omitempty"` - Host string `protobuf:"bytes,4,opt,name=host,proto3" json:"host,omitempty"` - Port uint32 `protobuf:"varint,5,opt,name=port,proto3" json:"port,omitempty"` - SecretRef string `protobuf:"bytes,6,opt,name=secret_ref,json=secretRef,proto3" json:"secret_ref,omitempty"` - Level CheckLevel `protobuf:"varint,7,opt,name=level,proto3,enum=proxy_pool.controlplane.v1.CheckLevel" json:"level,omitempty"` - RoutingName string `protobuf:"bytes,8,opt,name=routing_name,json=routingName,proto3" json:"routing_name,omitempty"` - TargetUrl string `protobuf:"bytes,9,opt,name=target_url,json=targetUrl,proto3" json:"target_url,omitempty"` - Timeout *durationpb.Duration `protobuf:"bytes,10,opt,name=timeout,proto3" json:"timeout,omitempty"` - Attempt uint32 `protobuf:"varint,11,opt,name=attempt,proto3" json:"attempt,omitempty"` - Deadline *timestamppb.Timestamp `protobuf:"bytes,12,opt,name=deadline,proto3" json:"deadline,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"open.v1"` + TaskId string `protobuf:"bytes,1,opt,name=task_id,json=taskId,proto3" json:"task_id,omitempty"` + ProxyId string `protobuf:"bytes,2,opt,name=proxy_id,json=proxyId,proto3" json:"proxy_id,omitempty"` + Protocol ProxyProtocol `protobuf:"varint,3,opt,name=protocol,proto3,enum=proxy_pool.controlplane.v1.ProxyProtocol" json:"protocol,omitempty"` + Host string `protobuf:"bytes,4,opt,name=host,proto3" json:"host,omitempty"` + Port uint32 `protobuf:"varint,5,opt,name=port,proto3" json:"port,omitempty"` + SecretRef string `protobuf:"bytes,6,opt,name=secret_ref,json=secretRef,proto3" json:"secret_ref,omitempty"` + Level CheckLevel `protobuf:"varint,7,opt,name=level,proto3,enum=proxy_pool.controlplane.v1.CheckLevel" json:"level,omitempty"` + RoutingName string `protobuf:"bytes,8,opt,name=routing_name,json=routingName,proto3" json:"routing_name,omitempty"` + TargetUrl string `protobuf:"bytes,9,opt,name=target_url,json=targetUrl,proto3" json:"target_url,omitempty"` + Timeout *durationpb.Duration `protobuf:"bytes,10,opt,name=timeout,proto3" json:"timeout,omitempty"` + Attempt uint32 `protobuf:"varint,11,opt,name=attempt,proto3" json:"attempt,omitempty"` + Deadline *timestamppb.Timestamp `protobuf:"bytes,12,opt,name=deadline,proto3" json:"deadline,omitempty"` + // Credential material is delivered only over the authenticated Checker mTLS + // stream and must be retained only for this task's lifetime. + Username string `protobuf:"bytes,13,opt,name=username,proto3" json:"username,omitempty"` + Password string `protobuf:"bytes,14,opt,name=password,proto3" json:"password,omitempty"` + CredentialVersion string `protobuf:"bytes,15,opt,name=credential_version,json=credentialVersion,proto3" json:"credential_version,omitempty"` + MaxAttempts uint32 `protobuf:"varint,16,opt,name=max_attempts,json=maxAttempts,proto3" json:"max_attempts,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *CheckTask) Reset() { @@ -1913,6 +1919,34 @@ func (x *CheckTask) GetDeadline() *timestamppb.Timestamp { return nil } +func (x *CheckTask) GetUsername() string { + if x != nil { + return x.Username + } + return "" +} + +func (x *CheckTask) GetPassword() string { + if x != nil { + return x.Password + } + return "" +} + +func (x *CheckTask) GetCredentialVersion() string { + if x != nil { + return x.CredentialVersion + } + return "" +} + +func (x *CheckTask) GetMaxAttempts() uint32 { + if x != nil { + return x.MaxAttempts + } + return 0 +} + type ObservationBatch struct { state protoimpl.MessageState `protogen:"open.v1"` CheckerId string `protobuf:"bytes,1,opt,name=checker_id,json=checkerId,proto3" json:"checker_id,omitempty"` @@ -2290,7 +2324,7 @@ const file_controlplane_v1_controlplane_proto_rawDesc = "" + "\vinstance_id\x18\x02 \x01(\tR\n" + "instanceId\x12\"\n" + "\rmax_in_flight\x18\x03 \x01(\rR\vmaxInFlight\x12Q\n" + - "\x10supported_levels\x18\x04 \x03(\x0e2&.proxy_pool.controlplane.v1.CheckLevelR\x0fsupportedLevels\"\xd4\x03\n" + + "\x10supported_levels\x18\x04 \x03(\x0e2&.proxy_pool.controlplane.v1.CheckLevelR\x0fsupportedLevels\"\xde\x04\n" + "\tCheckTask\x12\x17\n" + "\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x19\n" + "\bproxy_id\x18\x02 \x01(\tR\aproxyId\x12E\n" + @@ -2306,7 +2340,11 @@ const file_controlplane_v1_controlplane_proto_rawDesc = "" + "\atimeout\x18\n" + " \x01(\v2\x19.google.protobuf.DurationR\atimeout\x12\x18\n" + "\aattempt\x18\v \x01(\rR\aattempt\x126\n" + - "\bdeadline\x18\f \x01(\v2\x1a.google.protobuf.TimestampR\bdeadline\"\x84\x01\n" + + "\bdeadline\x18\f \x01(\v2\x1a.google.protobuf.TimestampR\bdeadline\x12\x1a\n" + + "\busername\x18\r \x01(\tR\busername\x12\x1a\n" + + "\bpassword\x18\x0e \x01(\tR\bpassword\x12-\n" + + "\x12credential_version\x18\x0f \x01(\tR\x11credentialVersion\x12!\n" + + "\fmax_attempts\x18\x10 \x01(\rR\vmaxAttempts\"\x84\x01\n" + "\x10ObservationBatch\x12\x1d\n" + "\n" + "checker_id\x18\x01 \x01(\tR\tcheckerId\x12Q\n" + diff --git a/internal/controller/health/grpc_handler.go b/internal/controller/health/grpc_handler.go index d89dda0..6714a98 100644 --- a/internal/controller/health/grpc_handler.go +++ b/internal/controller/health/grpc_handler.go @@ -3,13 +3,19 @@ package health import ( "context" "errors" + "math" + "strings" + "time" controlplanev1 "proxy-pool/gen/controlplane/v1" "proxy-pool/internal/domain/activitypool" healthDomain "proxy-pool/internal/domain/health" + proxyDomain "proxy-pool/internal/domain/proxy" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/durationpb" + "google.golang.org/protobuf/types/known/timestamppb" ) var ErrInvalidGRPCHandler = errors.New("invalid checker grpc handler") @@ -23,14 +29,22 @@ type CheckerIdentityAuthorizer interface { type GRPCHandlerOptions struct { MaxObservationsPerBatch int + MaxTasksPerClaim int + TaskBroker TaskBroker + Now func() time.Time } func DefaultGRPCHandlerOptions() GRPCHandlerOptions { - return GRPCHandlerOptions{MaxObservationsPerBatch: 1_000} + return GRPCHandlerOptions{ + MaxObservationsPerBatch: 1_000, + MaxTasksPerClaim: defaultMaxTasksPerClaim, + Now: time.Now, + } } -// GRPCHandler exposes fact reporting now. StreamCheckTasks remains inherited -// as unimplemented until the Controller's leased task scheduler is available. +// GRPCHandler exposes the Controller's checker task and fact boundaries. The +// Broker is optional while a deployment has no shared task store; reporting +// remains available for its existing external fact integrations. type GRPCHandler struct { controlplanev1.UnimplementedCheckerControlPlaneServer reducer *Reducer @@ -46,15 +60,72 @@ func NewGRPCHandler( if reducer == nil || nilInterface(identity) { return nil, ErrInvalidGRPCHandler } + defaults := DefaultGRPCHandlerOptions() if options.MaxObservationsPerBatch == 0 { - options = DefaultGRPCHandlerOptions() + options.MaxObservationsPerBatch = defaults.MaxObservationsPerBatch } - if options.MaxObservationsPerBatch < 0 { + if options.MaxTasksPerClaim == 0 { + options.MaxTasksPerClaim = defaults.MaxTasksPerClaim + } + if options.Now == nil { + options.Now = defaults.Now + } + if options.MaxObservationsPerBatch < 0 || options.MaxTasksPerClaim < 0 { return nil, ErrInvalidGRPCHandler } return &GRPCHandler{reducer: reducer, identity: identity, options: options}, nil } +func (handler *GRPCHandler) StreamCheckTasks( + request *controlplanev1.StreamCheckTasksRequest, + stream controlplanev1.CheckerControlPlane_StreamCheckTasksServer, +) error { + if handler == nil || handler.reducer == nil || nilInterface(handler.identity) || request == nil || stream == nil || + handler.options.Now == nil || request.GetCheckerId() == "" || request.GetInstanceId() == "" || + request.GetMaxInFlight() == 0 || request.GetMaxInFlight() > uint32(handler.options.MaxTasksPerClaim) { + return healthGRPCError(ErrInvalidGRPCHandler) + } + ctx := stream.Context() + if ctx == nil { + return healthGRPCError(ErrInvalidGRPCHandler) + } + if err := handler.identity.AuthorizeChecker(ctx, request.GetCheckerId()); err != nil { + return status.Error(codes.PermissionDenied, "checker identity is not authorized") + } + if nilInterface(handler.options.TaskBroker) { + return status.Error(codes.Unavailable, "checker task broker unavailable") + } + levels := make([]healthDomain.Level, 0, len(request.GetSupportedLevels())) + for _, value := range request.GetSupportedLevels() { + level, ok := decodeCheckLevel(value) + if !ok { + return healthGRPCError(ErrInvalidTaskClaim) + } + levels = append(levels, level) + } + tasks, err := handler.options.TaskBroker.Claim(ctx, TaskClaim{ + CheckerID: request.GetCheckerId(), InstanceID: request.GetInstanceId(), + MaxInFlight: int(request.GetMaxInFlight()), SupportedLevels: levels, + }) + if err != nil { + return healthGRPCError(err) + } + now := handler.options.Now().UTC() + if now.IsZero() { + return healthGRPCError(ErrInvalidGRPCHandler) + } + for _, task := range tasks { + wire, wireErr := wireCheckTask(task, now) + if wireErr != nil { + return healthGRPCError(wireErr) + } + if sendErr := stream.Send(wire); sendErr != nil { + return sendErr + } + } + return nil +} + func (handler *GRPCHandler) ReportObservations( ctx context.Context, request *controlplanev1.ObservationBatch, @@ -70,9 +141,23 @@ func (handler *GRPCHandler) ReportObservations( response := &controlplanev1.ReportObservationsResponse{} for _, item := range request.GetObservations() { observation, err := decodeHealthObservation(item) + if err == nil && !nilInterface(handler.options.TaskBroker) { + now := handler.options.Now().UTC() + if now.IsZero() { + return nil, healthGRPCError(ErrInvalidGRPCHandler) + } + err = handler.options.TaskBroker.AuthorizeObservation(ctx, request.GetCheckerId(), observation, now) + } if err == nil { _, err = handler.reducer.Apply(ctx, observation) } + if err == nil && !nilInterface(handler.options.TaskBroker) { + now := handler.options.Now().UTC() + if now.IsZero() { + return nil, healthGRPCError(ErrInvalidGRPCHandler) + } + err = handler.options.TaskBroker.CompleteObservation(ctx, request.GetCheckerId(), observation, now) + } if err == nil { response.Accepted++ continue @@ -120,7 +205,9 @@ func rejectedObservationError(err error) bool { return errors.Is(err, healthDomain.ErrInvalidObservation) || errors.Is(err, healthDomain.ErrInvalidFailureThreshold) || errors.Is(err, healthDomain.ErrStaleObservation) || errors.Is(err, healthDomain.ErrConflictingObservation) || errors.Is(err, activitypool.ErrInvalidHealthUpdate) || errors.Is(err, activitypool.ErrActivityNotFound) || - errors.Is(err, ErrInvalidThreshold) || errors.Is(err, ErrUnconfiguredUpstream) + errors.Is(err, ErrInvalidThreshold) || errors.Is(err, ErrUnconfiguredUpstream) || + errors.Is(err, ErrTaskNotFound) || errors.Is(err, ErrTaskLeaseExpired) || + errors.Is(err, ErrTaskLeaseNotOwned) || errors.Is(err, ErrTaskObservation) } func healthGRPCError(err error) error { @@ -131,9 +218,67 @@ func healthGRPCError(err error) error { return status.Error(codes.DeadlineExceeded, "checker control request deadline exceeded") case errors.Is(err, ErrInvalidGRPCHandler), errors.Is(err, healthDomain.ErrInvalidObservation), errors.Is(err, healthDomain.ErrInvalidFailureThreshold), errors.Is(err, activitypool.ErrInvalidHealthUpdate), - errors.Is(err, ErrInvalidThreshold): + errors.Is(err, ErrInvalidThreshold), errors.Is(err, ErrInvalidTaskClaim), errors.Is(err, ErrInvalidLeasedTask): return status.Error(codes.InvalidArgument, "invalid checker observation batch") default: return status.Error(codes.Unavailable, "checker control plane unavailable") } } + +func wireCheckTask(task LeasedTask, now time.Time) (*controlplanev1.CheckTask, error) { + if !validTaskIdentifier(task.TaskID) || !validTaskIdentifier(task.ProxyID) || task.Host == "" || + strings.TrimSpace(task.Host) != task.Host || task.Port == 0 || task.Attempts <= 0 || task.Attempts > math.MaxUint32 || + !task.Deadline.After(now) || (task.SecretRef == "") != (task.CredentialVersion == "") { + return nil, ErrInvalidLeasedTask + } + level, ok := wireCheckLevel(task.Level) + if !ok { + return nil, ErrInvalidLeasedTask + } + protocol, ok := wireTaskProtocol(task.Protocol) + if !ok { + return nil, ErrInvalidLeasedTask + } + if task.Level == healthDomain.LevelTarget { + if _, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{ + RoutingName: task.RoutingName, TargetURL: task.TargetURL, + }); err != nil { + return nil, ErrInvalidLeasedTask + } + } else if task.RoutingName != "" || task.TargetURL != "" { + return nil, ErrInvalidLeasedTask + } + return &controlplanev1.CheckTask{ + TaskId: task.TaskID, ProxyId: task.ProxyID, Protocol: protocol, Host: task.Host, Port: uint32(task.Port), + SecretRef: task.SecretRef, CredentialVersion: task.CredentialVersion, Username: task.Username, Password: task.Password, + Level: level, RoutingName: task.RoutingName, TargetUrl: task.TargetURL, + Timeout: durationpb.New(task.Deadline.Sub(now)), Attempt: 1, MaxAttempts: uint32(task.Attempts), + Deadline: timestamppb.New(task.Deadline.UTC()), + }, nil +} + +func wireCheckLevel(level healthDomain.Level) (controlplanev1.CheckLevel, bool) { + switch level { + case healthDomain.LevelBasic: + return controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, true + case healthDomain.LevelEgress: + return controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS, true + case healthDomain.LevelTarget: + return controlplanev1.CheckLevel_CHECK_LEVEL_TARGET, true + default: + return controlplanev1.CheckLevel_CHECK_LEVEL_UNSPECIFIED, false + } +} + +func wireTaskProtocol(protocol proxyDomain.Scheme) (controlplanev1.ProxyProtocol, bool) { + switch protocol { + case proxyDomain.SchemeHTTP: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP, true + case proxyDomain.SchemeHTTPS: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS, true + case proxyDomain.SchemeSOCKS5: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5, true + default: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_UNSPECIFIED, false + } +} diff --git a/internal/controller/health/grpc_handler_test.go b/internal/controller/health/grpc_handler_test.go index 86ad896..751ec61 100644 --- a/internal/controller/health/grpc_handler_test.go +++ b/internal/controller/health/grpc_handler_test.go @@ -3,14 +3,20 @@ package health import ( "context" "errors" + "io" + "net" "testing" "time" controlplanev1 "proxy-pool/gen/controlplane/v1" healthDomain "proxy-pool/internal/domain/health" + proxyDomain "proxy-pool/internal/domain/proxy" + "google.golang.org/grpc" "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" "google.golang.org/protobuf/types/known/durationpb" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -96,6 +102,94 @@ func TestGRPCHandlerReturnsUnavailableForStoreFailure(t *testing.T) { } } +func TestGRPCHandlerStreamsLeasedTasksAndFencesReportedFacts(t *testing.T) { + now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC) + broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{ + LeaseTTL: time.Minute, Now: func() time.Time { return now }, + Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) { + return TaskMaterial{ + Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080, + SecretRef: "cred-a", CredentialVersion: "v1", Username: "user", Password: "secret", + }, nil + }), + }) + if err != nil { + t.Fatalf("NewMemoryTaskBroker(): %v", err) + } + if offered, offerErr := broker.Offer(context.Background(), []PlannedTask{{ + Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now}, + Deadline: now.Add(10 * time.Second), Attempts: 2, + }}); offerErr != nil || offered != 1 { + t.Fatalf("Offer() = (%d, %v)", offered, offerErr) + } + global := &recordingGlobalStore{} + reducer, err := NewReducer(global, &recordingTargetStore{}, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil }) + if err != nil { + t.Fatalf("NewReducer(): %v", err) + } + handler, err := NewGRPCHandler(reducer, &recordingCheckerIdentity{}, GRPCHandlerOptions{ + MaxObservationsPerBatch: 2, MaxTasksPerClaim: 2, TaskBroker: broker, Now: func() time.Time { return now }, + }) + if err != nil { + t.Fatalf("NewGRPCHandler(): %v", err) + } + client, closeClient := newCheckerGRPCClient(t, handler) + defer closeClient() + + stream, err := client.StreamCheckTasks(context.Background(), &controlplanev1.StreamCheckTasksRequest{ + CheckerId: "checker-a", InstanceId: "instance-a", MaxInFlight: 2, + SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_BASIC}, + }) + if err != nil { + t.Fatalf("StreamCheckTasks(): %v", err) + } + task, err := stream.Recv() + if err != nil || task.GetProxyId() != "proxy-a" || task.GetPassword() != "secret" || task.GetCredentialVersion() != "v1" || + task.GetMaxAttempts() != 2 || task.GetAttempt() != 1 || task.GetTimeout().AsDuration() != 10*time.Second { + t.Fatalf("Recv() = (%+v, %v)", task, err) + } + if _, err := stream.Recv(); !errors.Is(err, io.EOF) { + t.Fatalf("second Recv() error = %v, want EOF", err) + } + response, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{ + CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{ + grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second)), + }, + }) + if err != nil || response.GetAccepted() != 1 || len(global.commands) != 1 { + t.Fatalf("ReportObservations() = (%+v, %v), global=%d", response, err, len(global.commands)) + } + rejected, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{ + CheckerId: "checker-b", Observations: []*controlplanev1.HealthObservation{ + grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second)), + }, + }) + if err != nil || rejected.GetRejected() != 1 || len(global.commands) != 1 { + t.Fatalf("ReportObservations(other checker) = (%+v, %v), global=%d", rejected, err, len(global.commands)) + } +} + +func newCheckerGRPCClient(t *testing.T, handler controlplanev1.CheckerControlPlaneServer) (controlplanev1.CheckerControlPlaneClient, func()) { + t.Helper() + listener := bufconn.Listen(1 << 20) + server := grpc.NewServer() + controlplanev1.RegisterCheckerControlPlaneServer(server, handler) + go func() { _ = server.Serve(listener) }() + connection, err := grpc.NewClient("passthrough:///checker-test", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return listener.Dial() + }), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + server.Stop() + _ = listener.Close() + t.Fatalf("grpc.NewClient(): %v", err) + } + return controlplanev1.NewCheckerControlPlaneClient(connection), func() { + _ = connection.Close() + server.Stop() + _ = listener.Close() + } +} + func grpcHealthObservation(taskID string, level controlplanev1.CheckLevel, observedAt time.Time) *controlplanev1.HealthObservation { item := &controlplanev1.HealthObservation{ TaskId: taskID, ProxyId: "proxy-a", Level: level, Success: true, diff --git a/internal/controller/health/task_broker.go b/internal/controller/health/task_broker.go new file mode 100644 index 0000000..57c78fd --- /dev/null +++ b/internal/controller/health/task_broker.go @@ -0,0 +1,487 @@ +package health + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "reflect" + "sort" + "strconv" + "strings" + "sync" + "time" + + healthDomain "proxy-pool/internal/domain/health" + proxyDomain "proxy-pool/internal/domain/proxy" +) + +var ( + ErrInvalidTaskBroker = errors.New("invalid health task broker") + ErrInvalidTaskClaim = errors.New("invalid health task claim") + ErrInvalidLeasedTask = errors.New("invalid leased health task") + ErrTaskNotFound = errors.New("health task not found") + ErrTaskLeaseExpired = errors.New("health task lease expired") + ErrTaskLeaseNotOwned = errors.New("health task lease is not owned by checker") + ErrTaskObservation = errors.New("health observation does not match task") +) + +const defaultMaxTasksPerClaim = 128 + +// TaskMaterial is the short-lived proxy connection material needed to execute +// one probe. It crosses only the authenticated Checker control stream and is +// deliberately redacted from formatted values and persistence APIs. +type TaskMaterial struct { + Protocol proxyDomain.Scheme + Host string + Port uint16 + SecretRef string + CredentialVersion string + Username string + Password string +} + +func (TaskMaterial) Format(state fmt.State, _ rune) { + _, _ = state.Write([]byte("health.TaskMaterial{Credentials:}")) +} + +// TaskMaterialResolver resolves a task's endpoint and credential material at +// claim time. A production implementation reads the Controller-local +// credential store; Checkers never read Redis or PostgreSQL. +type TaskMaterialResolver interface { + ResolveCheckTask(context.Context, Candidate) (TaskMaterial, error) +} + +type TaskMaterialResolverFunc func(context.Context, Candidate) (TaskMaterial, error) + +func (resolver TaskMaterialResolverFunc) ResolveCheckTask(ctx context.Context, candidate Candidate) (TaskMaterial, error) { + return resolver(ctx, candidate) +} + +// TaskClaim is one bounded pull request from a Checker process. MaxInFlight +// applies across all streams for the checker ID, so reconnecting cannot grow +// its local work window. +type TaskClaim struct { + CheckerID string + InstanceID string + MaxInFlight int + SupportedLevels []healthDomain.Level +} + +// LeasedTask is a Controller-assigned task. It holds material only in memory +// for the duration of a task lease and must not be logged. +type LeasedTask struct { + TaskID string + ProxyID string + Protocol proxyDomain.Scheme + Host string + Port uint16 + SecretRef string + CredentialVersion string + Username string + Password string + Level healthDomain.Level + RoutingName string + TargetURL string + Deadline time.Time + Attempts int +} + +func (LeasedTask) Format(state fmt.State, _ rune) { + _, _ = state.Write([]byte("health.LeasedTask{Credentials:}")) +} + +// TaskBroker is the shared task lease boundary. The Scheduler uses Offer, +// Checkers use Claim, and ReportObservations fences facts against the lease. +// The production Redis implementation will use this exact contract. +type TaskBroker interface { + TaskSink + Claim(context.Context, TaskClaim) ([]LeasedTask, error) + AuthorizeObservation(context.Context, string, healthDomain.Observation, time.Time) error + CompleteObservation(context.Context, string, healthDomain.Observation, time.Time) error +} + +type MemoryTaskBrokerOptions struct { + LeaseTTL time.Duration + MaxPerClaim int + Now func() time.Time + Material TaskMaterialResolver +} + +// MemoryTaskBroker is a deterministic contract implementation for unit tests +// and single-process development. Production Controller instances use Redis so +// leases and limits remain shared across replicas. +type MemoryTaskBroker struct { + mu sync.Mutex + options MemoryTaskBrokerOptions + tasks map[string]brokerTask + leases map[string]taskLease + finished map[string]taskCompletion +} + +type brokerTask struct { + plan PlannedTask +} + +type taskLease struct { + checkerID string + instanceID string + expiresAt time.Time +} + +type taskCompletion struct { + checkerID string + expiresAt time.Time +} + +func NewMemoryTaskBroker(options MemoryTaskBrokerOptions) (*MemoryTaskBroker, error) { + if options.LeaseTTL <= 0 || options.Now == nil || nilTaskMaterialResolver(options.Material) { + return nil, ErrInvalidTaskBroker + } + if options.MaxPerClaim == 0 { + options.MaxPerClaim = defaultMaxTasksPerClaim + } + if options.MaxPerClaim <= 0 { + return nil, ErrInvalidTaskBroker + } + return &MemoryTaskBroker{ + options: options, tasks: make(map[string]brokerTask), leases: make(map[string]taskLease), + finished: make(map[string]taskCompletion), + }, nil +} + +func (broker *MemoryTaskBroker) Offer(ctx context.Context, plans []PlannedTask) (int, error) { + if ctx == nil || broker == nil || broker.options.Now == nil { + return 0, ErrInvalidTaskBroker + } + if err := ctx.Err(); err != nil { + return 0, err + } + if len(plans) == 0 { + return 0, nil + } + now := broker.options.Now().UTC() + if now.IsZero() { + return 0, ErrInvalidTaskBroker + } + for _, plan := range plans { + if err := validatePlannedTask(plan, now); err != nil { + return 0, err + } + } + broker.mu.Lock() + defer broker.mu.Unlock() + broker.reapLocked(now) + offered := 0 + for _, plan := range plans { + taskID := deterministicTaskID(plan) + if _, exists := broker.tasks[taskID]; exists { + continue + } + broker.tasks[taskID] = brokerTask{plan: plan} + offered++ + } + return offered, nil +} + +func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]LeasedTask, error) { + if ctx == nil || broker == nil || broker.options.Now == nil || nilTaskMaterialResolver(broker.options.Material) { + return nil, ErrInvalidTaskBroker + } + if err := ctx.Err(); err != nil { + return nil, err + } + levels, err := normalizeTaskClaim(claim, broker.options.MaxPerClaim) + if err != nil { + return nil, err + } + now := broker.options.Now().UTC() + if now.IsZero() { + return nil, ErrInvalidTaskBroker + } + broker.mu.Lock() + defer broker.mu.Unlock() + broker.reapLocked(now) + capacity := claim.MaxInFlight - broker.activeForCheckerLocked(claim.CheckerID) + if capacity <= 0 { + return []LeasedTask{}, nil + } + if capacity > broker.options.MaxPerClaim { + capacity = broker.options.MaxPerClaim + } + taskIDs := broker.claimableTaskIDsLocked(levels) + if len(taskIDs) > capacity { + taskIDs = taskIDs[:capacity] + } + if len(taskIDs) == 0 { + return []LeasedTask{}, nil + } + materials := make([]TaskMaterial, len(taskIDs)) + for index, taskID := range taskIDs { + material, resolveErr := broker.options.Material.ResolveCheckTask(ctx, broker.tasks[taskID].plan.Candidate) + if resolveErr != nil { + return nil, resolveErr + } + if validateTaskMaterial(material) != nil { + return nil, ErrInvalidLeasedTask + } + materials[index] = material + } + result := make([]LeasedTask, len(taskIDs)) + for index, taskID := range taskIDs { + plan := broker.tasks[taskID].plan + leaseUntil := now.Add(broker.options.LeaseTTL) + if plan.Deadline.Before(leaseUntil) { + leaseUntil = plan.Deadline + } + broker.leases[taskID] = taskLease{checkerID: claim.CheckerID, instanceID: claim.InstanceID, expiresAt: leaseUntil} + result[index] = leasedTask(taskID, plan, materials[index]) + } + return result, nil +} + +func (broker *MemoryTaskBroker) AuthorizeObservation( + ctx context.Context, + checkerID string, + observation healthDomain.Observation, + now time.Time, +) error { + if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() { + return ErrInvalidTaskBroker + } + if err := ctx.Err(); err != nil { + return err + } + normalized, err := healthDomain.NormalizeObservation(observation) + if err != nil { + return err + } + broker.mu.Lock() + defer broker.mu.Unlock() + return broker.authorizeObservationLocked(checkerID, normalized, now.UTC()) +} + +func (broker *MemoryTaskBroker) CompleteObservation( + ctx context.Context, + checkerID string, + observation healthDomain.Observation, + now time.Time, +) error { + if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() { + return ErrInvalidTaskBroker + } + if err := ctx.Err(); err != nil { + return err + } + normalized, err := healthDomain.NormalizeObservation(observation) + if err != nil { + return err + } + now = now.UTC() + broker.mu.Lock() + defer broker.mu.Unlock() + if err := broker.authorizeObservationLocked(checkerID, normalized, now); err != nil { + return err + } + if _, complete := broker.finished[normalized.TaskID]; complete { + return nil + } + task := broker.tasks[normalized.TaskID] + delete(broker.leases, normalized.TaskID) + broker.finished[normalized.TaskID] = taskCompletion{checkerID: checkerID, expiresAt: task.plan.Deadline} + return nil +} + +func (broker *MemoryTaskBroker) authorizeObservationLocked( + checkerID string, + observation healthDomain.Observation, + now time.Time, +) error { + broker.reapLocked(now) + task, exists := broker.tasks[observation.TaskID] + if !exists { + return ErrTaskNotFound + } + if !taskMatchesObservation(task.plan, observation) { + return ErrTaskObservation + } + if finished, complete := broker.finished[observation.TaskID]; complete { + if finished.checkerID != checkerID { + return ErrTaskLeaseNotOwned + } + return nil + } + lease, leased := broker.leases[observation.TaskID] + if !leased || !lease.expiresAt.After(now) { + return ErrTaskLeaseExpired + } + if lease.checkerID != checkerID { + return ErrTaskLeaseNotOwned + } + return nil +} + +func (broker *MemoryTaskBroker) reapLocked(now time.Time) { + for taskID, task := range broker.tasks { + if !task.plan.Deadline.After(now) { + delete(broker.tasks, taskID) + delete(broker.leases, taskID) + delete(broker.finished, taskID) + } + } + for taskID, lease := range broker.leases { + if !lease.expiresAt.After(now) { + delete(broker.leases, taskID) + } + } + for taskID, complete := range broker.finished { + if !complete.expiresAt.After(now) { + delete(broker.finished, taskID) + } + } +} + +func (broker *MemoryTaskBroker) activeForCheckerLocked(checkerID string) int { + active := 0 + for _, lease := range broker.leases { + if lease.checkerID == checkerID { + active++ + } + } + return active +} + +func (broker *MemoryTaskBroker) claimableTaskIDsLocked(levels map[healthDomain.Level]struct{}) []string { + result := make([]string, 0, len(broker.tasks)) + for taskID, task := range broker.tasks { + if _, leased := broker.leases[taskID]; leased { + continue + } + if _, complete := broker.finished[taskID]; complete { + continue + } + if _, supported := levels[task.plan.Candidate.Level]; supported { + result = append(result, taskID) + } + } + sort.Slice(result, func(left, right int) bool { + leftTask, rightTask := broker.tasks[result[left]].plan, broker.tasks[result[right]].plan + if leftTask.Priority != rightTask.Priority { + return leftTask.Priority < rightTask.Priority + } + if !leftTask.Candidate.DueAt.Equal(rightTask.Candidate.DueAt) { + return leftTask.Candidate.DueAt.Before(rightTask.Candidate.DueAt) + } + return result[left] < result[right] + }) + return result +} + +func normalizeTaskClaim(claim TaskClaim, maximum int) (map[healthDomain.Level]struct{}, error) { + if !validTaskIdentifier(claim.CheckerID) || !validTaskIdentifier(claim.InstanceID) || claim.MaxInFlight <= 0 || + claim.MaxInFlight > maximum || len(claim.SupportedLevels) == 0 { + return nil, ErrInvalidTaskClaim + } + levels := make(map[healthDomain.Level]struct{}, len(claim.SupportedLevels)) + for _, level := range claim.SupportedLevels { + switch level { + case healthDomain.LevelBasic, healthDomain.LevelEgress, healthDomain.LevelTarget: + default: + return nil, ErrInvalidTaskClaim + } + if _, duplicate := levels[level]; duplicate { + return nil, ErrInvalidTaskClaim + } + levels[level] = struct{}{} + } + return levels, nil +} + +func validatePlannedTask(task PlannedTask, now time.Time) error { + if task.Deadline.IsZero() || !task.Deadline.After(now) || task.Attempts <= 0 || !validTaskIdentifier(task.Candidate.ProxyID) || + task.Candidate.DueAt.IsZero() { + return ErrInvalidLeasedTask + } + switch task.Candidate.State { + case proxyDomain.StateFetched, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy, proxyDomain.StateAvailable: + default: + return ErrInvalidLeasedTask + } + switch task.Candidate.Level { + case healthDomain.LevelBasic, healthDomain.LevelEgress: + if task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" { + return ErrInvalidLeasedTask + } + case healthDomain.LevelTarget: + if _, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{ + RoutingName: task.Candidate.RoutingName, TargetURL: task.Candidate.TargetURL, + }); err != nil { + return ErrInvalidLeasedTask + } + default: + return ErrInvalidLeasedTask + } + return nil +} + +func validateTaskMaterial(material TaskMaterial) error { + if strings.TrimSpace(material.Host) != material.Host || material.Host == "" || material.Port == 0 || + (material.SecretRef == "") != (material.CredentialVersion == "") { + return ErrInvalidLeasedTask + } + switch material.Protocol { + case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5: + return nil + default: + return ErrInvalidLeasedTask + } +} + +func deterministicTaskID(task PlannedTask) string { + payload := candidateIdentity(task.Candidate) + "\x00" + task.Deadline.UTC().Format(time.RFC3339Nano) + "\x00" + + strconv.Itoa(task.Attempts) + digest := sha256.Sum256([]byte(payload)) + return "check_" + hex.EncodeToString(digest[:]) +} + +func leasedTask(taskID string, plan PlannedTask, material TaskMaterial) LeasedTask { + return LeasedTask{ + TaskID: taskID, ProxyID: plan.Candidate.ProxyID, Protocol: material.Protocol, Host: material.Host, Port: material.Port, + SecretRef: material.SecretRef, CredentialVersion: material.CredentialVersion, Username: material.Username, Password: material.Password, + Level: plan.Candidate.Level, RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL, + Deadline: plan.Deadline.UTC(), Attempts: plan.Attempts, + } +} + +func taskMatchesObservation(task PlannedTask, observation healthDomain.Observation) bool { + return task.Candidate.ProxyID == observation.ProxyID && task.Candidate.Level == observation.Level && + task.Candidate.RoutingName == observation.RoutingName && task.Candidate.TargetURL == observation.TargetURL +} + +func validTaskIdentifier(value string) bool { + if value == "" || len(value) > 256 || strings.TrimSpace(value) != value { + return false + } + for _, character := range value { + if character <= ' ' || character == '\x7f' { + return false + } + } + return true +} + +func nilTaskMaterialResolver(value TaskMaterialResolver) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} + +var _ TaskBroker = (*MemoryTaskBroker)(nil) diff --git a/internal/controller/health/task_broker_test.go b/internal/controller/health/task_broker_test.go new file mode 100644 index 0000000..814721f --- /dev/null +++ b/internal/controller/health/task_broker_test.go @@ -0,0 +1,108 @@ +package health + +import ( + "context" + "errors" + "testing" + "time" + + healthDomain "proxy-pool/internal/domain/health" + proxyDomain "proxy-pool/internal/domain/proxy" +) + +func TestMemoryTaskBrokerClaimsSupportedTasksAndFencesObservations(t *testing.T) { + now := time.Date(2026, 8, 1, 9, 0, 0, 0, time.UTC) + broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{ + LeaseTTL: time.Minute, + Now: func() time.Time { return now }, + Material: TaskMaterialResolverFunc(func(_ context.Context, candidate Candidate) (TaskMaterial, error) { + return TaskMaterial{ + Protocol: proxyDomain.SchemeHTTP, Host: candidate.ProxyID + ".example", Port: 8080, + SecretRef: "cred-" + candidate.ProxyID, CredentialVersion: "v1", Username: "user", Password: "secret", + }, nil + }), + }) + if err != nil { + t.Fatalf("NewMemoryTaskBroker(): %v", err) + } + plans := []PlannedTask{ + {Candidate: Candidate{ProxyID: "basic", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now}, Deadline: now.Add(10 * time.Second), Attempts: 2}, + {Candidate: Candidate{ProxyID: "target", State: proxyDomain.StateAvailable, Level: healthDomain.LevelTarget, RoutingName: "route-a", TargetURL: "https://target.example/check", DueAt: now}, Deadline: now.Add(10 * time.Second), Attempts: 2}, + } + if offered, err := broker.Offer(context.Background(), plans); err != nil || offered != 2 { + t.Fatalf("Offer() = (%d, %v)", offered, err) + } + claimed, err := broker.Claim(context.Background(), TaskClaim{ + CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 2, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}, + }) + if err != nil || len(claimed) != 1 || claimed[0].ProxyID != "basic" || claimed[0].Password != "secret" { + t.Fatalf("Claim() = (%+v, %v)", claimed, err) + } + observation := healthDomain.Observation{ + TaskID: claimed[0].TaskID, ProxyID: "basic", Level: healthDomain.LevelBasic, + Success: true, Latency: time.Millisecond, ObservedAt: now.Add(time.Second), + } + if err := broker.AuthorizeObservation(context.Background(), "checker-b", observation, now.Add(time.Second)); !errors.Is(err, ErrTaskLeaseNotOwned) { + t.Fatalf("AuthorizeObservation(other checker) error = %v, want ErrTaskLeaseNotOwned", err) + } + if err := broker.AuthorizeObservation(context.Background(), "checker-a", observation, now.Add(time.Second)); err != nil { + t.Fatalf("AuthorizeObservation(owner): %v", err) + } + if err := broker.CompleteObservation(context.Background(), "checker-a", observation, now.Add(time.Second)); err != nil { + t.Fatalf("CompleteObservation(): %v", err) + } + if err := broker.AuthorizeObservation(context.Background(), "checker-a", observation, now.Add(2*time.Second)); err != nil { + t.Fatalf("AuthorizeObservation(replay): %v", err) + } +} + +func TestMemoryTaskBrokerReclaimsExpiredLeaseWithoutExceedingCheckerLimit(t *testing.T) { + now := time.Date(2026, 8, 1, 10, 0, 0, 0, time.UTC) + broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{ + LeaseTTL: time.Second, + Now: func() time.Time { return now }, + Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) { + return TaskMaterial{Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080}, nil + }), + }) + if err != nil { + t.Fatalf("NewMemoryTaskBroker(): %v", err) + } + plan := PlannedTask{Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now}, Deadline: now.Add(time.Minute), Attempts: 1} + if _, err := broker.Offer(context.Background(), []PlannedTask{plan}); err != nil { + t.Fatalf("Offer(): %v", err) + } + first, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}}) + if err != nil || len(first) != 1 { + t.Fatalf("first Claim() = (%+v, %v)", first, err) + } + second, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}}) + if err != nil || len(second) != 0 { + t.Fatalf("second Claim() = (%+v, %v)", second, err) + } + now = now.Add(2 * time.Second) + reclaimed, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-b", InstanceID: "instance-b", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}}) + if err != nil || len(reclaimed) != 1 || reclaimed[0].TaskID != first[0].TaskID { + t.Fatalf("reclaimed Claim() = (%+v, %v)", reclaimed, err) + } +} + +func TestMemoryTaskBrokerRejectsTerminalProxyState(t *testing.T) { + now := time.Date(2026, 8, 1, 11, 0, 0, 0, time.UTC) + broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{ + LeaseTTL: time.Second, Now: func() time.Time { return now }, + Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) { + return TaskMaterial{Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080}, nil + }), + }) + if err != nil { + t.Fatalf("NewMemoryTaskBroker(): %v", err) + } + _, err = broker.Offer(context.Background(), []PlannedTask{{ + Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateExtracted, Level: healthDomain.LevelBasic, DueAt: now}, + Deadline: now.Add(time.Second), Attempts: 1, + }}) + if !errors.Is(err, ErrInvalidLeasedTask) { + t.Fatalf("Offer(terminal state) error = %v, want ErrInvalidLeasedTask", err) + } +}