diff --git a/api/proto/controlplane/v1/controlplane.proto b/api/proto/controlplane/v1/controlplane.proto index f3804d8..c19ce48 100644 --- a/api/proto/controlplane/v1/controlplane.proto +++ b/api/proto/controlplane/v1/controlplane.proto @@ -182,6 +182,7 @@ message ReportRuntimeRequest { uint64 ownership_epoch = 4; repeated ProxyRuntime counters = 5; google.protobuf.Timestamp observed_at = 6; + uint64 report_sequence = 7; } message ProxyRuntime { diff --git a/internal/adapters/redisactivity/adapter.go b/internal/adapters/redisactivity/adapter.go index 48339ce..10ba4df 100644 --- a/internal/adapters/redisactivity/adapter.go +++ b/internal/adapters/redisactivity/adapter.go @@ -19,11 +19,13 @@ var ( ) type Options struct { - Namespace string - Credentials credentials.Store - OperationTTL time.Duration - MaxCandidateScan int - CleanupLimit int + Namespace string + Credentials credentials.Store + OperationTTL time.Duration + MaxCandidateScan int + MaxRuntimeCounters int + MaxInventoryScan int + CleanupLimit int } type Adapter struct { @@ -35,9 +37,16 @@ type Adapter struct { func New(client redis.Scripter, options Options) (*Adapter, error) { options.Namespace = strings.TrimSpace(options.Namespace) + if options.MaxRuntimeCounters == 0 { + options.MaxRuntimeCounters = options.MaxCandidateScan + } + if options.MaxInventoryScan == 0 { + options.MaxInventoryScan = options.MaxCandidateScan + } if nilInterface(client) || nilInterface(options.Credentials) || !namespacePattern.MatchString(options.Namespace) || options.OperationTTL <= 0 || - options.MaxCandidateScan <= 0 || options.CleanupLimit <= 0 { + options.MaxCandidateScan <= 0 || options.MaxRuntimeCounters <= 0 || + options.MaxInventoryScan <= 0 || options.CleanupLimit <= 0 { return nil, ErrInvalidOptions } return &Adapter{ diff --git a/internal/adapters/redisactivity/adapter_test.go b/internal/adapters/redisactivity/adapter_test.go index 17ab60d..6964dfe 100644 --- a/internal/adapters/redisactivity/adapter_test.go +++ b/internal/adapters/redisactivity/adapter_test.go @@ -58,6 +58,8 @@ func TestNewRejectsInvalidDependenciesAndOptions(t *testing.T) { {name: "colon in namespace", client: client, options: withNamespace(valid, "tenant:other")}, {name: "zero operation ttl", client: client, options: withOperationTTL(valid, 0)}, {name: "zero candidate scan", client: client, options: withMaxCandidateScan(valid, 0)}, + {name: "negative runtime counters", client: client, options: withMaxRuntimeCounters(valid, -1)}, + {name: "negative inventory scan", client: client, options: withMaxInventoryScan(valid, -1)}, {name: "negative cleanup limit", client: client, options: withCleanupLimit(valid, -1)}, } for _, tt := range tests { @@ -93,7 +95,9 @@ func TestNewBuildsClusterSafeKeyspaceAndHashesDynamicTokens(t *testing.T) { adapter.keys.records, adapter.keys.unique, adapter.keys.idkeys, adapter.keys.expiry, adapter.keys.available, adapter.keys.owners, adapter.keys.ownerExpiry, adapter.keys.epoch, adapter.keys.inventory, - adapter.keys.stateInventory, + adapter.keys.stateInventory, adapter.keys.workerSessions, + adapter.keys.workerSessionExpiry, adapter.keys.workerRuntime, + adapter.keys.workerRuntimeExpiry, } for _, key := range staticKeys { if strings.Count(key, "{activity}") != 1 || strings.Count(key, "{") != 1 || strings.Count(key, "}") != 1 { @@ -109,6 +113,7 @@ func TestNewBuildsClusterSafeKeyspaceAndHashesDynamicTokens(t *testing.T) { adapter.keys.region(raw), adapter.keys.carrier(raw), adapter.keys.upstream(raw), + adapter.keys.owned(raw), } for _, key := range dynamicKeys { if strings.Contains(key, raw) || strings.Count(key, "{activity}") != 1 || strings.Count(key, "{") != 1 { @@ -161,6 +166,7 @@ func TestProxyRecordCodecIsDeterministicStrictAndRedacted(t *testing.T) { UsableUntilMS: 58_000, LastCheckedAtMS: 2_000, LastSuccessAtMS: 2_000, LatencyNS: int64(25 * time.Millisecond), MaxConcurrency: 8, State: string(proxyDomain.StateAvailable), Tags: map[string]string{"region": "cn", "carrier": "ct"}, + OwnerIndexKey: "pp:{activity}:test:owned:index", } first, err := encodeProxyRecord(record) if err != nil { @@ -260,6 +266,16 @@ func withMaxCandidateScan(options Options, limit int) Options { return options } +func withMaxRuntimeCounters(options Options, limit int) Options { + options.MaxRuntimeCounters = limit + return options +} + +func withMaxInventoryScan(options Options, limit int) Options { + options.MaxInventoryScan = limit + return options +} + func withCleanupLimit(options Options, limit int) Options { options.CleanupLimit = limit return options diff --git a/internal/adapters/redisactivity/capacity.go b/internal/adapters/redisactivity/capacity.go new file mode 100644 index 0000000..1381c68 --- /dev/null +++ b/internal/adapters/redisactivity/capacity.go @@ -0,0 +1,45 @@ +package redisactivity + +import ( + "context" + "time" + + controllerPool "proxy-pool/internal/controller/pool" + "proxy-pool/internal/domain/activitypool" +) + +var _ controllerPool.InventoryReader = (*Adapter)(nil) + +func (a *Adapter) ReadInventory( + ctx context.Context, + upstreamID string, + safetyMargin time.Duration, +) (controllerPool.InventorySnapshot, error) { + if ctx == nil || a == nil || !runtimeClean(upstreamID) || safetyMargin < 0 { + return controllerPool.InventorySnapshot{}, activitypool.ErrInvalidInventory + } + if err := ctx.Err(); err != nil { + return controllerPool.InventorySnapshot{}, err + } + result, err := runScript(ctx, a.client, capacityScript, []string{ + a.keys.records, a.keys.inventory, a.keys.upstream(upstreamID), a.keys.owned(upstreamID), a.keys.owners, + a.keys.workerSessions, a.keys.workerSessionExpiry, + a.keys.workerRuntime, a.keys.workerRuntimeExpiry, + }, upstreamID, durationMillis(safetyMargin), a.options.MaxInventoryScan, a.options.CleanupLimit) + if err != nil { + return controllerPool.InventorySnapshot{}, err + } + var reply capacityScriptReply + if err := decodeScriptResult(result, &reply); err != nil { + return controllerPool.InventorySnapshot{}, err + } + if reply.Status == scriptInvalid { + return controllerPool.InventorySnapshot{}, activitypool.ErrInvalidInventory + } + if reply.Status != scriptOK || reply.Managed < 0 || reply.AvailableSlots < 0 { + return controllerPool.InventorySnapshot{}, invalidScriptReply("capacity inventory is unavailable") + } + return controllerPool.InventorySnapshot{ + Managed: reply.Managed, AvailableSlots: reply.AvailableSlots, + }, nil +} diff --git a/internal/adapters/redisactivity/capacity_integration_test.go b/internal/adapters/redisactivity/capacity_integration_test.go new file mode 100644 index 0000000..e607d29 --- /dev/null +++ b/internal/adapters/redisactivity/capacity_integration_test.go @@ -0,0 +1,105 @@ +//go:build integration + +package redisactivity + +import ( + "context" + "testing" + "time" + + controllerPool "proxy-pool/internal/controller/pool" + "proxy-pool/internal/domain/workerruntime" +) + +func TestRedisCapacityInventoryCombinesProxyAndWorkerRuntime(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + unowned := testProxy("proxy-unowned", "192.0.2.10") + unowned.MaxConcurrency = 10 + owned := testProxy("proxy-owned", "192.0.2.11") + owned.MaxConcurrency = 10 + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, unowned) + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, owned) + assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), + "proxy-owned", "worker-a", time.Minute) + if err != nil { + t.Fatalf("Assign(): %v", err) + } + if err := fixture.Adapter.ReplaceSession(context.Background(), workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", + AckedSnapshotVersion: 1, AckedOwnershipEpoch: assignment.Epoch, + }, time.Minute); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), workerruntime.Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 1, OwnershipEpoch: assignment.Epoch, ObservedAt: now, + Counters: []workerruntime.Counter{{ProxyID: "proxy-owned", Active: 3, Reserved: 2}}, + }, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(): %v", err) + } + + inventory, err := fixture.Adapter.ReadInventory(context.Background(), "provider-a", 0) + if err != nil || inventory.Managed != 2 || inventory.AvailableSlots != 15 { + t.Fatalf("ReadInventory() = %+v, %v; want managed=2 slots=15", inventory, err) + } + if _, ok := any(fixture.Adapter).(controllerPool.InventoryReader); !ok { + t.Fatal("Adapter does not implement pool.InventoryReader") + } + inventory, err = fixture.Adapter.ReadInventory(context.Background(), "provider-a", 2*time.Hour) + if err != nil || inventory.Managed != 2 || inventory.AvailableSlots != 0 { + t.Fatalf("ReadInventory(safety margin) = %+v, %v", inventory, err) + } + fixture.Adapter.options.MaxInventoryScan = 1 + if _, err := fixture.Adapter.ReadInventory(context.Background(), "provider-a", 0); err == nil { + t.Fatal("ReadInventory(over scan limit) error = nil") + } +} + +func TestRedisCapacityInventoryFailsClosedForExpiredRuntime(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + owned := testProxy("proxy-owned", "192.0.2.11") + owned.MaxConcurrency = 10 + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, owned) + assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), + "proxy-owned", "worker-a", time.Minute) + if err != nil { + t.Fatalf("Assign(): %v", err) + } + if err := fixture.Adapter.ReplaceSession(context.Background(), workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", + AckedSnapshotVersion: 1, AckedOwnershipEpoch: assignment.Epoch, + }, 100*time.Millisecond); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), workerruntime.Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 1, OwnershipEpoch: assignment.Epoch, ObservedAt: now, + }, 100*time.Millisecond); err != nil { + t.Fatalf("ReplaceRuntime(): %v", err) + } + time.Sleep(150 * time.Millisecond) + inventory, err := fixture.Adapter.ReadInventory(context.Background(), "provider-a", 0) + if err != nil || inventory.Managed != 1 || inventory.AvailableSlots != 0 { + t.Fatalf("ReadInventory(expired runtime) = %+v, %v", inventory, err) + } +} + +func TestRedisCapacityInventoryScanIsIsolatedPerUpstream(t *testing.T) { + fixture := newRedisTestFixture(t) + fixture.Adapter.options.MaxInventoryScan = 1 + now := redisTestNow() + target := testProxy("proxy-target", "192.0.2.10") + target.MaxConcurrency = 4 + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, target) + seedRedisAvailable(t, fixture.Adapter, "provider-b", now, now.Add(time.Second), 2*time.Minute, + testProxy("proxy-other-1", "192.0.2.11")) + seedRedisAvailable(t, fixture.Adapter, "provider-b", now, now.Add(2*time.Second), 2*time.Minute, + testProxy("proxy-other-2", "192.0.2.12")) + + inventory, err := fixture.Adapter.ReadInventory(context.Background(), "provider-a", 0) + if err != nil || inventory.Managed != 1 || inventory.AvailableSlots != 4 { + t.Fatalf("ReadInventory(provider-a) = %+v, %v; want isolated managed=1 slots=4", inventory, err) + } +} diff --git a/internal/adapters/redisactivity/codec.go b/internal/adapters/redisactivity/codec.go index 4f226ff..3ddbdf8 100644 --- a/internal/adapters/redisactivity/codec.go +++ b/internal/adapters/redisactivity/codec.go @@ -37,6 +37,7 @@ type proxyRecord struct { State string `json:"state"` Tags map[string]string `json:"tags,omitempty"` OwnerWorkerID string `json:"ownerWorkerId,omitempty"` + OwnerIndexKey string `json:"ownerIndexKey"` IndexKeys []string `json:"indexKeys,omitempty"` } @@ -203,7 +204,8 @@ func validateProxyRecord(record proxyRecord) error { record.CreatedAtMS <= 0 || record.ExpiresAtMS <= 0 || record.UsableUntilMS <= 0 || record.UsableUntilMS > record.ExpiresAtMS || record.LastCheckedAtMS < 0 || record.LastSuccessAtMS < 0 || record.LatencyNS < 0 || record.MaxConcurrency < 0 || - !validScheme(record.Scheme) || !validProxyState(record.State) { + !validScheme(record.Scheme) || !validProxyState(record.State) || + record.OwnerIndexKey == "" || !strings.Contains(record.OwnerIndexKey, "{activity}") { return ErrInvalidRecord } for _, key := range record.IndexKeys { diff --git a/internal/adapters/redisactivity/keys.go b/internal/adapters/redisactivity/keys.go index aded371..23e6a2b 100644 --- a/internal/adapters/redisactivity/keys.go +++ b/internal/adapters/redisactivity/keys.go @@ -9,33 +9,41 @@ import ( const redisKeyPrefix = "pp:{activity}:" type keyspace struct { - prefix string - records string - unique string - idkeys string - expiry string - available string - owners string - ownerExpiry string - epoch string - inventory string - stateInventory string + prefix string + records string + unique string + idkeys string + expiry string + available string + owners string + ownerExpiry string + epoch string + inventory string + stateInventory string + workerSessions string + workerSessionExpiry string + workerRuntime string + workerRuntimeExpiry string } func newKeyspace(namespace string) keyspace { prefix := redisKeyPrefix + namespace return keyspace{ - prefix: prefix, - records: prefix + ":records", - unique: prefix + ":unique", - idkeys: prefix + ":idkeys", - expiry: prefix + ":expiry", - available: prefix + ":available", - owners: prefix + ":owners", - ownerExpiry: prefix + ":owner-expiry", - epoch: prefix + ":epoch", - inventory: prefix + ":inventory", - stateInventory: prefix + ":state-inventory", + prefix: prefix, + records: prefix + ":records", + unique: prefix + ":unique", + idkeys: prefix + ":idkeys", + expiry: prefix + ":expiry", + available: prefix + ":available", + owners: prefix + ":owners", + ownerExpiry: prefix + ":owner-expiry", + epoch: prefix + ":epoch", + inventory: prefix + ":inventory", + stateInventory: prefix + ":state-inventory", + workerSessions: prefix + ":worker-sessions", + workerSessionExpiry: prefix + ":worker-session-expiry", + workerRuntime: prefix + ":worker-runtime", + workerRuntimeExpiry: prefix + ":worker-runtime-expiry", } } @@ -67,6 +75,10 @@ func (keys keyspace) upstream(value string) string { return keys.facet("upstream", value) } +func (keys keyspace) owned(value string) string { + return keys.facet("owned", value) +} + func (keys keyspace) facet(name, value string) string { return keys.prefix + ":" + name + ":" + digestToken(value) } diff --git a/internal/adapters/redisactivity/runtime.go b/internal/adapters/redisactivity/runtime.go new file mode 100644 index 0000000..26d1f28 --- /dev/null +++ b/internal/adapters/redisactivity/runtime.go @@ -0,0 +1,248 @@ +package redisactivity + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "sort" + "strconv" + "strings" + "time" + + "proxy-pool/internal/domain/workerruntime" +) + +const runtimeWireVersion = 1 + +const ( + runtimeReplaceSession = "replace_session" + runtimeReplaceReport = "replace_report" + runtimeRead = "read" +) + +type runtimeSessionWire struct { + Version int `json:"version"` + WorkerID string `json:"workerId"` + InstanceID string `json:"instanceId"` + SessionID string `json:"sessionId"` + AckedSnapshotVersion string `json:"ackedSnapshotVersion"` + AckedOwnershipEpoch string `json:"ackedOwnershipEpoch"` +} + +type runtimeCounterWire struct { + ProxyID string `json:"proxyId"` + Active int64 `json:"active"` + Reserved int64 `json:"reserved"` + Draining bool `json:"draining"` +} + +type runtimeReportWire struct { + Version int `json:"version"` + WorkerID string `json:"workerId"` + SessionID string `json:"sessionId"` + Sequence string `json:"sequence"` + SnapshotVersion string `json:"snapshotVersion"` + OwnershipEpoch string `json:"ownershipEpoch"` + ObservedAtMS int64 `json:"observedAtMs"` + Counters []runtimeCounterWire `json:"counters"` +} + +type runtimeOwnedProxyWire struct { + ProxyID string `json:"proxyId"` + WorkerID string `json:"workerId"` + OwnershipEpoch string `json:"ownershipEpoch"` +} + +type runtimeSnapshotWire struct { + ProxyID string `json:"proxyId"` + Active int64 `json:"active"` + Reserved int64 `json:"reserved"` + Draining bool `json:"draining"` + Fresh bool `json:"fresh"` +} + +var ( + _ workerruntime.SessionWriter = (*Adapter)(nil) + _ workerruntime.ReportWriter = (*Adapter)(nil) + _ workerruntime.RuntimeReader = (*Adapter)(nil) +) + +func (a *Adapter) ReplaceSession(ctx context.Context, session workerruntime.Session, ttl time.Duration) error { + if err := validateRuntimeCall(ctx, a); err != nil { + return err + } + if !runtimeClean(session.WorkerID) || !runtimeClean(session.InstanceID) || !runtimeClean(session.SessionID) || + session.AckedSnapshotVersion == 0 || session.AckedOwnershipEpoch == 0 || ttl <= 0 { + return workerruntime.ErrInvalidSession + } + payload, err := json.Marshal(runtimeSessionWire{ + Version: runtimeWireVersion, WorkerID: session.WorkerID, + InstanceID: session.InstanceID, SessionID: session.SessionID, + AckedSnapshotVersion: strconv.FormatUint(session.AckedSnapshotVersion, 10), + AckedOwnershipEpoch: strconv.FormatUint(session.AckedOwnershipEpoch, 10), + }) + if err != nil { + return workerruntime.ErrInvalidSession + } + reply, err := a.runRuntime(ctx, runtimeReplaceSession, durationMillis(ttl), payload, "") + if err != nil { + return err + } + switch reply.Status { + case scriptOK: + return nil + case scriptInvalid: + return workerruntime.ErrInvalidSession + case scriptStale: + return workerruntime.ErrStaleSession + default: + return invalidScriptReply("unexpected worker session reply") + } +} + +func (a *Adapter) ReplaceRuntime(ctx context.Context, report workerruntime.Report, ttl time.Duration) error { + if err := validateRuntimeCall(ctx, a); err != nil { + return err + } + payload, digest, err := a.encodeRuntimeReport(report, ttl) + if err != nil { + return err + } + reply, err := a.runRuntime(ctx, runtimeReplaceReport, durationMillis(ttl), payload, digest) + if err != nil { + return err + } + switch reply.Status { + case scriptOK: + return nil + case scriptInvalid: + return workerruntime.ErrInvalidReport + case scriptStale: + return workerruntime.ErrStaleReport + case scriptConflict: + return workerruntime.ErrConflictingReport + case scriptUnavailable: + return workerruntime.ErrStaleSession + default: + return invalidScriptReply("unexpected worker runtime reply") + } +} + +func (a *Adapter) ReadRuntime(ctx context.Context, proxies []workerruntime.OwnedProxy) ([]workerruntime.Snapshot, error) { + if err := validateRuntimeCall(ctx, a); err != nil { + return nil, err + } + if len(proxies) > a.options.MaxRuntimeCounters { + return nil, workerruntime.ErrInvalidQuery + } + wires := make([]runtimeOwnedProxyWire, len(proxies)) + seen := make(map[string]struct{}, len(proxies)) + for index, proxy := range proxies { + if !runtimeClean(proxy.ProxyID) || !runtimeClean(proxy.WorkerID) || proxy.OwnershipEpoch == 0 { + return nil, workerruntime.ErrInvalidQuery + } + key := proxy.WorkerID + "\x00" + proxy.ProxyID + if _, exists := seen[key]; exists { + return nil, workerruntime.ErrInvalidQuery + } + seen[key] = struct{}{} + wires[index] = runtimeOwnedProxyWire{ + ProxyID: proxy.ProxyID, WorkerID: proxy.WorkerID, + OwnershipEpoch: strconv.FormatUint(proxy.OwnershipEpoch, 10), + } + } + payload, err := json.Marshal(wires) + if err != nil { + return nil, workerruntime.ErrInvalidQuery + } + reply, err := a.runRuntime(ctx, runtimeRead, 0, payload, "") + if err != nil { + return nil, err + } + if reply.Status == scriptInvalid { + return nil, workerruntime.ErrInvalidQuery + } + if reply.Status != scriptOK || len(reply.Snapshots) != len(proxies) { + return nil, invalidScriptReply("unexpected worker runtime read reply") + } + result := make([]workerruntime.Snapshot, len(reply.Snapshots)) + for index, snapshot := range reply.Snapshots { + if snapshot.ProxyID != proxies[index].ProxyID || snapshot.Active < 0 || snapshot.Reserved < 0 { + return nil, invalidScriptReply("invalid worker runtime snapshot") + } + result[index] = workerruntime.Snapshot{ + ProxyID: snapshot.ProxyID, Active: snapshot.Active, Reserved: snapshot.Reserved, + Draining: snapshot.Draining, Fresh: snapshot.Fresh, + } + } + return result, nil +} + +func (a *Adapter) encodeRuntimeReport(report workerruntime.Report, ttl time.Duration) ([]byte, string, error) { + if ttl <= 0 || !runtimeClean(report.WorkerID) || !runtimeClean(report.SessionID) || + report.Sequence == 0 || report.SnapshotVersion == 0 || report.OwnershipEpoch == 0 || report.ObservedAt.IsZero() || + len(report.Counters) > a.options.MaxRuntimeCounters { + return nil, "", workerruntime.ErrInvalidReport + } + counters := append([]workerruntime.Counter(nil), report.Counters...) + sort.Slice(counters, func(left, right int) bool { return counters[left].ProxyID < counters[right].ProxyID }) + wires := make([]runtimeCounterWire, len(counters)) + for index, counter := range counters { + if !runtimeClean(counter.ProxyID) || counter.Active < 0 || counter.Reserved < 0 || + (index > 0 && counters[index-1].ProxyID == counter.ProxyID) { + return nil, "", workerruntime.ErrInvalidReport + } + wires[index] = runtimeCounterWire{ + ProxyID: counter.ProxyID, Active: counter.Active, + Reserved: counter.Reserved, Draining: counter.Draining, + } + } + payload, err := json.Marshal(runtimeReportWire{ + Version: runtimeWireVersion, WorkerID: report.WorkerID, SessionID: report.SessionID, + Sequence: strconv.FormatUint(report.Sequence, 10), + SnapshotVersion: strconv.FormatUint(report.SnapshotVersion, 10), + OwnershipEpoch: strconv.FormatUint(report.OwnershipEpoch, 10), + ObservedAtMS: report.ObservedAt.UTC().UnixMilli(), Counters: wires, + }) + if err != nil { + return nil, "", workerruntime.ErrInvalidReport + } + digest := sha256.Sum256(payload) + return payload, hex.EncodeToString(digest[:]), nil +} + +func (a *Adapter) runRuntime( + ctx context.Context, + operation string, + ttlMS int64, + payload []byte, + digest string, +) (runtimeScriptReply, error) { + result, err := runScript(ctx, a.client, runtimeScript, []string{ + a.keys.workerSessions, a.keys.workerSessionExpiry, + a.keys.workerRuntime, a.keys.workerRuntimeExpiry, a.keys.owners, + }, operation, ttlMS, a.options.CleanupLimit, string(payload), digest) + if err != nil { + return runtimeScriptReply{}, err + } + var reply runtimeScriptReply + if err := decodeScriptResult(result, &reply); err != nil { + return runtimeScriptReply{}, err + } + return reply, nil +} + +func validateRuntimeCall(ctx context.Context, adapter *Adapter) error { + if ctx == nil || adapter == nil { + return workerruntime.ErrInvalidStore + } + if err := ctx.Err(); err != nil { + return err + } + return nil +} + +func runtimeClean(value string) bool { + return value != "" && strings.TrimSpace(value) == value +} diff --git a/internal/adapters/redisactivity/runtime_integration_test.go b/internal/adapters/redisactivity/runtime_integration_test.go new file mode 100644 index 0000000..98032a9 --- /dev/null +++ b/internal/adapters/redisactivity/runtime_integration_test.go @@ -0,0 +1,147 @@ +//go:build integration + +package redisactivity + +import ( + "context" + "errors" + "testing" + "time" + + "proxy-pool/internal/domain/workerruntime" +) + +func TestRedisWorkerRuntimeReplacesSparseCountersAndFencesReports(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, + testProxy("proxy-a", "192.0.2.10")) + 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) + } + session := workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", + AckedSnapshotVersion: 3, AckedOwnershipEpoch: assignment.Epoch, + } + if err := fixture.Adapter.ReplaceSession(context.Background(), session, time.Minute); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } + report := workerruntime.Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 2, + SnapshotVersion: 3, OwnershipEpoch: assignment.Epoch, ObservedAt: now, + Counters: []workerruntime.Counter{{ProxyID: "proxy-a", Active: 2, Reserved: 1}}, + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(first): %v", err) + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(replay): %v", err) + } + conflict := report + conflict.Counters = []workerruntime.Counter{{ProxyID: "proxy-a", Active: 3}} + if err := fixture.Adapter.ReplaceRuntime(context.Background(), conflict, time.Minute); !errors.Is(err, workerruntime.ErrConflictingReport) { + t.Fatalf("ReplaceRuntime(conflict) error = %v", err) + } + stale := report + stale.Sequence = 1 + if err := fixture.Adapter.ReplaceRuntime(context.Background(), stale, time.Minute); !errors.Is(err, workerruntime.ErrStaleReport) { + t.Fatalf("ReplaceRuntime(stale) error = %v", err) + } + + query := []workerruntime.OwnedProxy{{ + ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: assignment.Epoch, + }} + got, err := fixture.Adapter.ReadRuntime(context.Background(), query) + if err != nil || len(got) != 1 || got[0] != (workerruntime.Snapshot{ + ProxyID: "proxy-a", Active: 2, Reserved: 1, Fresh: true, + }) { + t.Fatalf("ReadRuntime(first) = %+v, %v", got, err) + } + report.Sequence = 3 + report.Counters = nil + if err := fixture.Adapter.ReplaceRuntime(context.Background(), report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(empty): %v", err) + } + got, err = fixture.Adapter.ReadRuntime(context.Background(), query) + if err != nil || len(got) != 1 || got[0] != (workerruntime.Snapshot{ProxyID: "proxy-a", Fresh: true}) { + t.Fatalf("ReadRuntime(empty) = %+v, %v", got, err) + } + + if err := fixture.Adapter.ReplaceSession(context.Background(), workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-b", SessionID: "session-b", + AckedSnapshotVersion: 4, AckedOwnershipEpoch: assignment.Epoch + 1, + }, time.Minute); err != nil { + t.Fatalf("ReplaceSession(new): %v", err) + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, workerruntime.ErrStaleSession) { + t.Fatalf("ReplaceRuntime(old session) error = %v", err) + } +} + +func TestRedisWorkerRuntimeExpiresFailClosed(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute, + testProxy("proxy-a", "192.0.2.10")) + 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) + } + if err := fixture.Adapter.ReplaceSession(context.Background(), workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", + AckedSnapshotVersion: 1, AckedOwnershipEpoch: assignment.Epoch, + }, 100*time.Millisecond); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } + if err := fixture.Adapter.ReplaceRuntime(context.Background(), workerruntime.Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 1, OwnershipEpoch: assignment.Epoch, ObservedAt: now, + }, 100*time.Millisecond); err != nil { + t.Fatalf("ReplaceRuntime(): %v", err) + } + time.Sleep(150 * time.Millisecond) + got, err := fixture.Adapter.ReadRuntime(context.Background(), []workerruntime.OwnedProxy{{ + ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: assignment.Epoch, + }}) + if err != nil || len(got) != 1 || got[0].Fresh { + t.Fatalf("ReadRuntime(expired) = %+v, %v", got, err) + } +} + +func TestRedisWorkerRuntimeRejectsEmptyReportBeyondAcknowledgedSnapshot(t *testing.T) { + fixture := newRedisTestFixture(t) + now := redisTestNow() + if err := fixture.Adapter.ReplaceSession(context.Background(), workerruntime.Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", + AckedSnapshotVersion: 3, AckedOwnershipEpoch: 9, + }, time.Minute); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } + for name, report := range map[string]workerruntime.Report{ + "version": { + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 4, OwnershipEpoch: 9, ObservedAt: now, + }, + "epoch": { + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 3, OwnershipEpoch: 10, ObservedAt: now, + }, + } { + t.Run(name, func(t *testing.T) { + if err := fixture.Adapter.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, workerruntime.ErrStaleReport) { + t.Fatalf("ReplaceRuntime() error = %v, want ErrStaleReport", err) + } + }) + } +} + +func TestRedisWorkerRuntimeAcceptsEmptyRead(t *testing.T) { + fixture := newRedisTestFixture(t) + got, err := fixture.Adapter.ReadRuntime(context.Background(), nil) + if err != nil || got == nil || len(got) != 0 { + t.Fatalf("ReadRuntime(empty) = %#v, %v", got, err) + } +} diff --git a/internal/adapters/redisactivity/scripts.go b/internal/adapters/redisactivity/scripts.go index 5f41a28..e4007d3 100644 --- a/internal/adapters/redisactivity/scripts.go +++ b/internal/adapters/redisactivity/scripts.go @@ -74,6 +74,17 @@ type statusScriptInventory struct { Extracted int64 `json:"extracted"` } +type runtimeScriptReply struct { + Status scriptStatus `json:"status"` + Snapshots []runtimeSnapshotWire `json:"snapshots"` +} + +type capacityScriptReply struct { + Status scriptStatus `json:"status"` + Managed int `json:"managed"` + AvailableSlots int64 `json:"availableSlots,string"` +} + //go:embed scripts/upsert.lua var upsertSource string @@ -92,6 +103,12 @@ var sweepSource string //go:embed scripts/status.lua var statusSource string +//go:embed scripts/runtime.lua +var runtimeSource string + +//go:embed scripts/capacity.lua +var capacitySource string + var ( upsertScript = redis.NewScript(upsertSource) healthScript = redis.NewScript(healthSource) @@ -99,6 +116,8 @@ var ( ownershipScript = redis.NewScript(ownershipSource) sweepScript = redis.NewScript(sweepSource) statusScript = redis.NewScript(statusSource) + runtimeScript = redis.NewScript(runtimeSource) + capacityScript = redis.NewScript(capacitySource) ) func runScript(ctx context.Context, client redis.Scripter, script *redis.Script, keys []string, args ...any) (any, error) { diff --git a/internal/adapters/redisactivity/scripts/capacity.lua b/internal/adapters/redisactivity/scripts/capacity.lua new file mode 100644 index 0000000..83c1493 --- /dev/null +++ b/internal/adapters/redisactivity/scripts/capacity.lua @@ -0,0 +1,177 @@ +local records_key = KEYS[1] +local inventory_key = KEYS[2] +local available_upstream_key = KEYS[3] +local owned_upstream_key = KEYS[4] +local owners_key = KEYS[5] +local sessions_key = KEYS[6] +local session_expiry_key = KEYS[7] +local runtime_key = KEYS[8] +local runtime_expiry_key = KEYS[9] + +local upstream_id = ARGV[1] +local safety_margin_ms = tonumber(ARGV[2]) +local scan_limit = tonumber(ARGV[3]) +local cleanup_limit = tonumber(ARGV[4]) + +local function reply(status, managed, available_slots) + return cjson.encode({ + status = status, + managed = managed or 0, + availableSlots = tostring(available_slots or 0) + }) +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_uint(value) + return type(value) == 'string' and string.match(value, '^[0-9]+$') and + value ~= '0' and (string.len(value) == 1 or string.sub(value, 1, 1) ~= '0') +end + +local function compare_uint(left, right) + if string.len(left) ~= string.len(right) then + return string.len(left) < string.len(right) and -1 or 1 + end + if left == right then + return 0 + end + return left < right and -1 or 1 +end + +local function cleanup(now) + local expired_sessions = redis.call('ZRANGEBYSCORE', session_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit) + for _, worker_id in ipairs(expired_sessions) do + redis.call('HDEL', sessions_key, worker_id) + redis.call('HDEL', runtime_key, worker_id) + redis.call('ZREM', session_expiry_key, worker_id) + redis.call('ZREM', runtime_expiry_key, worker_id) + end + local expired_reports = redis.call('ZRANGEBYSCORE', runtime_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit) + for _, worker_id in ipairs(expired_reports) do + redis.call('HDEL', runtime_key, worker_id) + redis.call('ZREM', runtime_expiry_key, worker_id) + end +end + +if type(upstream_id) ~= 'string' or upstream_id == '' or not safety_margin_ms or safety_margin_ms < 0 or + not scan_limit or scan_limit <= 0 or not cleanup_limit or cleanup_limit <= 0 then + return reply('invalid', 0, 0) +end + +local now = now_ms() +cleanup(now) +local threshold = now + safety_margin_ms +local available_ids = redis.call('ZRANGEBYSCORE', available_upstream_key, '(' .. threshold, '+inf', + 'LIMIT', 0, scan_limit + 1) +if #available_ids > scan_limit then + return reply('unavailable', 0, 0) +end +local remaining = scan_limit - #available_ids +local owned_ids = redis.call('ZRANGEBYSCORE', owned_upstream_key, '(' .. threshold, '+inf', + 'LIMIT', 0, remaining + 1) +if #owned_ids > remaining then + return reply('unavailable', 0, 0) +end + +local managed = tonumber(redis.call('HGET', inventory_key, upstream_id) or '0') +if not managed or managed < 0 or managed ~= math.floor(managed) then + return reply('unavailable', 0, 0) +end +local available_slots = 0 +local worker_cache = {} +local seen = {} +for _, proxy_id in ipairs(available_ids) do + seen[proxy_id] = true + local record = decode_table(redis.call('HGET', records_key, proxy_id)) + if not record or type(record.sourceUpstream) ~= 'string' or type(record.state) ~= 'string' or + record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or + type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or + type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or + record.maxConcurrency ~= math.floor(record.maxConcurrency) or + (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HGET', owners_key, proxy_id) then + return reply('unavailable', 0, 0) + end + available_slots = available_slots + record.maxConcurrency +end + +for _, proxy_id in ipairs(owned_ids) do + if seen[proxy_id] then + return reply('unavailable', 0, 0) + end + local record = decode_table(redis.call('HGET', records_key, proxy_id)) + if not record or record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or + type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or + type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or + record.maxConcurrency ~= math.floor(record.maxConcurrency) or + type(record.ownerWorkerId) ~= 'string' or record.ownerWorkerId == '' then + return reply('unavailable', 0, 0) + end + local owner_worker_id = record.ownerWorkerId + local owner = decode_table(redis.call('HGET', owners_key, proxy_id)) + if not owner or owner.workerId ~= owner_worker_id or type(owner.epoch) ~= 'number' or + type(owner.expiresAtMs) ~= 'number' or owner.expiresAtMs <= now or + type(owner.draining) ~= 'boolean' then + return reply('unavailable', 0, 0) + end + local cached = worker_cache[owner_worker_id] + if not cached then + local session = decode_table(redis.call('HGET', sessions_key, owner_worker_id)) + local report = decode_table(redis.call('HGET', runtime_key, owner_worker_id)) + cached = {fresh = false, counters = {}} + if session and report and session.workerId == owner_worker_id and + report.workerId == owner_worker_id and session.sessionId == report.sessionId and + type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and + type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and + valid_uint(session.ackedSnapshotVersion) and valid_uint(session.ackedOwnershipEpoch) and + report.snapshotVersion == session.ackedSnapshotVersion and + report.ownershipEpoch == session.ackedOwnershipEpoch then + cached.fresh = true + cached.ownershipEpoch = report.ownershipEpoch + if type(report.counters) == 'table' then + for _, counter in pairs(report.counters) do + if type(counter) == 'table' and type(counter.proxyId) == 'string' then + cached.counters[counter.proxyId] = counter + end + end + end + end + worker_cache[owner_worker_id] = cached + end + local owner_epoch = tostring(owner.epoch) + if cached.fresh and valid_uint(owner_epoch) and + compare_uint(cached.ownershipEpoch, owner_epoch) >= 0 then + local counter = cached.counters[proxy_id] + local active = 0 + local reserved = 0 + local draining = false + if counter then + active = counter.active + reserved = counter.reserved + draining = counter.draining + end + if type(active) == 'number' and type(reserved) == 'number' and active >= 0 and reserved >= 0 and + active == math.floor(active) and reserved == math.floor(reserved) and + not draining and not owner.draining then + local slots = record.maxConcurrency - active - reserved + if slots > 0 then + available_slots = available_slots + slots + end + end + end +end + +return reply('ok', managed, available_slots) diff --git a/internal/adapters/redisactivity/scripts/extract.lua b/internal/adapters/redisactivity/scripts/extract.lua index 933d245..44c15e0 100644 --- a/internal/adapters/redisactivity/scripts/extract.lua +++ b/internal/adapters/redisactivity/scripts/extract.lua @@ -113,6 +113,12 @@ local function remove_available(proxy_id, record) end end +local function remove_owned(proxy_id, record) + if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then + redis.call('ZREM', record.ownerIndexKey, proxy_id) + end +end + local function remove_proxy(proxy_id) local raw = redis.call('HGET', records_key, proxy_id) local record = nil @@ -120,6 +126,7 @@ local function remove_proxy(proxy_id) local decoded decoded, record = pcall(cjson.decode, raw) remove_available(proxy_id, decoded and record or nil) + remove_owned(proxy_id, decoded and record or nil) if decoded and type(record) == 'table' and is_managed(record.state) then decrement_inventory(record.sourceUpstream) end diff --git a/internal/adapters/redisactivity/scripts/health.lua b/internal/adapters/redisactivity/scripts/health.lua index 8d2490e..30bfe1a 100644 --- a/internal/adapters/redisactivity/scripts/health.lua +++ b/internal/adapters/redisactivity/scripts/health.lua @@ -70,12 +70,19 @@ local function remove_available(id, record) end end +local function remove_owned(id, record) + if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then + redis.call('ZREM', record.ownerIndexKey, id) + end +end + local function remove_proxy(id) local raw = redis.call('HGET', records_key, id) local record = nil if raw then record = cjson.decode(raw) remove_available(id, record) + remove_owned(id, record) if is_managed(record.state) then decrement_inventory(record.sourceUpstream) end @@ -142,6 +149,9 @@ if not raw then return finish({status = 'not_found'}) end local record = cjson.decode(raw) +if type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' then + return finish({status = 'invalid'}) +end if tonumber(record.expiresAtMs) <= checked_at_ms then remove_proxy(proxy_id) return finish({status = 'not_found'}) @@ -184,11 +194,17 @@ end local encoded = cjson.encode(record) redis.call('HSET', records_key, proxy_id, encoded) local owned = (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1 -if next_state == 'AVAILABLE' and not owned and tonumber(record.usableUntilMs) > checked_at_ms then - redis.call('ZADD', available_key, record.usableUntilMs, proxy_id) - for _, index_key in ipairs(record.indexKeys or {}) do - redis.call('ZADD', index_key, record.usableUntilMs, proxy_id) - touch(index_key, tonumber(record.expiresAtMs)) +remove_owned(proxy_id, record) +if next_state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then + if owned then + redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id) + touch(record.ownerIndexKey, tonumber(record.expiresAtMs)) + else + redis.call('ZADD', available_key, record.usableUntilMs, proxy_id) + for _, index_key in ipairs(record.indexKeys or {}) do + redis.call('ZADD', index_key, record.usableUntilMs, proxy_id) + touch(index_key, tonumber(record.expiresAtMs)) + end end end touch(records_key, tonumber(record.expiresAtMs)) diff --git a/internal/adapters/redisactivity/scripts/ownership.lua b/internal/adapters/redisactivity/scripts/ownership.lua index 9312274..5106779 100644 --- a/internal/adapters/redisactivity/scripts/ownership.lua +++ b/internal/adapters/redisactivity/scripts/ownership.lua @@ -94,6 +94,12 @@ local function remove_available(id, record) end end +local function remove_owned(id, record) + if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then + redis.call('ZREM', record.ownerIndexKey, id) + end +end + local function add_available(id, record, at_ms) local usable_until_ms = record and tonumber(record.usableUntilMs) if not usable_until_ms or record.state ~= 'AVAILABLE' or usable_until_ms <= at_ms then @@ -118,6 +124,7 @@ local function remove_proxy(id) local decoded decoded, record = pcall(cjson.decode, raw) remove_available(id, decoded and record or nil) + remove_owned(id, decoded and record or nil) if decoded and type(record) == 'table' and is_managed(record.state) then decrement_inventory(record.sourceUpstream) end @@ -169,6 +176,7 @@ local function clear_owner(id, assignment, at_ms, restore) local raw_record = redis.call('HGET', records_key, id) local record = decode_table(raw_record) if record and (not assignment or record.ownerWorkerId == assignment.workerId) then + remove_owned(id, record) record.ownerWorkerId = nil redis.call('HSET', records_key, id, cjson.encode(record)) if restore then @@ -195,6 +203,7 @@ if operation == 'assign' then local record = decode_table(redis.call('HGET', records_key, proxy_id)) if not record or record.state ~= 'AVAILABLE' or + type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' or (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1 or not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then @@ -223,6 +232,8 @@ if operation == 'assign' then record.ownerWorkerId = worker_id redis.call('HSET', records_key, proxy_id, cjson.encode(record)) remove_available(proxy_id, record) + redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id) + touch(record.ownerIndexKey, tonumber(record.expiresAtMs)) return finish({status = 'ok', record = encoded}) end @@ -267,6 +278,8 @@ if operation == 'begin_drain' then current.assignmentVersion = tonumber(current.assignmentVersion) + 1 local encoded = cjson.encode(current) redis.call('HSET', owners_key, proxy_id, encoded) + local record = decode_table(redis.call('HGET', records_key, proxy_id)) + remove_owned(proxy_id, record) return finish({status = 'ok', record = encoded}) end return finish({status = 'ok', record = cjson.encode(current)}) diff --git a/internal/adapters/redisactivity/scripts/runtime.lua b/internal/adapters/redisactivity/scripts/runtime.lua new file mode 100644 index 0000000..e2f7c9e --- /dev/null +++ b/internal/adapters/redisactivity/scripts/runtime.lua @@ -0,0 +1,229 @@ +local sessions_key = KEYS[1] +local session_expiry_key = KEYS[2] +local runtime_key = KEYS[3] +local runtime_expiry_key = KEYS[4] +local owners_key = KEYS[5] + +local operation = ARGV[1] +local ttl_ms = tonumber(ARGV[2]) +local cleanup_limit = tonumber(ARGV[3]) +local payload = ARGV[4] +local digest = ARGV[5] + +local function reply(status, snapshots) + if snapshots then + return cjson.encode({status = status, snapshots = snapshots}) + end + return '{"status":' .. cjson.encode(status) .. ',"snapshots":[]}' +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_uint(value) + return type(value) == 'string' and string.match(value, '^[0-9]+$') and + value ~= '0' and (string.len(value) == 1 or string.sub(value, 1, 1) ~= '0') +end + +local function compare_uint(left, right) + if string.len(left) ~= string.len(right) then + return string.len(left) < string.len(right) and -1 or 1 + end + if left == right then + return 0 + end + return left < right and -1 or 1 +end + +local function cleanup(now) + local expired_sessions = redis.call('ZRANGEBYSCORE', session_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit) + for _, worker_id in ipairs(expired_sessions) do + redis.call('HDEL', sessions_key, worker_id) + redis.call('HDEL', runtime_key, worker_id) + redis.call('ZREM', session_expiry_key, worker_id) + redis.call('ZREM', runtime_expiry_key, worker_id) + end + local expired_reports = redis.call('ZRANGEBYSCORE', runtime_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit) + for _, worker_id in ipairs(expired_reports) do + redis.call('HDEL', runtime_key, worker_id) + redis.call('ZREM', runtime_expiry_key, worker_id) + end +end + +local function valid_session(value) + return value and value.version == 1 and type(value.workerId) == 'string' and value.workerId ~= '' and + type(value.instanceId) == 'string' and value.instanceId ~= '' and + type(value.sessionId) == 'string' and value.sessionId ~= '' and + valid_uint(value.ackedSnapshotVersion) and valid_uint(value.ackedOwnershipEpoch) +end + +local function valid_owner(value, worker_id, ownership_epoch, now) + return value and type(value.workerId) == 'string' and value.workerId == worker_id and + type(value.epoch) == 'number' and valid_uint(tostring(value.epoch)) and + compare_uint(ownership_epoch, tostring(value.epoch)) >= 0 and + type(value.expiresAtMs) == 'number' and value.expiresAtMs > now +end + +local now = now_ms() +cleanup(now) + +if operation == 'replace_session' then + if not ttl_ms or ttl_ms <= 0 then + return reply('invalid') + end + local session = decode_table(payload) + if not valid_session(session) then + return reply('invalid') + end + local current = decode_table(redis.call('HGET', sessions_key, session.workerId)) + if current and valid_session(current) and current.sessionId == session.sessionId and + current.instanceId == session.instanceId and type(current.expiresAtMs) == 'number' and + current.expiresAtMs > now then + local epoch_order = compare_uint(session.ackedOwnershipEpoch, current.ackedOwnershipEpoch) + local version_order = compare_uint(session.ackedSnapshotVersion, current.ackedSnapshotVersion) + if epoch_order < 0 or (epoch_order == 0 and version_order < 0) then + return reply('stale') + end + if epoch_order > 0 or version_order > 0 then + redis.call('HDEL', runtime_key, session.workerId) + redis.call('ZREM', runtime_expiry_key, session.workerId) + end + else + redis.call('HDEL', runtime_key, session.workerId) + redis.call('ZREM', runtime_expiry_key, session.workerId) + end + session.expiresAtMs = now + ttl_ms + redis.call('HSET', sessions_key, session.workerId, cjson.encode(session)) + redis.call('ZADD', session_expiry_key, session.expiresAtMs, session.workerId) + return reply('ok') +end + +if operation == 'replace_report' then + if not ttl_ms or ttl_ms <= 0 or type(digest) ~= 'string' or digest == '' then + return reply('invalid') + end + local report = decode_table(payload) + if not report or report.version ~= 1 or type(report.workerId) ~= 'string' or report.workerId == '' or + type(report.sessionId) ~= 'string' or report.sessionId == '' or not valid_uint(report.sequence) or + not valid_uint(report.snapshotVersion) or not valid_uint(report.ownershipEpoch) or + type(report.observedAtMs) ~= 'number' or type(report.counters) ~= 'table' then + return reply('invalid') + end + local session = decode_table(redis.call('HGET', sessions_key, report.workerId)) + if not valid_session(session) or session.workerId ~= report.workerId or session.sessionId ~= report.sessionId or + type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then + return reply('unavailable') + end + if report.snapshotVersion ~= session.ackedSnapshotVersion or + report.ownershipEpoch ~= session.ackedOwnershipEpoch then + return reply('stale') + end + local current = decode_table(redis.call('HGET', runtime_key, report.workerId)) + if current and current.sessionId == report.sessionId and valid_uint(current.sequence) then + local ordering = compare_uint(report.sequence, current.sequence) + if ordering < 0 then + return reply('stale') + end + if ordering == 0 then + if current.digest == digest then + return reply('ok') + end + return reply('conflict') + end + end + local seen = {} + for _, counter in pairs(report.counters) do + if type(counter) ~= 'table' or type(counter.proxyId) ~= 'string' or counter.proxyId == '' or + type(counter.active) ~= 'number' or counter.active < 0 or counter.active ~= math.floor(counter.active) or + type(counter.reserved) ~= 'number' or counter.reserved < 0 or counter.reserved ~= math.floor(counter.reserved) or + type(counter.draining) ~= 'boolean' or seen[counter.proxyId] then + return reply('invalid') + end + seen[counter.proxyId] = true + local owner = decode_table(redis.call('HGET', owners_key, counter.proxyId)) + if not valid_owner(owner, report.workerId, report.ownershipEpoch, now) then + return reply('stale') + end + end + report.digest = digest + report.expiresAtMs = now + ttl_ms + redis.call('HSET', runtime_key, report.workerId, cjson.encode(report)) + redis.call('ZADD', runtime_expiry_key, report.expiresAtMs, report.workerId) + session.expiresAtMs = now + ttl_ms + redis.call('HSET', sessions_key, report.workerId, cjson.encode(session)) + redis.call('ZADD', session_expiry_key, session.expiresAtMs, report.workerId) + return reply('ok') +end + +if operation == 'read' then + local queries = decode_table(payload) + if not queries then + return reply('invalid') + end + if next(queries) == nil then + return '{"status":"ok","snapshots":[]}' + end + local snapshots = cjson.decode('[]') + local cache = {} + for _, query in ipairs(queries) do + if type(query) ~= 'table' or type(query.proxyId) ~= 'string' or query.proxyId == '' or + type(query.workerId) ~= 'string' or query.workerId == '' or not valid_uint(query.ownershipEpoch) then + return reply('invalid') + end + local snapshot = {proxyId = query.proxyId, active = 0, reserved = 0, draining = false, fresh = false} + local owner = decode_table(redis.call('HGET', owners_key, query.proxyId)) + if valid_owner(owner, query.workerId, query.ownershipEpoch, now) and + compare_uint(query.ownershipEpoch, tostring(owner.epoch)) == 0 then + local cached = cache[query.workerId] + if not cached then + local session = decode_table(redis.call('HGET', sessions_key, query.workerId)) + local report = decode_table(redis.call('HGET', runtime_key, query.workerId)) + cached = {fresh = false, counters = {}} + if valid_session(session) and session.workerId == query.workerId and report and + report.workerId == query.workerId and report.sessionId == session.sessionId and + type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and + type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and + valid_uint(report.ownershipEpoch) and + report.snapshotVersion == session.ackedSnapshotVersion and + report.ownershipEpoch == session.ackedOwnershipEpoch then + cached.fresh = true + cached.ownershipEpoch = report.ownershipEpoch + if type(report.counters) == 'table' then + for _, counter in pairs(report.counters) do + if type(counter) == 'table' and type(counter.proxyId) == 'string' then + cached.counters[counter.proxyId] = counter + end + end + end + end + cache[query.workerId] = cached + end + if cached.fresh and compare_uint(cached.ownershipEpoch, query.ownershipEpoch) >= 0 then + snapshot.fresh = true + local counter = cached.counters[query.proxyId] + if counter then + snapshot.active = counter.active + snapshot.reserved = counter.reserved + snapshot.draining = counter.draining + end + end + end + snapshots[#snapshots + 1] = snapshot + end + return reply('ok', snapshots) +end + +return reply('invalid') diff --git a/internal/adapters/redisactivity/scripts/sweep.lua b/internal/adapters/redisactivity/scripts/sweep.lua index 280699c..fabb9ed 100644 --- a/internal/adapters/redisactivity/scripts/sweep.lua +++ b/internal/adapters/redisactivity/scripts/sweep.lua @@ -73,6 +73,12 @@ local function remove_available(proxy_id, record) end end +local function remove_owned(proxy_id, record) + if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then + redis.call('ZREM', record.ownerIndexKey, proxy_id) + end +end + local function remove_proxy(proxy_id) local raw = redis.call('HGET', records_key, proxy_id) local record = nil @@ -80,6 +86,7 @@ local function remove_proxy(proxy_id) local decoded decoded, record = pcall(cjson.decode, raw) remove_available(proxy_id, decoded and record or nil) + remove_owned(proxy_id, decoded and record or nil) if decoded and type(record) == 'table' and is_managed(record.state) then decrement_inventory(record.sourceUpstream) end diff --git a/internal/adapters/redisactivity/scripts/upsert.lua b/internal/adapters/redisactivity/scripts/upsert.lua index 15b9bb8..aa2d123 100644 --- a/internal/adapters/redisactivity/scripts/upsert.lua +++ b/internal/adapters/redisactivity/scripts/upsert.lua @@ -70,12 +70,19 @@ local function remove_available(proxy_id, record) end end +local function remove_owned(proxy_id, record) + if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then + redis.call('ZREM', record.ownerIndexKey, proxy_id) + end +end + local function remove_proxy(proxy_id) local raw = redis.call('HGET', records_key, proxy_id) local record = nil if raw then record = cjson.decode(raw) remove_available(proxy_id, record) + remove_owned(proxy_id, record) if is_managed(record.state) then decrement_inventory(record.sourceUpstream) end @@ -123,6 +130,16 @@ local function add_available(proxy_id, record) end end +local function sync_owned(proxy_id, record) + if record.state == 'AVAILABLE' and record.ownerWorkerId and record.ownerWorkerId ~= '' and + tonumber(record.usableUntilMs) > now_ms then + redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id) + touch(record.ownerIndexKey, tonumber(record.expiresAtMs)) + else + redis.call('ZREM', record.ownerIndexKey, proxy_id) + end +end + local function finish(reply) local encoded = cjson.encode(reply) redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms) @@ -132,6 +149,11 @@ end cleanup_expired() for _, candidate in ipairs(candidates) do + local decoded, incoming = pcall(cjson.decode, candidate.record) + if not decoded or type(incoming) ~= 'table' or + type(incoming.ownerIndexKey) ~= 'string' or incoming.ownerIndexKey == '' then + return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0}) + end local mapped = redis.call('HGET', idkeys_key, candidate.proxyId) if mapped and mapped ~= candidate.uniqueDigest then return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0}) @@ -174,6 +196,7 @@ for _, candidate in ipairs(candidates) do local encoded = cjson.encode(incoming) redis.call('HSET', records_key, incumbent_id, encoded) redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id) + sync_owned(incumbent_id, incoming) add_available(incumbent_id, incoming) if tonumber(incoming.expiresAtMs) > max_expiry_ms then max_expiry_ms = tonumber(incoming.expiresAtMs) @@ -192,6 +215,7 @@ for _, candidate in ipairs(candidates) do redis.call('HSET', unique_key, candidate.uniqueDigest, candidate.proxyId) redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest) redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId) + sync_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/upsert.go b/internal/adapters/redisactivity/upsert.go index 897f940..3071620 100644 --- a/internal/adapters/redisactivity/upsert.go +++ b/internal/adapters/redisactivity/upsert.go @@ -139,7 +139,8 @@ func (a *Adapter) prepareUpsertCandidate( CreatedAtMS: candidate.CreatedAt.UnixMilli(), ExpiresAtMS: expiresAt.UnixMilli(), UsableUntilMS: usableUntil.UnixMilli(), LatencyNS: int64(candidate.Latency), MaxConcurrency: candidate.MaxConcurrency, State: string(candidate.State), - Tags: cloneTags(candidate.Tags), IndexKeys: a.availableIndexKeys(candidate), + Tags: cloneTags(candidate.Tags), OwnerIndexKey: a.keys.owned(upstreamID), + IndexKeys: a.availableIndexKeys(candidate), } if candidate.LastCheckedAt != nil { record.LastCheckedAtMS = candidate.LastCheckedAt.UnixMilli() diff --git a/internal/controller/bootstrap/infrastructure.go b/internal/controller/bootstrap/infrastructure.go index 56d3a2e..3df49d3 100644 --- a/internal/controller/bootstrap/infrastructure.go +++ b/internal/controller/bootstrap/infrastructure.go @@ -100,11 +100,13 @@ func (*productionInfrastructure) Open( return ports{}, err } adapter, err := redisactivity.New(redisClient, redisactivity.Options{ - Namespace: redisNamespace, - Credentials: credentialStore, - OperationTTL: redisOperationTTL, - MaxCandidateScan: candidateScan(configuration), - CleanupLimit: redisCleanupLimit, + Namespace: redisNamespace, + Credentials: credentialStore, + OperationTTL: redisOperationTTL, + MaxCandidateScan: candidateScan(configuration), + MaxRuntimeCounters: credentialCapacity(configuration), + MaxInventoryScan: credentialCapacity(configuration), + CleanupLimit: redisCleanupLimit, }) if err != nil { return ports{}, err diff --git a/internal/controller/pool/inventory.go b/internal/controller/pool/inventory.go new file mode 100644 index 0000000..9988087 --- /dev/null +++ b/internal/controller/pool/inventory.go @@ -0,0 +1,17 @@ +package pool + +import ( + "context" + "time" +) + +type InventorySnapshot struct { + Managed int + AvailableSlots int64 +} + +type InventoryReader interface { + // ReadInventory returns a bounded, authoritative aggregate. Unknown Worker + // runtime must reduce capacity rather than being treated as idle. + ReadInventory(context.Context, string, time.Duration) (InventorySnapshot, error) +} diff --git a/internal/controller/pool/reconciler.go b/internal/controller/pool/reconciler.go index cc73821..3674f8c 100644 --- a/internal/controller/pool/reconciler.go +++ b/internal/controller/pool/reconciler.go @@ -62,15 +62,29 @@ func NewReconciler(policy ReconcilePolicy, budget *FetchBudget, notifier FetchNo // 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 { + return r.reconcileSnapshot( + InventorySnapshot{AvailableSlots: inventory.AvailableSlots(now, r.policy.SafetyMargin)}, + false, + ) +} + +func (r *Reconciler) ReconcileSnapshot(inventory InventorySnapshot) ReconcileDecision { + return r.reconcileSnapshot(inventory, true) +} + +func (r *Reconciler) reconcileSnapshot(inventory InventorySnapshot, synchronizeManaged bool) ReconcileDecision { usage := r.budget.Snapshot() - availableSlots := inventory.AvailableSlots(now, r.policy.SafetyMargin) + if synchronizeManaged && inventory.Managed >= 0 && usage.PendingExpected == 0 { + _ = r.budget.SynchronizeManaged(inventory.Managed) + usage = r.budget.Snapshot() + } pendingSlots := saturatingMultiply(int64(usage.PendingExpected), r.slotsPerProxy) decision := ReconcileDecision{ - AvailableSlots: availableSlots, + AvailableSlots: inventory.AvailableSlots, PendingExpected: usage.PendingExpected, FetchedTotal: usage.FetchedTotal, FetchAllowance: r.budget.FetchAllowance(), - EffectiveSlots: saturatingAdd(availableSlots, pendingSlots), + EffectiveSlots: saturatingAdd(inventory.AvailableSlots, pendingSlots), } r.mu.Lock() if r.refilling { diff --git a/internal/controller/pool/reconciler_test.go b/internal/controller/pool/reconciler_test.go index 2dd71e9..e836e2e 100644 --- a/internal/controller/pool/reconciler_test.go +++ b/internal/controller/pool/reconciler_test.go @@ -141,6 +141,53 @@ func TestPoolReconcilerPendingEstimatePausesWithoutEndingRefillEpisode(t *testin } } +func TestPoolReconcilerConsumesAuthoritativeInventorySnapshot(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "provider-a", MaxSize: 10, ExpectedPerFetch: 2, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + notifier := &recordingFetchNotifier{} + reconciler, err := NewReconciler(ReconcilePolicy{ + MinimumAvailableSlots: 3, TargetAvailableSlots: 8, + ExpectedPerFetch: 2, ExpectedSlotsPerFetch: 2, + }, budget, notifier) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + decision := reconciler.ReconcileSnapshot(InventorySnapshot{Managed: 9, AvailableSlots: 2}) + if decision.Triggered || decision.AvailableSlots != 2 || decision.FetchAllowance != 0 { + t.Fatalf("ReconcileSnapshot() = %+v, want authoritative managed inventory to close budget", decision) + } + if usage := budget.Snapshot(); usage.Managed != 9 { + t.Fatalf("FetchBudget.Managed = %d, want 9", usage.Managed) + } +} + +func TestPoolReconcilerLegacyInventoryDoesNotClearManagedBudget(t *testing.T) { + budget, err := NewFetchBudget(FetchBudgetConfig{ + UpstreamID: "provider-a", MaxSize: 10, ExpectedPerFetch: 2, Managed: 9, + }) + if err != nil { + t.Fatalf("NewFetchBudget(): %v", err) + } + reconciler, err := NewReconciler(ReconcilePolicy{ + MinimumAvailableSlots: 3, TargetAvailableSlots: 8, + ExpectedPerFetch: 2, ExpectedSlotsPerFetch: 2, + }, budget, &recordingFetchNotifier{}) + if err != nil { + t.Fatalf("NewReconciler(): %v", err) + } + decision := reconciler.Reconcile(time.Now(), upstream.Inventory{}) + if decision.Triggered || decision.FetchAllowance != 0 { + t.Fatalf("Reconcile(legacy) = %+v, want preserved managed budget", decision) + } + if usage := budget.Snapshot(); usage.Managed != 9 { + t.Fatalf("FetchBudget.Managed = %d, want 9", usage.Managed) + } +} + func inventoryWithSlots(now time.Time, slots int64) upstream.Inventory { return upstream.Inventory{Proxies: []upstream.ProxyCapacity{{ State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: slots, diff --git a/internal/domain/proxy/capacity.go b/internal/domain/proxy/capacity.go index 3d93c70..0369fef 100644 --- a/internal/domain/proxy/capacity.go +++ b/internal/domain/proxy/capacity.go @@ -13,19 +13,46 @@ var ( ) type Capacity struct { - max atomic.Uint32 - counters atomic.Uint64 + max atomic.Uint32 + counters atomic.Uint64 + configuredObserver *activityObserver + observer atomic.Pointer[activityObserver] } +type activityObserver struct{ notify func(bool) } + func NewCapacity(max int64) *Capacity { + return NewCapacityWithActivityObserver(max, nil) +} + +// NewCapacityWithActivityObserver reports successful zero-to-nonzero and +// nonzero-to-zero transitions. The observer must tolerate concurrent calls. +func NewCapacityWithActivityObserver(max int64, observer func(nonzero bool)) *Capacity { capacity := &Capacity{} if max < 0 || max > int64(counterMask) { max = 0 } capacity.max.Store(uint32(max)) + if observer != nil { + capacity.configuredObserver = &activityObserver{notify: observer} + capacity.observer.Store(capacity.configuredObserver) + } return capacity } +// SetActivityObservationEnabled lets snapshot ownership disable callbacks for +// current Proxies and enable them only while a runtime is retired and draining. +func (c *Capacity) SetActivityObservationEnabled(enabled bool) { + if c == nil || c.configuredObserver == nil { + return + } + if enabled { + c.observer.Store(c.configuredObserver) + return + } + c.observer.Store(nil) +} + func (c *Capacity) SetMax(max int64) bool { if max < 0 || max > int64(counterMask) { return false @@ -36,6 +63,15 @@ func (c *Capacity) SetMax(max int64) bool { func (c *Capacity) Max() int64 { return int64(c.max.Load()) } +func (c *Capacity) Counters() (active, reserved, maximum int64) { + if c == nil { + return 0, 0, 0 + } + packed := c.counters.Load() + activeCounter, reservedCounter := unpack(packed) + return int64(activeCounter), int64(reservedCounter), int64(c.max.Load()) +} + func (c *Capacity) Reserve() (*Reservation, bool) { for { current := c.counters.Load() @@ -45,6 +81,11 @@ func (c *Capacity) Reserve() (*Reservation, bool) { } next := pack(active, reserved+1) if c.counters.CompareAndSwap(current, next) { + if active+reserved == 0 { + if observer := c.observer.Load(); observer != nil { + observer.notify(true) + } + } return &Reservation{capacity: c}, true } } @@ -77,7 +118,16 @@ func (c *Capacity) cancel() { for { current := c.counters.Load() active, reserved := unpack(current) - if reserved == 0 || c.counters.CompareAndSwap(current, pack(active, reserved-1)) { + if reserved == 0 { + return + } + next := pack(active, reserved-1) + if c.counters.CompareAndSwap(current, next) { + if active+reserved == 1 { + if observer := c.observer.Load(); observer != nil { + observer.notify(false) + } + } return } } @@ -87,7 +137,16 @@ func (c *Capacity) release() { for { current := c.counters.Load() active, reserved := unpack(current) - if active == 0 || c.counters.CompareAndSwap(current, pack(active-1, reserved)) { + if active == 0 { + return + } + next := pack(active-1, reserved) + if c.counters.CompareAndSwap(current, next) { + if active+reserved == 1 { + if observer := c.observer.Load(); observer != nil { + observer.notify(false) + } + } return } } diff --git a/internal/domain/proxy/capacity_test.go b/internal/domain/proxy/capacity_test.go index 428aeff..fe36d91 100644 --- a/internal/domain/proxy/capacity_test.go +++ b/internal/domain/proxy/capacity_test.go @@ -2,6 +2,7 @@ package proxy import ( "errors" + "reflect" "sync" "sync/atomic" "testing" @@ -38,6 +39,56 @@ func TestReservationCancelReleasesReservedCapacity(t *testing.T) { assertCapacityCounters(t, capacity, 0, 0) } +func TestCapacityCountersReadsOnePackedSnapshot(t *testing.T) { + capacity := NewCapacity(3) + first, ok := capacity.Reserve() + if !ok { + t.Fatal("first Reserve() = false") + } + second, ok := capacity.Reserve() + if !ok { + t.Fatal("second Reserve() = false") + } + if err := first.Commit(); err != nil { + t.Fatalf("Commit(): %v", err) + } + + active, reserved, maximum := capacity.Counters() + if active != 1 || reserved != 1 || maximum != 3 { + t.Fatalf("Counters() = (%d, %d, %d), want (1, 1, 3)", active, reserved, maximum) + } + if err := first.Release(); err != nil { + t.Fatalf("Release(): %v", err) + } + if err := second.Cancel(); err != nil { + t.Fatalf("Cancel(): %v", err) + } +} + +func TestCapacityActivityObserverTracksOnlyNonzeroTransitions(t *testing.T) { + var transitions []bool + capacity := NewCapacityWithActivityObserver(2, func(nonzero bool) { + transitions = append(transitions, nonzero) + }) + first, ok := capacity.Reserve() + if !ok { + t.Fatal("Reserve(first) = false") + } + second, ok := capacity.Reserve() + if !ok { + t.Fatal("Reserve(second) = false") + } + if err := first.Cancel(); err != nil { + t.Fatalf("Cancel(first): %v", err) + } + if err := second.Cancel(); err != nil { + t.Fatalf("Cancel(second): %v", err) + } + if !reflect.DeepEqual(transitions, []bool{true, false}) { + t.Fatalf("transitions = %v, want [true false]", transitions) + } +} + func TestReservationCommitAndReleaseAreSingleUse(t *testing.T) { capacity := NewCapacity(1) reservation, ok := capacity.Reserve() diff --git a/internal/domain/workerruntime/memory.go b/internal/domain/workerruntime/memory.go new file mode 100644 index 0000000..bdcd89f --- /dev/null +++ b/internal/domain/workerruntime/memory.go @@ -0,0 +1,211 @@ +package workerruntime + +import ( + "context" + "crypto/sha256" + "encoding/json" + "sort" + "strings" + "sync" + "time" +) + +type MemoryStore struct { + mu sync.Mutex + now func() time.Time + sessions map[string]memorySession + reports map[string]memoryReport +} + +type memorySession struct { + value Session + expiresAt time.Time +} + +type memoryReport struct { + value Report + digest [sha256.Size]byte + expiresAt time.Time + counters map[string]Counter +} + +var ( + _ SessionWriter = (*MemoryStore)(nil) + _ ReportWriter = (*MemoryStore)(nil) + _ RuntimeReader = (*MemoryStore)(nil) +) + +func NewMemoryStore(now func() time.Time) (*MemoryStore, error) { + if now == nil { + return nil, ErrInvalidStore + } + return &MemoryStore{ + now: now, sessions: make(map[string]memorySession), reports: make(map[string]memoryReport), + }, nil +} + +func (store *MemoryStore) ReplaceSession(ctx context.Context, session Session, ttl time.Duration) error { + if ctx == nil || store == nil || !validSession(session) || ttl <= 0 { + return ErrInvalidSession + } + if err := ctx.Err(); err != nil { + return err + } + now := store.now().UTC() + if now.IsZero() { + return ErrInvalidStore + } + store.mu.Lock() + defer store.mu.Unlock() + current, exists := store.sessions[session.WorkerID] + identityChanged := exists && (current.value.SessionID != session.SessionID || current.value.InstanceID != session.InstanceID) + expired := exists && !current.expiresAt.After(now) + if exists && !identityChanged && !expired && sessionBefore(session, current.value) { + return ErrStaleSession + } + ackAdvanced := exists && !identityChanged && !expired && sessionAfter(session, current.value) + if identityChanged || expired || ackAdvanced { + delete(store.reports, session.WorkerID) + } + store.sessions[session.WorkerID] = memorySession{value: session, expiresAt: now.Add(ttl)} + return nil +} + +func (store *MemoryStore) ReplaceRuntime(ctx context.Context, report Report, ttl time.Duration) error { + if ctx == nil || store == nil || ttl <= 0 { + return ErrInvalidReport + } + if err := ctx.Err(); err != nil { + return err + } + normalized, counterIndex, err := normalizeReport(report) + if err != nil { + return err + } + payload, err := json.Marshal(normalized) + if err != nil { + return ErrInvalidReport + } + digest := sha256.Sum256(payload) + now := store.now().UTC() + if now.IsZero() { + return ErrInvalidStore + } + store.mu.Lock() + defer store.mu.Unlock() + session, exists := store.sessions[report.WorkerID] + if !exists || !session.expiresAt.After(now) || session.value.SessionID != report.SessionID { + return ErrStaleSession + } + if report.SnapshotVersion != session.value.AckedSnapshotVersion || + report.OwnershipEpoch != session.value.AckedOwnershipEpoch { + return ErrStaleReport + } + if current, exists := store.reports[report.WorkerID]; exists && current.value.SessionID == report.SessionID { + switch { + case normalized.Sequence < current.value.Sequence: + return ErrStaleReport + case normalized.Sequence == current.value.Sequence && digest != current.digest: + return ErrConflictingReport + case normalized.Sequence == current.value.Sequence: + return nil + } + } + store.reports[report.WorkerID] = memoryReport{ + value: normalized, digest: digest, expiresAt: now.Add(ttl), counters: counterIndex, + } + session.expiresAt = now.Add(ttl) + store.sessions[report.WorkerID] = session + return nil +} + +func (store *MemoryStore) ReadRuntime(ctx context.Context, proxies []OwnedProxy) ([]Snapshot, error) { + if ctx == nil || store == nil { + return nil, ErrInvalidQuery + } + if err := ctx.Err(); err != nil { + return nil, err + } + seen := make(map[string]struct{}, len(proxies)) + for _, proxy := range proxies { + if !clean(proxy.ProxyID) || !clean(proxy.WorkerID) || proxy.OwnershipEpoch == 0 { + return nil, ErrInvalidQuery + } + key := proxy.WorkerID + "\x00" + proxy.ProxyID + if _, exists := seen[key]; exists { + return nil, ErrInvalidQuery + } + seen[key] = struct{}{} + } + now := store.now().UTC() + if now.IsZero() { + return nil, ErrInvalidStore + } + store.mu.Lock() + defer store.mu.Unlock() + result := make([]Snapshot, len(proxies)) + for index, proxy := range proxies { + result[index].ProxyID = proxy.ProxyID + session, sessionExists := store.sessions[proxy.WorkerID] + report, reportExists := store.reports[proxy.WorkerID] + if !sessionExists || !reportExists || !session.expiresAt.After(now) || !report.expiresAt.After(now) || + report.value.SessionID != session.value.SessionID || + report.value.SnapshotVersion != session.value.AckedSnapshotVersion || + report.value.OwnershipEpoch != session.value.AckedOwnershipEpoch || + report.value.OwnershipEpoch < proxy.OwnershipEpoch { + continue + } + result[index].Fresh = true + if counter, exists := report.counters[proxy.ProxyID]; exists { + result[index].Active = counter.Active + result[index].Reserved = counter.Reserved + result[index].Draining = counter.Draining + } + } + return result, nil +} + +func normalizeReport(report Report) (Report, map[string]Counter, error) { + if !clean(report.WorkerID) || !clean(report.SessionID) || report.Sequence == 0 || + report.SnapshotVersion == 0 || report.OwnershipEpoch == 0 || report.ObservedAt.IsZero() { + return Report{}, nil, ErrInvalidReport + } + normalized := report + normalized.ObservedAt = report.ObservedAt.UTC() + normalized.Counters = append([]Counter(nil), report.Counters...) + sort.Slice(normalized.Counters, func(left, right int) bool { + return normalized.Counters[left].ProxyID < normalized.Counters[right].ProxyID + }) + index := make(map[string]Counter, len(normalized.Counters)) + for _, counter := range normalized.Counters { + if !clean(counter.ProxyID) || counter.Active < 0 || counter.Reserved < 0 { + return Report{}, nil, ErrInvalidReport + } + if _, exists := index[counter.ProxyID]; exists { + return Report{}, nil, ErrInvalidReport + } + index[counter.ProxyID] = counter + } + return normalized, index, nil +} + +func validSession(session Session) bool { + return clean(session.WorkerID) && clean(session.InstanceID) && clean(session.SessionID) && + session.AckedSnapshotVersion > 0 && session.AckedOwnershipEpoch > 0 +} + +func sessionBefore(left, right Session) bool { + return left.AckedOwnershipEpoch < right.AckedOwnershipEpoch || + (left.AckedOwnershipEpoch == right.AckedOwnershipEpoch && + left.AckedSnapshotVersion < right.AckedSnapshotVersion) +} + +func sessionAfter(left, right Session) bool { + return left.AckedOwnershipEpoch > right.AckedOwnershipEpoch || + (left.AckedOwnershipEpoch == right.AckedOwnershipEpoch && + left.AckedSnapshotVersion > right.AckedSnapshotVersion) +} + +func clean(value string) bool { + return value != "" && strings.TrimSpace(value) == value +} diff --git a/internal/domain/workerruntime/memory_test.go b/internal/domain/workerruntime/memory_test.go new file mode 100644 index 0000000..64760d4 --- /dev/null +++ b/internal/domain/workerruntime/memory_test.go @@ -0,0 +1,156 @@ +package workerruntime + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestMemoryStoreReplacesSparseRuntimeAndClearsMissingCounters(t *testing.T) { + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + store := newRuntimeStore(t, &now) + ctx := context.Background() + registerRuntimeSession(t, store, "session-a", time.Minute) + report := Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now, + Counters: []Counter{{ProxyID: "proxy-a", Active: 2, Reserved: 1}}, + } + if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(first): %v", err) + } + got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}}) + if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Active: 2, Reserved: 1, Fresh: true}) { + t.Fatalf("ReadRuntime(first) = %+v, %v", got, err) + } + + report.Sequence = 2 + report.ObservedAt = now.Add(time.Second) + report.Counters = nil + if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(empty): %v", err) + } + got, err = store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}}) + if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Fresh: true}) { + t.Fatalf("ReadRuntime(empty) = %+v, %v", got, err) + } +} + +func TestMemoryStoreFencesSessionsAndReportSequence(t *testing.T) { + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + store := newRuntimeStore(t, &now) + ctx := context.Background() + registerRuntimeSession(t, store, "session-a", time.Minute) + report := Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 2, + SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now, + Counters: []Counter{{ProxyID: "proxy-a", Active: 1}}, + } + if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(first): %v", err) + } + if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(idempotent): %v", err) + } + conflict := report + conflict.Counters = []Counter{{ProxyID: "proxy-a", Active: 2}} + if err := store.ReplaceRuntime(ctx, conflict, time.Minute); !errors.Is(err, ErrConflictingReport) { + t.Fatalf("ReplaceRuntime(conflict) error = %v", err) + } + stale := report + stale.Sequence = 1 + if err := store.ReplaceRuntime(ctx, stale, time.Minute); !errors.Is(err, ErrStaleReport) { + t.Fatalf("ReplaceRuntime(stale) error = %v", err) + } + registerRuntimeSession(t, store, "session-b", time.Minute) + if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrStaleSession) { + t.Fatalf("ReplaceRuntime(old session) error = %v", err) + } +} + +func TestMemoryStoreFailsClosedForExpiredOrOlderOwnershipReport(t *testing.T) { + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + store := newRuntimeStore(t, &now) + ctx := context.Background() + registerRuntimeSession(t, store, "session-a", time.Minute) + if err := store.ReplaceRuntime(ctx, Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now, + }, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(): %v", err) + } + queries := []OwnedProxy{ + {ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 10}, + {ProxyID: "proxy-b", WorkerID: "worker-a", OwnershipEpoch: 9}, + } + got, err := store.ReadRuntime(ctx, queries) + if err != nil || got[0].Fresh || !got[1].Fresh { + t.Fatalf("ReadRuntime(ownership fence) = %+v, %v", got, err) + } + now = now.Add(time.Minute) + got, err = store.ReadRuntime(ctx, queries[1:]) + if err != nil || len(got) != 1 || got[0].Fresh { + t.Fatalf("ReadRuntime(expired) = %+v, %v", got, err) + } +} + +func TestMemoryStoreRejectsRuntimeBeyondAcknowledgedSnapshot(t *testing.T) { + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + store := newRuntimeStore(t, &now) + registerRuntimeSession(t, store, "session-a", time.Minute) + report := Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 4, OwnershipEpoch: 9, ObservedAt: now, + } + if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrStaleReport) { + t.Fatalf("ReplaceRuntime(ahead snapshot) error = %v, want ErrStaleReport", err) + } + report.SnapshotVersion = 3 + report.OwnershipEpoch = 10 + if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrStaleReport) { + t.Fatalf("ReplaceRuntime(ahead epoch) error = %v, want ErrStaleReport", err) + } +} + +func TestMemoryStoreExpiredSameIdentitySessionDoesNotReactivateOldReport(t *testing.T) { + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + store := newRuntimeStore(t, &now) + registerRuntimeSession(t, store, "session-a", time.Second) + if err := store.ReplaceRuntime(context.Background(), Report{ + WorkerID: "worker-a", SessionID: "session-a", Sequence: 1, + SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now, + }, time.Minute); err != nil { + t.Fatalf("ReplaceRuntime(): %v", err) + } + session := store.sessions["worker-a"] + session.expiresAt = now.Add(time.Second) + store.sessions["worker-a"] = session + now = now.Add(2 * time.Second) + registerRuntimeSession(t, store, "session-a", time.Minute) + got, err := store.ReadRuntime(context.Background(), []OwnedProxy{{ + ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9, + }}) + if err != nil || len(got) != 1 || got[0].Fresh { + t.Fatalf("ReadRuntime(after re-register) = %+v, %v; want stale", got, err) + } +} + +func newRuntimeStore(t *testing.T, now *time.Time) *MemoryStore { + t.Helper() + store, err := NewMemoryStore(func() time.Time { return *now }) + if err != nil { + t.Fatalf("NewMemoryStore(): %v", err) + } + return store +} + +func registerRuntimeSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) { + t.Helper() + if err := store.ReplaceSession(context.Background(), Session{ + WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID, + AckedSnapshotVersion: 3, AckedOwnershipEpoch: 9, + }, ttl); err != nil { + t.Fatalf("ReplaceSession(): %v", err) + } +} diff --git a/internal/domain/workerruntime/runtime.go b/internal/domain/workerruntime/runtime.go new file mode 100644 index 0000000..f4861ef --- /dev/null +++ b/internal/domain/workerruntime/runtime.go @@ -0,0 +1,70 @@ +package workerruntime + +import ( + "context" + "errors" + "time" +) + +var ( + ErrInvalidStore = errors.New("invalid worker runtime store") + ErrInvalidSession = errors.New("invalid worker runtime session") + ErrInvalidReport = errors.New("invalid worker runtime report") + ErrInvalidQuery = errors.New("invalid worker runtime query") + ErrStaleSession = errors.New("stale worker runtime session") + ErrStaleReport = errors.New("stale worker runtime report") + ErrConflictingReport = errors.New("conflicting worker runtime report") +) + +type Session struct { + WorkerID string + InstanceID string + SessionID string + AckedSnapshotVersion uint64 + AckedOwnershipEpoch uint64 +} + +type Counter struct { + ProxyID string + Active int64 + Reserved int64 + Draining bool +} + +// Report is a complete sparse replacement. Missing counters are zero for the +// reported Worker snapshot; callers must increase Sequence for every update. +type Report struct { + WorkerID string + SessionID string + Sequence uint64 + SnapshotVersion uint64 + OwnershipEpoch uint64 + ObservedAt time.Time + Counters []Counter +} + +type OwnedProxy struct { + ProxyID string + WorkerID string + OwnershipEpoch uint64 +} + +type Snapshot struct { + ProxyID string + Active int64 + Reserved int64 + Draining bool + Fresh bool +} + +type SessionWriter interface { + ReplaceSession(context.Context, Session, time.Duration) error +} + +type ReportWriter interface { + ReplaceRuntime(context.Context, Report, time.Duration) error +} + +type RuntimeReader interface { + ReadRuntime(context.Context, []OwnedProxy) ([]Snapshot, error) +} diff --git a/internal/gateway/snapshot/store.go b/internal/gateway/snapshot/store.go index 6055a19..fe3b3ed 100644 --- a/internal/gateway/snapshot/store.go +++ b/internal/gateway/snapshot/store.go @@ -7,19 +7,90 @@ import ( "errors" "fmt" "sort" + "strings" "sync" "sync/atomic" "time" proxyDomain "proxy-pool/internal/domain/proxy" + "proxy-pool/internal/domain/workerruntime" ) var ( - ErrWrongTarget = errors.New("snapshot targets another cluster or worker") - ErrResyncRequired = errors.New("snapshot sequence requires a full resync") - ErrChecksumMismatch = errors.New("snapshot checksum mismatch") + ErrWrongTarget = errors.New("snapshot targets another cluster or worker") + ErrResyncRequired = errors.New("snapshot sequence requires a full resync") + ErrChecksumMismatch = errors.New("snapshot checksum mismatch") + ErrInvalidRuntimeReport = errors.New("invalid worker runtime report") + ErrInvalidRuntimeLimit = errors.New("invalid snapshot runtime limit") + ErrRuntimeLimitExceeded = errors.New("snapshot runtime limit exceeded") ) +const defaultRuntimeLimit = 1_000_000 + +const activeRuntimeShardCount = 64 + +type activeRuntimeShard struct { + mu sync.Mutex + entries map[string]*proxyDomain.Capacity +} + +type activeRuntimeIndex [activeRuntimeShardCount]activeRuntimeShard + +func (index *activeRuntimeIndex) track(proxyID string, runtime *proxyDomain.Capacity, nonzero bool) { + shard := &index[activeRuntimeShardIndex(proxyID)] + shard.mu.Lock() + defer shard.mu.Unlock() + if nonzero { + if shard.entries == nil { + shard.entries = make(map[string]*proxyDomain.Capacity) + } + shard.entries[proxyID] = runtime + return + } + active, reserved, _ := runtime.Counters() + if active == 0 && reserved == 0 && shard.entries[proxyID] == runtime { + delete(shard.entries, proxyID) + } +} + +func (index *activeRuntimeIndex) remove(proxyID string, runtime *proxyDomain.Capacity) { + shard := &index[activeRuntimeShardIndex(proxyID)] + shard.mu.Lock() + defer shard.mu.Unlock() + if shard.entries[proxyID] == runtime { + delete(shard.entries, proxyID) + } +} + +func (index *activeRuntimeIndex) rangeEntries(visit func(string, *proxyDomain.Capacity)) { + for shardIndex := range index { + shard := &index[shardIndex] + shard.mu.Lock() + for proxyID, runtime := range shard.entries { + visit(proxyID, runtime) + } + shard.mu.Unlock() + } +} + +func activeRuntimeShardIndex(proxyID string) uint64 { + const ( + offset = uint64(14695981039346656037) + prime = uint64(1099511628211) + ) + hash := offset + for index := 0; index < len(proxyID); index++ { + hash ^= uint64(proxyID[index]) + hash *= prime + } + return hash % activeRuntimeShardCount +} + +type runtimeRegistration struct { + capacity *proxyDomain.Capacity + current atomic.Bool +} + type Envelope struct { ClusterID string WorkerID string @@ -70,17 +141,29 @@ type Store struct { current atomic.Pointer[View] mu sync.Mutex - runtimes map[string]*proxyDomain.Capacity + runtimes map[string]*runtimeRegistration + active activeRuntimeIndex + limit int } func NewStore(clusterID, workerID string) *Store { return &Store{ clusterID: clusterID, workerID: workerID, - runtimes: make(map[string]*proxyDomain.Capacity), + runtimes: make(map[string]*runtimeRegistration), + limit: defaultRuntimeLimit, } } +func NewStoreWithRuntimeLimit(clusterID, workerID string, limit int) (*Store, error) { + if limit <= 0 { + return nil, ErrInvalidRuntimeLimit + } + store := NewStore(clusterID, workerID) + store.limit = limit + return store, nil +} + func (s *Store) Current() *View { if s == nil { return nil @@ -88,6 +171,49 @@ func (s *Store) Current() *View { return s.current.Load() } +func (s *Store) RuntimeReport(sessionID string, sequence uint64, observedAt time.Time) (workerruntime.Report, error) { + if s == nil || strings.TrimSpace(sessionID) != sessionID || sessionID == "" || sequence == 0 || observedAt.IsZero() { + return workerruntime.Report{}, ErrInvalidRuntimeReport + } + current := s.current.Load() + if current == nil { + return workerruntime.Report{}, ErrInvalidRuntimeReport + } + visible := make(map[string]struct{}, len(current.Entries)) + counters := make([]workerruntime.Counter, 0) + for _, entry := range current.Entries { + visible[entry.Proxy.ID] = struct{}{} + active, reserved, _ := entry.Runtime.Counters() + if active == 0 && reserved == 0 { + continue + } + counters = append(counters, workerruntime.Counter{ + ProxyID: entry.Proxy.ID, Active: active, Reserved: reserved, + Draining: entry.Proxy.State == proxyDomain.StateDraining, + }) + } + s.active.rangeEntries(func(proxyID string, runtime *proxyDomain.Capacity) { + if _, currentProxy := visible[proxyID]; currentProxy { + return + } + active, reserved, _ := runtime.Counters() + if active == 0 && reserved == 0 { + return + } + counters = append(counters, workerruntime.Counter{ + ProxyID: proxyID, Active: active, Reserved: reserved, Draining: true, + }) + }) + sort.Slice(counters, func(left, right int) bool { + return counters[left].ProxyID < counters[right].ProxyID + }) + return workerruntime.Report{ + WorkerID: s.workerID, SessionID: sessionID, Sequence: sequence, + SnapshotVersion: current.Version, OwnershipEpoch: current.Epoch, + ObservedAt: observedAt.UTC(), Counters: counters, + }, nil +} + func (s *Store) Apply(envelope Envelope) error { if s == nil { return fmt.Errorf("apply snapshot: nil store") @@ -118,19 +244,47 @@ func (s *Store) Apply(envelope Envelope) error { } proxies := cloneAndSort(envelope.Proxies) + newRuntimeCount := 0 + for _, descriptor := range proxies { + if s.runtimes[descriptor.ID] == nil { + newRuntimeCount++ + } + } + if len(s.runtimes)+newRuntimeCount > s.limit { + return ErrRuntimeLimitExceeded + } + if current != nil { + for _, entry := range current.Entries { + registration := s.runtimes[entry.Proxy.ID] + registration.capacity.SetActivityObservationEnabled(true) + registration.current.Store(false) + active, reserved, _ := registration.capacity.Counters() + s.active.track(entry.Proxy.ID, registration.capacity, active+reserved > 0) + } + } entries := make([]Entry, 0, len(proxies)) for _, descriptor := range proxies { - runtime := s.runtimes[descriptor.ID] - if runtime == nil { - runtime = proxyDomain.NewCapacity(descriptor.MaxConcurrency) + registration := s.runtimes[descriptor.ID] + if registration == nil { + proxyID := descriptor.ID + registration = &runtimeRegistration{} + registration.capacity = proxyDomain.NewCapacityWithActivityObserver(descriptor.MaxConcurrency, func(nonzero bool) { + if registration.current.Load() { + return + } + s.active.track(proxyID, registration.capacity, nonzero) + }) } else { - runtime.SetMax(descriptor.MaxConcurrency) + registration.capacity.SetMax(descriptor.MaxConcurrency) } + registration.current.Store(true) + registration.capacity.SetActivityObservationEnabled(false) + s.active.remove(descriptor.ID, registration.capacity) // 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}) + s.runtimes[descriptor.ID] = registration + entries = append(entries, Entry{Proxy: descriptor, Runtime: registration.capacity}) } next := &View{ diff --git a/internal/gateway/snapshot/store_test.go b/internal/gateway/snapshot/store_test.go index 2b6461a..bd6766d 100644 --- a/internal/gateway/snapshot/store_test.go +++ b/internal/gateway/snapshot/store_test.go @@ -8,6 +8,7 @@ import ( "time" proxyDomain "proxy-pool/internal/domain/proxy" + "proxy-pool/internal/domain/workerruntime" ) func TestStoreAppliesCompleteSnapshotsInOrder(t *testing.T) { @@ -266,6 +267,165 @@ func TestStoreReusesRuntimeWhenProxyDisappearsAndReappears(t *testing.T) { } } +func TestStoreRuntimeReportKeepsRemovedActiveProxyUntilRelease(t *testing.T) { + store := NewStore("cluster-a", "worker-a") + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + proxy := proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "provider-a", + State: proxyDomain.StateAvailable, MaxConcurrency: 1, + } + initial := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, + Proxies: []proxyDomain.Proxy{proxy}, + } + initial.Checksum = Checksum(initial.Proxies) + if err := store.Apply(initial); err != nil { + t.Fatalf("Apply(initial): %v", err) + } + runtime := store.Current().Entries[0].Runtime + reservation, ok := runtime.Reserve() + if !ok { + t.Fatal("Reserve() = false") + } + 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(nil) + if err := store.Apply(removed); err != nil { + t.Fatalf("Apply(remove): %v", err) + } + if got := store.activeRuntimeCount(); got != 1 { + t.Fatalf("activeRuntimeCount(removed) = %d, want 1", got) + } + + report, err := store.RuntimeReport("session-a", 7, now) + if err != nil { + t.Fatalf("RuntimeReport(): %v", err) + } + if report.WorkerID != "worker-a" || report.SessionID != "session-a" || report.Sequence != 7 || + report.SnapshotVersion != 2 || report.OwnershipEpoch != 1 || !report.ObservedAt.Equal(now) || + len(report.Counters) != 1 || report.Counters[0] != (workerruntime.Counter{ + ProxyID: "proxy-a", Active: 1, Draining: true, + }) { + t.Fatalf("RuntimeReport() = %+v", report) + } + if err := reservation.Release(); err != nil { + t.Fatalf("Release(): %v", err) + } + if got := store.activeRuntimeCount(); got != 0 { + t.Fatalf("activeRuntimeCount(released) = %d, want 0", got) + } + report, err = store.RuntimeReport("session-a", 8, now.Add(time.Second)) + if err != nil || len(report.Counters) != 0 { + t.Fatalf("RuntimeReport(after release) = %+v, %v", report, err) + } +} + +func TestStoreRuntimeReportMarksCurrentDrainingProxy(t *testing.T) { + store := NewStore("cluster-a", "worker-a") + now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC) + proxy := proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, + State: proxyDomain.StateDraining, MaxConcurrency: 1, + } + envelope := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, + Proxies: []proxyDomain.Proxy{proxy}, + } + envelope.Checksum = Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(): %v", err) + } + reservation, ok := store.Current().Entries[0].Runtime.Reserve() + if !ok { + t.Fatal("Reserve() = false") + } + report, err := store.RuntimeReport("session-a", 1, now) + if err != nil || len(report.Counters) != 1 || !report.Counters[0].Draining { + t.Fatalf("RuntimeReport() = %+v, %v", report, err) + } + if err := reservation.Cancel(); err != nil { + t.Fatalf("Cancel(): %v", err) + } +} + +func TestStoreBoundsHistoricalRuntimeRegistryAndKeepsApplyTransactional(t *testing.T) { + store, err := NewStoreWithRuntimeLimit("cluster-a", "worker-a", 1) + if err != nil { + t.Fatalf("NewStoreWithRuntimeLimit(): %v", err) + } + first := proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, + State: proxyDomain.StateAvailable, MaxConcurrency: 1, + } + envelope := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, + Full: true, Proxies: []proxyDomain.Proxy{first}, + } + envelope.Checksum = Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(first): %v", err) + } + removed := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true, + } + removed.Checksum = Checksum(nil) + if err := store.Apply(removed); err != nil { + t.Fatalf("Apply(removed): %v", err) + } + second := first + second.ID = "proxy-b" + overLimit := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 3, + Full: true, Proxies: []proxyDomain.Proxy{second}, + } + overLimit.Checksum = Checksum(overLimit.Proxies) + if err := store.Apply(overLimit); !errors.Is(err, ErrRuntimeLimitExceeded) { + t.Fatalf("Apply(over limit) error = %v, want ErrRuntimeLimitExceeded", err) + } + if current := store.Current(); current.Version != 2 || len(current.Entries) != 0 { + t.Fatalf("Current() after rejected apply = version %d entries %d", current.Version, len(current.Entries)) + } +} + +func TestStoreRuntimeActiveIndexDropsZeroCounters(t *testing.T) { + store := NewStore("cluster-a", "worker-a") + proxy := proxyDomain.Proxy{ + ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, + State: proxyDomain.StateAvailable, MaxConcurrency: 1, + } + envelope := Envelope{ + ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, + Full: true, Proxies: []proxyDomain.Proxy{proxy}, + } + envelope.Checksum = Checksum(envelope.Proxies) + if err := store.Apply(envelope); err != nil { + t.Fatalf("Apply(): %v", err) + } + reservation, ok := store.Current().Entries[0].Runtime.Reserve() + if !ok { + t.Fatal("Reserve() = false") + } + if got := store.activeRuntimeCount(); got != 0 { + t.Fatalf("activeRuntimeCount(current) = %d, want 0", got) + } + if err := reservation.Cancel(); err != nil { + t.Fatalf("Cancel(): %v", err) + } + if got := store.activeRuntimeCount(); got != 0 { + t.Fatalf("activeRuntimeCount() = %d, want 0", got) + } +} + +func (s *Store) activeRuntimeCount() int { + count := 0 + s.active.rangeEntries(func(_ string, _ *proxyDomain.Capacity) { + count++ + }) + return count +} + func collectSelectionIDs(selection Selection) []string { ids := make([]string, 0, selection.Len()) for index := 0; index < selection.Len(); index++ {