youfak/proxy-pool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add authoritative worker runtime capacity

Browse Source
This commit is contained in:
youfak 2026-07-30 15:46:18 +08:00
parent e9945d933f
commit 45ba6fb958
29 changed files with 2078 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/proto/controlplane/v1/controlplane.proto
View File

@ -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 {

11
internal/adapters/redisactivity/adapter.go
View File

@ -23,6 +23,8 @@ type Options struct {
Credentials credentials.Store
OperationTTL time.Duration
MaxCandidateScan int
MaxRuntimeCounters int
MaxInventoryScan int
CleanupLimit int
}
@ -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{

18
internal/adapters/redisactivity/adapter_test.go
View File

@ -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

45
internal/adapters/redisactivity/capacity.go Normal file
View File

@ -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
}

105
internal/adapters/redisactivity/capacity_integration_test.go Normal file
View File

@ -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)
}
}

4
internal/adapters/redisactivity/codec.go
View File

@ -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 {

12
internal/adapters/redisactivity/keys.go
View File

@ -20,6 +20,10 @@ type keyspace struct {
epoch string
inventory string
stateInventory string
workerSessions string
workerSessionExpiry string
workerRuntime string
workerRuntimeExpiry string
}
func newKeyspace(namespace string) keyspace {
@ -36,6 +40,10 @@ func newKeyspace(namespace string) keyspace {
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)
}

248
internal/adapters/redisactivity/runtime.go Normal file
View File

@ -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
}

147
internal/adapters/redisactivity/runtime_integration_test.go Normal file
View File

@ -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)
}
}

19
internal/adapters/redisactivity/scripts.go
View File

@ -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) {

177
internal/adapters/redisactivity/scripts/capacity.lua Normal file
View File

@ -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)

7
internal/adapters/redisactivity/scripts/extract.lua
View File

@ -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

18
internal/adapters/redisactivity/scripts/health.lua
View File

@ -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,12 +194,18 @@ 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
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))
touch(unique_key, tonumber(record.expiresAtMs))

13
internal/adapters/redisactivity/scripts/ownership.lua
View File

@ -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)})

229
internal/adapters/redisactivity/scripts/runtime.lua Normal file
View File

@ -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')

7
internal/adapters/redisactivity/scripts/sweep.lua
View File

@ -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

24
internal/adapters/redisactivity/scripts/upsert.lua
View File

@ -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

3
internal/adapters/redisactivity/upsert.go
View File

@ -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()

2
internal/controller/bootstrap/infrastructure.go
View File

@ -104,6 +104,8 @@ func (*productionInfrastructure) Open(
Credentials: credentialStore,
OperationTTL: redisOperationTTL,
MaxCandidateScan: candidateScan(configuration),
MaxRuntimeCounters: credentialCapacity(configuration),
MaxInventoryScan: credentialCapacity(configuration),
CleanupLimit: redisCleanupLimit,
})
if err != nil {

17
internal/controller/pool/inventory.go Normal file
View File

@ -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)
}

20
internal/controller/pool/reconciler.go
View File

@ -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 {

47
internal/controller/pool/reconciler_test.go
View File

@ -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,

63
internal/domain/proxy/capacity.go
View File

@ -15,17 +15,44 @@ var (
type Capacity struct {
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
}
}

51
internal/domain/proxy/capacity_test.go
View File

@ -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()

211
internal/domain/workerruntime/memory.go Normal file
View File

@ -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
}

156
internal/domain/workerruntime/memory_test.go Normal file
View File

@ -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)
}
}

70
internal/domain/workerruntime/runtime.go Normal file
View File

@ -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)
}

172
internal/gateway/snapshot/store.go
View File

@ -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")
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)
} else {
runtime.SetMax(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 {
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{

160
internal/gateway/snapshot/store_test.go
View File

@ -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++ {