feat: add authoritative worker runtime capacity

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

View File

@ -182,6 +182,7 @@ message ReportRuntimeRequest {
uint64 ownership_epoch = 4; uint64 ownership_epoch = 4;
repeated ProxyRuntime counters = 5; repeated ProxyRuntime counters = 5;
google.protobuf.Timestamp observed_at = 6; google.protobuf.Timestamp observed_at = 6;
uint64 report_sequence = 7;
} }
message ProxyRuntime { message ProxyRuntime {

View File

@ -19,11 +19,13 @@ var (
) )
type Options struct { type Options struct {
Namespace string Namespace string
Credentials credentials.Store Credentials credentials.Store
OperationTTL time.Duration OperationTTL time.Duration
MaxCandidateScan int MaxCandidateScan int
CleanupLimit int MaxRuntimeCounters int
MaxInventoryScan int
CleanupLimit int
} }
type Adapter struct { type Adapter struct {
@ -35,9 +37,16 @@ type Adapter struct {
func New(client redis.Scripter, options Options) (*Adapter, error) { func New(client redis.Scripter, options Options) (*Adapter, error) {
options.Namespace = strings.TrimSpace(options.Namespace) 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) || if nilInterface(client) || nilInterface(options.Credentials) ||
!namespacePattern.MatchString(options.Namespace) || options.OperationTTL <= 0 || !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 nil, ErrInvalidOptions
} }
return &Adapter{ return &Adapter{

View File

@ -58,6 +58,8 @@ func TestNewRejectsInvalidDependenciesAndOptions(t *testing.T) {
{name: "colon in namespace", client: client, options: withNamespace(valid, "tenant:other")}, {name: "colon in namespace", client: client, options: withNamespace(valid, "tenant:other")},
{name: "zero operation ttl", client: client, options: withOperationTTL(valid, 0)}, {name: "zero operation ttl", client: client, options: withOperationTTL(valid, 0)},
{name: "zero candidate scan", client: client, options: withMaxCandidateScan(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)}, {name: "negative cleanup limit", client: client, options: withCleanupLimit(valid, -1)},
} }
for _, tt := range tests { for _, tt := range tests {
@ -93,7 +95,9 @@ func TestNewBuildsClusterSafeKeyspaceAndHashesDynamicTokens(t *testing.T) {
adapter.keys.records, adapter.keys.unique, adapter.keys.idkeys, adapter.keys.records, adapter.keys.unique, adapter.keys.idkeys,
adapter.keys.expiry, adapter.keys.available, adapter.keys.owners, adapter.keys.expiry, adapter.keys.available, adapter.keys.owners,
adapter.keys.ownerExpiry, adapter.keys.epoch, adapter.keys.inventory, 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 { for _, key := range staticKeys {
if strings.Count(key, "{activity}") != 1 || strings.Count(key, "{") != 1 || strings.Count(key, "}") != 1 { 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.region(raw),
adapter.keys.carrier(raw), adapter.keys.carrier(raw),
adapter.keys.upstream(raw), adapter.keys.upstream(raw),
adapter.keys.owned(raw),
} }
for _, key := range dynamicKeys { for _, key := range dynamicKeys {
if strings.Contains(key, raw) || strings.Count(key, "{activity}") != 1 || strings.Count(key, "{") != 1 { 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, UsableUntilMS: 58_000, LastCheckedAtMS: 2_000, LastSuccessAtMS: 2_000,
LatencyNS: int64(25 * time.Millisecond), MaxConcurrency: 8, LatencyNS: int64(25 * time.Millisecond), MaxConcurrency: 8,
State: string(proxyDomain.StateAvailable), Tags: map[string]string{"region": "cn", "carrier": "ct"}, State: string(proxyDomain.StateAvailable), Tags: map[string]string{"region": "cn", "carrier": "ct"},
OwnerIndexKey: "pp:{activity}:test:owned:index",
} }
first, err := encodeProxyRecord(record) first, err := encodeProxyRecord(record)
if err != nil { if err != nil {
@ -260,6 +266,16 @@ func withMaxCandidateScan(options Options, limit int) Options {
return 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 { func withCleanupLimit(options Options, limit int) Options {
options.CleanupLimit = limit options.CleanupLimit = limit
return options return options

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
}

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

View File

@ -37,6 +37,7 @@ type proxyRecord struct {
State string `json:"state"` State string `json:"state"`
Tags map[string]string `json:"tags,omitempty"` Tags map[string]string `json:"tags,omitempty"`
OwnerWorkerID string `json:"ownerWorkerId,omitempty"` OwnerWorkerID string `json:"ownerWorkerId,omitempty"`
OwnerIndexKey string `json:"ownerIndexKey"`
IndexKeys []string `json:"indexKeys,omitempty"` IndexKeys []string `json:"indexKeys,omitempty"`
} }
@ -203,7 +204,8 @@ func validateProxyRecord(record proxyRecord) error {
record.CreatedAtMS <= 0 || record.ExpiresAtMS <= 0 || record.UsableUntilMS <= 0 || record.CreatedAtMS <= 0 || record.ExpiresAtMS <= 0 || record.UsableUntilMS <= 0 ||
record.UsableUntilMS > record.ExpiresAtMS || record.LastCheckedAtMS < 0 || record.UsableUntilMS > record.ExpiresAtMS || record.LastCheckedAtMS < 0 ||
record.LastSuccessAtMS < 0 || record.LatencyNS < 0 || record.MaxConcurrency < 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 return ErrInvalidRecord
} }
for _, key := range record.IndexKeys { for _, key := range record.IndexKeys {

View File

@ -9,33 +9,41 @@ import (
const redisKeyPrefix = "pp:{activity}:" const redisKeyPrefix = "pp:{activity}:"
type keyspace struct { type keyspace struct {
prefix string prefix string
records string records string
unique string unique string
idkeys string idkeys string
expiry string expiry string
available string available string
owners string owners string
ownerExpiry string ownerExpiry string
epoch string epoch string
inventory string inventory string
stateInventory string stateInventory string
workerSessions string
workerSessionExpiry string
workerRuntime string
workerRuntimeExpiry string
} }
func newKeyspace(namespace string) keyspace { func newKeyspace(namespace string) keyspace {
prefix := redisKeyPrefix + namespace prefix := redisKeyPrefix + namespace
return keyspace{ return keyspace{
prefix: prefix, prefix: prefix,
records: prefix + ":records", records: prefix + ":records",
unique: prefix + ":unique", unique: prefix + ":unique",
idkeys: prefix + ":idkeys", idkeys: prefix + ":idkeys",
expiry: prefix + ":expiry", expiry: prefix + ":expiry",
available: prefix + ":available", available: prefix + ":available",
owners: prefix + ":owners", owners: prefix + ":owners",
ownerExpiry: prefix + ":owner-expiry", ownerExpiry: prefix + ":owner-expiry",
epoch: prefix + ":epoch", epoch: prefix + ":epoch",
inventory: prefix + ":inventory", inventory: prefix + ":inventory",
stateInventory: prefix + ":state-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) 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 { func (keys keyspace) facet(name, value string) string {
return keys.prefix + ":" + name + ":" + digestToken(value) return keys.prefix + ":" + name + ":" + digestToken(value)
} }

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
}

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

View File

@ -74,6 +74,17 @@ type statusScriptInventory struct {
Extracted int64 `json:"extracted"` 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 //go:embed scripts/upsert.lua
var upsertSource string var upsertSource string
@ -92,6 +103,12 @@ var sweepSource string
//go:embed scripts/status.lua //go:embed scripts/status.lua
var statusSource string var statusSource string
//go:embed scripts/runtime.lua
var runtimeSource string
//go:embed scripts/capacity.lua
var capacitySource string
var ( var (
upsertScript = redis.NewScript(upsertSource) upsertScript = redis.NewScript(upsertSource)
healthScript = redis.NewScript(healthSource) healthScript = redis.NewScript(healthSource)
@ -99,6 +116,8 @@ var (
ownershipScript = redis.NewScript(ownershipSource) ownershipScript = redis.NewScript(ownershipSource)
sweepScript = redis.NewScript(sweepSource) sweepScript = redis.NewScript(sweepSource)
statusScript = redis.NewScript(statusSource) 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) { func runScript(ctx context.Context, client redis.Scripter, script *redis.Script, keys []string, args ...any) (any, error) {

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)

View File

@ -113,6 +113,12 @@ local function remove_available(proxy_id, record)
end end
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 function remove_proxy(proxy_id)
local raw = redis.call('HGET', records_key, proxy_id) local raw = redis.call('HGET', records_key, proxy_id)
local record = nil local record = nil
@ -120,6 +126,7 @@ local function remove_proxy(proxy_id)
local decoded local decoded
decoded, record = pcall(cjson.decode, raw) decoded, record = pcall(cjson.decode, raw)
remove_available(proxy_id, decoded and record or nil) 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 if decoded and type(record) == 'table' and is_managed(record.state) then
decrement_inventory(record.sourceUpstream) decrement_inventory(record.sourceUpstream)
end end

View File

@ -70,12 +70,19 @@ local function remove_available(id, record)
end end
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 function remove_proxy(id)
local raw = redis.call('HGET', records_key, id) local raw = redis.call('HGET', records_key, id)
local record = nil local record = nil
if raw then if raw then
record = cjson.decode(raw) record = cjson.decode(raw)
remove_available(id, record) remove_available(id, record)
remove_owned(id, record)
if is_managed(record.state) then if is_managed(record.state) then
decrement_inventory(record.sourceUpstream) decrement_inventory(record.sourceUpstream)
end end
@ -142,6 +149,9 @@ if not raw then
return finish({status = 'not_found'}) return finish({status = 'not_found'})
end end
local record = cjson.decode(raw) 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 if tonumber(record.expiresAtMs) <= checked_at_ms then
remove_proxy(proxy_id) remove_proxy(proxy_id)
return finish({status = 'not_found'}) return finish({status = 'not_found'})
@ -184,11 +194,17 @@ end
local encoded = cjson.encode(record) local encoded = cjson.encode(record)
redis.call('HSET', records_key, proxy_id, encoded) redis.call('HSET', records_key, proxy_id, encoded)
local owned = (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1 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)
redis.call('ZADD', available_key, record.usableUntilMs, proxy_id) if next_state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then
for _, index_key in ipairs(record.indexKeys or {}) do if owned then
redis.call('ZADD', index_key, record.usableUntilMs, proxy_id) redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
touch(index_key, tonumber(record.expiresAtMs)) 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
end end
touch(records_key, tonumber(record.expiresAtMs)) touch(records_key, tonumber(record.expiresAtMs))

View File

@ -94,6 +94,12 @@ local function remove_available(id, record)
end end
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 function add_available(id, record, at_ms)
local usable_until_ms = record and tonumber(record.usableUntilMs) 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 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 local decoded
decoded, record = pcall(cjson.decode, raw) decoded, record = pcall(cjson.decode, raw)
remove_available(id, decoded and record or nil) 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 if decoded and type(record) == 'table' and is_managed(record.state) then
decrement_inventory(record.sourceUpstream) decrement_inventory(record.sourceUpstream)
end end
@ -169,6 +176,7 @@ local function clear_owner(id, assignment, at_ms, restore)
local raw_record = redis.call('HGET', records_key, id) local raw_record = redis.call('HGET', records_key, id)
local record = decode_table(raw_record) local record = decode_table(raw_record)
if record and (not assignment or record.ownerWorkerId == assignment.workerId) then if record and (not assignment or record.ownerWorkerId == assignment.workerId) then
remove_owned(id, record)
record.ownerWorkerId = nil record.ownerWorkerId = nil
redis.call('HSET', records_key, id, cjson.encode(record)) redis.call('HSET', records_key, id, cjson.encode(record))
if restore then if restore then
@ -195,6 +203,7 @@ if operation == 'assign' then
local record = decode_table(redis.call('HGET', records_key, proxy_id)) local record = decode_table(redis.call('HGET', records_key, proxy_id))
if not record or record.state ~= 'AVAILABLE' or if not record or record.state ~= 'AVAILABLE' or
type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' or
(record.ownerWorkerId and record.ownerWorkerId ~= '') or (record.ownerWorkerId and record.ownerWorkerId ~= '') or
redis.call('HEXISTS', owners_key, proxy_id) == 1 or redis.call('HEXISTS', owners_key, proxy_id) == 1 or
not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
@ -223,6 +232,8 @@ if operation == 'assign' then
record.ownerWorkerId = worker_id record.ownerWorkerId = worker_id
redis.call('HSET', records_key, proxy_id, cjson.encode(record)) redis.call('HSET', records_key, proxy_id, cjson.encode(record))
remove_available(proxy_id, 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}) return finish({status = 'ok', record = encoded})
end end
@ -267,6 +278,8 @@ if operation == 'begin_drain' then
current.assignmentVersion = tonumber(current.assignmentVersion) + 1 current.assignmentVersion = tonumber(current.assignmentVersion) + 1
local encoded = cjson.encode(current) local encoded = cjson.encode(current)
redis.call('HSET', owners_key, proxy_id, encoded) 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}) return finish({status = 'ok', record = encoded})
end end
return finish({status = 'ok', record = cjson.encode(current)}) return finish({status = 'ok', record = cjson.encode(current)})

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

View File

@ -73,6 +73,12 @@ local function remove_available(proxy_id, record)
end end
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 function remove_proxy(proxy_id)
local raw = redis.call('HGET', records_key, proxy_id) local raw = redis.call('HGET', records_key, proxy_id)
local record = nil local record = nil
@ -80,6 +86,7 @@ local function remove_proxy(proxy_id)
local decoded local decoded
decoded, record = pcall(cjson.decode, raw) decoded, record = pcall(cjson.decode, raw)
remove_available(proxy_id, decoded and record or nil) 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 if decoded and type(record) == 'table' and is_managed(record.state) then
decrement_inventory(record.sourceUpstream) decrement_inventory(record.sourceUpstream)
end end

View File

@ -70,12 +70,19 @@ local function remove_available(proxy_id, record)
end end
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 function remove_proxy(proxy_id)
local raw = redis.call('HGET', records_key, proxy_id) local raw = redis.call('HGET', records_key, proxy_id)
local record = nil local record = nil
if raw then if raw then
record = cjson.decode(raw) record = cjson.decode(raw)
remove_available(proxy_id, record) remove_available(proxy_id, record)
remove_owned(proxy_id, record)
if is_managed(record.state) then if is_managed(record.state) then
decrement_inventory(record.sourceUpstream) decrement_inventory(record.sourceUpstream)
end end
@ -123,6 +130,16 @@ local function add_available(proxy_id, record)
end end
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 function finish(reply)
local encoded = cjson.encode(reply) local encoded = cjson.encode(reply)
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms) redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
@ -132,6 +149,11 @@ end
cleanup_expired() cleanup_expired()
for _, candidate in ipairs(candidates) do 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) local mapped = redis.call('HGET', idkeys_key, candidate.proxyId)
if mapped and mapped ~= candidate.uniqueDigest then if mapped and mapped ~= candidate.uniqueDigest then
return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0}) 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) local encoded = cjson.encode(incoming)
redis.call('HSET', records_key, incumbent_id, encoded) redis.call('HSET', records_key, incumbent_id, encoded)
redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id) redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id)
sync_owned(incumbent_id, incoming)
add_available(incumbent_id, incoming) add_available(incumbent_id, incoming)
if tonumber(incoming.expiresAtMs) > max_expiry_ms then if tonumber(incoming.expiresAtMs) > max_expiry_ms then
max_expiry_ms = tonumber(incoming.expiresAtMs) 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', unique_key, candidate.uniqueDigest, candidate.proxyId)
redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest) redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest)
redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId) redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId)
sync_owned(candidate.proxyId, incoming)
if is_managed(incoming.state) then if is_managed(incoming.state) then
redis.call('HINCRBY', inventory_key, candidate.upstream, 1) redis.call('HINCRBY', inventory_key, candidate.upstream, 1)
end end

View File

@ -139,7 +139,8 @@ func (a *Adapter) prepareUpsertCandidate(
CreatedAtMS: candidate.CreatedAt.UnixMilli(), ExpiresAtMS: expiresAt.UnixMilli(), CreatedAtMS: candidate.CreatedAt.UnixMilli(), ExpiresAtMS: expiresAt.UnixMilli(),
UsableUntilMS: usableUntil.UnixMilli(), LatencyNS: int64(candidate.Latency), UsableUntilMS: usableUntil.UnixMilli(), LatencyNS: int64(candidate.Latency),
MaxConcurrency: candidate.MaxConcurrency, State: string(candidate.State), 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 { if candidate.LastCheckedAt != nil {
record.LastCheckedAtMS = candidate.LastCheckedAt.UnixMilli() record.LastCheckedAtMS = candidate.LastCheckedAt.UnixMilli()

View File

@ -100,11 +100,13 @@ func (*productionInfrastructure) Open(
return ports{}, err return ports{}, err
} }
adapter, err := redisactivity.New(redisClient, redisactivity.Options{ adapter, err := redisactivity.New(redisClient, redisactivity.Options{
Namespace: redisNamespace, Namespace: redisNamespace,
Credentials: credentialStore, Credentials: credentialStore,
OperationTTL: redisOperationTTL, OperationTTL: redisOperationTTL,
MaxCandidateScan: candidateScan(configuration), MaxCandidateScan: candidateScan(configuration),
CleanupLimit: redisCleanupLimit, MaxRuntimeCounters: credentialCapacity(configuration),
MaxInventoryScan: credentialCapacity(configuration),
CleanupLimit: redisCleanupLimit,
}) })
if err != nil { if err != nil {
return ports{}, err return ports{}, err

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

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 // Reconcile centralizes the cold-path decision. The notifier may coalesce many
// calls; Provider Reconciler atomically reserves the budget before doing I/O. // calls; Provider Reconciler atomically reserves the budget before doing I/O.
func (r *Reconciler) Reconcile(now time.Time, inventory upstream.Inventory) ReconcileDecision { 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() 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) pendingSlots := saturatingMultiply(int64(usage.PendingExpected), r.slotsPerProxy)
decision := ReconcileDecision{ decision := ReconcileDecision{
AvailableSlots: availableSlots, AvailableSlots: inventory.AvailableSlots,
PendingExpected: usage.PendingExpected, PendingExpected: usage.PendingExpected,
FetchedTotal: usage.FetchedTotal, FetchedTotal: usage.FetchedTotal,
FetchAllowance: r.budget.FetchAllowance(), FetchAllowance: r.budget.FetchAllowance(),
EffectiveSlots: saturatingAdd(availableSlots, pendingSlots), EffectiveSlots: saturatingAdd(inventory.AvailableSlots, pendingSlots),
} }
r.mu.Lock() r.mu.Lock()
if r.refilling { if r.refilling {

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 { func inventoryWithSlots(now time.Time, slots int64) upstream.Inventory {
return upstream.Inventory{Proxies: []upstream.ProxyCapacity{{ return upstream.Inventory{Proxies: []upstream.ProxyCapacity{{
State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: slots, State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: slots,

View File

@ -13,19 +13,46 @@ var (
) )
type Capacity struct { type Capacity struct {
max atomic.Uint32 max atomic.Uint32
counters atomic.Uint64 counters atomic.Uint64
configuredObserver *activityObserver
observer atomic.Pointer[activityObserver]
} }
type activityObserver struct{ notify func(bool) }
func NewCapacity(max int64) *Capacity { 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{} capacity := &Capacity{}
if max < 0 || max > int64(counterMask) { if max < 0 || max > int64(counterMask) {
max = 0 max = 0
} }
capacity.max.Store(uint32(max)) capacity.max.Store(uint32(max))
if observer != nil {
capacity.configuredObserver = &activityObserver{notify: observer}
capacity.observer.Store(capacity.configuredObserver)
}
return capacity 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 { func (c *Capacity) SetMax(max int64) bool {
if max < 0 || max > int64(counterMask) { if max < 0 || max > int64(counterMask) {
return false 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) 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) { func (c *Capacity) Reserve() (*Reservation, bool) {
for { for {
current := c.counters.Load() current := c.counters.Load()
@ -45,6 +81,11 @@ func (c *Capacity) Reserve() (*Reservation, bool) {
} }
next := pack(active, reserved+1) next := pack(active, reserved+1)
if c.counters.CompareAndSwap(current, next) { 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 return &Reservation{capacity: c}, true
} }
} }
@ -77,7 +118,16 @@ func (c *Capacity) cancel() {
for { for {
current := c.counters.Load() current := c.counters.Load()
active, reserved := unpack(current) 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 return
} }
} }
@ -87,7 +137,16 @@ func (c *Capacity) release() {
for { for {
current := c.counters.Load() current := c.counters.Load()
active, reserved := unpack(current) 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 return
} }
} }

View File

@ -2,6 +2,7 @@ package proxy
import ( import (
"errors" "errors"
"reflect"
"sync" "sync"
"sync/atomic" "sync/atomic"
"testing" "testing"
@ -38,6 +39,56 @@ func TestReservationCancelReleasesReservedCapacity(t *testing.T) {
assertCapacityCounters(t, capacity, 0, 0) 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) { func TestReservationCommitAndReleaseAreSingleUse(t *testing.T) {
capacity := NewCapacity(1) capacity := NewCapacity(1)
reservation, ok := capacity.Reserve() reservation, ok := capacity.Reserve()

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
}

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

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

View File

@ -7,19 +7,90 @@ import (
"errors" "errors"
"fmt" "fmt"
"sort" "sort"
"strings"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
proxyDomain "proxy-pool/internal/domain/proxy" proxyDomain "proxy-pool/internal/domain/proxy"
"proxy-pool/internal/domain/workerruntime"
) )
var ( var (
ErrWrongTarget = errors.New("snapshot targets another cluster or worker") ErrWrongTarget = errors.New("snapshot targets another cluster or worker")
ErrResyncRequired = errors.New("snapshot sequence requires a full resync") ErrResyncRequired = errors.New("snapshot sequence requires a full resync")
ErrChecksumMismatch = errors.New("snapshot checksum mismatch") 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 { type Envelope struct {
ClusterID string ClusterID string
WorkerID string WorkerID string
@ -70,17 +141,29 @@ type Store struct {
current atomic.Pointer[View] current atomic.Pointer[View]
mu sync.Mutex mu sync.Mutex
runtimes map[string]*proxyDomain.Capacity runtimes map[string]*runtimeRegistration
active activeRuntimeIndex
limit int
} }
func NewStore(clusterID, workerID string) *Store { func NewStore(clusterID, workerID string) *Store {
return &Store{ return &Store{
clusterID: clusterID, clusterID: clusterID,
workerID: workerID, 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 { func (s *Store) Current() *View {
if s == nil { if s == nil {
return nil return nil
@ -88,6 +171,49 @@ func (s *Store) Current() *View {
return s.current.Load() 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 { func (s *Store) Apply(envelope Envelope) error {
if s == nil { if s == nil {
return fmt.Errorf("apply snapshot: nil store") return fmt.Errorf("apply snapshot: nil store")
@ -118,19 +244,47 @@ func (s *Store) Apply(envelope Envelope) error {
} }
proxies := cloneAndSort(envelope.Proxies) 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)) entries := make([]Entry, 0, len(proxies))
for _, descriptor := range proxies { for _, descriptor := range proxies {
runtime := s.runtimes[descriptor.ID] registration := s.runtimes[descriptor.ID]
if runtime == nil { if registration == nil {
runtime = proxyDomain.NewCapacity(descriptor.MaxConcurrency) 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 { } 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 // Keep runtimes for temporarily absent IDs. Old immutable views may still
// hold in-flight leases, so reclaiming here could reset active capacity if // hold in-flight leases, so reclaiming here could reset active capacity if
// the same Proxy reappears in a later snapshot. // the same Proxy reappears in a later snapshot.
s.runtimes[descriptor.ID] = runtime s.runtimes[descriptor.ID] = registration
entries = append(entries, Entry{Proxy: descriptor, Runtime: runtime}) entries = append(entries, Entry{Proxy: descriptor, Runtime: registration.capacity})
} }
next := &View{ next := &View{

View File

@ -8,6 +8,7 @@ import (
"time" "time"
proxyDomain "proxy-pool/internal/domain/proxy" proxyDomain "proxy-pool/internal/domain/proxy"
"proxy-pool/internal/domain/workerruntime"
) )
func TestStoreAppliesCompleteSnapshotsInOrder(t *testing.T) { 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 { func collectSelectionIDs(selection Selection) []string {
ids := make([]string, 0, selection.Len()) ids := make([]string, 0, selection.Len())
for index := 0; index < selection.Len(); index++ { for index := 0; index < selection.Len(); index++ {