feat: stream leased checker tasks
This commit is contained in:
youfak
parent
c9fa32bd41
commit
a427842954
10
README.md
10
README.md
@ -69,7 +69,7 @@ flowchart LR
|
||||
Controller -->|Fetch| Provider[Provider API]
|
||||
Controller --> Redis[(Redis)]
|
||||
Controller --> PostgreSQL[(PostgreSQL)]
|
||||
Checker[Checker<br/>健康链待闭环] -. Observation .-> Controller
|
||||
Checker[Checker<br/>任务协议已接入] -. 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)：依赖、探针、发布边界与故障处置
|
||||
|
||||
## 核心不变量
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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 启动参数
|
||||
|
||||
|
||||
@ -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}` 等由加载器从同名环境变量
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 | 配置校验测试 |
|
||||
|
||||
@ -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" +
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@ -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,
|
||||
|
||||
487
internal/controller/health/task_broker.go
Normal file
487
internal/controller/health/task_broker.go
Normal file
@ -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:<redacted>}"))
|
||||
}
|
||||
|
||||
// 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:<redacted>}"))
|
||||
}
|
||||
|
||||
// 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)
|
||||
108
internal/controller/health/task_broker_test.go
Normal file
108
internal/controller/health/task_broker_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user