diff --git a/README.md b/README.md index 36d48cc..3464847 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,8 @@ Proxy Pool 用 Controller 协调这些变化,并让 Gateway 数据面只消费 审计与 Outbox;不保存 Proxy 明细或逐次提取记录。 - **Gateway 组件**:HTTP 正向代理、HTTPS CONNECT、双向 Tunnel、重试、超时、 目的地址保护、本地快照存储、容量调度和 Worker 控制面 Register/Watch/ACK/ - Runtime 会话组件已有实现与定向测试;`proxy-gateway` 命令、凭据分发和权威 - Proxy/Routing Snapshot 发布仍待装配。 + Runtime 会话组件已有实现与定向测试。Controller 可向 Worker 下发无凭据引用的 + 已归属 Proxy 快照;`proxy-gateway` 命令、凭据分发和 Routing Snapshot 发布仍待装配。 - **安全边界**:Gateway、Distribution 与 Admin 使用各自的认证语义,并支持 CIDR、可信代理、严格请求解析和敏感信息最小化。 @@ -82,7 +82,7 @@ flowchart LR Redis 会话栅栏。 - **部分完成**:Gateway 传输与调度组件、Snapshot 本地存储、Worker ownership 与运行态领域组件、Docker Compose/Kubernetes 静态部署清单和 protobuf 契约。 -- **待完成**:Worker 权威 Proxy/Routing Snapshot 发布、Gateway 进程装配、Outcome +- **待完成**:带凭据 Proxy 与 Routing 的权威 Snapshot 发布、Gateway 进程装配、Outcome 上报、Checker 调度与健康状态链、完整 Routing 运行链,以及 loadgen 和代表性集群压测。 检查项数量不等于生产就绪度。静态部署清单与 protobuf descriptor 验证也不代表 diff --git a/docs/api/control-plane.md b/docs/api/control-plane.md index 5eb553f..fb39fc6 100644 --- a/docs/api/control-plane.md +++ b/docs/api/control-plane.md @@ -26,10 +26,11 @@ Controller 已实现并验证 `RegisterWorker`、`AcknowledgeSnapshot` 和 `controlPlane` 配置限定。 `WatchSnapshots` 已在 Register 后发送与当前 ownership epoch 对应的基础完整快照, -Gateway 校验后 ACK 并开始 Runtime 心跳。当前基础快照不包含权威 Proxy 或 Routing -内容,流会保持等待后续发布;权威快照发布器、增量、Gateway 进程装配、Outcome 与 -Checker 闭环尚未实现。`ReportOutcomes` 仍明确返回 `Unimplemented`;`100,000 QPS` -仍是未验证的设计目标。 +Gateway 校验后 ACK 并开始 Runtime 心跳。Controller 会从 Redis 的有界 Worker ownership +索引构建无凭据引用的已归属 Proxy 内容,并将租约到期收紧到 Proxy 的 +`usable_until`;带凭据 Proxy 会 fail-closed,直到凭据材料分发完成。Routing 内容、 +增量、Gateway 进程装配、Outcome 与 Checker 闭环尚未实现。`ReportOutcomes` 仍明确 +返回 `Unimplemented`;`100,000 QPS` 仍是未验证的设计目标。 ## 2. Worker 会话 diff --git a/docs/development/implementation-plan.md b/docs/development/implementation-plan.md index 2af91e6..0f8b452 100644 --- a/docs/development/implementation-plan.md +++ b/docs/development/implementation-plan.md @@ -210,10 +210,10 @@ Prometheus 运行时指标,三监听器隔离已通过测试;业务指标仍 WorkerControlPlane 现已接入 Controller 生命周期:Register、ACK 和 Runtime 报告均经 Redis 服务端 TTL 的 session/issued-snapshot/ACK 栅栏校验;mTLS SPIFFE -身份、消息/流限制和有界停机已实现。`WatchSnapshots` 会持续发送并保持当前 epoch -的基础完整 Snapshot;Gateway 已具备 Register/Watch/ACK/Runtime 会话协调组件。 -按 Worker 的可下发 ownership 索引已进入 Redis 原子脚本。权威 Proxy/Routing -payload 发布、凭据分发、Gateway 命令与 Outcome 上报仍未实现。 +身份、消息/流限制和有界停机已实现。`WatchSnapshots` 会发送当前 epoch 的基础完整 +Snapshot 并保持连接;Gateway 已具备 Register/Watch/ACK/Runtime 会话协调组件。 +按 Worker 的可下发 ownership 索引已进入 Redis 原子脚本,并可构建无凭据引用的 +已归属 Proxy payload。Routing payload、凭据分发、Gateway 命令与 Outcome 上报仍未实现。 已新增公用 `domain/activitypool` 契约及并发安全内存参考实现,Provider Reconciler 通过 `UpsertFetched` 写入带供应商 TTL 和分配安全余量的批次;已覆盖 diff --git a/docs/requirements/completion-audit.md b/docs/requirements/completion-audit.md index d0cc279..ee2e86c 100644 --- a/docs/requirements/completion-audit.md +++ b/docs/requirements/completion-audit.md @@ -51,8 +51,8 @@ 服务端 TTL,负向 ACK 会关闭 Runtime 写入栅栏。Controller 已装配 loopback plaintext fixture 与 SPIFFE mTLS 服务端;基础 Snapshot 流和 Gateway 的 Register/Watch/ACK/Runtime 会话协调已实现。Redis 以 Worker 可下发 ownership - 索引支持后续 payload 构建;权威 Proxy/Routing payload、凭据分发、Outcome 和 - Checker 尚未闭环。 + 索引构建无凭据引用的已归属 Proxy payload,并以租约收紧可用期;Routing payload、 + 凭据分发、Outcome 和 Checker 尚未闭环。 - `PostgreSQL 管理面`:已定义 `adminstate` 事务 seam、并发安全 MemoryStore、 公用契约、100 并发 Routing CAS、租约 Outbox 和只含六张管理表的 Schema;pgx Adapter 已在真实 PostgreSQL 18 上通过同一契约、迁移幂等、审计/Outbox @@ -94,9 +94,9 @@ CI 已配置 Linux race job。PostgreSQL 18 和 Redis 8.2 的隔离 Adapter fixt 5. Worker Register/ACK/Runtime gRPC 接收路径、Redis 服务端 TTL 会话生命周期、 SPIFFE mTLS 校验和 Controller 生命周期接线已完成;Redis Provider Leader、 分布式请求额度、Distribution Client 限制和 Provider 状态丢失重建已完成。 -6. Worker 基础网络快照流和 Gateway Snapshot 客户端已完成;权威 Proxy/Routing - payload、凭据分发与 Outcome 上报仍待完成。Redis ownership drain/ACK/过期回收 - 及按 Worker 的可下发索引已完成。 +6. Worker 基础网络快照流、无凭据引用 Proxy payload 和 Gateway Snapshot 客户端已 + 完成;Routing payload、凭据分发与 Outcome 上报仍待完成。Redis ownership + drain/ACK/过期回收及按 Worker 的可下发索引已完成。 7. Checker 调度、探测器和健康 reducer。 8. Admin/Distribution 细粒度授权和审计查询;Distribution 分布式限流已完成。 9. 真实 Compose/Kubernetes 集成、故障演练和代表性集群负载测试。 diff --git a/internal/adapters/redisactivity/ownership_integration_test.go b/internal/adapters/redisactivity/ownership_integration_test.go index 35cc0f0..26d6745 100644 --- a/internal/adapters/redisactivity/ownership_integration_test.go +++ b/internal/adapters/redisactivity/ownership_integration_test.go @@ -398,6 +398,51 @@ func TestRedisOwnershipAndMaintenanceValidateInputs(t *testing.T) { } } +func TestRedisReadWorkerSnapshotUsesWorkerOwnedIndex(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + candidate := testProxy("proxy-a", "192.0.2.10") + candidate.MaxConcurrency = 7 + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, candidate) + assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-a", time.Minute) + if err != nil { + t.Fatalf("Assign(): %v", err) + } + proxies, err := fixture.Adapter.ReadWorkerSnapshot(context.Background(), "worker-a", 1) + if err != nil || len(proxies) != 1 { + t.Fatalf("ReadWorkerSnapshot() = %+v, %v", proxies, err) + } + if proxies[0].Proxy.ID != "proxy-a" || proxies[0].OwnershipEpoch != assignment.Epoch || + !proxies[0].LeaseExpiresAt.Equal(assignment.ExpiresAt) { + t.Fatalf("snapshot proxy = %+v, want assignment %+v", proxies[0], assignment) + } + other, err := fixture.Adapter.ReadWorkerSnapshot(context.Background(), "worker-b", 1) + if err != nil || len(other) != 0 { + t.Fatalf("ReadWorkerSnapshot(other) = %+v, %v", other, err) + } + if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{ + ProxyID: "proxy-a", CheckedAt: now.Add(3 * time.Second), NextState: proxyDomain.StateSuspect, + }); err != nil { + t.Fatalf("ApplyHealth(suspect): %v", err) + } + proxies, err = fixture.Adapter.ReadWorkerSnapshot(context.Background(), "worker-a", 1) + if err != nil || len(proxies) != 0 { + t.Fatalf("ReadWorkerSnapshot(suspect) = %+v, %v", proxies, err) + } + if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{ + ProxyID: "proxy-a", CheckedAt: now.Add(4 * time.Second), NextState: proxyDomain.StateAvailable, + }); err != nil { + t.Fatalf("ApplyHealth(available): %v", err) + } + proxies, err = fixture.Adapter.ReadWorkerSnapshot(context.Background(), "worker-a", 1) + if err != nil || len(proxies) != 1 { + t.Fatalf("ReadWorkerSnapshot(available) = %+v, %v", proxies, err) + } + if _, err := fixture.Adapter.ReadWorkerSnapshot(context.Background(), "worker-a", 0); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) { + t.Fatalf("ReadWorkerSnapshot(zero limit) error = %v", err) + } +} + func assertRedisInventory(t *testing.T, adapter *Adapter, upstreamID string, now time.Time, want int) { t.Helper() inventory, err := adapter.Inventory(context.Background(), upstreamID, now) diff --git a/internal/adapters/redisactivity/scripts.go b/internal/adapters/redisactivity/scripts.go index 7996c0e..6453a94 100644 --- a/internal/adapters/redisactivity/scripts.go +++ b/internal/adapters/redisactivity/scripts.go @@ -88,6 +88,17 @@ type capacityScriptReply struct { AvailableSlots int64 `json:"availableSlots,string"` } +type workerSnapshotScriptReply struct { + Status scriptStatus `json:"status"` + Proxies []workerSnapshotProxyWire `json:"proxies"` +} + +type workerSnapshotProxyWire struct { + Record string `json:"record"` + OwnershipEpoch string `json:"ownershipEpoch"` + LeaseExpiresAtMS int64 `json:"leaseExpiresAtMs"` +} + //go:embed scripts/upsert.lua var upsertSource string @@ -112,15 +123,19 @@ var runtimeSource string //go:embed scripts/capacity.lua var capacitySource string +//go:embed scripts/worker_snapshot.lua +var workerSnapshotSource string + var ( - upsertScript = redis.NewScript(upsertSource) - healthScript = redis.NewScript(healthSource) - extractScript = redis.NewScript(extractSource) - ownershipScript = redis.NewScript(ownershipSource) - sweepScript = redis.NewScript(sweepSource) - statusScript = redis.NewScript(statusSource) - runtimeScript = redis.NewScript(runtimeSource) - capacityScript = redis.NewScript(capacitySource) + upsertScript = redis.NewScript(upsertSource) + healthScript = redis.NewScript(healthSource) + extractScript = redis.NewScript(extractSource) + ownershipScript = redis.NewScript(ownershipSource) + sweepScript = redis.NewScript(sweepSource) + statusScript = redis.NewScript(statusSource) + runtimeScript = redis.NewScript(runtimeSource) + capacityScript = redis.NewScript(capacitySource) + workerSnapshotScript = redis.NewScript(workerSnapshotSource) ) func runScript(ctx context.Context, client redis.Scripter, script *redis.Script, keys []string, args ...any) (any, error) { diff --git a/internal/adapters/redisactivity/scripts/health.lua b/internal/adapters/redisactivity/scripts/health.lua index f25e727..f3c6ed7 100644 --- a/internal/adapters/redisactivity/scripts/health.lua +++ b/internal/adapters/redisactivity/scripts/health.lua @@ -88,6 +88,29 @@ local function remove_worker_owned(id) end end +local function sync_worker_owned(id, record) + local raw = redis.call('HGET', owners_key, id) + if not raw then + return + end + local decoded, assignment = pcall(cjson.decode, raw) + if not decoded or type(assignment) ~= 'table' or + type(assignment.workerIndexKey) ~= 'string' or assignment.workerIndexKey == '' then + return + end + if assignment.workerId == record.ownerWorkerId and not assignment.draining and + type(assignment.expiresAtMs) == 'number' and assignment.expiresAtMs > checked_at_ms and + record.state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then + redis.call('ZADD', assignment.workerIndexKey, assignment.expiresAtMs, id) + local current = redis.call('PEXPIRETIME', assignment.workerIndexKey) + if current < tonumber(record.expiresAtMs) then + redis.call('PEXPIREAT', assignment.workerIndexKey, tonumber(record.expiresAtMs)) + end + return + end + redis.call('ZREM', assignment.workerIndexKey, id) +end + local function remove_proxy(id) local raw = redis.call('HGET', records_key, id) local record = nil @@ -220,6 +243,7 @@ if next_state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms end end end +sync_worker_owned(proxy_id, record) touch(records_key, tonumber(record.expiresAtMs)) touch(unique_key, tonumber(record.expiresAtMs)) touch(idkeys_key, tonumber(record.expiresAtMs)) diff --git a/internal/adapters/redisactivity/scripts/upsert.lua b/internal/adapters/redisactivity/scripts/upsert.lua index dd62a3c..bd8283a 100644 --- a/internal/adapters/redisactivity/scripts/upsert.lua +++ b/internal/adapters/redisactivity/scripts/upsert.lua @@ -88,6 +88,29 @@ local function remove_worker_owned(proxy_id) end end +local function sync_worker_owned(proxy_id, record) + local raw = redis.call('HGET', owners_key, proxy_id) + if not raw then + return + end + local decoded, assignment = pcall(cjson.decode, raw) + if not decoded or type(assignment) ~= 'table' or + type(assignment.workerIndexKey) ~= 'string' or assignment.workerIndexKey == '' then + return + end + if assignment.workerId == record.ownerWorkerId and not assignment.draining and + type(assignment.expiresAtMs) == 'number' and assignment.expiresAtMs > now_ms and + record.state == 'AVAILABLE' and tonumber(record.usableUntilMs) > now_ms then + redis.call('ZADD', assignment.workerIndexKey, assignment.expiresAtMs, proxy_id) + local current = redis.call('PEXPIRETIME', assignment.workerIndexKey) + if current < tonumber(record.expiresAtMs) then + redis.call('PEXPIREAT', assignment.workerIndexKey, tonumber(record.expiresAtMs)) + end + return + end + redis.call('ZREM', assignment.workerIndexKey, proxy_id) +end + local function remove_proxy(proxy_id) local raw = redis.call('HGET', records_key, proxy_id) local record = nil @@ -210,6 +233,7 @@ for _, candidate in ipairs(candidates) do redis.call('HSET', records_key, incumbent_id, encoded) redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id) sync_owned(incumbent_id, incoming) + sync_worker_owned(incumbent_id, incoming) add_available(incumbent_id, incoming) if tonumber(incoming.expiresAtMs) > max_expiry_ms then max_expiry_ms = tonumber(incoming.expiresAtMs) @@ -229,6 +253,7 @@ for _, candidate in ipairs(candidates) do redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest) redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId) sync_owned(candidate.proxyId, incoming) + sync_worker_owned(candidate.proxyId, incoming) if is_managed(incoming.state) then redis.call('HINCRBY', inventory_key, candidate.upstream, 1) end diff --git a/internal/adapters/redisactivity/scripts/worker_snapshot.lua b/internal/adapters/redisactivity/scripts/worker_snapshot.lua new file mode 100644 index 0000000..a58da1a --- /dev/null +++ b/internal/adapters/redisactivity/scripts/worker_snapshot.lua @@ -0,0 +1,75 @@ +local records_key = KEYS[1] +local owners_key = KEYS[2] +local worker_owned_key = KEYS[3] + +local worker_id = ARGV[1] +local limit = tonumber(ARGV[2]) + +local function reply(status, proxies) + if not proxies or #proxies == 0 then + return '{"status":' .. cjson.encode(status) .. ',"proxies":[]}' + end + return cjson.encode({status = status, proxies = proxies}) +end + +local function now_ms() + local value = redis.call('TIME') + return tonumber(value[1]) * 1000 + math.floor(tonumber(value[2]) / 1000) +end + +local function decode_table(value) + if not value then + return nil + end + local ok, decoded = pcall(cjson.decode, value) + if not ok or type(decoded) ~= 'table' then + return nil + end + return decoded +end + +local function valid_assignment(value, now) + return value and value.version == 1 and type(value.proxyId) == 'string' and value.proxyId ~= '' and + type(value.workerId) == 'string' and value.workerId == worker_id and type(value.epoch) == 'number' and + value.epoch > 0 and type(value.expiresAtMs) == 'number' and value.expiresAtMs > now and + type(value.draining) == 'boolean' and not value.draining and + type(value.workerIndexKey) == 'string' and value.workerIndexKey == worker_owned_key +end + +local function valid_record(value, proxy_id, now) + return value and value.version == 1 and value.id == proxy_id and value.state == 'AVAILABLE' and + value.ownerWorkerId == worker_id and type(value.scheme) == 'string' and + (value.scheme == 'http' or value.scheme == 'https' or value.scheme == 'socks5') and + type(value.host) == 'string' and value.host ~= '' and type(value.port) == 'number' and + value.port > 0 and value.port <= 65535 and type(value.sourceUpstream) == 'string' and + value.sourceUpstream ~= '' and type(value.expiresAtMs) == 'number' and value.expiresAtMs > now and + type(value.usableUntilMs) == 'number' and value.usableUntilMs > now and + type(value.maxConcurrency) == 'number' and value.maxConcurrency > 0 +end + +if type(worker_id) ~= 'string' or worker_id == '' or not limit or limit <= 0 then + return reply('invalid') +end + +if redis.call('ZCARD', worker_owned_key) > limit then + return reply('unavailable') +end + +local now = now_ms() +local ids = redis.call('ZRANGEBYSCORE', worker_owned_key, now + 1, '+inf', 'LIMIT', 0, limit) +local proxies = cjson.decode('[]') +for _, proxy_id in ipairs(ids) do + local assignment = decode_table(redis.call('HGET', owners_key, proxy_id)) + local raw_record = redis.call('HGET', records_key, proxy_id) + local record = decode_table(raw_record) + if valid_assignment(assignment, now) and valid_record(record, proxy_id, now) then + proxies[#proxies + 1] = { + record = raw_record, + ownershipEpoch = tostring(assignment.epoch), + leaseExpiresAtMs = assignment.expiresAtMs, + } + else + redis.call('ZREM', worker_owned_key, proxy_id) + end +end +return reply('ok', proxies) diff --git a/internal/adapters/redisactivity/worker_snapshot.go b/internal/adapters/redisactivity/worker_snapshot.go new file mode 100644 index 0000000..377db6f --- /dev/null +++ b/internal/adapters/redisactivity/worker_snapshot.go @@ -0,0 +1,69 @@ +package redisactivity + +import ( + "context" + "errors" + "strconv" + "time" + + ownershipDomain "proxy-pool/internal/domain/ownership" + proxyDomain "proxy-pool/internal/domain/proxy" +) + +var _ ownershipDomain.SnapshotReader = (*Adapter)(nil) + +// ReadWorkerSnapshot returns every currently assignable proxy owned by one +// Worker. The Lua script rejects an oversized view rather than returning a +// partial snapshot that could silently withdraw still-owned proxies. +func (a *Adapter) ReadWorkerSnapshot( + ctx context.Context, + workerID string, + limit int, +) ([]ownershipDomain.SnapshotProxy, error) { + if ctx == nil || a == nil || !runtimeClean(workerID) || limit <= 0 { + return nil, ownershipDomain.ErrInvalidOwnership + } + if err := ctx.Err(); err != nil { + return nil, err + } + result, err := runScript(ctx, a.client, workerSnapshotScript, []string{ + a.keys.records, a.keys.owners, a.keys.workerOwned(workerID), + }, workerID, limit) + if err != nil { + return nil, err + } + var reply workerSnapshotScriptReply + if err := decodeScriptResult(result, &reply); err != nil { + return nil, err + } + switch reply.Status { + case scriptOK: + case scriptInvalid: + return nil, ownershipDomain.ErrInvalidOwnership + case scriptUnavailable: + return nil, errors.Join(ownershipDomain.ErrOwnershipUnavailable, errors.New("worker snapshot limit exceeded")) + default: + return nil, invalidScriptReply("unexpected worker snapshot reply") + } + proxies := make([]ownershipDomain.SnapshotProxy, 0, len(reply.Proxies)) + for _, item := range reply.Proxies { + record, err := decodeProxyRecord(item.Record) + if err != nil { + return nil, invalidScriptReply("worker snapshot contained invalid proxy record") + } + epoch, err := strconv.ParseUint(item.OwnershipEpoch, 10, 64) + if err != nil || epoch == 0 || item.LeaseExpiresAtMS <= 0 { + return nil, invalidScriptReply("worker snapshot contained invalid ownership") + } + proxy := proxyRecordEntry(record).Proxy + if proxy.State != proxyDomain.StateAvailable || proxy.MaxConcurrency <= 0 || proxy.ExpiresAt == nil || + proxy.UsableUntil == nil { + return nil, invalidScriptReply("worker snapshot contained unavailable proxy") + } + proxies = append(proxies, ownershipDomain.SnapshotProxy{ + Proxy: proxy, OwnershipEpoch: epoch, + LeaseExpiresAt: time.UnixMilli(item.LeaseExpiresAtMS).UTC(), + }) + } + return proxies, nil +} diff --git a/internal/controller/bootstrap/bootstrap.go b/internal/controller/bootstrap/bootstrap.go index 7da6f92..a26cf89 100644 --- a/internal/controller/bootstrap/bootstrap.go +++ b/internal/controller/bootstrap/bootstrap.go @@ -22,6 +22,7 @@ import ( "proxy-pool/internal/controller/worker" "proxy-pool/internal/domain/activitypool" extractionDomain "proxy-pool/internal/domain/extraction" + ownershipDomain "proxy-pool/internal/domain/ownership" "proxy-pool/internal/domain/workerruntime" "proxy-pool/internal/platform/admission" "proxy-pool/internal/platform/credentials" @@ -232,12 +233,18 @@ func runWithWorkerFactory( if nilInterface(opened.workerStore) { return errors.Join(ErrStartup, ErrInvalidOptions) } + var snapshotReader ownershipDomain.SnapshotReader + if reader, ok := opened.workerStore.(ownershipDomain.SnapshotReader); ok { + snapshotReader = reader + } service, serviceErr := worker.NewService(opened.workerStore, worker.Options{ ProtocolVersion: loaded.Value.ControlPlane.ProtocolVersion, HeartbeatInterval: loaded.Value.ControlPlane.HeartbeatInterval.Value(), SessionTTL: loaded.Value.ControlPlane.SessionTTL.Value(), MaxStaleAge: loaded.Value.ControlPlane.MaxStaleAge.Value(), MaxRuntimeCounters: loaded.Value.ControlPlane.MaxRuntimeCounters, + MaxSnapshotBytes: loaded.Value.ControlPlane.MaxMessageBytes, + SnapshotReader: snapshotReader, }) if serviceErr != nil { return fmt.Errorf("%w: build Worker control service: %w", ErrStartup, serviceErr) diff --git a/internal/controller/worker/owned_snapshot_source.go b/internal/controller/worker/owned_snapshot_source.go new file mode 100644 index 0000000..3ac01f2 --- /dev/null +++ b/internal/controller/worker/owned_snapshot_source.go @@ -0,0 +1,142 @@ +package worker + +import ( + "context" + "errors" + "math" + "time" + + controlplanev1 "proxy-pool/gen/controlplane/v1" + "proxy-pool/internal/controlplane/snapshotwire" + ownershipDomain "proxy-pool/internal/domain/ownership" + proxyDomain "proxy-pool/internal/domain/proxy" + "proxy-pool/internal/domain/workerruntime" + + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/timestamppb" +) + +var ErrSnapshotCredentialsUnavailable = errors.New("worker snapshot credential material is unavailable") + +// OwnedSnapshotSource builds a complete Worker view from the Redis ownership +// index. It only accepts a complete bounded result from the reader. +type OwnedSnapshotSource struct { + epochs OwnershipEpochReader + reader ownershipDomain.SnapshotReader + validFor time.Duration + maxProxies int + maxBytes int + now func() time.Time +} + +func NewOwnedSnapshotSource( + epochs OwnershipEpochReader, + reader ownershipDomain.SnapshotReader, + validFor time.Duration, + maxProxies int, + maxBytes int, + now func() time.Time, +) (*OwnedSnapshotSource, error) { + if epochs == nil || reader == nil || validFor <= 0 || maxProxies <= 0 || maxBytes <= 0 || now == nil { + return nil, ErrSnapshotsUnavailable + } + return &OwnedSnapshotSource{ + epochs: epochs, reader: reader, validFor: validFor, maxProxies: maxProxies, maxBytes: maxBytes, now: now, + }, nil +} + +func (source *OwnedSnapshotSource) Watch(ctx context.Context, request SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) { + if source == nil || ctx == nil || !workerruntime.ValidIdentifier(request.WorkerID) || + !workerruntime.ValidIdentifier(request.SessionID) || request.LastAppliedVersion == math.MaxUint64 { + return nil, ErrSnapshotsUnavailable + } + epoch, err := source.epochs.CurrentOwnershipEpoch(ctx) + if err != nil { + return nil, err + } + if epoch == 0 { + return nil, ErrSnapshotsUnavailable + } + now := source.now().UTC() + if now.IsZero() { + return nil, ErrSnapshotsUnavailable + } + proxies, err := source.reader.ReadWorkerSnapshot(ctx, request.WorkerID, source.maxProxies) + if err != nil { + return nil, err + } + full := &controlplanev1.WorkerSnapshot{ + Version: request.LastAppliedVersion + 1, OwnershipEpoch: epoch, + GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(source.validFor)), + Proxies: make([]*controlplanev1.OwnedProxy, 0, len(proxies)), + } + for _, item := range proxies { + owned, validUntil, err := wireOwnedProxy(item, now) + if err != nil { + return nil, err + } + if validUntil.Before(full.ValidUntil.AsTime()) { + full.ValidUntil = timestamppb.New(validUntil) + } + full.Proxies = append(full.Proxies, owned) + } + checksum, err := snapshotwire.Checksum(full) + if err != nil { + return nil, err + } + full.Checksum = append([]byte(nil), checksum[:]...) + if proto.Size(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}) > source.maxBytes { + return nil, ErrSnapshotsUnavailable + } + updates := make(chan *controlplanev1.WorkerSnapshot, 1) + updates <- full + return updates, nil +} + +func wireOwnedProxy(item ownershipDomain.SnapshotProxy, now time.Time) (*controlplanev1.OwnedProxy, time.Time, error) { + proxy := item.Proxy + if item.OwnershipEpoch == 0 || item.LeaseExpiresAt.IsZero() || proxy.CredentialVersion != "" || proxy.SecretRef != "" { + if proxy.CredentialVersion != "" || proxy.SecretRef != "" { + return nil, time.Time{}, ErrSnapshotCredentialsUnavailable + } + return nil, time.Time{}, ErrSnapshotsUnavailable + } + protocol, ok := wireProtocol(proxy.Scheme) + if !ok || !workerruntime.ValidIdentifier(proxy.ID) || !workerruntime.ValidIdentifier(proxy.SourceUpstream) || + proxy.Host == "" || proxy.Port == 0 || proxy.MaxConcurrency <= 0 || proxy.MaxConcurrency > math.MaxUint32 || + proxy.ExpiresAt == nil || proxy.UsableUntil == nil { + return nil, time.Time{}, ErrSnapshotsUnavailable + } + validUntil := proxy.UsableUntil.UTC() + if item.LeaseExpiresAt.Before(validUntil) { + validUntil = item.LeaseExpiresAt.UTC() + } + if !validUntil.After(now) || !proxy.ExpiresAt.After(now) { + return nil, time.Time{}, ErrSnapshotsUnavailable + } + tags := make(map[string]string, len(proxy.Tags)) + for key, value := range proxy.Tags { + tags[key] = value + } + return &controlplanev1.OwnedProxy{ + Id: proxy.ID, Upstream: proxy.SourceUpstream, Protocol: protocol, Host: proxy.Host, Port: uint32(proxy.Port), + Username: proxy.Username, ExpiresAt: timestamppb.New(proxy.ExpiresAt.UTC()), + MaxConcurrency: uint32(proxy.MaxConcurrency), Tags: tags, OwnershipEpoch: item.OwnershipEpoch, + UsableUntil: timestamppb.New(validUntil), + }, validUntil, nil +} + +func wireProtocol(scheme proxyDomain.Scheme) (controlplanev1.ProxyProtocol, bool) { + switch scheme { + case proxyDomain.SchemeHTTP: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP, true + case proxyDomain.SchemeHTTPS: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS, true + case proxyDomain.SchemeSOCKS5: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5, true + default: + return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_UNSPECIFIED, false + } +} + +var _ SnapshotSource = (*OwnedSnapshotSource)(nil) diff --git a/internal/controller/worker/server.go b/internal/controller/worker/server.go index abb2e12..07bfe0c 100644 --- a/internal/controller/worker/server.go +++ b/internal/controller/worker/server.go @@ -52,6 +52,11 @@ func NewServer(controlPlane config.ControlPlane, service Service, options Server return nil, err } snapshots := options.Snapshots + if snapshots == nil { + if provider, ok := service.(interface{ SnapshotSource() SnapshotSource }); ok { + snapshots = provider.SnapshotSource() + } + } if snapshots == nil { snapshots, err = NewInitialSnapshotSource(service, controlPlane.MaxStaleAge.Value(), time.Now) if err != nil { diff --git a/internal/controller/worker/service.go b/internal/controller/worker/service.go index 0ecd9dc..db63ea6 100644 --- a/internal/controller/worker/service.go +++ b/internal/controller/worker/service.go @@ -8,6 +8,7 @@ import ( "errors" "time" + ownershipDomain "proxy-pool/internal/domain/ownership" "proxy-pool/internal/domain/workerruntime" ) @@ -55,6 +56,8 @@ type Options struct { SessionTTL time.Duration MaxStaleAge time.Duration MaxRuntimeCounters int + MaxSnapshotBytes int + SnapshotReader ownershipDomain.SnapshotReader SessionID func() (string, error) } @@ -98,8 +101,9 @@ func (service *service) IssueSnapshot(ctx context.Context, reference workerrunti } type service struct { - store workerruntime.ControlStore - options Options + store workerruntime.ControlStore + options Options + snapshots SnapshotSource } func NewService(store workerruntime.ControlStore, options Options) (Service, error) { @@ -111,7 +115,24 @@ func NewService(store workerruntime.ControlStore, options Options) (Service, err if options.SessionID == nil { options.SessionID = randomSessionID } - return &service{store: store, options: options}, nil + result := &service{store: store, options: options} + if options.SnapshotReader != nil { + source, err := NewOwnedSnapshotSource( + result, options.SnapshotReader, options.MaxStaleAge, options.MaxRuntimeCounters, options.MaxSnapshotBytes, time.Now, + ) + if err != nil { + return nil, errors.Join(ErrInvalidCommand, err) + } + result.snapshots = source + } + return result, nil +} + +func (service *service) SnapshotSource() SnapshotSource { + if service == nil { + return nil + } + return service.snapshots } func (service *service) Register(ctx context.Context, command RegisterCommand) (Registration, error) { diff --git a/internal/controller/worker/snapshot_source_test.go b/internal/controller/worker/snapshot_source_test.go index 928a6ef..ca6b279 100644 --- a/internal/controller/worker/snapshot_source_test.go +++ b/internal/controller/worker/snapshot_source_test.go @@ -2,10 +2,13 @@ package worker import ( "context" + "errors" "testing" "time" "proxy-pool/internal/controlplane/snapshotwire" + ownershipDomain "proxy-pool/internal/domain/ownership" + proxyDomain "proxy-pool/internal/domain/proxy" ) func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) { @@ -30,6 +33,61 @@ func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) { } } +func TestOwnedSnapshotSourceBuildsBoundedProxySnapshot(t *testing.T) { + now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC) + expiresAt := now.Add(10 * time.Minute) + usableUntil := now.Add(5 * time.Minute) + source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{ + Proxy: proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, + SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil, + MaxConcurrency: 7, State: proxyDomain.StateAvailable, Tags: map[string]string{"region": "cn"}, + }, + OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute), + }}}, time.Minute*2, 10, 4096, func() time.Time { return now }) + if err != nil { + t.Fatalf("NewOwnedSnapshotSource(): %v", err) + } + updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"}) + if err != nil { + t.Fatalf("Watch(): %v", err) + } + full := <-updates + if full.GetVersion() != 1 || full.GetOwnershipEpoch() != 9 || len(full.GetProxies()) != 1 { + t.Fatalf("snapshot = %+v", full) + } + proxy := full.GetProxies()[0] + if proxy.GetOwnershipEpoch() != 4 || !proxy.GetUsableUntil().AsTime().Equal(now.Add(time.Minute)) || + !full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) { + t.Fatalf("wire proxy = %+v, valid until = %s", proxy, full.GetValidUntil().AsTime()) + } + checksum, err := snapshotwire.Checksum(full) + if err != nil || string(checksum[:]) != string(full.GetChecksum()) { + t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum) + } +} + +func TestOwnedSnapshotSourceRejectsCredentialReferenceUntilMaterialIsAvailable(t *testing.T) { + now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC) + expiresAt := now.Add(time.Minute) + usableUntil := now.Add(time.Minute) + source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{ + Proxy: proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, + SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil, + MaxConcurrency: 1, State: proxyDomain.StateAvailable, CredentialVersion: "v1", + }, + OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute), + }}}, time.Minute, 10, 4096, func() time.Time { return now }) + if err != nil { + t.Fatalf("NewOwnedSnapshotSource(): %v", err) + } + _, err = source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"}) + if !errors.Is(err, ErrSnapshotCredentialsUnavailable) { + t.Fatalf("Watch() error = %v, want ErrSnapshotCredentialsUnavailable", err) + } +} + type epochReaderStub struct { epoch uint64 err error @@ -38,3 +96,12 @@ type epochReaderStub struct { func (reader epochReaderStub) CurrentOwnershipEpoch(context.Context) (uint64, error) { return reader.epoch, reader.err } + +type snapshotReaderStub struct { + proxies []ownershipDomain.SnapshotProxy + err error +} + +func (reader snapshotReaderStub) ReadWorkerSnapshot(context.Context, string, int) ([]ownershipDomain.SnapshotProxy, error) { + return reader.proxies, reader.err +} diff --git a/internal/domain/ownership/snapshot.go b/internal/domain/ownership/snapshot.go new file mode 100644 index 0000000..eb25b0a --- /dev/null +++ b/internal/domain/ownership/snapshot.go @@ -0,0 +1,22 @@ +package ownership + +import ( + "context" + "time" + + proxyDomain "proxy-pool/internal/domain/proxy" +) + +// SnapshotProxy is one currently assignable proxy in a Worker-specific view. +// LeaseExpiresAt bounds how long the view may remain valid without refresh. +type SnapshotProxy struct { + Proxy proxyDomain.Proxy + OwnershipEpoch uint64 + LeaseExpiresAt time.Time +} + +// SnapshotReader returns a complete bounded Worker view. Implementations must +// never silently truncate the result when the requested limit is exceeded. +type SnapshotReader interface { + ReadWorkerSnapshot(context.Context, string, int) ([]SnapshotProxy, error) +}