feat: implement proxy pool control plane core
This commit is contained in:
parent
7ce9778bdf
commit
139ff46a13
@ -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.
|
||||
|
||||
|
||||
@ -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 测试清单 |
|
||||
|
||||
|
||||
@ -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 端到端验收结论。
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
136
internal/config/redact.go
Normal file
136
internal/config/redact.go
Normal file
@ -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 "<nil>"
|
||||
}
|
||||
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
|
||||
}
|
||||
147
internal/config/resolve.go
Normal file
147
internal/config/resolve.go
Normal file
@ -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
|
||||
}
|
||||
@ -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":
|
||||
|
||||
178
internal/controller/extraction/service.go
Normal file
178
internal/controller/extraction/service.go
Normal file
@ -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
|
||||
}
|
||||
191
internal/controller/extraction/service_test.go
Normal file
191
internal/controller/extraction/service_test.go
Normal file
@ -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
|
||||
}
|
||||
174
internal/controller/pool/fetch_budget.go
Normal file
174
internal/controller/pool/fetch_budget.go
Normal file
@ -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
|
||||
}
|
||||
155
internal/controller/pool/fetch_budget_test.go
Normal file
155
internal/controller/pool/fetch_budget_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
73
internal/controller/pool/ownership.go
Normal file
73
internal/controller/pool/ownership.go
Normal file
@ -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)
|
||||
}
|
||||
173
internal/controller/pool/ownership_test.go
Normal file
173
internal/controller/pool/ownership_test.go
Normal file
@ -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
|
||||
}
|
||||
64
internal/controller/pool/reconciler.go
Normal file
64
internal/controller/pool/reconciler.go
Normal file
@ -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
|
||||
}
|
||||
76
internal/controller/pool/reconciler_test.go
Normal file
76
internal/controller/pool/reconciler_test.go
Normal file
@ -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++ }
|
||||
61
internal/controller/provider/ports.go
Normal file
61
internal/controller/provider/ports.go
Normal file
@ -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
|
||||
}
|
||||
304
internal/controller/provider/reconciler.go
Normal file
304
internal/controller/provider/reconciler.go
Normal file
@ -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() }
|
||||
840
internal/controller/provider/reconciler_test.go
Normal file
840
internal/controller/provider/reconciler_test.go
Normal file
@ -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) }
|
||||
@ -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 {
|
||||
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
35
internal/domain/ownership/ownership.go
Normal file
35
internal/domain/ownership/ownership.go
Normal file
@ -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
|
||||
}
|
||||
@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
168
internal/domain/routing/strategy.go
Normal file
168
internal/domain/routing/strategy.go
Normal file
@ -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}
|
||||
}
|
||||
232
internal/domain/routing/strategy_test.go
Normal file
232
internal/domain/routing/strategy_test.go
Normal file
@ -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
|
||||
}
|
||||
13
internal/domain/upstream/fetch_capacity.go
Normal file
13
internal/domain/upstream/fetch_capacity.go
Normal file
@ -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
|
||||
}
|
||||
81
internal/domain/upstream/pool.go
Normal file
81
internal/domain/upstream/pool.go
Normal file
@ -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
|
||||
}
|
||||
}
|
||||
51
internal/domain/upstream/pool_test.go
Normal file
51
internal/domain/upstream/pool_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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"
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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"
|
||||
}
|
||||
|
||||
88
internal/platform/admission/fixed_window.go
Normal file
88
internal/platform/admission/fixed_window.go
Normal file
@ -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
|
||||
}
|
||||
76
internal/platform/admission/fixed_window_test.go
Normal file
76
internal/platform/admission/fixed_window_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
28
internal/platform/coalesce/signal.go
Normal file
28
internal/platform/coalesce/signal.go
Normal file
@ -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()
|
||||
}
|
||||
}
|
||||
32
internal/platform/coalesce/signal_test.go
Normal file
32
internal/platform/coalesce/signal_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
13
progress.md
13
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 验证。
|
||||
|
||||
Loading…
Reference in New Issue
Block a user