From 139ff46a132c322f08983d5e10b1f97621909177 Mon Sep 17 00:00:00 2001 From: youfak Date: Tue, 28 Jul 2026 21:31:24 +0800 Subject: [PATCH] feat: implement proxy pool control plane core --- docs/development/implementation-plan.md | 57 +- docs/requirements/traceability.md | 37 +- docs/testing/strategy.md | 15 + internal/config/config.go | 30 +- internal/config/config_test.go | 348 ++++++++ internal/config/load.go | 19 +- internal/config/redact.go | 136 +++ internal/config/resolve.go | 147 +++ internal/config/validate.go | 364 +++++++- internal/controller/extraction/service.go | 178 ++++ .../controller/extraction/service_test.go | 191 ++++ internal/controller/pool/fetch_budget.go | 174 ++++ internal/controller/pool/fetch_budget_test.go | 155 ++++ internal/controller/pool/ownership.go | 73 ++ internal/controller/pool/ownership_test.go | 173 ++++ internal/controller/pool/reconciler.go | 64 ++ internal/controller/pool/reconciler_test.go | 76 ++ internal/controller/provider/ports.go | 61 ++ internal/controller/provider/reconciler.go | 304 +++++++ .../controller/provider/reconciler_test.go | 840 ++++++++++++++++++ internal/domain/extraction/extraction.go | 241 ++++- internal/domain/extraction/extraction_test.go | 163 ++++ internal/domain/ownership/ownership.go | 35 + internal/domain/routing/routing_test.go | 91 ++ internal/domain/routing/rule.go | 22 +- internal/domain/routing/sequential.go | 188 +++- internal/domain/routing/strategy.go | 168 ++++ internal/domain/routing/strategy_test.go | 232 +++++ internal/domain/upstream/fetch_capacity.go | 13 + internal/domain/upstream/pool.go | 81 ++ internal/domain/upstream/pool_test.go | 51 ++ internal/gateway/dispatch/dispatcher.go | 56 +- internal/gateway/dispatch/dispatcher_test.go | 197 ++++ internal/gateway/snapshot/store.go | 154 +++- internal/gateway/snapshot/store_test.go | 253 ++++++ internal/platform/admission/fixed_window.go | 88 ++ .../platform/admission/fixed_window_test.go | 76 ++ internal/platform/coalesce/signal.go | 28 + internal/platform/coalesce/signal_test.go | 32 + progress.md | 13 + 40 files changed, 5446 insertions(+), 178 deletions(-) create mode 100644 internal/config/redact.go create mode 100644 internal/config/resolve.go create mode 100644 internal/controller/extraction/service.go create mode 100644 internal/controller/extraction/service_test.go create mode 100644 internal/controller/pool/fetch_budget.go create mode 100644 internal/controller/pool/fetch_budget_test.go create mode 100644 internal/controller/pool/ownership.go create mode 100644 internal/controller/pool/ownership_test.go create mode 100644 internal/controller/pool/reconciler.go create mode 100644 internal/controller/pool/reconciler_test.go create mode 100644 internal/controller/provider/ports.go create mode 100644 internal/controller/provider/reconciler.go create mode 100644 internal/controller/provider/reconciler_test.go create mode 100644 internal/domain/ownership/ownership.go create mode 100644 internal/domain/routing/strategy.go create mode 100644 internal/domain/routing/strategy_test.go create mode 100644 internal/domain/upstream/fetch_capacity.go create mode 100644 internal/domain/upstream/pool.go create mode 100644 internal/domain/upstream/pool_test.go create mode 100644 internal/platform/admission/fixed_window.go create mode 100644 internal/platform/admission/fixed_window_test.go create mode 100644 internal/platform/coalesce/signal.go create mode 100644 internal/platform/coalesce/signal_test.go diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index 467237b..c1586e2 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -64,32 +64,32 @@ test/{fixtures,integration,e2e,load}/ **Files:** `internal/config/config.go`, `load.go`, `validate.go`, corresponding tests, `configs/default.yaml`, `docs/configuration/reference.md` -- [ ] Define versioned types for security, gateway, distribution, admin, metrics, +- [x] Define versioned types for security, gateway, distribution, admin, metrics, storage, routing, upstream/provider/api/proxyAuth/pool/capacity/lifecycle/fetch/check. -- [ ] Decode one YAML document with known fields enabled and resolve `${ENV}` plus +- [x] Decode one YAML document with known fields enabled and resolve `${ENV}` plus secret file references without logging values. -- [ ] Validate listener protection, routing references/order, regexes, strategy fields, +- [x] Validate listener protection, routing references/order, regexes, strategy fields, positive limits, TTL margins, pool/fetch limits, auth modes, and exposure modes. -- [ ] Add table tests for every invalid condition in CFG requirements. +- [x] Add table tests for every invalid condition in CFG requirements. ## Task 3: Proxy Domain and Capacity **Files:** `internal/domain/proxy/*.go`, corresponding tests -- [ ] Implement Proxy fields, UTC TTL precedence, canonical host/port, and unique key. -- [ ] Implement state transitions and reject illegal transitions. +- [x] Implement Proxy fields, UTC TTL precedence, canonical host/port, and unique key. +- [x] Implement state transitions and reject illegal transitions. - [ ] Implement sharded runtime counters with CAS Reserve, Commit, Cancel, Release. -- [ ] Prove with 1,000 concurrent goroutines that effective capacity is never exceeded. +- [x] Prove with 1,000 concurrent goroutines that effective capacity is never exceeded. - [ ] Add race coverage and duplicate-release invariant metrics hook. ## Task 4: Routing and Sequential Switching **Files:** `internal/domain/routing/*.go`, corresponding tests -- [ ] Compile first-match host/method/path/header rules into an immutable RuleSet. -- [ ] Implement random, round-robin, weighted, least-connections, and sequential. -- [ ] Model upstream empty counters separately from per-routing current indexes. -- [ ] Implement versioned CAS switch so simultaneous threshold observers advance once. +- [x] Compile first-match host/method/path/header rules into an immutable RuleSet. +- [x] Implement random, round-robin, weighted, least-connections, and sequential. +- [x] Model upstream empty counters separately from per-routing current indexes. +- [x] Implement versioned CAS switch so simultaneous threshold observers advance once. - [ ] Cover four-empty-then-success, five-empty, A-to-B-only, disabled references, end behavior, and explicit onUnavailable. @@ -97,13 +97,13 @@ test/{fixtures,integration,e2e,load}/ **Files:** `internal/controller/provider/*.go`, `internal/domain/upstream/*.go`, tests -- [ ] Implement Valid, Empty, DuplicateOnly, and Error result classes exactly as the +- [x] Implement Valid, Empty, DuplicateOnly, and Error result classes exactly as the traceability matrix defines. -- [ ] Implement one coalesced reconcile signal per Upstream using singleflight. -- [ ] Enforce requestInterval, maxInFlight, maxSize, maxTotal, timeout, retry, +- [x] Implement one coalesced reconcile signal per Upstream using singleflight. +- [x] Enforce requestInterval, maxInFlight, maxSize, maxTotal, timeout, retry, exponential backoff, jitter, and Retry-After. - [ ] Define ProviderAdapter and safe TemplateParser ports; add fixture adapters. -- [ ] Test that 100 concurrent capacity signals do not fan out 100 Provider calls. +- [x] Test that 100 concurrent capacity signals do not fan out 100 Provider calls. ## Task 6: Pool Reconciliation and Ownership @@ -111,32 +111,32 @@ test/{fixtures,integration,e2e,load}/ - [ ] Compute Available Slots from eligible Proxy capacity, Active, Reserved, TTL, health, ownership, pending expected fetch, and gateway reserve. -- [ ] Implement pool.maxSize and fetch.maxTotal as distinct counters. -- [ ] Allocate each Proxy to one Worker with epoch/version/expiry ownership. -- [ ] Implement revoke -> drain -> ACK -> unowned transition. -- [ ] Test Worker crash expiry and prevent simultaneous dual ownership. +- [x] Implement pool.maxSize and fetch.maxTotal as distinct counters. +- [x] Allocate each Proxy to one Worker with epoch/version/expiry ownership. +- [x] Implement revoke -> drain -> ACK -> unowned transition. +- [x] Test Worker crash expiry and prevent simultaneous dual ownership. ## Task 7: Exclusive Extraction **Files:** `internal/domain/extraction/*.go`, `internal/controller/extraction/*.go`, `internal/adapters/memory/extraction.go`, tests -- [ ] Implement POST extraction command with protocol/region/carrier/upstream filters. -- [ ] Enforce minRemainingTTL, maxHealthCheckAge, maxCount, client limits, and +- [x] Implement POST extraction command with protocol/region/carrier/upstream filters. +- [x] Enforce minRemainingTTL, maxHealthCheckAge, maxCount, client limits, and reserveForGateway. -- [ ] Atomically transition AVAILABLE to EXTRACTED and append audit records. -- [ ] Implement partial and allOrNothing without Lease, release, or renewal concepts. -- [ ] Run 1,000 concurrent claim attempts and prove every Proxy ID appears at most once. +- [x] Atomically transition AVAILABLE to EXTRACTED and append audit records. +- [x] Implement partial and allOrNothing without Lease, release, or renewal concepts. +- [x] Run 1,000 concurrent claim attempts and prove every Proxy ID appears at most once. ## Task 8: Immutable Snapshot and Dispatch **Files:** `internal/gateway/snapshot/*.go`, `internal/gateway/dispatch/*.go`, tests -- [ ] Define cluster/worker/epoch/version/checksum snapshot envelopes. +- [x] Define cluster/worker/epoch/version/checksum snapshot envelopes. - [ ] Build indexes in the background and atomically swap complete snapshots. -- [ ] Reject version gaps and wrong epochs; request full resync. -- [ ] Implement Dispatch Acquire/Commit/Release over local owned Proxy runtime. -- [ ] Benchmark 100k Proxy snapshots and record allocations and latency. +- [x] Reject version gaps and wrong epochs; request full resync. +- [x] Implement Dispatch Acquire/Commit/Release over local owned Proxy runtime. +- [x] Benchmark 100k Proxy snapshots and record allocations and latency. ## Task 9: Gateway Transport @@ -214,4 +214,3 @@ test/{fixtures,integration,e2e,load}/ representative cluster load run exists. - [ ] Confirm no TODO/TBD/placeholders, secrets, unbounded queues, high-cardinality metric labels, extraction Lease APIs, or conflicting maxSize semantics remain. - diff --git a/docs/requirements/traceability.md b/docs/requirements/traceability.md index a86d5e7..4639257 100644 --- a/docs/requirements/traceability.md +++ b/docs/requirements/traceability.md @@ -17,11 +17,11 @@ | ID | 最终需求 | 来源 | 验证证据 | |---|---|---|---| -| ROUTE-001 | Routing 自上而下匹配，首条命中停止 | 3534-3798, 5825-6467 | 路由单测 | +| ROUTE-001 | Routing 自上而下匹配，首条命中停止 | 3534-3798, 5825-6467 | `rule.go` 与不可变/首命中单测 | | ROUTE-002 | Routing 与 Upstream 生命周期解耦 | 3534-3798 | 包依赖与配置模型 | -| ROUTE-003 | 支持 sequential、random、roundRobin、weighted、leastConnections | 5825-6467 | 策略契约测试 | -| ROUTE-004 | Sequential 连续空结果达到阈值后原子切换一次 | 5295-5824, 6520-6617 | 并发切换测试 | -| ROUTE-005 | 空计数属于 Upstream，当前选择属于 Routing | 8442-8529 | 状态模型与多 Routing 测试 | +| ROUTE-003 | 支持 sequential、random、roundRobin、weighted、leastConnections | 5825-6467 | `strategy_test.go`、`routing_test.go` | +| ROUTE-004 | Sequential 连续空结果达到阈值后原子切换一次 | 5295-5824, 6520-6617 | `RoutingCursor` 版本 CAS 与 100 并发测试 | +| ROUTE-005 | 空计数属于 Upstream，当前选择属于 Routing | 8442-8529 | 共享 `UpstreamEmptyState` 双 Routing 测试 | | ROUTE-006 | 旧 Upstream 已有 Proxy 继续耗尽，不因切换直接丢弃 | 6618-6641 | Drain 测试 | | ROUTE-007 | 无可用 Upstream 时显式 reject、wait 或 direct，默认 reject | 5075-5294, 6743-6760 | 配置默认值与端到端测试 | @@ -29,14 +29,14 @@ | ID | 最终需求 | 来源 | 验证证据 | |---|---|---|---| -| FETCH-001 | 每个 Provider 有独立 requestInterval、maxInFlight、timeout 和 retry | 968-2394 | Fetcher 单测 | -| FETCH-002 | 大量缺池信号合并为 singleflight/容量 1 通知 | 2067-2136, 8808-8849 | 100 并发请求测试 | -| FETCH-003 | 错误使用指数退避和抖动，429 尊重 Retry-After | 1601-1831, 8808-8856 | 时钟驱动测试 | +| FETCH-001 | 每个 Provider 有独立 requestInterval、maxInFlight、timeout 和 retry | 968-2394 | `provider/reconciler_test.go` | +| FETCH-002 | 大量缺池信号合并为 singleflight/容量 1 通知 | 2067-2136, 8808-8849 | `coalesce.Signal` 与 100 并发通知测试 | +| FETCH-003 | 错误使用指数退避和抖动，429 尊重 Retry-After | 1601-1831, 8808-8856 | 注入时钟/随机数/Retry-After 测试 | | FETCH-004 | Provider 获取由单逻辑 Leader 执行 | 1403-1580 | 多实例锁测试 | -| FETCH-005 | Empty 与 Error 分开；只有合法候选为零时 Empty++ | 8442-8529 | 分类表驱动测试 | -| FETCH-006 | 重复候选不当作 Empty，记录独立指标 | 8442-8480 | 去重测试 | +| FETCH-005 | Empty 与 Error 分开；只有合法候选为零时 Empty++ | 8442-8529 | `fetch_result_test.go` 分类矩阵 | +| FETCH-006 | 重复候选不当作 Empty，记录独立指标 | 8442-8480 | DuplicateOnly 分类与 Provider 测试 | | FETCH-007 | 模板限制响应大小、执行时间、函数集和外部访问 | 8808-8856 | 安全测试 | -| FETCH-008 | pool.maxSize 与 fetch.maxTotal 语义分离 | 9190-9280 | 配置校验与计数测试 | +| FETCH-008 | pool.maxSize 与 fetch.maxTotal 语义分离 | 9190-9280 | `FetchBudget` 并发预占/释放测试 | ## Proxy 生命周期与容量 @@ -46,8 +46,8 @@ | PROXY-002 | 唯一键包含 scheme、host、port、username、credentialVersion | 6655-6727, 8605-8678 | 去重单测 | | PROXY-003 | TTL 来源优先级明确并统一 UTC | 681-747, 8655-8678 | TTL 表驱动测试 | | CAP-001 | Gateway 分配使用 Reserved -> Active 原子转换 | 1203-1467, 8530-8597 | 高并发竞态测试 | -| CAP-002 | 补池依据 Available Slots，不只看 Proxy 数量 | 1203-1402, 8530-8597 | 容量单测 | -| CAP-003 | pool.maxSize 包括 FETCHED/CHECKING/AVAILABLE/SUSPECT/DRAINING 与 pending expected | 3001-3533, 6642-6680 | 并发 fetch 上限测试 | +| CAP-002 | 补池依据 Available Slots，不只看 Proxy 数量 | 1203-1402, 8530-8597 | `Inventory.AvailableSlots` 与 Pool Reconciler 测试 | +| CAP-003 | pool.maxSize 包括 FETCHED/CHECKING/AVAILABLE/SUSPECT/DRAINING 与 pending expected | 3001-3533, 6642-6680 | `FetchBudget` 100 并发额度预占测试 | | CAP-004 | TTL safety margin 内禁止新分配 | 173-220, 6728-6741 | 时钟测试 | | CAP-005 | 多 Worker 不在热路径访问 Redis 计数 | 1403-1467 | 依赖审计与压测 | @@ -66,13 +66,13 @@ | ID | 最终需求 | 来源 | 验证证据 | |---|---|---|---| | DIST-001 | API 提取固定为一次性独占发放，不使用 Lease | 9083-9404 | Domain 状态机与 API 测试 | -| DIST-002 | AVAILABLE -> EXTRACTED 必须原子完成后才能返回 | 9083-9189 | 并发提取测试 | +| DIST-002 | AVAILABLE -> EXTRACTED 必须原子完成后才能返回 | 9083-9189 | 共享 Repository 所有权/提取 100 轮竞态测试 | | DIST-003 | 支持 partial 与 allOrNothing，默认 partial | 9190-9215 | API 契约测试 | -| DIST-004 | 保存审计记录，不提供释放接口 | 9216-9252 | Repository 测试与 OpenAPI | -| DIST-005 | 返回 expiresAt 与 remainingTtlSeconds | 9334-9360 | 响应测试 | +| DIST-004 | 保存审计记录，不提供释放接口 | 9216-9252 | 原子审计、幂等测试与 OpenAPI | +| DIST-005 | 返回 expiresAt 与 remainingTtlSeconds | 9334-9360 | `extraction/service_test.go` | | DIST-006 | 提取前校验 minRemainingTTL 与 maxHealthCheckAge | 9334-9369 | 过滤测试 | | DIST-007 | reserveForGateway 防止 Extract 清空共享池 | 9281-9333 | 共享池测试 | -| DIST-008 | 提取认证可关闭，关闭后仍有来源识别与全局限制 | 8112-8441 | 安全配置测试 | +| DIST-008 | 提取认证可关闭，关闭后仍有来源识别与全局限制 | 8112-8441 | 来源身份准入与 `FixedWindow` 并发测试 | ## 健康、安全、运维与测试 @@ -81,10 +81,9 @@ | HEALTH-001 | 全局健康与 Routing/目标健康分离 | 221-270, 8679-8708 | 健康 reducer 测试 | | HEALTH-002 | 健康调度有 jitter、maxInFlight 和分级频率 | 8679-8736 | 调度测试 | | HEALTH-003 | 失败分级 SUSPECT -> UNHEALTHY -> REMOVE | 8679-8736 | 状态机测试 | -| SEC-001 | API 认证与 Proxy 认证分离，Secret 统一脱敏 | 7528-8111, 8904-8945 | 配置类型与日志测试 | +| SEC-001 | API 认证与 Proxy 认证分离，Secret 统一脱敏 | 7528-8111, 8904-8945 | `Config.Redacted/String/GoString` 泄漏回归测试 | | SEC-002 | 非回环监听无保护时严格模式启动失败 | 8112-8441 | 配置校验测试 | -| OPS-001 | 配置校验后构建不可变快照并原子替换 | 8959-8999 | 热更新并发测试 | +| OPS-001 | 配置校验后构建不可变快照并原子替换 | 8959-8999 | 100k 索引、版本/epoch 与并发 Apply/Acquire 测试 | | OPS-002 | 优雅停机停止新请求/Fetch，等待现有流量后超时关闭 | 8981-9000 | 进程测试 | | OBS-001 | 指标禁止 Proxy IP、session、Client、完整 URL 高基数标签 | 9001-9029 | 指标描述符测试 | | TEST-001 | 覆盖对话中列出的 11 个关键并发与故障场景 | 9030-9082 | CI 测试清单 | - diff --git a/docs/testing/strategy.md b/docs/testing/strategy.md index bea241d..c775021 100644 --- a/docs/testing/strategy.md +++ b/docs/testing/strategy.md @@ -47,3 +47,18 @@ soak 测试单独标记，不混入快速单测。 只有在代表性环境持续达到目标且满足错误率和延迟门槛后，才能把“设计目标” 改为“已验证容量”。 + +## 5. 当前本地微基准 + +2026-07-28，Windows/amd64、Intel Core Ultra 7 155H： + +```text +BenchmarkAcquire100kIndexed-22 3553592 640.0 ns/op 256 B/op 2 allocs/op +BenchmarkStoreApply100k-22 1 472.7 ms/op 654 MB/op 2700642 allocs/op +``` + +`Acquire` 已使用 scheme/upstream/tag 索引，结果只代表本地选择和容量预留。 +`Store.Apply` 属于冷路径且当前内存开销较高；运行态为防止旧快照在途连接超配， +暂不自动回收曾出现过的 Proxy ID。后续需要基于 RCU/引用计数定义安全回收点。 +这些数据不包含网络、认证、Provider、存储或多 Worker 协调，不能作为 +100k QPS 端到端验收结论。 diff --git a/internal/config/config.go b/internal/config/config.go index f78ee7e..7446c77 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -62,22 +62,26 @@ type Access struct { } type Auth struct { - Mode string `yaml:"mode"` - Username string `yaml:"username"` - Password string `yaml:"password"` - Token string `yaml:"token"` - Header string `yaml:"header"` - CIDRs []string `yaml:"cidrs"` - Methods []AuthMethod `yaml:"methods"` + Mode string `yaml:"mode"` + Username string `yaml:"username"` + Password string `yaml:"password"` + PasswordFile string `yaml:"passwordFile"` + Token string `yaml:"token"` + TokenFile string `yaml:"tokenFile"` + Header string `yaml:"header"` + CIDRs []string `yaml:"cidrs"` + Methods []AuthMethod `yaml:"methods"` } type AuthMethod struct { - Mode string `yaml:"mode"` - Username string `yaml:"username"` - Password string `yaml:"password"` - Header string `yaml:"header"` - Value string `yaml:"value"` - CIDRs []string `yaml:"cidrs"` + Mode string `yaml:"mode"` + Username string `yaml:"username"` + Password string `yaml:"password"` + PasswordFile string `yaml:"passwordFile"` + Header string `yaml:"header"` + Value string `yaml:"value"` + ValueFile string `yaml:"valueFile"` + CIDRs []string `yaml:"cidrs"` } type Limits struct { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index dfbb1bb..02efb1e 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,6 +1,7 @@ package config import ( + "fmt" "os" "path/filepath" "strings" @@ -35,6 +36,7 @@ routing: strategy: type: sequential switchAfterEmptyFetch: 5 + endBehavior: stayLast onUnavailable: action: reject upstreams: @@ -143,3 +145,349 @@ func TestShippedConfigurationsAreValid(t *testing.T) { }) } } + +func TestLoadResolvedExpandsEnvironmentWithoutChangingTemplateVariables(t *testing.T) { + configured := strings.Replace(validConfig, ` auth: + type: none`, ` auth: + type: basic + username: "${PROVIDER_USER}" + password: "${PROVIDER_PASSWORD}"`, 1) + configured = strings.Replace(configured, "template: '{{.}}'", "template: '{{$x := .}}{{$x}}'", 1) + resolver := fixtureResolver{environment: map[string]string{ + "PROVIDER_USER": "alice", + "PROVIDER_PASSWORD": "secret", + }} + + cfg, err := LoadResolved(strings.NewReader(configured), resolver) + if err != nil { + t.Fatalf("LoadResolved(): %v", err) + } + auth := cfg.Upstreams["provider-a"].API.Auth + if auth.Username != "alice" || auth.Password != "secret" { + t.Fatalf("resolved auth = %+v", auth) + } + if got := cfg.Upstreams["provider-a"].API.Template; got != "{{$x := .}}{{$x}}" { + t.Fatalf("template = %q, want template variable unchanged", got) + } +} + +func TestLoadResolvedReadsSecretFileAndClearsReference(t *testing.T) { + configured := strings.Replace(validConfig, ` auth: + type: none`, ` auth: + type: basic + username: alice + passwordFile: /run/secrets/provider-password`, 1) + resolver := fixtureResolver{files: map[string]string{ + "/run/secrets/provider-password": "file-secret\r\n", + }} + + cfg, err := LoadResolved(strings.NewReader(configured), resolver) + if err != nil { + t.Fatalf("LoadResolved(): %v", err) + } + auth := cfg.Upstreams["provider-a"].API.Auth + if auth.Password != "file-secret" || auth.PasswordFile != "" { + t.Fatalf("resolved auth = %+v", auth) + } +} + +func TestLoadRejectsMultipleYAMLDocuments(t *testing.T) { + _, err := Load(strings.NewReader(validConfig + "\n---\nversion: 1\n")) + if err == nil || !strings.Contains(err.Error(), "multiple YAML documents") { + t.Fatalf("Load() error = %v, want multiple document error", err) + } +} + +func TestValidateRejectsUnsupportedListenerAuthMode(t *testing.T) { + broken := strings.Replace(validConfig, "mode: none", "mode: custom", 1) + _, err := Load(strings.NewReader(broken)) + if err == nil || !strings.Contains(err.Error(), "auth.mode") { + t.Fatalf("Load() error = %v, want auth.mode error", err) + } +} + +func TestValidateRejectsInvalidConfigurationMatrix(t *testing.T) { + tests := []struct { + name string + mutate func(*Config) + want string + }{ + { + name: "unsupported strategy type", + mutate: func(cfg *Config) { + cfg.Routing[0].Strategy.Type = "custom" + }, + want: "strategy.type", + }, + { + name: "unsupported unavailable action", + mutate: func(cfg *Config) { + cfg.Routing[0].OnUnavailable.Action = "fallback" + }, + want: "onUnavailable.action", + }, + { + name: "negative cumulative fetch limit", + mutate: func(cfg *Config) { + upstream := cfg.Upstreams["provider-a"] + upstream.Fetch.MaxTotal = -1 + cfg.Upstreams["provider-a"] = upstream + }, + want: "fetch.maxTotal", + }, + { + name: "negative gateway reserve", + mutate: func(cfg *Config) { + cfg.Distribution.Extraction.ReserveForGateway = -1 + }, + want: "reserveForGateway", + }, + { + name: "no enabled upstream", + mutate: func(cfg *Config) { + for name, upstream := range cfg.Upstreams { + upstream.Enabled = false + cfg.Upstreams[name] = upstream + } + }, + want: "enabled upstream", + }, + { + name: "invalid trusted proxy CIDR", + mutate: func(cfg *Config) { + cfg.Gateway.Access.TrustedProxies = []string{"not-a-cidr"} + }, + want: "trustedProxies", + }, + { + name: "invalid destination deny CIDR", + mutate: func(cfg *Config) { + cfg.Gateway.DestinationPolicy.DenyCIDRs = []string{"not-a-cidr"} + }, + want: "denyCIDRs", + }, + { + name: "unsupported routing purpose", + mutate: func(cfg *Config) { + cfg.Routing[0].Purpose = "custom" + }, + want: "purpose", + }, + { + name: "unsupported sequential end behavior", + mutate: func(cfg *Config) { + cfg.Routing[0].Strategy.EndBehavior = "restart" + }, + want: "endBehavior", + }, + { + name: "weighted strategy missing weight", + mutate: func(cfg *Config) { + cfg.Routing[0].Strategy = Strategy{Type: "weighted", Weights: map[string]int{}} + }, + want: "weights", + }, + { + name: "weighted strategy unknown upstream", + mutate: func(cfg *Config) { + cfg.Routing[0].Strategy = Strategy{Type: "weighted", Weights: map[string]int{ + "provider-a": 1, + "provider-b": 1, + }} + }, + want: "provider-b", + }, + { + name: "weighted strategy nonpositive weight", + mutate: func(cfg *Config) { + cfg.Routing[0].Strategy = Strategy{Type: "weighted", Weights: map[string]int{"provider-a": 0}} + }, + want: "weight", + }, + { + name: "negative fetch request interval", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.RequestInterval = Duration(-1) }) + }, + want: "requestInterval", + }, + { + name: "zero fetch timeout", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.Timeout = 0 }) + }, + want: "fetch.timeout", + }, + { + name: "zero fetch attempts", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.MaxAttempts = 0 }) + }, + want: "fetch.maxAttempts", + }, + { + name: "zero fetch concurrency", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.MaxInFlight = 0 }) + }, + want: "fetch.maxInFlight", + }, + { + name: "negative fetch response limit", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.MaxResponseBytes = -1 }) + }, + want: "maxResponseBytes", + }, + { + name: "negative template timeout", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.TemplateTimeout = Duration(-1) }) + }, + want: "templateTimeout", + }, + { + name: "fetch jitter over one hundred", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Fetch.Retry.Jitter = 101 }) + }, + want: "fetch.retry.jitter", + }, + { + name: "zero check interval", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.Interval = 0 }) + }, + want: "check.interval", + }, + { + name: "check jitter over one hundred", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.Jitter = 101 }) + }, + want: "check.jitter", + }, + { + name: "zero check concurrency", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.MaxInFlight = 0 }) + }, + want: "check.maxInFlight", + }, + { + name: "zero check timeout", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.Timeout = 0 }) + }, + want: "check.timeout", + }, + { + name: "zero check attempts", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.MaxAttempts = 0 }) + }, + want: "check.maxAttempts", + }, + { + name: "zero check failure threshold", + mutate: func(cfg *Config) { + updateUpstream(cfg, func(upstream *Upstream) { upstream.Check.MaxConsecutiveFailures = 0 }) + }, + want: "maxConsecutiveFailures", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := mustLoadValidConfig(t) + test.mutate(cfg) + err := Validate(cfg) + if err == nil || !strings.Contains(err.Error(), test.want) { + t.Fatalf("Validate() error = %v, want substring %q", err, test.want) + } + }) + } +} + +func TestResolvedConfigFormattingRedactsSecrets(t *testing.T) { + configured := strings.Replace(validConfig, ` auth: + mode: none`, ` auth: + mode: usernamePassword + username: gateway-user + password: "${GATEWAY_PASSWORD}"`, 1) + configured = strings.Replace(configured, ` auth: + mode: none`, ` auth: + mode: apiKey + header: X-API-Key + token: "${DISTRIBUTION_TOKEN}"`, 1) + configured = strings.Replace(configured, ` auth: + type: none`, ` auth: + type: apiKey + location: header + name: X-Provider-Key + value: "${PROVIDER_API_KEY}"`, 1) + resolver := fixtureResolver{environment: map[string]string{ + "GATEWAY_PASSWORD": "password-marker", + "DISTRIBUTION_TOKEN": "token-marker", + "PROVIDER_API_KEY": "api-key-marker", + }} + + cfg, err := LoadResolved(strings.NewReader(configured), resolver) + if err != nil { + t.Fatalf("LoadResolved(): %v", err) + } + for _, format := range []string{"%v", "%+v", "%#v"} { + formatted := fmt.Sprintf(format, cfg) + for _, secret := range []string{"password-marker", "token-marker", "api-key-marker"} { + if strings.Contains(formatted, secret) { + t.Fatalf("format %q leaked secret %q: %s", format, secret, formatted) + } + } + if !strings.Contains(formatted, "[REDACTED]") { + t.Fatalf("format %q did not contain a redaction marker: %s", format, formatted) + } + } + + redacted := cfg.Redacted() + if redacted.Gateway.Auth.Password != "[REDACTED]" || + redacted.Distribution.Auth.Token != "[REDACTED]" || + redacted.Upstreams["provider-a"].API.Auth.Value != "[REDACTED]" { + t.Fatalf("Redacted() retained a secret: %+v", &redacted) + } + if cfg.Gateway.Auth.Password != "password-marker" { + t.Fatal("Redacted() mutated the source configuration") + } +} + +func mustLoadValidConfig(t *testing.T) *Config { + t.Helper() + cfg, err := Load(strings.NewReader(validConfig)) + if err != nil { + t.Fatalf("Load(validConfig): %v", err) + } + return cfg +} + +func updateUpstream(cfg *Config, update func(*Upstream)) { + upstream := cfg.Upstreams["provider-a"] + update(&upstream) + cfg.Upstreams["provider-a"] = upstream +} + +type fixtureResolver struct { + environment map[string]string + files map[string]string +} + +func (r fixtureResolver) LookupEnv(name string) (string, bool) { + value, ok := r.environment[name] + return value, ok +} + +func (r fixtureResolver) ReadFile(path string) ([]byte, error) { + value, ok := r.files[path] + if !ok { + return nil, fmt.Errorf("fixture file %q not found", path) + } + return []byte(value), nil +} diff --git a/internal/config/load.go b/internal/config/load.go index d282e93..a8f4579 100644 --- a/internal/config/load.go +++ b/internal/config/load.go @@ -8,6 +8,17 @@ import ( ) func Load(reader io.Reader) (*Config, error) { + cfg, err := decode(reader) + if err != nil { + return nil, err + } + if err := Validate(cfg); err != nil { + return nil, err + } + return cfg, nil +} + +func decode(reader io.Reader) (*Config, error) { decoder := yaml.NewDecoder(reader) decoder.KnownFields(true) @@ -15,8 +26,12 @@ func Load(reader io.Reader) (*Config, error) { if err := decoder.Decode(&cfg); err != nil { return nil, fmt.Errorf("decode configuration: %w", err) } - if err := Validate(&cfg); err != nil { - return nil, err + var extra any + if err := decoder.Decode(&extra); err != io.EOF { + if err != nil { + return nil, fmt.Errorf("decode trailing configuration: %w", err) + } + return nil, fmt.Errorf("decode configuration: multiple YAML documents are not allowed") } return &cfg, nil } diff --git a/internal/config/redact.go b/internal/config/redact.go new file mode 100644 index 0000000..17c26c4 --- /dev/null +++ b/internal/config/redact.go @@ -0,0 +1,136 @@ +package config + +import "encoding/json" + +const redactedSecret = "[REDACTED]" + +// Redacted returns a detached configuration view suitable for diagnostics. +func (c *Config) Redacted() Config { + if c == nil { + return Config{} + } + redacted := cloneConfig(*c) + redactAuth(&redacted.Gateway.Auth) + redactAuth(&redacted.Distribution.Auth) + redactAuth(&redacted.Admin.Auth) + for name, upstream := range redacted.Upstreams { + redactProviderAuth(&upstream.API.Auth) + upstream.ProxyAuth.Password = redact(upstream.ProxyAuth.Password) + redacted.Upstreams[name] = upstream + } + redacted.Storage.PostgresURL = redact(redacted.Storage.PostgresURL) + redacted.Storage.RedisURL = redact(redacted.Storage.RedisURL) + return redacted +} + +func (c *Config) String() string { + if c == nil { + return "" + } + encoded, err := json.Marshal(c.Redacted()) + if err != nil { + return `{"config":"[REDACTED]"}` + } + return string(encoded) +} + +func (c *Config) GoString() string { return c.String() } + +func redactAuth(auth *Auth) { + auth.Password = redact(auth.Password) + auth.Token = redact(auth.Token) + for index := range auth.Methods { + auth.Methods[index].Password = redact(auth.Methods[index].Password) + auth.Methods[index].Value = redact(auth.Methods[index].Value) + } +} + +func redactProviderAuth(auth *ProviderAuth) { + auth.Password = redact(auth.Password) + auth.Token = redact(auth.Token) + auth.Value = redact(auth.Value) +} + +func redact(value string) string { + if value == "" { + return "" + } + return redactedSecret +} + +func cloneConfig(source Config) Config { + cloned := source + cloned.Defaults.Check.URLs = cloneStrings(source.Defaults.Check.URLs) + cloned.Gateway = cloneListener(source.Gateway) + cloned.Distribution.Listener = cloneListener(source.Distribution.Listener) + cloned.Admin = cloneListener(source.Admin) + cloned.Routing = make([]Routing, len(source.Routing)) + for index, route := range source.Routing { + cloned.Routing[index] = cloneRouting(route) + } + cloned.Upstreams = make(map[string]Upstream, len(source.Upstreams)) + for name, upstream := range source.Upstreams { + cloned.Upstreams[name] = cloneUpstream(upstream) + } + return cloned +} + +func cloneListener(source Listener) Listener { + cloned := source + cloned.Access.AllowCIDRs = cloneStrings(source.Access.AllowCIDRs) + cloned.Access.TrustedProxies = cloneStrings(source.Access.TrustedProxies) + cloned.Auth.CIDRs = cloneStrings(source.Auth.CIDRs) + cloned.Auth.Methods = append([]AuthMethod(nil), source.Auth.Methods...) + for index := range cloned.Auth.Methods { + cloned.Auth.Methods[index].CIDRs = cloneStrings(source.Auth.Methods[index].CIDRs) + } + cloned.Retry.RetryMethods = cloneStrings(source.Retry.RetryMethods) + cloned.DestinationPolicy.DenyCIDRs = cloneStrings(source.DestinationPolicy.DenyCIDRs) + return cloned +} + +func cloneRouting(source Routing) Routing { + cloned := source + cloned.Match.Methods = cloneStrings(source.Match.Methods) + cloned.Match.Headers = cloneStringMap(source.Match.Headers) + cloned.Upstreams = cloneStrings(source.Upstreams) + cloned.Strategy.Weights = cloneIntMap(source.Strategy.Weights) + return cloned +} + +func cloneUpstream(source Upstream) Upstream { + cloned := source + cloned.Exposure = cloneStrings(source.Exposure) + cloned.Provider.Protocols = cloneStrings(source.Provider.Protocols) + cloned.API.Headers = cloneStringMap(source.API.Headers) + cloned.API.Query = cloneStringMap(source.API.Query) + cloned.API.Body.Value = cloneStringMap(source.API.Body.Value) + cloned.Check.URLs = cloneStrings(source.Check.URLs) + return cloned +} + +func cloneStrings(source []string) []string { + return append([]string(nil), source...) +} + +func cloneStringMap(source map[string]string) map[string]string { + if source == nil { + return nil + } + cloned := make(map[string]string, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +func cloneIntMap(source map[string]int) map[string]int { + if source == nil { + return nil + } + cloned := make(map[string]int, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} diff --git a/internal/config/resolve.go b/internal/config/resolve.go new file mode 100644 index 0000000..1aa5909 --- /dev/null +++ b/internal/config/resolve.go @@ -0,0 +1,147 @@ +package config + +import ( + "fmt" + "io" + "os" + "reflect" + "regexp" + "strings" +) + +var environmentPattern = regexp.MustCompile(`\$\{([A-Z_][A-Z0-9_]*)\}`) + +type Resolver interface { + LookupEnv(string) (string, bool) + ReadFile(string) ([]byte, error) +} + +type OSResolver struct{} + +func (OSResolver) LookupEnv(name string) (string, bool) { return os.LookupEnv(name) } +func (OSResolver) ReadFile(path string) ([]byte, error) { return os.ReadFile(path) } + +func LoadResolved(reader io.Reader, resolver Resolver) (*Config, error) { + if resolver == nil { + return nil, fmt.Errorf("resolve configuration: resolver is required") + } + cfg, err := decode(reader) + if err != nil { + return nil, err + } + if err := expandValue(reflect.ValueOf(cfg).Elem(), resolver.LookupEnv); err != nil { + return nil, err + } + if err := resolveSecretFiles(cfg, resolver); err != nil { + return nil, err + } + if err := Validate(cfg); err != nil { + return nil, err + } + return cfg, nil +} + +func resolveSecretFiles(cfg *Config, resolver Resolver) error { + for _, auth := range []*Auth{&cfg.Gateway.Auth, &cfg.Distribution.Auth, &cfg.Admin.Auth} { + if err := resolveSecret(&auth.Password, &auth.PasswordFile, resolver); err != nil { + return err + } + if err := resolveSecret(&auth.Token, &auth.TokenFile, resolver); err != nil { + return err + } + for index := range auth.Methods { + method := &auth.Methods[index] + if err := resolveSecret(&method.Password, &method.PasswordFile, resolver); err != nil { + return err + } + if err := resolveSecret(&method.Value, &method.ValueFile, resolver); err != nil { + return err + } + } + } + for name, upstream := range cfg.Upstreams { + if err := resolveSecret(&upstream.API.Auth.Password, &upstream.API.Auth.PasswordFile, resolver); err != nil { + return fmt.Errorf("resolve upstream %q api password: %w", name, err) + } + if err := resolveSecret(&upstream.API.Auth.Token, &upstream.API.Auth.TokenFile, resolver); err != nil { + return fmt.Errorf("resolve upstream %q api token: %w", name, err) + } + if err := resolveSecret(&upstream.API.Auth.Value, &upstream.API.Auth.ValueFile, resolver); err != nil { + return fmt.Errorf("resolve upstream %q api key: %w", name, err) + } + if err := resolveSecret(&upstream.ProxyAuth.Password, &upstream.ProxyAuth.PasswordFile, resolver); err != nil { + return fmt.Errorf("resolve upstream %q proxy password: %w", name, err) + } + cfg.Upstreams[name] = upstream + } + return nil +} + +func resolveSecret(value, fileRef *string, resolver Resolver) error { + if *value != "" { + *fileRef = "" + return nil + } + if *fileRef == "" { + return nil + } + content, err := resolver.ReadFile(*fileRef) + if err != nil { + return err + } + resolved := strings.TrimRight(string(content), "\r\n") + if resolved == "" { + return fmt.Errorf("secret file %q is empty", *fileRef) + } + *value = resolved + *fileRef = "" + return nil +} + +func expandValue(value reflect.Value, lookup func(string) (string, bool)) error { + switch value.Kind() { + case reflect.String: + expanded, err := expandEnvironment(value.String(), lookup) + if err != nil { + return err + } + value.SetString(expanded) + case reflect.Struct: + for index := 0; index < value.NumField(); index++ { + if err := expandValue(value.Field(index), lookup); err != nil { + return err + } + } + case reflect.Slice: + for index := 0; index < value.Len(); index++ { + if err := expandValue(value.Index(index), lookup); err != nil { + return err + } + } + case reflect.Map: + iterator := value.MapRange() + for iterator.Next() { + item := reflect.New(value.Type().Elem()).Elem() + item.Set(iterator.Value()) + if err := expandValue(item, lookup); err != nil { + return err + } + value.SetMapIndex(iterator.Key(), item) + } + } + return nil +} + +func expandEnvironment(value string, lookup func(string) (string, bool)) (string, error) { + var resolveErr error + expanded := environmentPattern.ReplaceAllStringFunc(value, func(match string) string { + name := environmentPattern.FindStringSubmatch(match)[1] + resolved, ok := lookup(name) + if !ok { + resolveErr = fmt.Errorf("resolve configuration: environment variable %s is not set", name) + return match + } + return resolved + }) + return expanded, resolveErr +} diff --git a/internal/config/validate.go b/internal/config/validate.go index bc47681..efb2aa3 100644 --- a/internal/config/validate.go +++ b/internal/config/validate.go @@ -28,48 +28,47 @@ func Validate(cfg *Config) error { return err } } + if fetchConfigured(cfg.Defaults.Fetch) { + if err := validateFetch("defaults.fetch", cfg.Defaults.Fetch); err != nil { + return err + } + } + if err := validateCheck("defaults.check", cfg.Defaults.Check); err != nil { + return err + } + enabledUpstreams := 0 for name, upstream := range cfg.Upstreams { + if upstream.Enabled { + enabledUpstreams++ + } if err := validateUpstream(name, upstream); err != nil { return err } } + if enabledUpstreams == 0 { + return fmt.Errorf("validate configuration: at least one enabled upstream is required") + } seen := make(map[string]struct{}, len(cfg.Routing)) for index, route := range cfg.Routing { - if route.Name == "" { - return fmt.Errorf("validate routing[%d]: name is required", index) - } - if _, ok := seen[route.Name]; ok { - return fmt.Errorf("validate routing %q: duplicate name", route.Name) - } - seen[route.Name] = struct{}{} - if route.Match.HostRegex != "" { - if _, err := regexp.Compile(route.Match.HostRegex); err != nil { - return fmt.Errorf("validate routing %q hostRegex: %w", route.Name, err) - } - } - if route.Match.PathRegex != "" { - if _, err := regexp.Compile(route.Match.PathRegex); err != nil { - return fmt.Errorf("validate routing %q pathRegex: %w", route.Name, err) - } - } - for _, upstream := range route.Upstreams { - if _, ok := cfg.Upstreams[upstream]; !ok { - return fmt.Errorf("validate routing %q: upstream %q does not exist", route.Name, upstream) - } - } - if route.Strategy.Type == "sequential" && route.Strategy.SwitchAfterEmptyFetch <= 0 { - return fmt.Errorf("validate routing %q: switchAfterEmptyFetch must be greater than zero", route.Name) - } - if route.OnUnavailable.Action == "" { - return fmt.Errorf("validate routing %q: onUnavailable.action is required", route.Name) + if err := validateRouting(index, route, cfg.Upstreams, seen); err != nil { + return err } } if cfg.Distribution.Enabled { - if cfg.Distribution.Extraction.MaxCountPerRequest <= 0 { - return fmt.Errorf("validate distribution: maxCountPerRequest must be greater than zero") + if err := requirePositive("distribution.maxCountPerRequest", cfg.Distribution.Extraction.MaxCountPerRequest); err != nil { + return err } - if cfg.Distribution.Extraction.Fulfillment != "partial" && cfg.Distribution.Extraction.Fulfillment != "allOrNothing" { - return fmt.Errorf("validate distribution: fulfillment must be partial or allOrNothing") + if err := validateEnum("distribution.fulfillment", cfg.Distribution.Extraction.Fulfillment, "partial", "allOrNothing"); err != nil { + return err + } + if err := requireNonNegative("distribution.minRemainingTTL", cfg.Distribution.Extraction.MinRemainingTTL); err != nil { + return err + } + if err := requireNonNegative("distribution.maxHealthCheckAge", cfg.Distribution.Extraction.MaxHealthCheckAge); err != nil { + return err + } + if err := requireNonNegative("distribution.reserveForGateway", cfg.Distribution.Extraction.ReserveForGateway); err != nil { + return err } } return nil @@ -82,6 +81,9 @@ func validateListener(name string, listener Listener, security Security) error { if listener.Listen == "" { return fmt.Errorf("validate %s: listen is required", name) } + if err := validateListenerAuth(name, listener.Auth); err != nil { + return err + } host, _, err := net.SplitHostPort(listener.Listen) if err != nil { return fmt.Errorf("validate %s listen: %w", name, err) @@ -89,9 +91,164 @@ func validateListener(name string, listener Listener, security Security) error { if security.RequireProtectionOnPublicListen && isPublicHost(host) && listener.Auth.Mode == "none" && len(listener.Access.AllowCIDRs) == 0 { return fmt.Errorf("validate %s: unprotected public listener is forbidden", name) } - for _, cidr := range listener.Access.AllowCIDRs { + if err := validateCIDRs(name+" allowCIDRs", listener.Access.AllowCIDRs); err != nil { + return err + } + if err := validateCIDRs(name+" trustedProxies", listener.Access.TrustedProxies); err != nil { + return err + } + if err := validateCIDRs(name+" destinationPolicy.denyCIDRs", listener.DestinationPolicy.DenyCIDRs); err != nil { + return err + } + return nil +} + +func validateRouting(index int, route Routing, upstreams map[string]Upstream, seen map[string]struct{}) error { + if route.Name == "" { + return fmt.Errorf("validate routing[%d]: name is required", index) + } + if _, ok := seen[route.Name]; ok { + return fmt.Errorf("validate routing %q: duplicate name", route.Name) + } + seen[route.Name] = struct{}{} + scope := fmt.Sprintf("routing %q", route.Name) + if err := validateEnum(scope+" purpose", route.Purpose, "gateway", "extract"); err != nil { + return err + } + if route.Match.HostRegex != "" { + if _, err := regexp.Compile(route.Match.HostRegex); err != nil { + return fmt.Errorf("validate %s hostRegex: %w", scope, err) + } + } + if route.Match.PathRegex != "" { + if _, err := regexp.Compile(route.Match.PathRegex); err != nil { + return fmt.Errorf("validate %s pathRegex: %w", scope, err) + } + } + for _, upstream := range route.Upstreams { + if _, ok := upstreams[upstream]; !ok { + return fmt.Errorf("validate %s: upstream %q does not exist", scope, upstream) + } + } + if err := validateStrategy(scope, route.Upstreams, route.Strategy); err != nil { + return err + } + if err := validateEnum(scope+" onUnavailable.action", route.OnUnavailable.Action, "reject", "wait", "direct"); err != nil { + return err + } + if route.OnUnavailable.WaitTimeout < 0 { + return fmt.Errorf("validate %s onUnavailable.waitTimeout: must be non-negative", scope) + } + if route.OnUnavailable.Action == "wait" && route.OnUnavailable.WaitTimeout <= 0 { + return fmt.Errorf("validate %s onUnavailable.waitTimeout: must be greater than zero for wait", scope) + } + return nil +} + +func validateStrategy(scope string, upstreams []string, strategy Strategy) error { + if err := validateEnum(scope+" strategy.type", strategy.Type, + "sequential", "random", "roundRobin", "weighted", "leastConnections"); err != nil { + return err + } + if strategy.Type == "sequential" { + if err := requirePositive(scope+" strategy.switchAfterEmptyFetch", strategy.SwitchAfterEmptyFetch); err != nil { + return err + } + if strategy.EndBehavior != "" { + if err := validateEnum(scope+" strategy.endBehavior", strategy.EndBehavior, "stop", "loop", "stayLast"); err != nil { + return err + } + } + } + if strategy.Type != "weighted" { + if len(strategy.Weights) > 0 { + return fmt.Errorf("validate %s strategy.weights: only weighted strategy accepts weights", scope) + } + return nil + } + if len(strategy.Weights) == 0 { + return fmt.Errorf("validate %s strategy.weights: one positive weight per upstream is required", scope) + } + upstreamSet := make(map[string]struct{}, len(upstreams)) + for _, upstream := range upstreams { + upstreamSet[upstream] = struct{}{} + weight, ok := strategy.Weights[upstream] + if !ok { + return fmt.Errorf("validate %s strategy.weights: upstream %q has no weight", scope, upstream) + } + if err := requirePositive(scope+" strategy weight for "+upstream, weight); err != nil { + return err + } + } + for upstream := range strategy.Weights { + if _, ok := upstreamSet[upstream]; !ok { + return fmt.Errorf("validate %s strategy.weights: upstream %q is not referenced", scope, upstream) + } + } + return nil +} + +func validateListenerAuth(listener string, auth Auth) error { + switch auth.Mode { + case "none": + return nil + case "usernamePassword": + if auth.Username == "" || (auth.Password == "" && auth.PasswordFile == "") { + return fmt.Errorf("validate %s auth.mode usernamePassword: username and password are required", listener) + } + case "apiKey": + if auth.Header == "" || (auth.Token == "" && auth.TokenFile == "") { + return fmt.Errorf("validate %s auth.mode apiKey: header and token are required", listener) + } + case "ipWhitelist": + if len(auth.CIDRs) == 0 { + return fmt.Errorf("validate %s auth.mode ipWhitelist: cidrs are required", listener) + } + if err := validateCIDRs(listener+" auth", auth.CIDRs); err != nil { + return err + } + case "any": + if len(auth.Methods) == 0 { + return fmt.Errorf("validate %s auth.mode any: methods are required", listener) + } + for index, method := range auth.Methods { + if err := validateAuthMethod(listener, index, method); err != nil { + return err + } + } + default: + return fmt.Errorf("validate %s auth.mode: unsupported value %q", listener, auth.Mode) + } + return nil +} + +func validateAuthMethod(listener string, index int, method AuthMethod) error { + switch method.Mode { + case "usernamePassword": + if method.Username == "" || (method.Password == "" && method.PasswordFile == "") { + return fmt.Errorf("validate %s auth.methods[%d]: username and password are required", listener, index) + } + case "apiKey": + if method.Header == "" || (method.Value == "" && method.ValueFile == "") { + return fmt.Errorf("validate %s auth.methods[%d]: header and value are required", listener, index) + } + case "ipWhitelist": + if len(method.CIDRs) == 0 { + return fmt.Errorf("validate %s auth.methods[%d]: cidrs are required", listener, index) + } + if err := validateCIDRs(listener+" auth method", method.CIDRs); err != nil { + return err + } + default: + return fmt.Errorf("validate %s auth.methods[%d]: unsupported mode %q", listener, index, method.Mode) + } + return nil +} + +func validateCIDRs(name string, cidrs []string) error { + for _, cidr := range cidrs { if _, _, err := net.ParseCIDR(cidr); err != nil { - return fmt.Errorf("validate %s allowCIDRs %q: %w", name, cidr, err) + return fmt.Errorf("validate %s CIDR %q: %w", name, cidr, err) } } return nil @@ -101,20 +258,33 @@ func validateUpstream(name string, upstream Upstream) error { if !upstream.Enabled { return nil } - if upstream.Pool.MaxSize <= 0 { - return fmt.Errorf("validate upstream %q: pool.maxSize must be greater than zero", name) + scope := fmt.Sprintf("upstream %q", name) + if err := requirePositive(scope+" pool.maxSize", upstream.Pool.MaxSize); err != nil { + return err + } + if err := requireNonNegative(scope+" fetch.maxTotal", upstream.Fetch.MaxTotal); err != nil { + return err } if upstream.Fetch.MaxTotal > 0 && upstream.Fetch.MaxTotal < upstream.Pool.MaxSize { - return fmt.Errorf("validate upstream %q: fetch.maxTotal cannot be lower than pool.maxSize", name) + return fmt.Errorf("validate %s fetch.maxTotal: cannot be lower than pool.maxSize", scope) } - if upstream.Capacity.MaxConcurrencyPerProxy <= 0 { - return fmt.Errorf("validate upstream %q: maxConcurrencyPerProxy must be greater than zero", name) + if err := requirePositive(scope+" capacity.maxConcurrencyPerProxy", upstream.Capacity.MaxConcurrencyPerProxy); err != nil { + return err + } + if err := requireNonNegative(scope+" lifecycle.ttl", upstream.Lifecycle.TTL); err != nil { + return err + } + if err := requireNonNegative(scope+" lifecycle.allocationSafetyMargin", upstream.Lifecycle.AllocationSafetyMargin); err != nil { + return err } if upstream.Lifecycle.TTL > 0 && upstream.Lifecycle.AllocationSafetyMargin >= upstream.Lifecycle.TTL { - return fmt.Errorf("validate upstream %q: allocationSafetyMargin must be lower than ttl", name) + return fmt.Errorf("validate %s lifecycle.allocationSafetyMargin: must be lower than ttl", scope) } - if upstream.Fetch.RequestInterval < 0 || upstream.Fetch.MaxInFlight <= 0 || upstream.Fetch.MaxAttempts <= 0 { - return fmt.Errorf("validate upstream %q: fetch limits must be positive", name) + if err := validateFetch(scope+" fetch", upstream.Fetch); err != nil { + return err + } + if err := validateCheck(scope+" check", upstream.Check); err != nil { + return err } if upstream.API.URL != "" { parsed, err := url.Parse(upstream.API.URL) @@ -139,6 +309,118 @@ func validateUpstream(name string, upstream Upstream) error { return nil } +func validateFetch(scope string, fetch Fetch) error { + if err := requireNonNegative(scope+".requestInterval", fetch.RequestInterval); err != nil { + return err + } + if err := requirePositive(scope+".timeout", fetch.Timeout); err != nil { + return err + } + if err := requirePositive(scope+".maxAttempts", fetch.MaxAttempts); err != nil { + return err + } + if err := requirePositive(scope+".maxInFlight", fetch.MaxInFlight); err != nil { + return err + } + if err := requireNonNegative(scope+".maxTotal", fetch.MaxTotal); err != nil { + return err + } + if err := requireNonNegative(scope+".maxResponseBytes", fetch.MaxResponseBytes); err != nil { + return err + } + if err := requireNonNegative(scope+".templateTimeout", fetch.TemplateTimeout); err != nil { + return err + } + if err := requirePercentage(scope+".retry.jitter", fetch.Retry.Jitter); err != nil { + return err + } + if err := requireNonNegative(scope+".retry.initial", fetch.Retry.Initial); err != nil { + return err + } + if err := requireNonNegative(scope+".retry.max", fetch.Retry.Max); err != nil { + return err + } + if (fetch.Retry.Initial == 0) != (fetch.Retry.Max == 0) { + return fmt.Errorf("validate %s.retry: initial and max must be configured together", scope) + } + if fetch.Retry.Max > 0 && fetch.Retry.Max < fetch.Retry.Initial { + return fmt.Errorf("validate %s.retry.max: must not be lower than initial", scope) + } + return nil +} + +func validateCheck(scope string, check Check) error { + if !checkConfigured(check) { + return nil + } + if err := requirePositive(scope+".interval", check.Interval); err != nil { + return err + } + if err := requirePercentage(scope+".jitter", check.Jitter); err != nil { + return err + } + if err := requirePositive(scope+".maxInFlight", check.MaxInFlight); err != nil { + return err + } + if err := requirePositive(scope+".timeout", check.Timeout); err != nil { + return err + } + if err := requirePositive(scope+".maxAttempts", check.MaxAttempts); err != nil { + return err + } + if err := requirePositive(scope+".maxConsecutiveFailures", check.MaxConsecutiveFailures); err != nil { + return err + } + return nil +} + +func fetchConfigured(fetch Fetch) bool { + return fetch.RequestInterval != 0 || fetch.Timeout != 0 || fetch.MaxAttempts != 0 || + fetch.MaxInFlight != 0 || fetch.MaxTotal != 0 || fetch.MaxResponseBytes != 0 || + fetch.TemplateTimeout != 0 || fetch.Retry.Initial != 0 || fetch.Retry.Max != 0 || + fetch.Retry.Jitter != 0 +} + +func checkConfigured(check Check) bool { + return check.Interval != 0 || check.Jitter != 0 || check.MaxInFlight != 0 || + check.Timeout != 0 || check.MaxAttempts != 0 || check.MaxConsecutiveFailures != 0 || + len(check.URLs) != 0 +} + +type validationNumber interface { + ~int | ~int64 +} + +func requirePositive[T validationNumber](field string, value T) error { + if value <= 0 { + return fmt.Errorf("validate %s: must be greater than zero", field) + } + return nil +} + +func requireNonNegative[T validationNumber](field string, value T) error { + if value < 0 { + return fmt.Errorf("validate %s: must be non-negative", field) + } + return nil +} + +func requirePercentage(field string, value int) error { + if value < 0 || value > 100 { + return fmt.Errorf("validate %s: must be between 0 and 100", field) + } + return nil +} + +func validateEnum(field, value string, allowed ...string) error { + for _, candidate := range allowed { + if value == candidate { + return nil + } + } + return fmt.Errorf("validate %s: unsupported value %q; allowed values: %s", field, value, strings.Join(allowed, ", ")) +} + func validateProviderAuth(upstream string, auth ProviderAuth) error { switch auth.Type { case "", "none": diff --git a/internal/controller/extraction/service.go b/internal/controller/extraction/service.go new file mode 100644 index 0000000..3e2cf0f --- /dev/null +++ b/internal/controller/extraction/service.go @@ -0,0 +1,178 @@ +package extraction + +import ( + "context" + "errors" + "fmt" + "time" + + domain "github.com/proxy-pool/proxy-pool/internal/domain/extraction" +) + +var ( + ErrInvalidRequest = errors.New("invalid extraction request") + ErrCountExceeded = errors.New("extraction count exceeds policy") + ErrInvalidFulfillment = errors.New("invalid extraction fulfillment") + ErrInvalidServicePolicy = errors.New("invalid extraction service policy") + ErrAdmissionRejected = errors.New("extraction admission rejected") +) + +type Policy struct { + MaxCountPerRequest int + DefaultFulfillment domain.Fulfillment + MinRemainingTTL time.Duration + MaxHealthCheckAge time.Duration + ReserveForGateway int +} + +type Filters struct { + Protocols []string + Regions []string + Carriers []string + Upstreams []string +} + +type Request struct { + RequestID string + ClientID string + SourceIP string + IdempotencyKey string + Count int + Fulfillment domain.Fulfillment + Filters Filters +} + +type ExtractedProxy struct { + ID string + Protocol string + Host string + Port uint16 + Username string + Password string + URL string + Region string + Carrier string + Upstream string + ExpiresAt time.Time + RemainingTTLSeconds int64 + ExtractedAt time.Time +} + +type Response struct { + RequestID string + Requested int + Returned int + Proxies []ExtractedProxy +} + +type Service struct { + store domain.Store + policy Policy + admission Admission + now func() time.Time +} + +type Admission interface { + Admit(context.Context, string) error +} + +func NewService(store domain.Store, policy Policy, admission Admission, now func() time.Time) (*Service, error) { + if admission == nil { + return nil, fmt.Errorf("%w: admission is required", ErrInvalidServicePolicy) + } + if store == nil { + return nil, fmt.Errorf("%w: store is required", ErrInvalidServicePolicy) + } + if policy.MaxCountPerRequest <= 0 || policy.MinRemainingTTL < 0 || + policy.MaxHealthCheckAge < 0 || policy.ReserveForGateway < 0 { + return nil, ErrInvalidServicePolicy + } + if policy.DefaultFulfillment != domain.Partial && policy.DefaultFulfillment != domain.AllOrNothing { + return nil, ErrInvalidServicePolicy + } + if now == nil { + now = time.Now + } + return &Service{store: store, policy: policy, admission: admission, now: now}, nil +} + +func (s *Service) Extract(ctx context.Context, request Request) (Response, error) { + response := Response{RequestID: request.RequestID, Requested: request.Count} + if request.RequestID == "" || (request.ClientID == "" && request.SourceIP == "") || request.Count <= 0 { + return response, ErrInvalidRequest + } + if request.Count > s.policy.MaxCountPerRequest { + return response, ErrCountExceeded + } + fulfillment := request.Fulfillment + if fulfillment == "" { + fulfillment = s.policy.DefaultFulfillment + } + if fulfillment != domain.Partial && fulfillment != domain.AllOrNothing { + return response, ErrInvalidFulfillment + } + if err := s.admission.Admit(ctx, admissionKey(request)); err != nil { + return response, errors.Join(ErrAdmissionRejected, err) + } + + now := s.now().UTC() + clientID := request.ClientID + if clientID == "" { + clientID = "anonymous" + } + result, err := s.store.Extract(ctx, domain.Command{ + RequestID: request.RequestID, + ClientID: clientID, + SourceIP: request.SourceIP, + IdempotencyKey: request.IdempotencyKey, + Requested: request.Count, + Fulfillment: fulfillment, + Now: now, + MinRemainingTTL: s.policy.MinRemainingTTL, + MaxHealthCheckAge: s.policy.MaxHealthCheckAge, + ReserveForGateway: s.policy.ReserveForGateway, + Protocols: append([]string(nil), request.Filters.Protocols...), + Regions: append([]string(nil), request.Filters.Regions...), + Carriers: append([]string(nil), request.Filters.Carriers...), + Upstreams: append([]string(nil), request.Filters.Upstreams...), + }) + if err != nil { + return response, err + } + + response.Proxies = make([]ExtractedProxy, 0, len(result.Items)) + extractedAt := result.ExtractedAt + if extractedAt.IsZero() { + extractedAt = now + } + for _, candidate := range result.Items { + remaining := int64(0) + if !candidate.ExpiresAt.IsZero() && candidate.ExpiresAt.After(extractedAt) { + remaining = int64(candidate.ExpiresAt.Sub(extractedAt) / time.Second) + } + response.Proxies = append(response.Proxies, ExtractedProxy{ + ID: candidate.ID, + Protocol: candidate.Protocol, + Host: candidate.Host, + Port: candidate.Port, + Username: candidate.Username, + Password: candidate.Password, + URL: candidate.URL, + Region: candidate.Region, + Carrier: candidate.Carrier, + Upstream: candidate.Upstream, + ExpiresAt: candidate.ExpiresAt, + RemainingTTLSeconds: remaining, + ExtractedAt: extractedAt, + }) + } + response.Returned = len(response.Proxies) + return response, nil +} + +func admissionKey(request Request) string { + if request.ClientID != "" { + return "client:" + request.ClientID + } + return "source:" + request.SourceIP +} diff --git a/internal/controller/extraction/service_test.go b/internal/controller/extraction/service_test.go new file mode 100644 index 0000000..16c1fb2 --- /dev/null +++ b/internal/controller/extraction/service_test.go @@ -0,0 +1,191 @@ +package extraction + +import ( + "context" + "errors" + "testing" + "time" + + domain "github.com/proxy-pool/proxy-pool/internal/domain/extraction" +) + +func TestServiceAppliesPolicyAndBuildsResponse(t *testing.T) { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + store := &recordingStore{result: domain.Result{ + Requested: 2, + Returned: 1, + Items: []domain.Candidate{{ + ID: "p1", + Protocol: "http", + Host: "192.0.2.10", + Port: 8080, + URL: "http://user:pass@192.0.2.10:8080", + Upstream: "provider-a", + ExpiresAt: now.Add(95 * time.Second), + }}, + }} + service, err := NewService(store, Policy{ + MaxCountPerRequest: 20, + DefaultFulfillment: domain.Partial, + MinRemainingTTL: 30 * time.Second, + MaxHealthCheckAge: 15 * time.Second, + ReserveForGateway: 5, + }, allowAllAdmission{}, func() time.Time { return now }) + if err != nil { + t.Fatalf("NewService(): %v", err) + } + + response, err := service.Extract(context.Background(), Request{ + RequestID: "req-1", + ClientID: "client-1", + SourceIP: "192.0.2.30", + Count: 2, + Filters: Filters{ + Protocols: []string{"http"}, + Regions: []string{"shanghai"}, + }, + }) + if err != nil { + t.Fatalf("Extract(): %v", err) + } + + if store.command.Fulfillment != domain.Partial || store.command.ReserveForGateway != 5 { + t.Fatalf("store command policy = %+v", store.command) + } + if store.command.RequestID != "req-1" || store.command.ClientID != "client-1" { + t.Fatalf("store command audit context = %+v", store.command) + } + if response.RequestID != "req-1" || response.Requested != 2 || response.Returned != 1 { + t.Fatalf("response = %+v", response) + } + if len(response.Proxies) != 1 || response.Proxies[0].RemainingTTLSeconds != 95 { + t.Fatalf("response proxies = %+v", response.Proxies) + } + if !response.Proxies[0].ExtractedAt.Equal(now) { + t.Fatalf("extractedAt = %s, want %s", response.Proxies[0].ExtractedAt, now) + } +} + +func TestServiceRejectsCountAbovePolicyBeforeStore(t *testing.T) { + store := &recordingStore{} + service, err := NewService(store, Policy{ + MaxCountPerRequest: 1, + DefaultFulfillment: domain.Partial, + }, allowAllAdmission{}, time.Now) + if err != nil { + t.Fatalf("NewService(): %v", err) + } + + _, err = service.Extract(context.Background(), Request{ + RequestID: "req-1", + ClientID: "client-1", + Count: 2, + }) + if !errors.Is(err, ErrCountExceeded) { + t.Fatalf("Extract() error = %v, want ErrCountExceeded", err) + } + if store.calls != 0 { + t.Fatalf("store calls = %d, want 0", store.calls) + } +} + +func TestServiceAppliesAdmissionBeforeStoreUsingStableIdentity(t *testing.T) { + store := &recordingStore{} + admission := &recordingAdmission{err: ErrAdmissionRejected} + service, err := NewService(store, Policy{ + MaxCountPerRequest: 1, + DefaultFulfillment: domain.Partial, + }, admission, time.Now) + if err != nil { + t.Fatalf("NewService(): %v", err) + } + + _, err = service.Extract(context.Background(), Request{ + RequestID: "req-1", + SourceIP: "192.0.2.30", + Count: 1, + }) + if !errors.Is(err, ErrAdmissionRejected) { + t.Fatalf("Extract() error = %v, want ErrAdmissionRejected", err) + } + if admission.key != "source:192.0.2.30" || admission.calls != 1 { + t.Fatalf("admission = %+v, want source identity", admission) + } + if store.calls != 0 { + t.Fatalf("store calls = %d, want 0", store.calls) + } +} + +func TestServiceIdempotentReplayKeepsOriginalExtractionTime(t *testing.T) { + firstTime := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + secondTime := firstTime.Add(time.Minute) + store := domain.NewMemoryStore([]domain.Candidate{{ + ID: "p1", + State: domain.Available, + ExpiresAt: firstTime.Add(2 * time.Minute), + LastCheckedAt: firstTime, + }}) + times := []time.Time{firstTime, secondTime} + service, err := NewService(store, Policy{ + MaxCountPerRequest: 1, + DefaultFulfillment: domain.Partial, + }, allowAllAdmission{}, func() time.Time { + value := times[0] + times = times[1:] + return value + }) + if err != nil { + t.Fatalf("NewService(): %v", err) + } + request := Request{ + RequestID: "req-1", + ClientID: "client-1", + IdempotencyKey: "idem-12345678", + Count: 1, + } + + first, err := service.Extract(context.Background(), request) + if err != nil { + t.Fatalf("first Extract(): %v", err) + } + request.RequestID = "req-2" + second, err := service.Extract(context.Background(), request) + if err != nil { + t.Fatalf("second Extract(): %v", err) + } + if !second.Proxies[0].ExtractedAt.Equal(first.Proxies[0].ExtractedAt) { + t.Fatalf("replayed extractedAt = %s, want %s", second.Proxies[0].ExtractedAt, first.Proxies[0].ExtractedAt) + } + if second.Proxies[0].RemainingTTLSeconds != first.Proxies[0].RemainingTTLSeconds { + t.Fatalf("replayed remaining TTL = %d, want %d", second.Proxies[0].RemainingTTLSeconds, first.Proxies[0].RemainingTTLSeconds) + } +} + +type recordingStore struct { + command domain.Command + result domain.Result + err error + calls int +} + +type recordingAdmission struct { + key string + err error + calls int +} + +type allowAllAdmission struct{} + +func (allowAllAdmission) Admit(context.Context, string) error { return nil } + +func (a *recordingAdmission) Admit(_ context.Context, key string) error { + a.calls++ + a.key = key + return a.err +} + +func (s *recordingStore) Extract(_ context.Context, command domain.Command) (domain.Result, error) { + s.calls++ + s.command = command + return s.result, s.err +} diff --git a/internal/controller/pool/fetch_budget.go b/internal/controller/pool/fetch_budget.go new file mode 100644 index 0000000..466b84d --- /dev/null +++ b/internal/controller/pool/fetch_budget.go @@ -0,0 +1,174 @@ +package pool + +import ( + "errors" + "fmt" + "sync" + + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" +) + +var ( + ErrInvalidFetchBudget = errors.New("invalid fetch budget") + ErrInvalidFetchCompletion = errors.New("invalid fetch completion") + ErrFetchPermitFinished = errors.New("fetch permit is already finished") + ErrInvalidManagedRelease = errors.New("invalid managed proxy release") +) + +type FetchBudgetConfig struct { + UpstreamID string + MaxSize int + MaxTotal int64 + ExpectedPerFetch int + Managed int + FetchedTotal int64 +} + +type FetchBudgetSnapshot struct { + Managed int + PendingExpected int + FetchedTotal int64 +} + +// FetchBudget owns both current-inventory and cumulative-fetch accounting for +// one upstream. Reserving the expected response before I/O closes the race +// between concurrent provider calls. +type FetchBudget struct { + mu sync.Mutex + + upstreamID string + maxSize int + maxTotal int64 + expected int + usage FetchBudgetSnapshot +} + +func NewFetchBudget(config FetchBudgetConfig) (*FetchBudget, error) { + if config.UpstreamID == "" || config.MaxSize <= 0 || config.ExpectedPerFetch <= 0 || + config.ExpectedPerFetch > config.MaxSize || config.MaxTotal < 0 || + config.Managed < 0 || config.FetchedTotal < 0 { + return nil, ErrInvalidFetchBudget + } + if config.MaxTotal > 0 && int64(config.ExpectedPerFetch) > config.MaxTotal { + return nil, ErrInvalidFetchBudget + } + return &FetchBudget{ + upstreamID: config.UpstreamID, + maxSize: config.MaxSize, + maxTotal: config.MaxTotal, + expected: config.ExpectedPerFetch, + usage: FetchBudgetSnapshot{ + Managed: config.Managed, + FetchedTotal: config.FetchedTotal, + }, + }, nil +} + +func (b *FetchBudget) ReserveFetch(upstreamID string) (upstream.FetchPermit, bool, error) { + if b == nil || upstreamID == "" || upstreamID != b.upstreamID { + return nil, false, fmt.Errorf("%w: upstream %q", ErrInvalidFetchBudget, upstreamID) + } + b.mu.Lock() + defer b.mu.Unlock() + if !b.canReserveLocked() { + return nil, false, nil + } + b.usage.PendingExpected += b.expected + return &fetchPermit{budget: b, expected: b.expected}, true, nil +} + +// FetchAllowance is an advisory snapshot for reconciliation. ReserveFetch is +// still the atomic authority immediately before provider I/O. +func (b *FetchBudget) FetchAllowance() int { + if b == nil { + return 0 + } + b.mu.Lock() + defer b.mu.Unlock() + if b.canReserveLocked() { + return b.expected + } + return 0 +} + +func (b *FetchBudget) canReserveLocked() bool { + poolRoom := b.maxSize - b.usage.Managed - b.usage.PendingExpected + if poolRoom < b.expected { + return false + } + if b.maxTotal > 0 { + totalRoom := b.maxTotal - b.usage.FetchedTotal - int64(b.usage.PendingExpected) + return totalRoom >= int64(b.expected) + } + return true +} + +func (b *FetchBudget) Snapshot() FetchBudgetSnapshot { + if b == nil { + return FetchBudgetSnapshot{} + } + b.mu.Lock() + defer b.mu.Unlock() + return b.usage +} + +// ReleaseManaged returns current-inventory capacity after extraction, expiry, +// or removal. It deliberately does not restore the cumulative fetch quota. +func (b *FetchBudget) ReleaseManaged(count int) error { + if b == nil || count < 0 { + return ErrInvalidManagedRelease + } + b.mu.Lock() + defer b.mu.Unlock() + if count > b.usage.Managed { + return ErrInvalidManagedRelease + } + b.usage.Managed -= count + return nil +} + +type fetchPermit struct { + budget *FetchBudget + expected int + finished bool +} + +func (p *fetchPermit) Expected() int { + if p == nil { + return 0 + } + return p.expected +} + +func (p *fetchPermit) Complete(fetched, retained int) error { + if p == nil || p.budget == nil { + return ErrFetchPermitFinished + } + if fetched < 0 || retained < 0 || retained > fetched || retained > p.expected { + return ErrInvalidFetchCompletion + } + p.budget.mu.Lock() + defer p.budget.mu.Unlock() + if p.finished { + return ErrFetchPermitFinished + } + p.finished = true + p.budget.usage.PendingExpected -= p.expected + p.budget.usage.Managed += retained + p.budget.usage.FetchedTotal += int64(fetched) + return nil +} + +func (p *fetchPermit) Cancel() error { + if p == nil || p.budget == nil { + return ErrFetchPermitFinished + } + p.budget.mu.Lock() + defer p.budget.mu.Unlock() + if p.finished { + return ErrFetchPermitFinished + } + p.finished = true + p.budget.usage.PendingExpected -= p.expected + return nil +} diff --git a/internal/controller/pool/fetch_budget_test.go b/internal/controller/pool/fetch_budget_test.go new file mode 100644 index 0000000..c6ddb9b --- /dev/null +++ b/internal/controller/pool/fetch_budget_test.go @@ -0,0 +1,155 @@ +package pool + +import ( + "errors" + "sync" + "testing" +) + +func TestFetchBudgetReservesExpectedCapacityAndSeparatesCounters(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "provider-a", + MaxSize: 10, + MaxTotal: 20, + ExpectedPerFetch: 4, + Managed: 2, + FetchedTotal: 3, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + + first, ok, err := budget.ReserveFetch("provider-a") + if err != nil || !ok { + t.Fatalf("first ReserveFetch() = (_, %v, %v), want permit", ok, err) + } + second, ok, err := budget.ReserveFetch("provider-a") + if err != nil || !ok { + t.Fatalf("second ReserveFetch() = (_, %v, %v), want permit", ok, err) + } + if _, ok, err := budget.ReserveFetch("provider-a"); err != nil || ok { + t.Fatalf("third ReserveFetch() = (_, %v, %v), want no capacity", ok, err) + } + + if err := first.Complete(4, 3); err != nil { + t.Fatalf("first.Complete(): %v", err) + } + if err := second.Cancel(); err != nil { + t.Fatalf("second.Cancel(): %v", err) + } + usage := budget.Snapshot() + if usage.Managed != 5 || usage.PendingExpected != 0 || usage.FetchedTotal != 7 { + t.Fatalf("Snapshot() = %+v, want managed=5 pending=0 fetched=7", usage) + } + if err := budget.ReleaseManaged(2); err != nil { + t.Fatalf("ReleaseManaged(): %v", err) + } + if usage := budget.Snapshot(); usage.Managed != 3 || usage.FetchedTotal != 7 { + t.Fatalf("Snapshot() after release = %+v, want managed=3 fetched=7", usage) + } +} + +func TestFetchBudgetRejectsManagedCounterUnderflow(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 1, Managed: 1, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + if err := budget.ReleaseManaged(2); !errors.Is(err, ErrInvalidManagedRelease) { + t.Fatalf("ReleaseManaged() error = %v, want ErrInvalidManagedRelease", err) + } +} + +func TestFetchBudgetRequiresWholeExpectedBatchToFitLimits(t *testing.T) { + tests := []struct { + name string + config FetchBudgetConfig + }{ + { + name: "pool size", + config: FetchBudgetConfig{ + UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 4, Managed: 7, + }, + }, + { + name: "cumulative total", + config: FetchBudgetConfig{ + UpstreamID: "a", MaxSize: 10, MaxTotal: 5, ExpectedPerFetch: 4, FetchedTotal: 2, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + budget, err := NewFetchBudget(tt.config) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + if _, ok, err := budget.ReserveFetch("a"); err != nil || ok { + t.Fatalf("ReserveFetch() = (_, %v, %v), want no capacity", ok, err) + } + }) + } +} + +func TestFetchBudgetConcurrentReservationsNeverExceedMaxSize(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 1, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + + var wg sync.WaitGroup + permits := make(chan interface{ Cancel() error }, 100) + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + permit, ok, reserveErr := budget.ReserveFetch("a") + if reserveErr != nil { + t.Errorf("ReserveFetch(): %v", reserveErr) + return + } + if ok { + permits <- permit + } + }() + } + wg.Wait() + close(permits) + + count := 0 + for permit := range permits { + count++ + if err := permit.Cancel(); err != nil { + t.Errorf("Cancel(): %v", err) + } + } + if count != 10 { + t.Fatalf("reserved permits = %d, want 10", count) + } +} + +func TestFetchPermitRejectsDoubleFinishAndInvalidCounts(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 2, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + permit, ok, err := budget.ReserveFetch("a") + if err != nil || !ok { + t.Fatalf("ReserveFetch() = (_, %v, %v), want permit", ok, err) + } + if err := permit.Complete(1, 2); !errors.Is(err, ErrInvalidFetchCompletion) { + t.Fatalf("Complete() error = %v, want ErrInvalidFetchCompletion", err) + } + if err := permit.Cancel(); err != nil { + t.Fatalf("Cancel(): %v", err) + } + if err := permit.Cancel(); !errors.Is(err, ErrFetchPermitFinished) { + t.Fatalf("second Cancel() error = %v, want ErrFetchPermitFinished", err) + } +} diff --git a/internal/controller/pool/ownership.go b/internal/controller/pool/ownership.go new file mode 100644 index 0000000..499db51 --- /dev/null +++ b/internal/controller/pool/ownership.go @@ -0,0 +1,73 @@ +package pool + +import ( + "time" + + ownershipDomain "github.com/proxy-pool/proxy-pool/internal/domain/ownership" +) + +var ( + ErrInvalidOwnership = ownershipDomain.ErrInvalidOwnership + ErrOwnershipUnavailable = ownershipDomain.ErrOwnershipUnavailable + ErrAlreadyOwned = ownershipDomain.ErrAlreadyOwned + ErrStaleAssignment = ownershipDomain.ErrStaleAssignment + ErrNotDraining = ownershipDomain.ErrNotDraining + ErrDrainNotReady = ownershipDomain.ErrDrainNotReady +) + +type Assignment = ownershipDomain.Assignment + +type OwnershipManager struct { + repository ownershipDomain.Repository +} + +// NewOwnershipManager requires the same authoritative repository used by +// extraction, so Worker assignment and AVAILABLE -> EXTRACTED cannot race. +func NewOwnershipManager(repository ownershipDomain.Repository) (*OwnershipManager, error) { + if repository == nil { + return nil, ErrInvalidOwnership + } + return &OwnershipManager{repository: repository}, nil +} + +func (m *OwnershipManager) Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) { + if m == nil || m.repository == nil { + return Assignment{}, ErrInvalidOwnership + } + return m.repository.Assign(now, proxyID, workerID, ttl) +} + +func (m *OwnershipManager) Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) { + if m == nil || m.repository == nil { + return Assignment{}, ErrInvalidOwnership + } + return m.repository.Renew(now, proxyID, workerID, epoch, ttl) +} + +func (m *OwnershipManager) BeginDrain(proxyID, workerID string, epoch uint64) (Assignment, error) { + if m == nil || m.repository == nil { + return Assignment{}, ErrInvalidOwnership + } + return m.repository.BeginDrain(proxyID, workerID, epoch) +} + +func (m *OwnershipManager) AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error { + if m == nil || m.repository == nil { + return ErrInvalidOwnership + } + return m.repository.AcknowledgeDrain(proxyID, workerID, epoch, active, reserved) +} + +func (m *OwnershipManager) Get(proxyID string) (Assignment, bool) { + if m == nil || m.repository == nil { + return Assignment{}, false + } + return m.repository.Get(proxyID) +} + +func (m *OwnershipManager) Expire(now time.Time) []Assignment { + if m == nil || m.repository == nil { + return nil + } + return m.repository.Expire(now) +} diff --git a/internal/controller/pool/ownership_test.go b/internal/controller/pool/ownership_test.go new file mode 100644 index 0000000..d93fcbc --- /dev/null +++ b/internal/controller/pool/ownership_test.go @@ -0,0 +1,173 @@ +package pool + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + extractionDomain "github.com/proxy-pool/proxy-pool/internal/domain/extraction" +) + +func TestOwnershipManagerPreventsDualAssignment(t *testing.T) { + manager := newTestOwnershipManager(t, "proxy-1") + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + var succeeded atomic.Int64 + var wg sync.WaitGroup + + for index := range 100 { + wg.Add(1) + go func() { + defer wg.Done() + _, err := manager.Assign(now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute) + if err == nil { + succeeded.Add(1) + return + } + if !errors.Is(err, ErrAlreadyOwned) { + t.Errorf("Assign(): %v", err) + } + }() + } + wg.Wait() + + if got := succeeded.Load(); got != 1 { + t.Fatalf("successful assignments = %d, want 1", got) + } +} + +func TestOwnershipManagerRenewsOnlyCurrentAssignment(t *testing.T) { + manager := newTestOwnershipManager(t, "proxy-1") + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + assigned, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + if err != nil { + t.Fatalf("Assign(): %v", err) + } + renewed, err := manager.Renew(now.Add(30*time.Second), "proxy-1", "worker-1", assigned.Epoch, time.Minute) + if err != nil { + t.Fatalf("Renew(): %v", err) + } + if renewed.Version != assigned.Version+1 || !renewed.ExpiresAt.Equal(now.Add(90*time.Second)) { + t.Fatalf("renewed assignment = %+v", renewed) + } + if _, err := manager.Renew(now, "proxy-1", "worker-2", assigned.Epoch, time.Minute); !errors.Is(err, ErrStaleAssignment) { + t.Fatalf("Renew(stale) error = %v, want ErrStaleAssignment", err) + } + if expired := manager.Expire(now.Add(time.Minute)); len(expired) != 0 { + t.Fatalf("renewed assignment expired at old deadline: %+v", expired) + } +} + +func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) { + for iteration := range 100 { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + store := extractionDomain.NewMemoryStore([]extractionDomain.Candidate{{ + ID: "proxy-1", State: extractionDomain.Available, + }}) + manager, err := NewOwnershipManager(store) + if err != nil { + t.Fatalf("iteration %d NewOwnershipManager(): %v", iteration, err) + } + + start := make(chan struct{}) + assigned := make(chan bool, 1) + extracted := make(chan bool, 1) + go func() { + <-start + _, assignErr := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + if assignErr != nil && !errors.Is(assignErr, ErrOwnershipUnavailable) { + t.Errorf("iteration %d Assign(): %v", iteration, assignErr) + } + assigned <- assignErr == nil + }() + go func() { + <-start + result, extractErr := store.Extract(context.Background(), extractionDomain.Command{ + ClientID: "client-1", Requested: 1, Fulfillment: extractionDomain.Partial, Now: now, + }) + if extractErr != nil { + t.Errorf("iteration %d Extract(): %v", iteration, extractErr) + } + extracted <- result.Returned == 1 + }() + close(start) + + wins := 0 + if <-assigned { + wins++ + } + if <-extracted { + wins++ + } + if wins != 1 { + t.Fatalf("iteration %d successful ownership/extraction operations = %d, want 1", iteration, wins) + } + } +} + +func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) { + manager := newTestOwnershipManager(t, "proxy-1") + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + assignment, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + if err != nil { + t.Fatalf("Assign(): %v", err) + } + draining, err := manager.BeginDrain("proxy-1", "worker-1", assignment.Epoch) + if err != nil { + t.Fatalf("BeginDrain(): %v", err) + } + if !draining.Draining || draining.Version != assignment.Version+1 { + t.Fatalf("draining assignment = %+v", draining) + } + + if err := manager.AcknowledgeDrain("proxy-1", "worker-1", assignment.Epoch, 1, 0); !errors.Is(err, ErrDrainNotReady) { + t.Fatalf("AcknowledgeDrain(active) error = %v, want ErrDrainNotReady", err) + } + if err := manager.AcknowledgeDrain("proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil { + t.Fatalf("AcknowledgeDrain(zero): %v", err) + } + if _, ok := manager.Get("proxy-1"); ok { + t.Fatal("assignment still exists after drain acknowledgement") + } +} + +func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) { + manager := newTestOwnershipManager(t, "proxy-1") + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + first, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute) + if err != nil { + t.Fatalf("Assign(first): %v", err) + } + if expired := manager.Expire(now.Add(59 * time.Second)); len(expired) != 0 { + t.Fatalf("expired early: %+v", expired) + } + if expired := manager.Expire(now.Add(time.Minute)); len(expired) != 1 || expired[0].ProxyID != "proxy-1" { + t.Fatalf("Expire() = %+v, want proxy-1", expired) + } + + second, err := manager.Assign(now.Add(time.Minute), "proxy-1", "worker-2", time.Minute) + if err != nil { + t.Fatalf("Assign(second): %v", err) + } + if second.Epoch <= first.Epoch { + t.Fatalf("second epoch = %d, want greater than %d", second.Epoch, first.Epoch) + } +} + +func newTestOwnershipManager(t *testing.T, proxyIDs ...string) *OwnershipManager { + t.Helper() + candidates := make([]extractionDomain.Candidate, 0, len(proxyIDs)) + for _, proxyID := range proxyIDs { + candidates = append(candidates, extractionDomain.Candidate{ + ID: proxyID, State: extractionDomain.Available, + }) + } + manager, err := NewOwnershipManager(extractionDomain.NewMemoryStore(candidates)) + if err != nil { + t.Fatalf("NewOwnershipManager(): %v", err) + } + return manager +} diff --git a/internal/controller/pool/reconciler.go b/internal/controller/pool/reconciler.go new file mode 100644 index 0000000..3ab72a5 --- /dev/null +++ b/internal/controller/pool/reconciler.go @@ -0,0 +1,64 @@ +package pool + +import ( + "errors" + "time" + + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" +) + +var ErrInvalidReconcilePolicy = errors.New("invalid pool reconcile policy") + +type ReconcilePolicy struct { + MinimumAvailableSlots int64 + ExpectedPerFetch int + SafetyMargin time.Duration +} + +type FetchNotifier interface { + Notify() +} + +type ReconcileDecision struct { + AvailableSlots int64 + PendingExpected int + FetchedTotal int64 + FetchAllowance int + Triggered bool +} + +type Reconciler struct { + policy ReconcilePolicy + budget *FetchBudget + notifier FetchNotifier +} + +func NewReconciler(policy ReconcilePolicy, budget *FetchBudget, notifier FetchNotifier) (*Reconciler, error) { + if policy.MinimumAvailableSlots <= 0 || policy.ExpectedPerFetch <= 0 || + policy.SafetyMargin < 0 || budget == nil || notifier == nil { + return nil, ErrInvalidReconcilePolicy + } + if budget.expected != policy.ExpectedPerFetch { + return nil, ErrInvalidReconcilePolicy + } + return &Reconciler{policy: policy, budget: budget, notifier: notifier}, nil +} + +// Reconcile centralizes the cold-path decision. The notifier may coalesce many +// calls; Provider Reconciler atomically reserves the budget before doing I/O. +func (r *Reconciler) Reconcile(now time.Time, inventory upstream.Inventory) ReconcileDecision { + usage := r.budget.Snapshot() + decision := ReconcileDecision{ + AvailableSlots: inventory.AvailableSlots(now, r.policy.SafetyMargin), + PendingExpected: usage.PendingExpected, + FetchedTotal: usage.FetchedTotal, + FetchAllowance: r.budget.FetchAllowance(), + } + if decision.AvailableSlots >= r.policy.MinimumAvailableSlots || + decision.FetchAllowance < r.policy.ExpectedPerFetch { + return decision + } + r.notifier.Notify() + decision.Triggered = true + return decision +} diff --git a/internal/controller/pool/reconciler_test.go b/internal/controller/pool/reconciler_test.go new file mode 100644 index 0000000..8da9aab --- /dev/null +++ b/internal/controller/pool/reconciler_test.go @@ -0,0 +1,76 @@ +package pool + +import ( + "testing" + "time" + + proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" +) + +func TestPoolReconcilerSignalsOnlyWhenSlotsAreLowAndWholeFetchFits(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "provider-a", MaxSize: 10, MaxTotal: 20, ExpectedPerFetch: 2, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + notifier := &recordingFetchNotifier{} + reconciler, err := NewReconciler(ReconcilePolicy{ + MinimumAvailableSlots: 5, + ExpectedPerFetch: 2, + SafetyMargin: 10 * time.Second, + }, budget, notifier) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + inventory := upstream.Inventory{ + Proxies: []upstream.ProxyCapacity{{ + State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: 4, Active: 3, + }}, + MaxSize: 10, + MaxTotal: 20, + } + + decision := reconciler.Reconcile(now, inventory) + if !decision.Triggered || decision.AvailableSlots != 1 || decision.FetchAllowance != 2 { + t.Fatalf("Reconcile() = %+v, want triggered with one available slot", decision) + } + if notifier.calls != 1 { + t.Fatalf("Notify() calls = %d, want 1", notifier.calls) + } +} + +func TestPoolReconcilerUsesBudgetPendingAndCumulativeCounters(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "provider-a", MaxSize: 2, MaxTotal: 2, ExpectedPerFetch: 2, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + permit, ok, err := budget.ReserveFetch("provider-a") + if err != nil || !ok { + t.Fatalf("ReserveFetch() = (_, %v, %v), want permit", ok, err) + } + defer permit.Cancel() + notifier := &recordingFetchNotifier{} + reconciler, err := NewReconciler(ReconcilePolicy{ + MinimumAvailableSlots: 1, ExpectedPerFetch: 2, + }, budget, notifier) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + decision := reconciler.Reconcile(time.Now(), upstream.Inventory{MaxSize: 2, MaxTotal: 2}) + if decision.Triggered || decision.PendingExpected != 2 || decision.FetchAllowance != 0 { + t.Fatalf("Reconcile() = %+v, want pending fetch to suppress signal", decision) + } + if notifier.calls != 0 { + t.Fatalf("Notify() calls = %d, want 0", notifier.calls) + } +} + +type recordingFetchNotifier struct{ calls int } + +func (n *recordingFetchNotifier) Notify() { n.calls++ } diff --git a/internal/controller/provider/ports.go b/internal/controller/provider/ports.go new file mode 100644 index 0000000..0c88eb3 --- /dev/null +++ b/internal/controller/provider/ports.go @@ -0,0 +1,61 @@ +package provider + +import ( + "context" + "time" + + proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" +) + +type FetchResponse struct { + // Body belongs to the caller until Parse returns and must not be reused earlier. + Body []byte + RetryAfter time.Duration +} + +// ProviderAdapter may be called concurrently up to Config.MaxInFlight times. +type ProviderAdapter interface { + Fetch(context.Context) (FetchResponse, error) +} + +// RetryableError lets an adapter stop retries for permanent provider failures. +// Plain errors are treated as transient for backward-compatible fail-safe retries. +type RetryableError interface { + error + Retryable() bool +} + +// Parser must be safe for concurrent calls and must honor context cancellation. +type Parser interface { + Parse(context.Context, []byte) ([]proxyDomain.Proxy, error) +} + +// CandidateSink owns deduplication, is concurrency-safe, and returns the new +// count. On error it must not retain candidates; a successful count must be +// between zero and len(candidates). +type CandidateSink interface { + Add(context.Context, string, []proxyDomain.Proxy) (int, error) +} + +type Result struct { + UpstreamID string + Class upstream.FetchClass + ValidCount int + NewCount int + Err error + Attempt int +} + +type ResultRecorder interface { + // Record may be called concurrently and must not retain mutable result data. + Record(Result) +} + +type Ports struct { + Adapter ProviderAdapter + Parser Parser + Candidates CandidateSink + Results ResultRecorder + Capacity upstream.FetchCapacity +} diff --git a/internal/controller/provider/reconciler.go b/internal/controller/provider/reconciler.go new file mode 100644 index 0000000..b3ee354 --- /dev/null +++ b/internal/controller/provider/reconciler.go @@ -0,0 +1,304 @@ +package provider + +import ( + "context" + "errors" + "fmt" + "math/rand" + "sync" + "time" + + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" + "github.com/proxy-pool/proxy-pool/internal/platform/coalesce" +) + +type Config struct { + UpstreamID string + RequestInterval time.Duration + Timeout time.Duration + MaxAttempts int + MaxInFlight int + Retry RetryConfig +} + +type RetryConfig struct { + Initial time.Duration + Max time.Duration + Jitter int +} + +type Reconciler struct { + config Config + ports Ports + runtime Runtime + signal *coalesce.Signal + inFlight chan struct{} + rateMu sync.Mutex + nextRequest time.Time +} + +type Clock interface { + Now() time.Time +} + +type Sleeper interface { + Sleep(context.Context, time.Duration) error +} + +type Random interface { + Float64() float64 +} + +type Runtime struct { + Clock Clock + Sleeper Sleeper + Random Random +} + +func NewReconciler(config Config, ports Ports, runtimes ...Runtime) (*Reconciler, error) { + if config.UpstreamID == "" { + return nil, fmt.Errorf("new provider reconciler: upstream ID is required") + } + if config.RequestInterval < 0 || config.Timeout <= 0 || config.MaxAttempts <= 0 || config.MaxInFlight <= 0 { + return nil, fmt.Errorf("new provider reconciler: fetch limits must be positive") + } + if config.Retry.Initial < 0 || config.Retry.Max < 0 || config.Retry.Jitter < 0 || config.Retry.Jitter > 100 { + return nil, fmt.Errorf("new provider reconciler: retry settings are invalid") + } + if (config.Retry.Initial == 0) != (config.Retry.Max == 0) { + return nil, fmt.Errorf("new provider reconciler: retry initial and max must be configured together") + } + if config.Retry.Max > 0 && config.Retry.Initial > config.Retry.Max { + return nil, fmt.Errorf("new provider reconciler: retry initial delay exceeds maximum") + } + if ports.Adapter == nil || ports.Parser == nil || ports.Candidates == nil || + ports.Results == nil || ports.Capacity == nil { + return nil, fmt.Errorf("new provider reconciler: all ports are required") + } + runtime := Runtime{Clock: systemClock{}, Sleeper: timerSleeper{}, Random: globalRandom{}} + if len(runtimes) > 1 { + return nil, fmt.Errorf("new provider reconciler: at most one runtime is allowed") + } + if len(runtimes) == 1 { + if runtimes[0].Clock != nil { + runtime.Clock = runtimes[0].Clock + } + if runtimes[0].Sleeper != nil { + runtime.Sleeper = runtimes[0].Sleeper + } + if runtimes[0].Random != nil { + runtime.Random = runtimes[0].Random + } + } + return &Reconciler{ + config: config, + ports: ports, + runtime: runtime, + signal: coalesce.NewSignal(), + inFlight: make(chan struct{}, config.MaxInFlight), + }, nil +} + +func (r *Reconciler) Notify() { + r.signal.Notify() +} + +func (r *Reconciler) Run(ctx context.Context) error { + var workers sync.WaitGroup + defer workers.Wait() + for { + if ctx.Err() != nil { + return nil + } + if err := r.signal.Wait(ctx); err != nil { + return nil + } + select { + case r.inFlight <- struct{}{}: + case <-ctx.Done(): + return nil + } + workers.Add(1) + go func() { + defer workers.Done() + defer func() { <-r.inFlight }() + r.reconcile(ctx) + }() + } +} + +func (r *Reconciler) reconcile(ctx context.Context) { + for attempt := 1; attempt <= r.config.MaxAttempts; attempt++ { + response, result, retryable, ok := r.fetchAttempt(ctx, attempt) + if !ok { + return + } + r.ports.Results.Record(result) + if result.Class != upstream.FetchError || !retryable || attempt == r.config.MaxAttempts { + return + } + if delay := r.retryDelay(attempt, response.RetryAfter); delay > 0 { + if err := r.runtime.Sleeper.Sleep(ctx, delay); err != nil { + return + } + } + if ctx.Err() != nil { + return + } + } +} + +func (r *Reconciler) fetchAttempt(ctx context.Context, attempt int) (FetchResponse, Result, bool, bool) { + if err := r.waitForRequestSlot(ctx); err != nil { + return FetchResponse{}, Result{}, false, false + } + permit, available, err := r.ports.Capacity.ReserveFetch(r.config.UpstreamID) + if err != nil { + resultErr := fmt.Errorf("reserve fetch capacity: %w", err) + return FetchResponse{}, Result{ + UpstreamID: r.config.UpstreamID, + Class: upstream.FetchError, + Err: resultErr, + Attempt: attempt, + }, false, true + } + if !available { + return FetchResponse{}, Result{}, false, false + } + if permit == nil || permit.Expected() <= 0 { + if permit != nil { + _ = permit.Cancel() + } + return FetchResponse{}, Result{ + UpstreamID: r.config.UpstreamID, + Class: upstream.FetchError, + Err: fmt.Errorf("reserve fetch capacity: invalid permit"), + Attempt: attempt, + }, false, true + } + permitFinished := false + defer func() { + if !permitFinished { + _ = permit.Cancel() + } + }() + + callCtx := ctx + cancel := func() {} + if r.config.Timeout > 0 { + callCtx, cancel = context.WithTimeout(ctx, r.config.Timeout) + } + defer cancel() + response, callErr := r.ports.Adapter.Fetch(callCtx) + + var parseErr, candidateErr, capacityErr error + var validCount, newCount int + if callErr == nil { + candidates, err := r.ports.Parser.Parse(callCtx, response.Body) + parseErr = err + validCount = len(candidates) + if parseErr == nil && validCount > 0 { + retained := candidates + if expected := permit.Expected(); expected < len(retained) { + retained = retained[:expected] + } + newCount, candidateErr = r.ports.Candidates.Add(callCtx, r.config.UpstreamID, retained) + } + } + if callErr == nil && parseErr == nil && candidateErr == nil { + capacityErr = permit.Complete(validCount, newCount) + permitFinished = capacityErr == nil + } + resultErr := errors.Join(callErr, parseErr, candidateErr, capacityErr) + class := upstream.ClassifyFetchResult(callErr, errors.Join(parseErr, candidateErr, capacityErr), validCount, newCount) + return response, Result{ + UpstreamID: r.config.UpstreamID, + Class: class, + ValidCount: validCount, + NewCount: newCount, + Err: resultErr, + Attempt: attempt, + }, isRetryable(callErr) || parseErr != nil, true +} + +func isRetryable(err error) bool { + if err == nil { + return false + } + var classified RetryableError + if errors.As(err, &classified) { + return classified.Retryable() + } + return true +} + +func (r *Reconciler) waitForRequestSlot(ctx context.Context) error { + r.rateMu.Lock() + now := r.runtime.Clock.Now() + requestAt := now + if r.nextRequest.After(requestAt) { + requestAt = r.nextRequest + } + r.nextRequest = requestAt.Add(r.config.RequestInterval) + r.rateMu.Unlock() + + if delay := requestAt.Sub(now); delay > 0 { + return r.runtime.Sleeper.Sleep(ctx, delay) + } + return nil +} + +func (r *Reconciler) retryDelay(failedAttempt int, retryAfter time.Duration) time.Duration { + if retryAfter > 0 { + if r.config.Retry.Max > 0 && retryAfter > r.config.Retry.Max { + return r.config.Retry.Max + } + return retryAfter + } + + delay := r.config.Retry.Initial + for attempt := 1; attempt < failedAttempt; attempt++ { + if r.config.Retry.Max > 0 && delay >= r.config.Retry.Max/2 { + delay = r.config.Retry.Max + break + } + delay *= 2 + } + if delay <= 0 || r.config.Retry.Jitter == 0 { + return delay + } + + random := r.runtime.Random.Float64() + if random < 0 { + random = 0 + } else if random > 1 { + random = 1 + } + spread := float64(r.config.Retry.Jitter) / 100 + delay = time.Duration(float64(delay) * (1 + (2*random-1)*spread)) + if r.config.Retry.Max > 0 && delay > r.config.Retry.Max { + return r.config.Retry.Max + } + return delay +} + +type systemClock struct{} + +func (systemClock) Now() time.Time { return time.Now() } + +type timerSleeper struct{} + +func (timerSleeper) Sleep(ctx context.Context, duration time.Duration) error { + timer := time.NewTimer(duration) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +type globalRandom struct{} + +func (globalRandom) Float64() float64 { return rand.Float64() } diff --git a/internal/controller/provider/reconciler_test.go b/internal/controller/provider/reconciler_test.go new file mode 100644 index 0000000..9314b43 --- /dev/null +++ b/internal/controller/provider/reconciler_test.go @@ -0,0 +1,840 @@ +package provider + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" + "github.com/proxy-pool/proxy-pool/internal/domain/upstream" +) + +func TestReconcilerCoalescesConcurrentNotifications(t *testing.T) { + var calls atomic.Int64 + result := make(chan Result, 1) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 1, + }, Ports{ + Adapter: adapterFunc(func(context.Context) (FetchResponse, error) { + calls.Add(1) + return FetchResponse{Body: []byte("fixture")}, nil + }), + Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) { + return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil + }), + Candidates: candidateSinkFunc(func(context.Context, string, []proxyDomain.Proxy) (int, error) { + return 1, nil + }), + Results: resultRecorderFunc(func(got Result) { result <- got }), + Capacity: unlimitedFetchCapacity{}, + }) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + var wg sync.WaitGroup + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + reconciler.Notify() + }() + } + wg.Wait() + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + + select { + case got := <-result: + if got.Class != upstream.FetchValid { + t.Fatalf("result class = %q, want %q", got.Class, upstream.FetchValid) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for reconcile result") + } + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + if got := calls.Load(); got != 1 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 1", got) + } +} + +func TestReconcilerEnforcesRequestInterval(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + calledAt := make(chan time.Time, 2) + results := make(chan Result, 2) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + RequestInterval: 250 * time.Millisecond, + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 1, + }, successfulPorts(func() { calledAt <- clock.Now() }, results), Runtime{ + Clock: clock, + Sleeper: sleeper, + }) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + + reconciler.Notify() + <-results + reconciler.Notify() + <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + + first, second := <-calledAt, <-calledAt + if got := second.Sub(first); got != 250*time.Millisecond { + t.Fatalf("Provider calls separated by %s, want 250ms", got) + } + if got := sleeper.Durations(); len(got) != 1 || got[0] != 250*time.Millisecond { + t.Fatalf("Sleep durations = %v, want [250ms]", got) + } +} + +func TestReconcilerRetriesErrorsWithExponentialBackoffAndJitter(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + results := make(chan Result, 3) + var calls atomic.Int64 + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + if calls.Add(1) < 3 { + return FetchResponse{}, errors.New("provider unavailable") + } + return FetchResponse{Body: []byte("fixture")}, nil + }) + ports.Candidates = candidateSinkFunc(func(context.Context, string, []proxyDomain.Proxy) (int, error) { + return 0, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 3, + MaxInFlight: 1, + Retry: RetryConfig{ + Initial: 100 * time.Millisecond, + Max: time.Second, + Jitter: 20, + }, + }, ports, Runtime{Clock: clock, Sleeper: sleeper, Random: fixedRandom(0.75)}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + + wantClasses := []upstream.FetchClass{ + upstream.FetchError, + upstream.FetchError, + upstream.FetchDuplicateOnly, + } + for attempt, want := range wantClasses { + select { + case got := <-results: + if got.Class != want || got.Attempt != attempt+1 { + t.Fatalf("result %d = {class:%q attempt:%d}, want {class:%q attempt:%d}", attempt, got.Class, got.Attempt, want, attempt+1) + } + case <-time.After(time.Second): + t.Fatalf("timed out waiting for result %d", attempt) + } + } + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + + if got := sleeper.Durations(); len(got) != 2 || got[0] != 110*time.Millisecond || got[1] != 220*time.Millisecond { + t.Fatalf("Sleep durations = %v, want [110ms 220ms]", got) + } +} + +func TestReconcilerClassifiesFetchResults(t *testing.T) { + tests := []struct { + name string + callErr error + parseErr error + candidates int + newCount int + want upstream.FetchClass + }{ + {name: "valid", candidates: 2, newCount: 1, want: upstream.FetchValid}, + {name: "empty", want: upstream.FetchEmpty}, + {name: "duplicate only", candidates: 2, want: upstream.FetchDuplicateOnly}, + {name: "provider error", callErr: errors.New("HTTP 500"), want: upstream.FetchError}, + {name: "parser error", parseErr: errors.New("invalid template output"), want: upstream.FetchError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + results := make(chan Result, 1) + ports := Ports{ + Adapter: adapterFunc(func(context.Context) (FetchResponse, error) { + return FetchResponse{Body: []byte("fixture")}, tt.callErr + }), + Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) { + return make([]proxyDomain.Proxy, tt.candidates), tt.parseErr + }), + Candidates: candidateSinkFunc(func(context.Context, string, []proxyDomain.Proxy) (int, error) { + return tt.newCount, nil + }), + Results: resultRecorderFunc(func(got Result) { results <- got }), + Capacity: unlimitedFetchCapacity{}, + } + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + got := runSingleReconcile(t, reconciler, results) + if got.Class != tt.want { + t.Fatalf("result class = %q, want %q", got.Class, tt.want) + } + }) + } +} + +func TestReconcilerHonorsRetryAfterBeforeBackoff(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + calledAt := make(chan time.Time, 2) + results := make(chan Result, 2) + var calls atomic.Int64 + ports := successfulPorts(func() { calledAt <- clock.Now() }, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + calledAt <- clock.Now() + if calls.Add(1) == 1 { + return FetchResponse{RetryAfter: 700 * time.Millisecond}, errors.New("rate limited") + } + return FetchResponse{Body: []byte("fixture")}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + RequestInterval: 100 * time.Millisecond, + Timeout: time.Second, + MaxAttempts: 2, + MaxInFlight: 1, + Retry: RetryConfig{Initial: 50 * time.Millisecond, Max: time.Second, Jitter: 20}, + }, ports, Runtime{Clock: clock, Sleeper: sleeper, Random: fixedRandom(1)}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + <-results + <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + + first, second := <-calledAt, <-calledAt + if got := second.Sub(first); got != 700*time.Millisecond { + t.Fatalf("Provider calls separated by %s, want Retry-After 700ms", got) + } +} + +func TestReconcilerCapsRetryAfter(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + results := make(chan Result, 2) + var calls atomic.Int64 + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + if calls.Add(1) == 1 { + return FetchResponse{RetryAfter: time.Minute}, errors.New("rate limited") + } + return FetchResponse{Body: []byte("fixture")}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 2, + MaxInFlight: 1, + Retry: RetryConfig{Initial: 50 * time.Millisecond, Max: 2 * time.Second}, + }, ports, Runtime{Clock: clock, Sleeper: sleeper}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + <-results + <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + if got := sleeper.Durations(); len(got) != 1 || got[0] != 2*time.Second { + t.Fatalf("Sleep durations = %v, want capped Retry-After [2s]", got) + } +} + +func TestReconcilerAppliesAttemptTimeoutToEveryPort(t *testing.T) { + results := make(chan Result, 1) + ports := successfulPorts(func() {}, results) + var stages atomic.Int64 + assertDeadline := func(ctx context.Context, stage string) { + deadline, ok := ctx.Deadline() + if !ok { + t.Errorf("%s context has no deadline", stage) + return + } + remaining := time.Until(deadline) + if remaining <= 0 || remaining > 250*time.Millisecond { + t.Errorf("%s deadline remaining = %s, want (0, 250ms]", stage, remaining) + } + stages.Add(1) + } + ports.Adapter = adapterFunc(func(ctx context.Context) (FetchResponse, error) { + assertDeadline(ctx, "ProviderAdapter.Fetch") + return FetchResponse{Body: []byte("fixture")}, nil + }) + ports.Parser = parserFunc(func(ctx context.Context, _ []byte) ([]proxyDomain.Proxy, error) { + assertDeadline(ctx, "Parser.Parse") + return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil + }) + ports.Candidates = candidateSinkFunc(func(ctx context.Context, _ string, _ []proxyDomain.Proxy) (int, error) { + assertDeadline(ctx, "CandidateSink.Add") + return 1, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: 250 * time.Millisecond, + MaxAttempts: 1, + MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + if got := runSingleReconcile(t, reconciler, results); got.Class != upstream.FetchValid { + t.Fatalf("result class = %q, want %q", got.Class, upstream.FetchValid) + } + if got := stages.Load(); got != 3 { + t.Fatalf("ports observing timeout = %d, want 3", got) + } +} + +func TestReconcilerDropsNotificationFanoutWhileFetchIsInFlight(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + results := make(chan Result, 2) + var calls atomic.Int64 + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + calls.Add(1) + close(started) + <-release + return FetchResponse{Body: []byte("fixture")}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + <-started + + var wg sync.WaitGroup + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + reconciler.Notify() + }() + } + wg.Wait() + cancel() + close(release) + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + if got := calls.Load(); got != 1 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 1", got) + } +} + +func TestReconcilerEnforcesMaxInFlightAcrossRunConsumers(t *testing.T) { + started := make(chan struct{}, 2) + release := make(chan struct{}, 2) + results := make(chan Result, 2) + var active atomic.Int64 + var maximum atomic.Int64 + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + current := active.Add(1) + for { + observed := maximum.Load() + if current <= observed || maximum.CompareAndSwap(observed, current) { + break + } + } + started <- struct{}{} + <-release + active.Add(-1) + return FetchResponse{Body: []byte("fixture")}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 2) + go func() { done <- reconciler.Run(ctx) }() + go func() { done <- reconciler.Run(ctx) }() + + reconciler.Notify() + <-started + reconciler.Notify() + release <- struct{}{} + <-started + release <- struct{}{} + <-results + <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("first Run(): %v", err) + } + if err := <-done; err != nil { + t.Fatalf("second Run(): %v", err) + } + if got := maximum.Load(); got != 1 { + t.Fatalf("maximum ProviderAdapter.Fetch() in flight = %d, want 1", got) + } +} + +func TestReconcilerUsesConfiguredMaxInFlight(t *testing.T) { + started := make(chan struct{}, 2) + release := make(chan struct{}, 2) + results := make(chan Result, 2) + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + started <- struct{}{} + <-release + return FetchResponse{Body: []byte("fixture")}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 1, + MaxInFlight: 2, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + <-started + reconciler.Notify() + select { + case <-started: + case <-time.After(time.Second): + release <- struct{}{} + cancel() + <-done + t.Fatal("second ProviderAdapter.Fetch did not use available in-flight slot") + } + release <- struct{}{} + release <- struct{}{} + <-results + <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } +} + +func TestReconcilerDoesNotRefetchWhenCandidateSinkFails(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + results := make(chan Result, 3) + var calls atomic.Int64 + ports := successfulPorts(func() { calls.Add(1) }, results) + ports.Candidates = candidateSinkFunc(func(context.Context, string, []proxyDomain.Proxy) (int, error) { + return 0, errors.New("candidate store unavailable") + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 3, + MaxInFlight: 1, + Retry: RetryConfig{Initial: 100 * time.Millisecond, Max: time.Second}, + }, ports, Runtime{Clock: clock, Sleeper: sleeper}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + result := runSingleReconcile(t, reconciler, results) + if result.Class != upstream.FetchError { + t.Fatalf("result class = %q, want %q", result.Class, upstream.FetchError) + } + if got := calls.Load(); got != 1 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 1", got) + } + if got := sleeper.Durations(); len(got) != 0 { + t.Fatalf("Sleep durations = %v, want no retry backoff", got) + } +} + +func TestReconcilerDoesNotRetryPermanentAdapterError(t *testing.T) { + results := make(chan Result, 1) + sleeper := &errorSleeper{} + var calls atomic.Int64 + ports := successfulPorts(func() {}, results) + ports.Adapter = adapterFunc(func(context.Context) (FetchResponse, error) { + calls.Add(1) + return FetchResponse{}, permanentFetchError("authentication rejected") + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 3, + MaxInFlight: 1, + Retry: RetryConfig{Initial: time.Second, Max: time.Second}, + }, ports, Runtime{Sleeper: sleeper}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + result := runSingleReconcile(t, reconciler, results) + if result.Class != upstream.FetchError { + t.Fatalf("result class = %q, want %q", result.Class, upstream.FetchError) + } + if got := calls.Load(); got != 1 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 1", got) + } + if got := sleeper.calls.Load(); got != 0 { + t.Fatalf("Sleeper.Sleep() calls = %d, want 0", got) + } +} + +func TestReconcilerDoesNotCallProviderWithoutFetchCapacity(t *testing.T) { + results := make(chan Result, 1) + var calls atomic.Int64 + ports := successfulPorts(func() { calls.Add(1) }, results) + ports.Capacity = fetchCapacityFunc(func(string) (upstream.FetchPermit, bool, error) { + return nil, false, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + time.Sleep(20 * time.Millisecond) + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + if got := calls.Load(); got != 0 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 0", got) + } + select { + case result := <-results: + t.Fatalf("unexpected fetch result: %+v", result) + default: + } +} + +func TestReconcilerCompletesFetchPermitWithFetchedAndRetainedCounts(t *testing.T) { + results := make(chan Result, 1) + completed := make(chan fetchCompletion, 1) + ports := successfulPorts(func() {}, results) + ports.Parser = parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) { + return []proxyDomain.Proxy{{ID: "p1"}, {ID: "p2"}, {ID: "p3"}}, nil + }) + ports.Candidates = candidateSinkFunc(func(_ context.Context, _ string, candidates []proxyDomain.Proxy) (int, error) { + if len(candidates) != 2 { + t.Errorf("CandidateSink candidates = %d, want permit limit 2", len(candidates)) + } + return 1, nil + }) + ports.Capacity = fetchCapacityFunc(func(string) (upstream.FetchPermit, bool, error) { + return &recordingFetchPermit{expected: 2, completed: completed}, true, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1, + }, ports) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + result := runSingleReconcile(t, reconciler, results) + if result.ValidCount != 3 || result.NewCount != 1 { + t.Fatalf("result = %+v, want valid=3 new=1", result) + } + if got := <-completed; got.fetched != 3 || got.retained != 1 { + t.Fatalf("fetch completion = %+v, want fetched=3 retained=1", got) + } +} + +func TestReconcilerRetriesParserErrorsAsFetchErrors(t *testing.T) { + clock := newFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + sleeper := &fakeSleeper{clock: clock} + results := make(chan Result, 2) + var calls atomic.Int64 + var parses atomic.Int64 + ports := successfulPorts(func() { calls.Add(1) }, results) + ports.Parser = parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) { + if parses.Add(1) == 1 { + return nil, errors.New("temporary parser failure") + } + return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil + }) + reconciler, err := NewReconciler(Config{ + UpstreamID: "provider-a", + Timeout: time.Second, + MaxAttempts: 2, + MaxInFlight: 1, + Retry: RetryConfig{Initial: 100 * time.Millisecond, Max: time.Second}, + }, ports, Runtime{Clock: clock, Sleeper: sleeper}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + first, second := <-results, <-results + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + if first.Class != upstream.FetchError || second.Class != upstream.FetchValid { + t.Fatalf("result classes = [%q %q], want [%q %q]", first.Class, second.Class, upstream.FetchError, upstream.FetchValid) + } + if got := calls.Load(); got != 2 { + t.Fatalf("ProviderAdapter.Fetch() calls = %d, want 2", got) + } + if got := sleeper.Durations(); len(got) != 1 || got[0] != 100*time.Millisecond { + t.Fatalf("Sleep durations = %v, want [100ms]", got) + } +} + +func TestNewReconcilerRejectsInvalidSchedulingConfiguration(t *testing.T) { + results := make(chan Result, 1) + ports := successfulPorts(func() {}, results) + tests := []struct { + name string + config Config + }{ + {name: "missing upstream", config: Config{Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1}}, + {name: "negative interval", config: Config{UpstreamID: "a", RequestInterval: -1, Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1}}, + {name: "missing timeout", config: Config{UpstreamID: "a", MaxAttempts: 1, MaxInFlight: 1}}, + {name: "missing attempts", config: Config{UpstreamID: "a", Timeout: time.Second, MaxInFlight: 1}}, + {name: "missing in flight", config: Config{UpstreamID: "a", Timeout: time.Second, MaxAttempts: 1}}, + {name: "invalid jitter", config: Config{UpstreamID: "a", Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1, Retry: RetryConfig{Jitter: 101}}}, + {name: "initial exceeds max", config: Config{UpstreamID: "a", Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1, Retry: RetryConfig{Initial: time.Second, Max: time.Millisecond}}}, + {name: "incomplete retry pair", config: Config{UpstreamID: "a", Timeout: time.Second, MaxAttempts: 1, MaxInFlight: 1, Retry: RetryConfig{Max: time.Second}}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if _, err := NewReconciler(tt.config, ports); err == nil { + t.Fatal("NewReconciler() error = nil, want invalid configuration error") + } + }) + } +} + +func runSingleReconcile(t *testing.T, reconciler *Reconciler, results <-chan Result) Result { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- reconciler.Run(ctx) }() + reconciler.Notify() + var result Result + select { + case result = <-results: + case <-time.After(time.Second): + t.Fatal("timed out waiting for reconcile result") + } + cancel() + if err := <-done; err != nil { + t.Fatalf("Run(): %v", err) + } + return result +} + +func successfulPorts(onFetch func(), results chan<- Result) Ports { + return Ports{ + Adapter: adapterFunc(func(context.Context) (FetchResponse, error) { + onFetch() + return FetchResponse{Body: []byte("fixture")}, nil + }), + Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) { + return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil + }), + Candidates: candidateSinkFunc(func(context.Context, string, []proxyDomain.Proxy) (int, error) { + return 1, nil + }), + Results: resultRecorderFunc(func(got Result) { results <- got }), + Capacity: unlimitedFetchCapacity{}, + } +} + +type fakeClock struct { + mu sync.Mutex + now time.Time +} + +func newFakeClock(now time.Time) *fakeClock { return &fakeClock{now: now} } + +func (c *fakeClock) Now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.now +} + +func (c *fakeClock) Advance(duration time.Duration) { + c.mu.Lock() + c.now = c.now.Add(duration) + c.mu.Unlock() +} + +type fakeSleeper struct { + mu sync.Mutex + clock *fakeClock + durations []time.Duration +} + +type fixedRandom float64 + +func (r fixedRandom) Float64() float64 { return float64(r) } + +type permanentFetchError string + +func (e permanentFetchError) Error() string { return string(e) } +func (permanentFetchError) Retryable() bool { return false } + +type errorSleeper struct{ calls atomic.Int64 } + +func (s *errorSleeper) Sleep(context.Context, time.Duration) error { + s.calls.Add(1) + return errors.New("unexpected sleep") +} + +func (s *fakeSleeper) Sleep(ctx context.Context, duration time.Duration) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + s.durations = append(s.durations, duration) + s.mu.Unlock() + s.clock.Advance(duration) + return nil +} + +func (s *fakeSleeper) Durations() []time.Duration { + s.mu.Lock() + defer s.mu.Unlock() + return append([]time.Duration(nil), s.durations...) +} + +type adapterFunc func(context.Context) (FetchResponse, error) + +func (f adapterFunc) Fetch(ctx context.Context) (FetchResponse, error) { return f(ctx) } + +type parserFunc func(context.Context, []byte) ([]proxyDomain.Proxy, error) + +func (f parserFunc) Parse(ctx context.Context, body []byte) ([]proxyDomain.Proxy, error) { + return f(ctx, body) +} + +type candidateSinkFunc func(context.Context, string, []proxyDomain.Proxy) (int, error) + +func (f candidateSinkFunc) Add(ctx context.Context, upstreamID string, candidates []proxyDomain.Proxy) (int, error) { + return f(ctx, upstreamID, candidates) +} + +type fetchCapacityFunc func(string) (upstream.FetchPermit, bool, error) + +func (f fetchCapacityFunc) ReserveFetch(upstreamID string) (upstream.FetchPermit, bool, error) { + return f(upstreamID) +} + +type unlimitedFetchCapacity struct{} + +func (unlimitedFetchCapacity) ReserveFetch(string) (upstream.FetchPermit, bool, error) { + return &recordingFetchPermit{expected: int(^uint(0) >> 1)}, true, nil +} + +type fetchCompletion struct { + fetched int + retained int +} + +type recordingFetchPermit struct { + expected int + completed chan<- fetchCompletion +} + +func (p *recordingFetchPermit) Expected() int { return p.expected } + +func (p *recordingFetchPermit) Complete(fetched, retained int) error { + if p.completed != nil { + p.completed <- fetchCompletion{fetched: fetched, retained: retained} + } + return nil +} + +func (*recordingFetchPermit) Cancel() error { return nil } + +type resultRecorderFunc func(Result) + +func (f resultRecorderFunc) Record(result Result) { f(result) } diff --git a/internal/domain/extraction/extraction.go b/internal/domain/extraction/extraction.go index 033790c..1aea894 100644 --- a/internal/domain/extraction/extraction.go +++ b/internal/domain/extraction/extraction.go @@ -6,6 +6,8 @@ import ( "sort" "sync" "time" + + ownershipDomain "github.com/proxy-pool/proxy-pool/internal/domain/ownership" ) type Fulfillment string @@ -22,14 +24,23 @@ const ( Extracted State = "EXTRACTED" ) -var ErrInsufficientProxies = errors.New("insufficient proxies") +var ( + ErrInsufficientProxies = errors.New("insufficient proxies") + ErrIdempotencyConflict = errors.New("idempotency key was reused with a different extraction request") + ErrInvalidCommand = errors.New("invalid extraction command") +) type Candidate struct { ID string Protocol string + Host string + Port uint16 + Username string + Password string Region string Carrier string Upstream string + OwnerWorkerID string URL string State State ExpiresAt time.Time @@ -37,6 +48,10 @@ type Candidate struct { } type Command struct { + RequestID string + ClientID string + SourceIP string + IdempotencyKey string Requested int Fulfillment Fulfillment Now time.Time @@ -60,9 +75,10 @@ type Record struct { } type Result struct { - Requested int - Returned int - Items []Candidate + Requested int + Returned int + ExtractedAt time.Time + Items []Candidate } type Store interface { @@ -72,6 +88,17 @@ type Store interface { type MemoryStore struct { mu sync.Mutex candidates map[string]Candidate + records []Record + idempotent map[string]idempotencyEntry + nextEpoch uint64 + ownership map[string]ownershipDomain.Assignment +} + +var _ ownershipDomain.Repository = (*MemoryStore)(nil) + +type idempotencyEntry struct { + command Command + result Result } func NewMemoryStore(candidates []Candidate) *MemoryStore { @@ -79,14 +106,38 @@ func NewMemoryStore(candidates []Candidate) *MemoryStore { for _, candidate := range candidates { items[candidate.ID] = candidate } - return &MemoryStore{candidates: items} + return &MemoryStore{ + candidates: items, + idempotent: make(map[string]idempotencyEntry), + ownership: make(map[string]ownershipDomain.Assignment), + } } -func (s *MemoryStore) Extract(_ context.Context, command Command) (Result, error) { +func (s *MemoryStore) Extract(ctx context.Context, command Command) (Result, error) { + result := Result{Requested: command.Requested} + if err := ctx.Err(); err != nil { + return result, err + } + if command.Requested < 0 || command.ReserveForGateway < 0 || + command.MinRemainingTTL < 0 || command.MaxHealthCheckAge < 0 || + (command.Fulfillment != Partial && command.Fulfillment != AllOrNothing) { + return result, ErrInvalidCommand + } s.mu.Lock() defer s.mu.Unlock() + if err := ctx.Err(); err != nil { + return result, err + } - result := Result{Requested: command.Requested} + idempotencyKey := command.ClientID + "\x00" + command.IdempotencyKey + if command.IdempotencyKey != "" { + if committed, ok := s.idempotent[idempotencyKey]; ok { + if !sameIdempotentRequest(committed.command, command) { + return result, ErrIdempotencyConflict + } + return cloneResult(committed.result), nil + } + } if command.Requested <= 0 { return result, nil } @@ -115,13 +166,187 @@ func (s *MemoryStore) Extract(_ context.Context, command Command) (Result, error candidate.State = Extracted s.candidates[candidate.ID] = candidate result.Items = append(result.Items, candidate) + s.records = append(s.records, Record{ + ProxyID: candidate.ID, + ClientID: command.ClientID, + SourceIP: command.SourceIP, + RequestID: command.RequestID, + Upstream: candidate.Upstream, + ExtractedAt: command.Now, + ExpiresAt: candidate.ExpiresAt, + }) } result.Returned = len(result.Items) + if result.Returned > 0 { + result.ExtractedAt = command.Now + } + if command.IdempotencyKey != "" { + s.idempotent[idempotencyKey] = idempotencyEntry{ + command: cloneCommand(command), + result: cloneResult(result), + } + } return result, nil } +func (s *MemoryStore) Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (ownershipDomain.Assignment, error) { + if proxyID == "" || workerID == "" || ttl <= 0 { + return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership + } + s.mu.Lock() + defer s.mu.Unlock() + candidate, ok := s.candidates[proxyID] + if current, exists := s.ownership[proxyID]; exists && current.ExpiresAt.After(now) { + return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned + } else if exists { + if candidate.OwnerWorkerID == current.WorkerID { + candidate.OwnerWorkerID = "" + } + delete(s.ownership, proxyID) + } + if !ok || candidate.State != Available || candidate.OwnerWorkerID != "" { + return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable + } + s.nextEpoch++ + assignment := ownershipDomain.Assignment{ + ProxyID: proxyID, WorkerID: workerID, Epoch: s.nextEpoch, Version: 1, + ExpiresAt: now.UTC().Add(ttl), + } + candidate.OwnerWorkerID = workerID + s.candidates[proxyID] = candidate + s.ownership[proxyID] = assignment + return assignment, nil +} + +func (s *MemoryStore) Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (ownershipDomain.Assignment, error) { + if ttl <= 0 { + return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership + } + s.mu.Lock() + defer s.mu.Unlock() + assignment, ok := s.ownership[proxyID] + if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch || !assignment.ExpiresAt.After(now) { + return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment + } + assignment.ExpiresAt = now.UTC().Add(ttl) + assignment.Version++ + s.ownership[proxyID] = assignment + return assignment, nil +} + +func (s *MemoryStore) BeginDrain(proxyID, workerID string, epoch uint64) (ownershipDomain.Assignment, error) { + s.mu.Lock() + defer s.mu.Unlock() + assignment, ok := s.ownership[proxyID] + if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch { + return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment + } + if assignment.Draining { + return assignment, nil + } + assignment.Draining = true + assignment.Version++ + s.ownership[proxyID] = assignment + return assignment, nil +} + +func (s *MemoryStore) AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error { + if active < 0 || reserved < 0 { + return ownershipDomain.ErrInvalidOwnership + } + s.mu.Lock() + defer s.mu.Unlock() + assignment, ok := s.ownership[proxyID] + if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch { + return ownershipDomain.ErrStaleAssignment + } + if !assignment.Draining { + return ownershipDomain.ErrNotDraining + } + if active > 0 || reserved > 0 { + return ownershipDomain.ErrDrainNotReady + } + if candidate, exists := s.candidates[proxyID]; exists && candidate.OwnerWorkerID == workerID { + candidate.OwnerWorkerID = "" + s.candidates[proxyID] = candidate + } + delete(s.ownership, proxyID) + return nil +} + +func (s *MemoryStore) Get(proxyID string) (ownershipDomain.Assignment, bool) { + s.mu.Lock() + defer s.mu.Unlock() + assignment, ok := s.ownership[proxyID] + return assignment, ok +} + +func (s *MemoryStore) Expire(now time.Time) []ownershipDomain.Assignment { + s.mu.Lock() + defer s.mu.Unlock() + expired := make([]ownershipDomain.Assignment, 0) + for proxyID, assignment := range s.ownership { + if assignment.ExpiresAt.After(now) { + continue + } + if candidate, ok := s.candidates[proxyID]; ok && candidate.OwnerWorkerID == assignment.WorkerID { + candidate.OwnerWorkerID = "" + s.candidates[proxyID] = candidate + } + expired = append(expired, assignment) + delete(s.ownership, proxyID) + } + sort.Slice(expired, func(i, j int) bool { return expired[i].ProxyID < expired[j].ProxyID }) + return expired +} + +func (s *MemoryStore) Records() []Record { + s.mu.Lock() + defer s.mu.Unlock() + return append([]Record(nil), s.records...) +} + +func cloneResult(result Result) Result { + result.Items = append([]Candidate(nil), result.Items...) + return result +} + +func cloneCommand(command Command) Command { + command.Protocols = append([]string(nil), command.Protocols...) + command.Regions = append([]string(nil), command.Regions...) + command.Carriers = append([]string(nil), command.Carriers...) + command.Upstreams = append([]string(nil), command.Upstreams...) + return command +} + +func sameIdempotentRequest(left, right Command) bool { + return left.Requested == right.Requested && + left.Fulfillment == right.Fulfillment && + equalSet(left.Protocols, right.Protocols) && + equalSet(left.Regions, right.Regions) && + equalSet(left.Carriers, right.Carriers) && + equalSet(left.Upstreams, right.Upstreams) +} + +func equalSet(left, right []string) bool { + if len(left) != len(right) { + return false + } + counts := make(map[string]int, len(left)) + for _, value := range left { + counts[value]++ + } + for _, value := range right { + counts[value]-- + if counts[value] < 0 { + return false + } + } + return true +} + func eligibleForExtraction(candidate Candidate, command Command) bool { - if candidate.State != Available { + if candidate.State != Available || candidate.OwnerWorkerID != "" { return false } if !candidate.ExpiresAt.IsZero() && candidate.ExpiresAt.Sub(command.Now) < command.MinRemainingTTL { diff --git a/internal/domain/extraction/extraction_test.go b/internal/domain/extraction/extraction_test.go index 8791d3d..27cfde1 100644 --- a/internal/domain/extraction/extraction_test.go +++ b/internal/domain/extraction/extraction_test.go @@ -2,6 +2,7 @@ package extraction import ( "context" + "errors" "sync" "testing" "time" @@ -81,3 +82,165 @@ func TestAllOrNothingDoesNotConsumePartialInventory(t *testing.T) { t.Fatalf("inventory was consumed by failed all-or-nothing: result=%+v err=%v", partial, err) } } + +func TestMemoryStoreCommitsAuditWithExtraction(t *testing.T) { + now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC) + store := NewMemoryStore([]Candidate{{ + ID: "p1", + Upstream: "provider-a", + State: Available, + ExpiresAt: now.Add(time.Minute), + LastCheckedAt: now, + }}) + + _, err := store.Extract(context.Background(), Command{ + RequestID: "req-1", + ClientID: "client-1", + SourceIP: "192.0.2.30", + Requested: 1, + Fulfillment: Partial, + Now: now, + MinRemainingTTL: 30 * time.Second, + MaxHealthCheckAge: 10 * time.Second, + }) + if err != nil { + t.Fatalf("Extract(): %v", err) + } + + records := store.Records() + if len(records) != 1 { + t.Fatalf("audit records = %d, want 1", len(records)) + } + record := records[0] + if record.ProxyID != "p1" || record.ClientID != "client-1" || record.RequestID != "req-1" { + t.Fatalf("audit record = %+v", record) + } + if !record.ExtractedAt.Equal(now) || !record.ExpiresAt.Equal(now.Add(time.Minute)) { + t.Fatalf("audit timestamps = %+v", record) + } +} + +func TestMemoryStoreReplaysCommittedIdempotentResult(t *testing.T) { + now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC) + store := NewMemoryStore([]Candidate{{ + ID: "p1", + State: Available, + ExpiresAt: now.Add(time.Minute), + LastCheckedAt: now, + }}) + command := Command{ + RequestID: "req-1", + ClientID: "client-1", + IdempotencyKey: "idem-12345678", + Requested: 1, + Fulfillment: Partial, + Now: now, + MinRemainingTTL: 30 * time.Second, + MaxHealthCheckAge: 10 * time.Second, + } + + first, err := store.Extract(context.Background(), command) + if err != nil { + t.Fatalf("first Extract(): %v", err) + } + command.RequestID = "req-2" + second, err := store.Extract(context.Background(), command) + if err != nil { + t.Fatalf("second Extract(): %v", err) + } + + if len(first.Items) != 1 || len(second.Items) != 1 || second.Items[0].ID != first.Items[0].ID { + t.Fatalf("idempotent results: first=%+v second=%+v", first, second) + } + if got := len(store.Records()); got != 1 { + t.Fatalf("audit records = %d, want 1", got) + } +} + +func TestMemoryStoreRejectsIdempotencyKeyReuseWithDifferentRequest(t *testing.T) { + now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC) + store := NewMemoryStore([]Candidate{{ + ID: "p1", + State: Available, + ExpiresAt: now.Add(time.Minute), + LastCheckedAt: now, + }}) + command := Command{ + RequestID: "req-1", + ClientID: "client-1", + IdempotencyKey: "idem-12345678", + Requested: 1, + Fulfillment: Partial, + Now: now, + MinRemainingTTL: 30 * time.Second, + MaxHealthCheckAge: 10 * time.Second, + } + if _, err := store.Extract(context.Background(), command); err != nil { + t.Fatalf("first Extract(): %v", err) + } + + command.Requested = 2 + if _, err := store.Extract(context.Background(), command); err != ErrIdempotencyConflict { + t.Fatalf("second Extract() error = %v, want ErrIdempotencyConflict", err) + } +} + +func TestMemoryStoreExtractsOnlyUnownedProxy(t *testing.T) { + now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC) + store := NewMemoryStore([]Candidate{ + {ID: "owned", OwnerWorkerID: "worker-1", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now}, + {ID: "unowned", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now}, + }) + + result, err := store.Extract(context.Background(), Command{ + RequestID: "req-1", + ClientID: "client-1", + Requested: 2, + Fulfillment: Partial, + Now: now, + MinRemainingTTL: 30 * time.Second, + MaxHealthCheckAge: 10 * time.Second, + }) + if err != nil { + t.Fatalf("Extract(): %v", err) + } + if len(result.Items) != 1 || result.Items[0].ID != "unowned" { + t.Fatalf("extracted items = %+v, want only unowned", result.Items) + } +} + +func TestMemoryStoreValidatesCommandAndHonorsCancellation(t *testing.T) { + store := NewMemoryStore([]Candidate{{ID: "p1", State: Available}}) + if _, err := store.Extract(context.Background(), Command{ + Requested: 1, Fulfillment: Partial, ReserveForGateway: -1, + }); !errors.Is(err, ErrInvalidCommand) { + t.Fatalf("Extract(negative reserve) error = %v, want ErrInvalidCommand", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := store.Extract(ctx, Command{Requested: 1, Fulfillment: Partial}); !errors.Is(err, context.Canceled) { + t.Fatalf("Extract(canceled) error = %v, want context.Canceled", err) + } + result, err := store.Extract(context.Background(), Command{Requested: 1, Fulfillment: Partial}) + if err != nil || result.Returned != 1 { + t.Fatalf("candidate changed after canceled request: result=%+v err=%v", result, err) + } +} + +func TestMemoryStoreClonesIdempotencyCommandFilters(t *testing.T) { + store := NewMemoryStore([]Candidate{{ID: "p1", Protocol: "http", State: Available}}) + protocols := []string{"http"} + command := Command{ + ClientID: "client-1", IdempotencyKey: "idem-1", Requested: 1, + Fulfillment: Partial, Protocols: protocols, + } + if _, err := store.Extract(context.Background(), command); err != nil { + t.Fatalf("first Extract(): %v", err) + } + protocols[0] = "socks5" + command.Protocols = []string{"http"} + if _, err := store.Extract(context.Background(), command); err != nil { + t.Fatalf("idempotent replay after caller mutation: %v", err) + } +} diff --git a/internal/domain/ownership/ownership.go b/internal/domain/ownership/ownership.go new file mode 100644 index 0000000..f9f2342 --- /dev/null +++ b/internal/domain/ownership/ownership.go @@ -0,0 +1,35 @@ +package ownership + +import ( + "errors" + "time" +) + +var ( + ErrInvalidOwnership = errors.New("invalid proxy ownership request") + ErrOwnershipUnavailable = errors.New("proxy is unavailable for ownership") + ErrAlreadyOwned = errors.New("proxy is already owned") + ErrStaleAssignment = errors.New("proxy ownership assignment is stale") + ErrNotDraining = errors.New("proxy ownership is not draining") + ErrDrainNotReady = errors.New("proxy still has active or reserved runtime") +) + +type Assignment struct { + ProxyID string + WorkerID string + Epoch uint64 + Version uint64 + ExpiresAt time.Time + Draining bool +} + +// Repository is the shared authority for ownership changes. Implementations +// that also support extraction must serialize both operations transactionally. +type Repository interface { + Assign(now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) + Renew(now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) + BeginDrain(proxyID, workerID string, epoch uint64) (Assignment, error) + AcknowledgeDrain(proxyID, workerID string, epoch uint64, active, reserved int64) error + Get(proxyID string) (Assignment, bool) + Expire(now time.Time) []Assignment +} diff --git a/internal/domain/routing/routing_test.go b/internal/domain/routing/routing_test.go index 9459757..be82142 100644 --- a/internal/domain/routing/routing_test.go +++ b/internal/domain/routing/routing_test.go @@ -20,6 +20,41 @@ func TestRuleSetUsesFirstMatchingRule(t *testing.T) { } } +func TestRuleSetIsDetachedFromInputAndReturnedRules(t *testing.T) { + input := []Rule{{ + Name: "route-a", + Match: Match{ + HostRegex: `.*`, + Methods: []string{"GET"}, + Headers: map[string]string{"X-Tenant": "a"}, + }, + Upstreams: []string{"provider-a"}, + }} + rules, err := Compile(input) + if err != nil { + t.Fatalf("Compile(): %v", err) + } + input[0].Match.Methods[0] = "POST" + input[0].Match.Headers["X-Tenant"] = "changed" + input[0].Upstreams[0] = "changed" + + matched, ok := rules.Match(Request{ + Host: "example.com", Method: "GET", Headers: map[string]string{"X-Tenant": "a"}, + }) + if !ok || matched.Upstreams[0] != "provider-a" { + t.Fatalf("Match() after input mutation = %+v, %v", matched, ok) + } + matched.Match.Methods[0] = "DELETE" + matched.Match.Headers["X-Tenant"] = "returned-change" + matched.Upstreams[0] = "returned-change" + second, ok := rules.Match(Request{ + Host: "example.com", Method: "GET", Headers: map[string]string{"X-Tenant": "a"}, + }) + if !ok || second.Upstreams[0] != "provider-a" || second.Match.Methods[0] != "GET" { + t.Fatalf("Match() after returned-rule mutation = %+v, %v", second, ok) + } +} + func TestSequentialSwitchesOnceAtThreshold(t *testing.T) { sequence, err := NewSequential([]string{"a", "b", "c"}, 5) if err != nil { @@ -60,3 +95,59 @@ func TestSequentialValidFetchResetsEmptyCount(t *testing.T) { t.Fatalf("Current() = %q, want a after reset", got) } } + +func TestSequentialSharesUpstreamEmptyStateAcrossRoutingCursors(t *testing.T) { + empty := NewUpstreamEmptyState() + first, err := NewSequentialWithState([]string{"a", "b"}, 5, EndStayLast, empty) + if err != nil { + t.Fatalf("NewSequentialWithState(first): %v", err) + } + second, err := NewSequentialWithState([]string{"a", "c"}, 5, EndStayLast, empty) + if err != nil { + t.Fatalf("NewSequentialWithState(second): %v", err) + } + + for range 5 { + first.ObserveEmpty("a") + } + if got := first.Current(); got != "b" { + t.Fatalf("first.Current() = %q, want b", got) + } + if got := second.EmptyCount("a"); got != 5 { + t.Fatalf("second.EmptyCount(a) = %d, want shared count 5", got) + } + if !second.ObserveEmpty("a") || second.Current() != "c" { + t.Fatalf("second did not advance from shared empty state: current=%q", second.Current()) + } +} + +func TestSequentialStopEndBehaviorHasNoCurrentSelection(t *testing.T) { + sequence, err := NewSequentialWithState([]string{"a"}, 1, EndStop, NewUpstreamEmptyState()) + if err != nil { + t.Fatalf("NewSequentialWithState(): %v", err) + } + if !sequence.ObserveEmpty("a") { + t.Fatal("ObserveEmpty() = false, want transition to stopped") + } + if current, ok, version := sequence.CurrentSelection(); ok || current != "" || version != 2 { + t.Fatalf("CurrentSelection() = %q, %v, %d; want stopped version 2", current, ok, version) + } +} + +func TestSequentialLoopDoesNotReuseSameEmptyEpisode(t *testing.T) { + empty := NewUpstreamEmptyState() + sequence, err := NewSequentialWithState([]string{"a", "b"}, 1, EndLoop, empty) + if err != nil { + t.Fatalf("NewSequentialWithState(): %v", err) + } + if !sequence.ObserveEmpty("a") || !sequence.ObserveEmpty("b") || sequence.Current() != "a" { + t.Fatalf("sequence did not loop to a: current=%q", sequence.Current()) + } + if sequence.ObserveEmpty("a") || sequence.Current() != "a" { + t.Fatal("sequence reused the same a empty episode") + } + sequence.ObserveValid("a") + if !sequence.ObserveEmpty("a") || sequence.Current() != "b" { + t.Fatal("sequence did not advance after a new empty episode") + } +} diff --git a/internal/domain/routing/rule.go b/internal/domain/routing/rule.go index 5561739..fdfc290 100644 --- a/internal/domain/routing/rule.go +++ b/internal/domain/routing/rule.go @@ -63,10 +63,11 @@ func Compile(rules []Rule) (*RuleSet, error) { return nil, fmt.Errorf("compile routing %q path: %w", rule.Name, err) } } - if rule.Action == "" && len(rule.Upstreams) > 0 { - rule.Action = ActionProxy + owned := cloneRule(rule) + if owned.Action == "" && len(owned.Upstreams) > 0 { + owned.Action = ActionProxy } - compiled = append(compiled, compiledRule{rule: rule, host: host, path: path}) + compiled = append(compiled, compiledRule{rule: owned, host: host, path: path}) } return &RuleSet{rules: compiled}, nil } @@ -90,11 +91,24 @@ func (r *RuleSet) Match(request Request) (Rule, bool) { if !headersMatch(candidate.rule.Match.Headers, request.Headers) { continue } - return candidate.rule, true + return cloneRule(candidate.rule), true } return Rule{}, false } +func cloneRule(source Rule) Rule { + cloned := source + cloned.Match.Methods = append([]string(nil), source.Match.Methods...) + cloned.Upstreams = append([]string(nil), source.Upstreams...) + if source.Match.Headers != nil { + cloned.Match.Headers = make(map[string]string, len(source.Match.Headers)) + for name, value := range source.Match.Headers { + cloned.Match.Headers[name] = value + } + } + return cloned +} + func containsFold(values []string, target string) bool { return slices.ContainsFunc(values, func(value string) bool { return strings.EqualFold(value, target) diff --git a/internal/domain/routing/sequential.go b/internal/domain/routing/sequential.go index 64ed648..28afe9c 100644 --- a/internal/domain/routing/sequential.go +++ b/internal/domain/routing/sequential.go @@ -5,57 +5,203 @@ import ( "sync" ) -type Sequential struct { +type EndBehavior string + +const ( + EndStayLast EndBehavior = "stayLast" + EndStop EndBehavior = "stop" + EndLoop EndBehavior = "loop" +) + +type EmptyObservation struct { + Count int + Generation uint64 +} + +type upstreamEmpty struct { + count int + generation uint64 +} + +// UpstreamEmptyState is shared by every Routing that references an Upstream. +// A generation represents one uninterrupted empty-result episode. +type UpstreamEmptyState struct { + mu sync.RWMutex + nextGeneration uint64 + states map[string]upstreamEmpty +} + +func NewUpstreamEmptyState() *UpstreamEmptyState { + return &UpstreamEmptyState{states: make(map[string]upstreamEmpty)} +} + +func (s *UpstreamEmptyState) ObserveEmpty(upstream string) EmptyObservation { + s.mu.Lock() + defer s.mu.Unlock() + state := s.states[upstream] + if state.count == 0 { + s.nextGeneration++ + state.generation = s.nextGeneration + } + state.count++ + s.states[upstream] = state + return EmptyObservation{Count: state.count, Generation: state.generation} +} + +func (s *UpstreamEmptyState) ObserveValid(upstream string) { + s.mu.Lock() + defer s.mu.Unlock() + state := s.states[upstream] + state.count = 0 + s.states[upstream] = state +} + +func (s *UpstreamEmptyState) Count(upstream string) int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.states[upstream].count +} + +type CursorSnapshot struct { + Index int + Version uint64 + Stopped bool +} + +// RoutingCursor owns only per-Routing selection state. Advance uses an +// expected version so simultaneous threshold observers can change it once. +type RoutingCursor struct { mu sync.RWMutex + state CursorSnapshot + processed map[string]uint64 +} + +func NewRoutingCursor() *RoutingCursor { + return &RoutingCursor{ + state: CursorSnapshot{Version: 1}, + processed: make(map[string]uint64), + } +} + +func (c *RoutingCursor) Snapshot() CursorSnapshot { + c.mu.RLock() + defer c.mu.RUnlock() + return c.state +} + +func (c *RoutingCursor) Advance( + expectedVersion uint64, + currentUpstream string, + emptyGeneration uint64, + upstreamCount int, + end EndBehavior, +) bool { + c.mu.Lock() + defer c.mu.Unlock() + if c.state.Stopped || c.state.Version != expectedVersion || emptyGeneration == 0 || + c.processed[currentUpstream] >= emptyGeneration { + return false + } + c.processed[currentUpstream] = emptyGeneration + if c.state.Index+1 < upstreamCount { + c.state.Index++ + c.state.Version++ + return true + } + switch end { + case EndStop: + c.state.Stopped = true + c.state.Version++ + return true + case EndLoop: + if upstreamCount > 1 { + c.state.Index = 0 + c.state.Version++ + return true + } + } + return false +} + +type Sequential struct { upstreams []string threshold int - current int - empty map[string]int + end EndBehavior + empty *UpstreamEmptyState + cursor *RoutingCursor } func NewSequential(upstreams []string, threshold int) (*Sequential, error) { + return NewSequentialWithState(upstreams, threshold, EndStayLast, NewUpstreamEmptyState()) +} + +func NewSequentialWithState( + upstreams []string, + threshold int, + end EndBehavior, + empty *UpstreamEmptyState, +) (*Sequential, error) { if len(upstreams) == 0 { return nil, fmt.Errorf("sequential strategy requires at least one upstream") } if threshold <= 0 { return nil, fmt.Errorf("sequential threshold must be greater than zero") } + if end == "" { + end = EndStayLast + } + if end != EndStayLast && end != EndStop && end != EndLoop { + return nil, fmt.Errorf("sequential end behavior %q is invalid", end) + } + if empty == nil { + return nil, fmt.Errorf("sequential empty state is required") + } + for _, upstream := range upstreams { + if upstream == "" { + return nil, fmt.Errorf("sequential upstream is required") + } + } return &Sequential{ upstreams: append([]string(nil), upstreams...), threshold: threshold, - empty: make(map[string]int, len(upstreams)), + end: end, + empty: empty, + cursor: NewRoutingCursor(), }, nil } func (s *Sequential) Current() string { - s.mu.RLock() - defer s.mu.RUnlock() - return s.upstreams[s.current] + current, ok, _ := s.CurrentSelection() + if !ok { + return "" + } + return current +} + +func (s *Sequential) CurrentSelection() (upstream string, available bool, version uint64) { + state := s.cursor.Snapshot() + if state.Stopped || state.Index < 0 || state.Index >= len(s.upstreams) { + return "", false, state.Version + } + return s.upstreams[state.Index], true, state.Version } func (s *Sequential) ObserveEmpty(upstream string) bool { - s.mu.Lock() - defer s.mu.Unlock() - - s.empty[upstream]++ - if s.upstreams[s.current] != upstream || s.empty[upstream] < s.threshold { + observation := s.empty.ObserveEmpty(upstream) + if observation.Count < s.threshold { return false } - if s.current+1 >= len(s.upstreams) { + current, available, version := s.CurrentSelection() + if !available || current != upstream { return false } - s.current++ - return true + return s.cursor.Advance(version, upstream, observation.Generation, len(s.upstreams), s.end) } func (s *Sequential) ObserveValid(upstream string) { - s.mu.Lock() - defer s.mu.Unlock() - s.empty[upstream] = 0 + s.empty.ObserveValid(upstream) } func (s *Sequential) EmptyCount(upstream string) int { - s.mu.RLock() - defer s.mu.RUnlock() - return s.empty[upstream] + return s.empty.Count(upstream) } diff --git a/internal/domain/routing/strategy.go b/internal/domain/routing/strategy.go new file mode 100644 index 0000000..a2cd2ac --- /dev/null +++ b/internal/domain/routing/strategy.go @@ -0,0 +1,168 @@ +package routing + +import ( + "errors" + "math/rand" + "sync" +) + +var ErrNoCandidate = errors.New("no eligible routing candidate") + +type Candidate struct { + Name string + Weight int + Active int64 + Eligible bool +} + +type Selector interface { + Select([]Candidate) (Candidate, error) +} + +type RandomSource interface { + Intn(int) int +} + +type randomSelector struct { + source *synchronizedRandomSource +} + +type globalRandomSource struct{} + +func (globalRandomSource) Intn(n int) int { + return rand.Intn(n) +} + +type synchronizedRandomSource struct { + mu sync.Mutex + source RandomSource +} + +func (s *synchronizedRandomSource) Intn(n int) int { + s.mu.Lock() + defer s.mu.Unlock() + return s.source.Intn(n) +} + +func NewRandom(source ...RandomSource) Selector { + return &randomSelector{source: newSynchronizedRandomSource(source)} +} + +func (s *randomSelector) Select(candidates []Candidate) (Candidate, error) { + eligibleCount := 0 + for _, candidate := range candidates { + if candidate.Eligible { + eligibleCount++ + } + } + if eligibleCount == 0 { + return Candidate{}, ErrNoCandidate + } + + selected := s.source.Intn(eligibleCount) + for _, candidate := range candidates { + if !candidate.Eligible { + continue + } + if selected == 0 { + return candidate, nil + } + selected-- + } + return Candidate{}, ErrNoCandidate +} + +type roundRobinSelector struct { + mu sync.Mutex + next int +} + +func NewRoundRobin() Selector { + return &roundRobinSelector{} +} + +func (s *roundRobinSelector) Select(candidates []Candidate) (Candidate, error) { + if len(candidates) == 0 { + return Candidate{}, ErrNoCandidate + } + + s.mu.Lock() + defer s.mu.Unlock() + start := s.next % len(candidates) + for offset := range len(candidates) { + selected := (start + offset) % len(candidates) + if !candidates[selected].Eligible { + continue + } + s.next = (selected + 1) % len(candidates) + return candidates[selected], nil + } + return Candidate{}, ErrNoCandidate +} + +type weightedSelector struct { + source *synchronizedRandomSource +} + +func NewWeighted(source ...RandomSource) Selector { + return &weightedSelector{source: newSynchronizedRandomSource(source)} +} + +func (s *weightedSelector) Select(candidates []Candidate) (Candidate, error) { + totalWeight := 0 + for _, candidate := range candidates { + if candidate.Eligible && candidate.Weight > 0 { + if candidate.Weight > int(^uint(0)>>1)-totalWeight { + return Candidate{}, ErrNoCandidate + } + totalWeight += candidate.Weight + } + } + if totalWeight == 0 { + return Candidate{}, ErrNoCandidate + } + + target := s.source.Intn(totalWeight) + for _, candidate := range candidates { + if !candidate.Eligible || candidate.Weight <= 0 { + continue + } + if target < candidate.Weight { + return candidate, nil + } + target -= candidate.Weight + } + return Candidate{}, ErrNoCandidate +} + +type leastConnectionsSelector struct{} + +func NewLeastConnections() Selector { + return leastConnectionsSelector{} +} + +func (leastConnectionsSelector) Select(candidates []Candidate) (Candidate, error) { + var selected Candidate + found := false + for _, candidate := range candidates { + if !candidate.Eligible { + continue + } + if !found || candidate.Active < selected.Active { + selected = candidate + found = true + } + } + if !found { + return Candidate{}, ErrNoCandidate + } + return selected, nil +} + +func newSynchronizedRandomSource(sources []RandomSource) *synchronizedRandomSource { + var source RandomSource = globalRandomSource{} + if len(sources) > 0 && sources[0] != nil { + source = sources[0] + } + return &synchronizedRandomSource{source: source} +} diff --git a/internal/domain/routing/strategy_test.go b/internal/domain/routing/strategy_test.go new file mode 100644 index 0000000..5ff4982 --- /dev/null +++ b/internal/domain/routing/strategy_test.go @@ -0,0 +1,232 @@ +package routing + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +type fixedRandomSource struct { + values []int + next int +} + +func (s *fixedRandomSource) Intn(n int) int { + value := s.values[s.next] + s.next++ + return value % n +} + +func TestSelectorsReturnStableErrorWhenNoCandidateIsEligible(t *testing.T) { + selectors := map[string]Selector{ + "random": NewRandom(), + "round robin": NewRoundRobin(), + "weighted": NewWeighted(), + "least connections": NewLeastConnections(), + } + for name, selector := range selectors { + t.Run(name, func(t *testing.T) { + _, err := selector.Select([]Candidate{{Name: "disabled"}}) + if !errors.Is(err, ErrNoCandidate) { + t.Fatalf("Select() error = %v, want ErrNoCandidate", err) + } + }) + } +} + +func TestRandomSelectsOnlyFromEligibleCandidates(t *testing.T) { + selector := NewRandom(&fixedRandomSource{values: []int{1}}) + candidates := []Candidate{ + {Name: "disabled", Eligible: false}, + {Name: "a", Eligible: true}, + {Name: "b", Eligible: true}, + } + + got, err := selector.Select(candidates) + if err != nil { + t.Fatalf("Select(): %v", err) + } + if got.Name != "b" { + t.Fatalf("Select() = %q, want b", got.Name) + } +} + +func TestRoundRobinCyclesThroughEligibleCandidates(t *testing.T) { + selector := NewRoundRobin() + candidates := []Candidate{ + {Name: "disabled", Eligible: false}, + {Name: "a", Eligible: true}, + {Name: "b", Eligible: true}, + } + + for call, want := range []string{"a", "b", "a"} { + got, err := selector.Select(candidates) + if err != nil { + t.Fatalf("Select() call %d: %v", call+1, err) + } + if got.Name != want { + t.Fatalf("Select() call %d = %q, want %q", call+1, got.Name, want) + } + } +} + +func TestRoundRobinKeepsInputOrderWhenEligibilityChanges(t *testing.T) { + selector := NewRoundRobin() + candidates := []Candidate{ + {Name: "a", Eligible: true}, + {Name: "b", Eligible: true}, + {Name: "c", Eligible: true}, + } + + first, err := selector.Select(candidates) + if err != nil { + t.Fatalf("first Select(): %v", err) + } + if first.Name != "a" { + t.Fatalf("first Select() = %q, want a", first.Name) + } + + candidates[0].Eligible = false + second, err := selector.Select(candidates) + if err != nil { + t.Fatalf("second Select(): %v", err) + } + if second.Name != "b" { + t.Fatalf("second Select() = %q, want b", second.Name) + } +} + +func TestWeightedSelectsByEligibleCandidateWeight(t *testing.T) { + selector := NewWeighted(&fixedRandomSource{values: []int{0, 1, 2, 4}}) + candidates := []Candidate{ + {Name: "disabled", Weight: 100, Eligible: false}, + {Name: "a", Weight: 2, Eligible: true}, + {Name: "b", Weight: 3, Eligible: true}, + } + + for call, want := range []string{"a", "a", "b", "b"} { + got, err := selector.Select(candidates) + if err != nil { + t.Fatalf("Select() call %d: %v", call+1, err) + } + if got.Name != want { + t.Fatalf("Select() call %d = %q, want %q", call+1, got.Name, want) + } + } +} + +func TestWeightedReturnsStableErrorWhenWeightSumOverflows(t *testing.T) { + selector := NewWeighted(panicRandomSource{}) + maxInt := int(^uint(0) >> 1) + + defer func() { + if recovered := recover(); recovered != nil { + t.Fatalf("Select() panicked: %v", recovered) + } + }() + _, err := selector.Select([]Candidate{ + {Name: "a", Weight: maxInt, Eligible: true}, + {Name: "b", Weight: 1, Eligible: true}, + }) + if !errors.Is(err, ErrNoCandidate) { + t.Fatalf("Select() error = %v, want ErrNoCandidate", err) + } +} + +type panicRandomSource struct{} + +func (panicRandomSource) Intn(int) int { + panic("random source should not be called") +} + +func TestLeastConnectionsChoosesFirstEligibleMinimum(t *testing.T) { + selector := NewLeastConnections() + candidates := []Candidate{ + {Name: "disabled", Active: 0, Eligible: false}, + {Name: "busy", Active: 8, Eligible: true}, + {Name: "first-idle", Active: 2, Eligible: true}, + {Name: "second-idle", Active: 2, Eligible: true}, + } + + got, err := selector.Select(candidates) + if err != nil { + t.Fatalf("Select(): %v", err) + } + if got.Name != "first-idle" { + t.Fatalf("Select() = %q, want first-idle", got.Name) + } +} + +func TestRoundRobinIsSafeForConcurrentCalls(t *testing.T) { + selector := NewRoundRobin() + candidates := []Candidate{ + {Name: "a", Eligible: true}, + {Name: "b", Eligible: true}, + } + + results := make(chan string, 1000) + var wg sync.WaitGroup + for range 1000 { + wg.Add(1) + go func() { + defer wg.Done() + candidate, err := selector.Select(candidates) + if err != nil { + t.Errorf("Select(): %v", err) + return + } + results <- candidate.Name + }() + } + wg.Wait() + close(results) + + counts := map[string]int{} + for name := range results { + counts[name]++ + } + if counts["a"] != 500 || counts["b"] != 500 { + t.Fatalf("concurrent selections = %v, want a:500 b:500", counts) + } +} + +func TestInjectedRandomSourcesAreSerialized(t *testing.T) { + for name, selector := range map[string]Selector{ + "random": NewRandom(&concurrencyDetectingSource{}), + "weighted": NewWeighted(&concurrencyDetectingSource{}), + } { + t.Run(name, func(t *testing.T) { + candidates := []Candidate{ + {Name: "a", Weight: 1, Eligible: true}, + {Name: "b", Weight: 1, Eligible: true}, + } + + var wg sync.WaitGroup + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := selector.Select(candidates); err != nil { + t.Errorf("Select(): %v", err) + } + }() + } + wg.Wait() + }) + } +} + +type concurrencyDetectingSource struct { + active atomic.Bool +} + +func (s *concurrencyDetectingSource) Intn(int) int { + if !s.active.CompareAndSwap(false, true) { + panic("concurrent RandomSource call") + } + time.Sleep(100 * time.Microsecond) + s.active.Store(false) + return 0 +} diff --git a/internal/domain/upstream/fetch_capacity.go b/internal/domain/upstream/fetch_capacity.go new file mode 100644 index 0000000..ab32523 --- /dev/null +++ b/internal/domain/upstream/fetch_capacity.go @@ -0,0 +1,13 @@ +package upstream + +// FetchCapacity atomically accounts for expected provider responses before a +// request starts, preventing concurrent reconciliation from exceeding pool limits. +type FetchCapacity interface { + ReserveFetch(upstreamID string) (FetchPermit, bool, error) +} + +type FetchPermit interface { + Expected() int + Complete(fetched, retained int) error + Cancel() error +} diff --git a/internal/domain/upstream/pool.go b/internal/domain/upstream/pool.go new file mode 100644 index 0000000..f1bb65c --- /dev/null +++ b/internal/domain/upstream/pool.go @@ -0,0 +1,81 @@ +package upstream + +import ( + "time" + + proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" +) + +type ProxyCapacity struct { + State proxyDomain.State + ExpiresAt time.Time + Max int64 + Active int64 + Reserved int64 +} + +type Inventory struct { + Proxies []ProxyCapacity + PendingExpected int + FetchedTotal int64 + MaxSize int + MaxTotal int64 +} + +func (i Inventory) AvailableSlots(now time.Time, safetyMargin time.Duration) int64 { + var available int64 + deadline := now.Add(safetyMargin) + for _, candidate := range i.Proxies { + if candidate.State != proxyDomain.StateAvailable { + continue + } + if !candidate.ExpiresAt.IsZero() && !candidate.ExpiresAt.After(deadline) { + continue + } + slots := candidate.Max - candidate.Active - candidate.Reserved + if slots > 0 { + available += slots + } + } + return available +} + +func (i Inventory) ManagedCount() int { + count := max(i.PendingExpected, 0) + for _, candidate := range i.Proxies { + if managedState(candidate.State) { + count++ + } + } + return count +} + +func (i Inventory) FetchAllowance(requested int) int { + if requested <= 0 || i.MaxSize <= 0 { + return 0 + } + allowed := min(requested, max(i.MaxSize-i.ManagedCount(), 0)) + if i.MaxTotal > 0 { + remaining := i.MaxTotal - i.FetchedTotal + if remaining <= 0 { + return 0 + } + if int64(allowed) > remaining { + allowed = int(remaining) + } + } + return allowed +} + +func managedState(state proxyDomain.State) bool { + switch state { + case proxyDomain.StateFetched, + proxyDomain.StateChecking, + proxyDomain.StateAvailable, + proxyDomain.StateSuspect, + proxyDomain.StateDraining: + return true + default: + return false + } +} diff --git a/internal/domain/upstream/pool_test.go b/internal/domain/upstream/pool_test.go new file mode 100644 index 0000000..d4875fc --- /dev/null +++ b/internal/domain/upstream/pool_test.go @@ -0,0 +1,51 @@ +package upstream + +import ( + "testing" + "time" + + proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" +) + +func TestInventoryAvailableSlotsUsesOnlyAllocatableCapacity(t *testing.T) { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + inventory := Inventory{Proxies: []ProxyCapacity{ + {State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: 10, Active: 3, Reserved: 2}, + {State: proxyDomain.StateAvailable, ExpiresAt: now.Add(5 * time.Second), Max: 10}, + {State: proxyDomain.StateChecking, ExpiresAt: now.Add(time.Minute), Max: 10}, + {State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: 2, Active: 3}, + }} + + if got := inventory.AvailableSlots(now, 10*time.Second); got != 5 { + t.Fatalf("AvailableSlots() = %d, want 5", got) + } +} + +func TestInventoryFetchAllowanceSeparatesPoolAndCumulativeLimits(t *testing.T) { + inventory := Inventory{ + Proxies: []ProxyCapacity{ + {State: proxyDomain.StateFetched}, + {State: proxyDomain.StateChecking}, + {State: proxyDomain.StateAvailable}, + {State: proxyDomain.StateSuspect}, + {State: proxyDomain.StateDraining}, + {State: proxyDomain.StateExtracted}, + }, + PendingExpected: 2, + FetchedTotal: 98, + MaxSize: 10, + MaxTotal: 100, + } + + if got := inventory.ManagedCount(); got != 7 { + t.Fatalf("ManagedCount() = %d, want 7", got) + } + if got := inventory.FetchAllowance(10); got != 2 { + t.Fatalf("FetchAllowance() = %d, want 2 from cumulative quota", got) + } + + inventory.MaxTotal = 0 + if got := inventory.FetchAllowance(10); got != 3 { + t.Fatalf("FetchAllowance() with unlimited cumulative quota = %d, want 3 from pool size", got) + } +} diff --git a/internal/gateway/dispatch/dispatcher.go b/internal/gateway/dispatch/dispatcher.go index af0674f..0332cd6 100644 --- a/internal/gateway/dispatch/dispatcher.go +++ b/internal/gateway/dispatch/dispatcher.go @@ -52,10 +52,22 @@ func (d *Dispatcher) Acquire(request Request) (*Lease, error) { request.Now = time.Now().UTC() } - start := int((d.cursor.Add(1) - 1) % uint64(len(view.Entries))) - for offset := 0; offset < len(view.Entries); offset++ { - entry := view.Entries[(start+offset)%len(view.Entries)] - if !eligible(entry.Proxy, request) { + selection := view.Select(snapshot.Query{ + Now: request.Now, + Scheme: request.Scheme, + Upstreams: request.Upstreams, + RequiredTags: request.RequiredTags, + Exclude: request.Exclude, + SafetyMargin: request.SafetyMargin, + }) + if selection.Len() == 0 { + return nil, ErrNoCandidate + } + + start := int((d.cursor.Add(1) - 1) % uint64(selection.Len())) + for offset := 0; offset < selection.Len(); offset++ { + entry, ok := selection.EntryAt((start + offset) % selection.Len()) + if !ok { continue } reservation, ok := entry.Runtime.Reserve() @@ -71,39 +83,3 @@ func (d *Dispatcher) Acquire(request Request) (*Lease, error) { } return nil, ErrNoCandidate } - -func eligible(candidate proxyDomain.Proxy, request Request) bool { - if candidate.State != proxyDomain.StateAvailable { - return false - } - if request.Scheme != "" && candidate.Scheme != request.Scheme { - return false - } - if _, excluded := request.Exclude[candidate.ID]; excluded { - return false - } - if candidate.ExpiresAt != nil && !candidate.ExpiresAt.After(request.Now.Add(request.SafetyMargin)) { - return false - } - if !contains(request.Upstreams, candidate.SourceUpstream) { - return false - } - for key, value := range request.RequiredTags { - if candidate.Tags[key] != value { - return false - } - } - return true -} - -func contains(allowed []string, value string) bool { - if len(allowed) == 0 { - return true - } - for _, candidate := range allowed { - if candidate == value { - return true - } - } - return false -} diff --git a/internal/gateway/dispatch/dispatcher_test.go b/internal/gateway/dispatch/dispatcher_test.go index dd05377..09bf079 100644 --- a/internal/gateway/dispatch/dispatcher_test.go +++ b/internal/gateway/dispatch/dispatcher_test.go @@ -2,7 +2,9 @@ package dispatch import ( "errors" + "fmt" "sync" + "sync/atomic" "testing" "time" @@ -88,3 +90,198 @@ func TestAcquireNeverOversubscribesSnapshotProxy(t *testing.T) { t.Fatalf("reserved = %d, want 8", count) } } + +func TestAcquireSurvivesConcurrentSnapshotApply(t *testing.T) { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + store := snapshot.NewStore("cluster-a", "worker-a") + base := makeTestProxies(256) + envelope := snapshot.Envelope{ + ClusterID: "cluster-a", + WorkerID: "worker-a", + Epoch: 1, + Version: 1, + Full: true, + Proxies: base, + } + envelope.Checksum = snapshot.Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(base): %v", err) + } + + dispatcher := New(store) + var successes atomic.Int64 + var applyErr atomic.Value + var wg sync.WaitGroup + + wg.Add(1) + go func() { + defer wg.Done() + steps := []struct { + epoch uint64 + version uint64 + }{ + {epoch: 1, version: 2}, + {epoch: 1, version: 3}, + {epoch: 1, version: 4}, + {epoch: 2, version: 1}, + {epoch: 2, version: 2}, + {epoch: 2, version: 3}, + {epoch: 3, version: 1}, + {epoch: 3, version: 2}, + } + for _, step := range steps { + next := snapshot.Envelope{ + ClusterID: "cluster-a", + WorkerID: "worker-a", + Epoch: step.epoch, + Version: step.version, + Full: true, + Proxies: makeTestProxies(256), + } + next.Checksum = snapshot.Checksum(next.Proxies) + if err := store.Apply(next); err != nil { + applyErr.Store(err) + return + } + } + }() + + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + for range 512 { + lease, err := dispatcher.Acquire(Request{ + Now: now, + Scheme: proxyDomain.SchemeHTTP, + Upstreams: []string{"upstream-a"}, + RequiredTags: map[string]string{ + "region": "cn-east", + }, + }) + if err != nil { + if !errors.Is(err, ErrNoCandidate) { + t.Errorf("Acquire(): %v", err) + return + } + continue + } + successes.Add(1) + if err := lease.Cancel(); err != nil { + t.Errorf("Cancel(): %v", err) + return + } + } + }() + } + wg.Wait() + + if err, _ := applyErr.Load().(error); err != nil { + t.Fatalf("Apply(): %v", err) + } + if successes.Load() == 0 { + t.Fatal("Acquire() never succeeded during concurrent Apply") + } +} + +func BenchmarkAcquire100kIndexed(b *testing.B) { + store := snapshot.NewStore("cluster-a", "worker-a") + proxies := makeTestProxies(100_000) + envelope := snapshot.Envelope{ + ClusterID: "cluster-a", + WorkerID: "worker-a", + Epoch: 1, + Version: 1, + Full: true, + Proxies: proxies, + } + envelope.Checksum = snapshot.Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + b.Fatalf("Apply(): %v", err) + } + + dispatcher := New(store) + request := Request{ + Now: time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC), + Scheme: proxyDomain.SchemeHTTP, + Upstreams: []string{"upstream-a"}, + RequiredTags: map[string]string{ + "region": "cn-east", + "tier": "gold", + }, + SafetyMargin: 5 * time.Second, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + lease, err := dispatcher.Acquire(request) + if err != nil { + b.Fatalf("Acquire(): %v", err) + } + if err := lease.Cancel(); err != nil { + b.Fatalf("Cancel(): %v", err) + } + } +} + +func makeTestProxies(count int) []proxyDomain.Proxy { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + expiresAt := now.Add(10 * time.Minute) + proxies := make([]proxyDomain.Proxy, 0, count) + for index := range count { + proxies = append(proxies, proxyDomain.Proxy{ + ID: fmt.Sprintf("proxy-%06d", index), + Scheme: schemeForIndex(index), + Host: fmt.Sprintf("10.0.%d.%d", index/256, index%256), + Port: uint16(20000 + index%1000), + SourceUpstream: upstreamForIndex(index), + State: proxyDomain.StateAvailable, + MaxConcurrency: 8, + ExpiresAt: &expiresAt, + Tags: map[string]string{ + "region": regionForIndex(index), + "tier": tierForIndex(index), + }, + }) + } + return proxies +} + +func schemeForIndex(index int) proxyDomain.Scheme { + switch index % 3 { + case 0: + return proxyDomain.SchemeHTTP + case 1: + return proxyDomain.SchemeHTTPS + default: + return proxyDomain.SchemeSOCKS5 + } +} + +func upstreamForIndex(index int) string { + if index%2 == 0 { + return "upstream-a" + } + return "upstream-b" +} + +func regionForIndex(index int) string { + switch index % 4 { + case 0: + return "cn-east" + case 1: + return "us-west" + case 2: + return "eu-central" + default: + return "ap-south" + } +} + +func tierForIndex(index int) string { + if index%5 == 0 { + return "gold" + } + return "silver" +} diff --git a/internal/gateway/snapshot/store.go b/internal/gateway/snapshot/store.go index 54dc94b..546e42e 100644 --- a/internal/gateway/snapshot/store.go +++ b/internal/gateway/snapshot/store.go @@ -9,6 +9,7 @@ import ( "sort" "sync" "sync/atomic" + "time" proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" ) @@ -41,6 +42,26 @@ type View struct { Version uint64 Checksum string Entries []Entry + + all []int + byScheme map[proxyDomain.Scheme][]int + byUpstream map[string][]int + byTag map[string][]int +} + +type Query struct { + Now time.Time + Scheme proxyDomain.Scheme + Upstreams []string + RequiredTags map[string]string + Exclude map[string]struct{} + SafetyMargin time.Duration +} + +type Selection struct { + view *View + base []int + query Query } type Store struct { @@ -102,10 +123,13 @@ func (s *Store) Apply(envelope Envelope) error { runtime := s.runtimes[descriptor.ID] if runtime == nil { runtime = proxyDomain.NewCapacity(descriptor.MaxConcurrency) - s.runtimes[descriptor.ID] = runtime } else { runtime.SetMax(descriptor.MaxConcurrency) } + // Keep runtimes for temporarily absent IDs. Old immutable views may still + // hold in-flight leases, so reclaiming here could reset active capacity if + // the same Proxy reappears in a later snapshot. + s.runtimes[descriptor.ID] = runtime entries = append(entries, Entry{Proxy: descriptor, Runtime: runtime}) } @@ -117,10 +141,51 @@ func (s *Store) Apply(envelope Envelope) error { Checksum: envelope.Checksum, Entries: entries, } + next.buildIndexes() s.current.Store(next) return nil } +func (v *View) Select(query Query) Selection { + if query.Now.IsZero() { + query.Now = time.Now().UTC() + } + if v == nil { + return Selection{} + } + + base := v.all + if query.Scheme != "" { + base = chooseSmaller(base, v.byScheme[query.Scheme]) + } + if len(query.Upstreams) == 1 { + base = chooseSmaller(base, v.byUpstream[query.Upstreams[0]]) + } else if len(query.Upstreams) > 1 { + if merged := v.unionUpstreams(query.Upstreams); len(merged) > 0 { + base = chooseSmaller(base, merged) + } + } + for key, value := range query.RequiredTags { + base = chooseSmaller(base, v.byTag[tagKey(key, value)]) + } + return Selection{view: v, base: base, query: query} +} + +func (s Selection) Len() int { + return len(s.base) +} + +func (s Selection) EntryAt(index int) (Entry, bool) { + if s.view == nil || index < 0 || index >= len(s.base) { + return Entry{}, false + } + entry := s.view.Entries[s.base[index]] + if !matchesQuery(entry.Proxy, s.query) { + return Entry{}, false + } + return entry, true +} + func Checksum(proxies []proxyDomain.Proxy) string { canonical := cloneAndSort(proxies) encoded, err := json.Marshal(canonical) @@ -147,3 +212,90 @@ func cloneAndSort(source []proxyDomain.Proxy) []proxyDomain.Proxy { }) return cloned } + +func (v *View) buildIndexes() { + if v == nil { + return + } + v.all = make([]int, len(v.Entries)) + v.byScheme = make(map[proxyDomain.Scheme][]int) + v.byUpstream = make(map[string][]int) + v.byTag = make(map[string][]int) + for index, entry := range v.Entries { + v.all[index] = index + v.byScheme[entry.Proxy.Scheme] = append(v.byScheme[entry.Proxy.Scheme], index) + v.byUpstream[entry.Proxy.SourceUpstream] = append(v.byUpstream[entry.Proxy.SourceUpstream], index) + for key, value := range entry.Proxy.Tags { + v.byTag[tagKey(key, value)] = append(v.byTag[tagKey(key, value)], index) + } + } +} + +func (v *View) unionUpstreams(upstreams []string) []int { + total := 0 + for _, upstream := range upstreams { + total += len(v.byUpstream[upstream]) + } + if total == 0 { + return nil + } + merged := make([]int, 0, total) + for _, upstream := range upstreams { + merged = append(merged, v.byUpstream[upstream]...) + } + sort.Ints(merged) + return merged +} + +func chooseSmaller(current, candidate []int) []int { + if len(current) == 0 { + return candidate + } + if len(candidate) == 0 { + return candidate + } + if len(candidate) < len(current) { + return candidate + } + return current +} + +func matchesQuery(candidate proxyDomain.Proxy, query Query) bool { + if candidate.State != proxyDomain.StateAvailable { + return false + } + if query.Scheme != "" && candidate.Scheme != query.Scheme { + return false + } + if _, excluded := query.Exclude[candidate.ID]; excluded { + return false + } + if candidate.ExpiresAt != nil && !candidate.ExpiresAt.After(query.Now.Add(query.SafetyMargin)) { + return false + } + if !contains(query.Upstreams, candidate.SourceUpstream) { + return false + } + for key, value := range query.RequiredTags { + if candidate.Tags[key] != value { + return false + } + } + return true +} + +func contains(allowed []string, value string) bool { + if len(allowed) == 0 { + return true + } + for _, candidate := range allowed { + if candidate == value { + return true + } + } + return false +} + +func tagKey(key, value string) string { + return key + "\x00" + value +} diff --git a/internal/gateway/snapshot/store_test.go b/internal/gateway/snapshot/store_test.go index cdd120f..74cce14 100644 --- a/internal/gateway/snapshot/store_test.go +++ b/internal/gateway/snapshot/store_test.go @@ -2,7 +2,10 @@ package snapshot import ( "errors" + "fmt" + "reflect" "testing" + "time" proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy" ) @@ -78,3 +81,253 @@ func TestStoreRejectsWrongWorkerVersionGapAndChecksum(t *testing.T) { t.Fatalf("checksum error = %v, want ErrChecksumMismatch", err) } } + +func TestViewSelectFiltersBySchemeUpstreamTagAndExclude(t *testing.T) { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + expiresSoon := now.Add(5 * time.Second) + expiresLater := now.Add(time.Minute) + + store := NewStore("cluster-a", "worker-a") + proxies := []proxyDomain.Proxy{ + {ID: "http-a-cn", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east", "tier": "gold"}}, + {ID: "http-b-cn", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "b", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}}, + {ID: "socks-a-cn", Scheme: proxyDomain.SchemeSOCKS5, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}}, + {ID: "http-a-us", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "us-west"}}, + {ID: "expiring", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresSoon, Tags: map[string]string{"region": "cn-east"}}, + } + envelope := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies} + envelope.Checksum = Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(): %v", err) + } + + view := store.Current() + if view == nil { + t.Fatal("Current() returned nil view") + } + + selection := view.Select(Query{ + Now: now, + Scheme: proxyDomain.SchemeHTTP, + Upstreams: []string{"a", "c"}, + RequiredTags: map[string]string{"region": "cn-east"}, + Exclude: map[string]struct{}{"expiring": {}}, + SafetyMargin: 10 * time.Second, + }) + if got, want := collectSelectionIDs(selection), []string{"http-a-cn"}; !reflect.DeepEqual(got, want) { + t.Fatalf("selection ids = %v, want %v", got, want) + } +} + +func TestStoreApplyReusesCapacityAcrossVersionsAndEpochs(t *testing.T) { + store := NewStore("cluster-a", "worker-a") + firstProxies := []proxyDomain.Proxy{{ + ID: "stable", + Scheme: proxyDomain.SchemeHTTP, + Host: "127.0.0.1", + Port: 18080, + MaxConcurrency: 2, + State: proxyDomain.StateAvailable, + }} + first := Envelope{ + ClusterID: "cluster-a", + WorkerID: "worker-a", + Epoch: 1, + Version: 1, + Full: true, + Proxies: firstProxies, + } + first.Checksum = Checksum(first.Proxies) + if err := store.Apply(first); err != nil { + t.Fatalf("Apply(first): %v", err) + } + + initialView := store.Current() + initialRuntime := initialView.Entries[0].Runtime + reservation, ok := initialRuntime.Reserve() + if !ok { + t.Fatal("Reserve() = false, want true") + } + if err := reservation.Commit(); err != nil { + t.Fatalf("Commit(): %v", err) + } + + second := first + second.Version = 2 + second.Proxies = []proxyDomain.Proxy{{ + ID: "stable", + Scheme: proxyDomain.SchemeHTTP, + Host: "localhost", + Port: 18080, + MaxConcurrency: 5, + State: proxyDomain.StateAvailable, + }} + second.Checksum = Checksum(second.Proxies) + if err := store.Apply(second); err != nil { + t.Fatalf("Apply(second): %v", err) + } + + third := second + third.Epoch = 2 + third.Version = 1 + third.Checksum = Checksum(third.Proxies) + if err := store.Apply(third); err != nil { + t.Fatalf("Apply(third): %v", err) + } + + current := store.Current() + if got := current.Entries[0].Runtime; got != initialRuntime { + t.Fatal("runtime pointer was replaced, want capacity reuse") + } + if got := current.Entries[0].Runtime.Max(); got != 5 { + t.Fatalf("runtime max = %d, want 5", got) + } + if got := current.Entries[0].Runtime.Active(); got != 1 { + t.Fatalf("runtime active = %d, want 1", got) + } + if err := reservation.Release(); err != nil { + t.Fatalf("Release(): %v", err) + } +} + +func TestStoreReusesRuntimeWhenProxyDisappearsAndReappears(t *testing.T) { + store := NewStore("cluster-a", "worker-a") + proxy := proxyDomain.Proxy{ + ID: "p1", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable, MaxConcurrency: 1, + } + first := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, + Proxies: []proxyDomain.Proxy{proxy}, + } + first.Checksum = Checksum(first.Proxies) + if err := store.Apply(first); err != nil { + t.Fatalf("Apply(first): %v", err) + } + runtime := store.Current().Entries[0].Runtime + reservation, ok := runtime.Reserve() + if !ok { + t.Fatal("Reserve() = false, want true") + } + if err := reservation.Commit(); err != nil { + t.Fatalf("Commit(): %v", err) + } + + removed := Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true} + removed.Checksum = Checksum(removed.Proxies) + if err := store.Apply(removed); err != nil { + t.Fatalf("Apply(removed): %v", err) + } + reappeared := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 3, Full: true, + Proxies: []proxyDomain.Proxy{proxy}, + } + reappeared.Checksum = Checksum(reappeared.Proxies) + if err := store.Apply(reappeared); err != nil { + t.Fatalf("Apply(reappeared): %v", err) + } + + reused := store.Current().Entries[0].Runtime + if reused != runtime || reused.Active() != 1 { + t.Fatalf("reappeared runtime = %p active=%d, want %p active=1", reused, reused.Active(), runtime) + } + if _, ok := reused.Reserve(); ok { + t.Fatal("Reserve() succeeded despite inherited active capacity") + } + if err := reservation.Release(); err != nil { + t.Fatalf("Release(): %v", err) + } +} + +func collectSelectionIDs(selection Selection) []string { + ids := make([]string, 0, selection.Len()) + for index := 0; index < selection.Len(); index++ { + entry, ok := selection.EntryAt(index) + if !ok { + continue + } + ids = append(ids, entry.Proxy.ID) + } + return ids +} + +func BenchmarkStoreApply100k(b *testing.B) { + store := NewStore("cluster-a", "worker-a") + proxies := makeBenchmarkProxies(100_000) + + b.ReportAllocs() + for index := 0; index < b.N; index++ { + envelope := Envelope{ + ClusterID: "cluster-a", + WorkerID: "worker-a", + Epoch: 1, + Version: uint64(index + 1), + Full: true, + Proxies: proxies, + } + envelope.Checksum = Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + b.Fatalf("Apply(): %v", err) + } + } +} + +func makeBenchmarkProxies(count int) []proxyDomain.Proxy { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + expiresAt := now.Add(30 * time.Minute) + proxies := make([]proxyDomain.Proxy, 0, count) + for index := range count { + proxies = append(proxies, proxyDomain.Proxy{ + ID: fmt.Sprintf("proxy-%06d", index), + Scheme: schemeForBenchmarkIndex(index), + Host: "127.0.0.1", + Port: uint16(20000 + index%1000), + SourceUpstream: upstreamForBenchmarkIndex(index), + State: proxyDomain.StateAvailable, + MaxConcurrency: 8, + ExpiresAt: &expiresAt, + Tags: map[string]string{ + "region": regionForBenchmarkIndex(index), + "tier": tierForBenchmarkIndex(index), + }, + }) + } + return proxies +} + +func schemeForBenchmarkIndex(index int) proxyDomain.Scheme { + switch index % 3 { + case 0: + return proxyDomain.SchemeHTTP + case 1: + return proxyDomain.SchemeHTTPS + default: + return proxyDomain.SchemeSOCKS5 + } +} + +func upstreamForBenchmarkIndex(index int) string { + if index%2 == 0 { + return "upstream-a" + } + return "upstream-b" +} + +func regionForBenchmarkIndex(index int) string { + switch index % 4 { + case 0: + return "cn-east" + case 1: + return "us-west" + case 2: + return "eu-central" + default: + return "ap-south" + } +} + +func tierForBenchmarkIndex(index int) string { + if index%5 == 0 { + return "gold" + } + return "silver" +} diff --git a/internal/platform/admission/fixed_window.go b/internal/platform/admission/fixed_window.go new file mode 100644 index 0000000..983f709 --- /dev/null +++ b/internal/platform/admission/fixed_window.go @@ -0,0 +1,88 @@ +package admission + +import ( + "context" + "errors" + "sync" + "time" +) + +var ( + ErrInvalidConfig = errors.New("invalid admission limiter configuration") + ErrInvalidIdentity = errors.New("invalid admission identity") + ErrGlobalLimit = errors.New("global admission limit exceeded") + ErrPerKeyLimit = errors.New("per-key admission limit exceeded") +) + +type FixedWindowConfig struct { + Window time.Duration + Global int + PerKey int + Now func() time.Time +} + +// FixedWindow provides one concurrency-safe admission primitive for listener +// global and stable-client request limits. +type FixedWindow struct { + mu sync.Mutex + + window time.Duration + global int + perKey int + now func() time.Time + + windowID int64 + globalUsed int + keyUsed map[string]int +} + +func NewFixedWindow(config FixedWindowConfig) (*FixedWindow, error) { + if config.Window <= 0 || config.Global < 0 || config.PerKey < 0 || + (config.Global == 0 && config.PerKey == 0) { + return nil, ErrInvalidConfig + } + if config.Now == nil { + config.Now = time.Now + } + return &FixedWindow{ + window: config.Window, + global: config.Global, + perKey: config.PerKey, + now: config.Now, + keyUsed: make(map[string]int), + }, nil +} + +func (l *FixedWindow) Admit(ctx context.Context, key string) error { + if err := ctx.Err(); err != nil { + return err + } + if l == nil || key == "" { + return ErrInvalidIdentity + } + + l.mu.Lock() + defer l.mu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + windowID := l.now().UnixNano() / int64(l.window) + if windowID != l.windowID { + l.windowID = windowID + l.globalUsed = 0 + clear(l.keyUsed) + } + if l.global > 0 && l.globalUsed >= l.global { + return ErrGlobalLimit + } + if l.perKey > 0 && l.keyUsed[key] >= l.perKey { + return ErrPerKeyLimit + } + if l.global > 0 { + l.globalUsed++ + } + if l.perKey > 0 { + l.keyUsed[key]++ + } + return nil +} diff --git a/internal/platform/admission/fixed_window_test.go b/internal/platform/admission/fixed_window_test.go new file mode 100644 index 0000000..01818f1 --- /dev/null +++ b/internal/platform/admission/fixed_window_test.go @@ -0,0 +1,76 @@ +package admission + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestFixedWindowEnforcesGlobalAndPerKeyLimits(t *testing.T) { + now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + limiter, err := NewFixedWindow(FixedWindowConfig{ + Window: time.Minute, + Global: 3, + PerKey: 2, + Now: func() time.Time { return now }, + }) + if err != nil { + t.Fatalf("NewFixedWindow(): %v", err) + } + + if err := limiter.Admit(context.Background(), "client-a"); err != nil { + t.Fatalf("Admit(a1): %v", err) + } + if err := limiter.Admit(context.Background(), "client-a"); err != nil { + t.Fatalf("Admit(a2): %v", err) + } + if err := limiter.Admit(context.Background(), "client-a"); !errors.Is(err, ErrPerKeyLimit) { + t.Fatalf("Admit(a3) error = %v, want ErrPerKeyLimit", err) + } + if err := limiter.Admit(context.Background(), "client-b"); err != nil { + t.Fatalf("Admit(b1): %v", err) + } + if err := limiter.Admit(context.Background(), "client-c"); !errors.Is(err, ErrGlobalLimit) { + t.Fatalf("Admit(c1) error = %v, want ErrGlobalLimit", err) + } +} + +func TestFixedWindowResetsAndIsConcurrencySafe(t *testing.T) { + var current atomic.Int64 + base := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC) + limiter, err := NewFixedWindow(FixedWindowConfig{ + Window: time.Minute, + Global: 10, + PerKey: 10, + Now: func() time.Time { + return base.Add(time.Duration(current.Load())) + }, + }) + if err != nil { + t.Fatalf("NewFixedWindow(): %v", err) + } + + var accepted atomic.Int64 + var wg sync.WaitGroup + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + if limiter.Admit(context.Background(), "client-a") == nil { + accepted.Add(1) + } + }() + } + wg.Wait() + if got := accepted.Load(); got != 10 { + t.Fatalf("accepted = %d, want 10", got) + } + + current.Store(int64(time.Minute)) + if err := limiter.Admit(context.Background(), "client-a"); err != nil { + t.Fatalf("Admit(after reset): %v", err) + } +} diff --git a/internal/platform/coalesce/signal.go b/internal/platform/coalesce/signal.go new file mode 100644 index 0000000..7595762 --- /dev/null +++ b/internal/platform/coalesce/signal.go @@ -0,0 +1,28 @@ +package coalesce + +import "context" + +// Signal keeps at most one pending notification and never blocks producers. +type Signal struct { + ready chan struct{} +} + +func NewSignal() *Signal { + return &Signal{ready: make(chan struct{}, 1)} +} + +func (s *Signal) Notify() { + select { + case s.ready <- struct{}{}: + default: + } +} + +func (s *Signal) Wait(ctx context.Context) error { + select { + case <-s.ready: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/internal/platform/coalesce/signal_test.go b/internal/platform/coalesce/signal_test.go new file mode 100644 index 0000000..c43c613 --- /dev/null +++ b/internal/platform/coalesce/signal_test.go @@ -0,0 +1,32 @@ +package coalesce + +import ( + "context" + "errors" + "sync" + "testing" +) + +func TestSignalCoalescesConcurrentNotifications(t *testing.T) { + signal := NewSignal() + + var wg sync.WaitGroup + for range 100 { + wg.Add(1) + go func() { + defer wg.Done() + signal.Notify() + }() + } + wg.Wait() + + if err := signal.Wait(context.Background()); err != nil { + t.Fatalf("Wait(): %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := signal.Wait(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("second Wait() error = %v, want context.Canceled", err) + } +} diff --git a/progress.md b/progress.md index 3aa707b..61ec5c9 100644 --- a/progress.md +++ b/progress.md @@ -21,3 +21,16 @@ 代表性集群压测尚未实施，已在完成审计中明确列出。 - 已生成 `proxy-pool-docs-v1.0.zip`，包含 50 个条目，SHA-256 为 `A6882B71196210CE3594A10992C2A0A73EA3A1CF31D8A570709CA999133CE104`。 +- 已实现 random、roundRobin、weighted、leastConnections，并将 Sequential + 拆成共享 Upstream 空结果状态与每 Routing 版本化游标。 +- 已实现 Provider 合并通知、完整 attempt 超时、永久错误契约、指数退避、 + Retry-After、原子 FetchBudget 以及 Pool Reconciler。 +- 已强制 Worker 所有权与 Exclusive Extraction 注入同一权威 Repository， + 覆盖 100 轮并发竞争、续期、Drain ACK 和 Worker 过期回收。 +- 已实现提取服务层策略映射、幂等审计、来源身份准入和公共 FixedWindow + 全局/Client 限流器。 +- 已为 100k Proxy 快照建立 scheme/upstream/tag 索引；本机持续基准中 + `Acquire` 为 640 ns/op、256 B/op、2 allocs/op。该数据仅证明本地选择热路径， + 不代表 100k QPS 集群端到端容量。 +- 本轮 `go test -count=1 ./...`、`go vet ./...`、`go build ./...` 和 + `git diff --check` 通过；Windows `CGO_ENABLED=0`，race 仍由 Linux CI 验证。