feat: index active proxies by worker

internal/adapters/redisactivity/codec.go

@@ -45,6 +45,7 @@ type ownershipRecord struct {
 	Version           int    `json:"version"`
 	ProxyID           string `json:"proxyId"`
 	WorkerID          string `json:"workerId"`
+	WorkerIndexKey    string `json:"workerIndexKey,omitempty"`
 	Epoch             uint64 `json:"epoch"`
 	AssignmentVersion uint64 `json:"assignmentVersion"`
 	ExpiresAtMS       int64  `json:"expiresAtMs"`
@@ -221,6 +222,9 @@ func validateOwnershipRecord(record ownershipRecord) error {
 		record.Epoch == 0 || record.AssignmentVersion == 0 || record.ExpiresAtMS <= 0 {
 		return ErrInvalidRecord
 	}
+	if record.WorkerIndexKey != "" && !strings.Contains(record.WorkerIndexKey, "{activity}") {
+		return ErrInvalidRecord
+	}
 	return nil
 }
 

internal/adapters/redisactivity/keys.go

@@ -83,6 +83,10 @@ func (keys keyspace) owned(value string) string {
 	return keys.facet("owned", value)
 }
 
+func (keys keyspace) workerOwned(workerID string) string {
+	return keys.facet("worker-owned", workerID)
+}
+
 func (keys keyspace) facet(name, value string) string {
 	return keys.prefix + ":" + name + ":" + digestToken(value)
 }

internal/adapters/redisactivity/ownership.go

@@ -226,7 +226,7 @@ func (a *Adapter) runOwnership(
 		a.keys.inventory, a.keys.stateInventory, a.keys.owners, a.keys.ownerExpiry,
 		a.keys.epoch, operationKey,
 	}, operation, operationTTLMillis(a.options.OperationTTL), a.options.CleanupLimit,
-		nowMS, proxyID, workerID, epoch, value, active, reserved)
+		nowMS, proxyID, workerID, epoch, value, active, reserved, a.keys.workerOwned(workerID))
 	if err != nil {
 		return ownershipScriptReply{}, err
 	}

internal/adapters/redisactivity/ownership_integration_test.go

@@ -36,6 +36,7 @@ func TestRedisOwnershipLifecycle(t *testing.T) {
 	if current, ok, err := fixture.Adapter.Get(context.Background(), "proxy-a"); err != nil || !ok || current != assigned {
 		t.Fatalf("Get(assigned) = %+v, %t, %v", current, ok, err)
 	}
+	assertRedisWorkerOwned(t, fixture, "worker-a", "proxy-a")
 	blocked, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
 		RequestID: "req-owned", ClientID: "client-a", Requested: 1,
 		Fulfillment: extractionDomain.Partial, Now: now.Add(3 * time.Second),
@@ -50,6 +51,9 @@ func TestRedisOwnershipLifecycle(t *testing.T) {
 		!renewed.ExpiresAt.Equal(now.Add(2*time.Minute)) {
 		t.Fatalf("Renew() = %+v, %v", renewed, err)
 	}
+	if score, err := fixture.Client.ZScore(context.Background(), fixture.Adapter.keys.workerOwned("worker-a"), "proxy-a").Result(); err != nil || int64(score) != renewed.ExpiresAt.UnixMilli() {
+		t.Fatalf("worker-owned score = %f, %v; want %d", score, err, renewed.ExpiresAt.UnixMilli())
+	}
 	if _, err := fixture.Adapter.Renew(context.Background(), now.Add(31*time.Second),
 		"proxy-a", "worker-a", assigned.Epoch+1, time.Minute); !errors.Is(err, ownershipDomain.ErrStaleAssignment) {
 		t.Fatalf("Renew(stale epoch) error = %v", err)
@@ -59,6 +63,7 @@ func TestRedisOwnershipLifecycle(t *testing.T) {
 	if err != nil || !draining.Draining || draining.Version != renewed.Version+1 {
 		t.Fatalf("BeginDrain() = %+v, %v", draining, err)
 	}
+	assertRedisWorkerOwned(t, fixture, "worker-a")
 	replayed, err := fixture.Adapter.BeginDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch)
 	if err != nil || replayed != draining {
 		t.Fatalf("BeginDrain(replay) = %+v, %v", replayed, err)
@@ -72,6 +77,7 @@ func TestRedisOwnershipLifecycle(t *testing.T) {
 	if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 0, 0); err != nil {
 		t.Fatalf("AcknowledgeDrain(): %v", err)
 	}
+	assertRedisWorkerOwned(t, fixture, "worker-a")
 	assertRedisKeysHaveTTL(t, fixture,
 		fixture.Adapter.keys.available,
 		fixture.Adapter.keys.protocol("http"),
@@ -355,6 +361,7 @@ func TestRedisSweepExpiredIsLimitedAndCleansOwnership(t *testing.T) {
 		ownedFixture.Adapter.keys.region("cn"),
 		ownedFixture.Adapter.keys.carrier("ct"),
 		ownedFixture.Adapter.keys.upstream("provider-a"),
+		ownedFixture.Adapter.keys.workerOwned("worker-a"),
 	} {
 		if count, err := ownedFixture.Client.ZCard(context.Background(), sortedSet).Result(); err != nil || count != 0 {
 			t.Fatalf("sorted set %s entries after sweep = %d, %v", sortedSet, count, err)
@@ -399,6 +406,22 @@ func assertRedisInventory(t *testing.T, adapter *Adapter, upstreamID string, now
 	}
 }
 
+func assertRedisWorkerOwned(t *testing.T, fixture redisTestFixture, workerID string, want ...string) {
+	t.Helper()
+	actual, err := fixture.Client.ZRange(context.Background(), fixture.Adapter.keys.workerOwned(workerID), 0, -1).Result()
+	if err != nil {
+		t.Fatalf("ZRange(worker-owned %s): %v", workerID, err)
+	}
+	if len(actual) != len(want) {
+		t.Fatalf("worker-owned %s = %v, want %v", workerID, actual, want)
+	}
+	for index, proxyID := range want {
+		if actual[index] != proxyID {
+			t.Fatalf("worker-owned %s = %v, want %v", workerID, actual, want)
+		}
+	}
+}
+
 func assertRedisKeysHaveTTL(t *testing.T, fixture redisTestFixture, keys ...string) {
 	t.Helper()
 	for _, key := range keys {

internal/adapters/redisactivity/scripts/extract.lua

@@ -119,6 +119,18 @@ local function remove_owned(proxy_id, record)
     end
 end
 
+local function remove_worker_owned(proxy_id)
+    local raw = redis.call('HGET', owners_key, proxy_id)
+    if not raw then
+        return
+    end
+    local decoded, assignment = pcall(cjson.decode, raw)
+    if decoded and type(assignment) == 'table' and
+        type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
+        redis.call('ZREM', assignment.workerIndexKey, proxy_id)
+    end
+end
+
 local function remove_proxy(proxy_id)
     local raw = redis.call('HGET', records_key, proxy_id)
     local record = nil
@@ -143,6 +155,7 @@ local function remove_proxy(proxy_id)
     redis.call('HDEL', idkeys_key, proxy_id)
     redis.call('HDEL', records_key, proxy_id)
     redis.call('ZREM', expiry_key, proxy_id)
+    remove_worker_owned(proxy_id)
     redis.call('HDEL', owners_key, proxy_id)
     redis.call('ZREM', owner_expiry_key, proxy_id)
 end

internal/adapters/redisactivity/scripts/health.lua

@@ -76,6 +76,18 @@ local function remove_owned(id, record)
     end
 end
 
+local function remove_worker_owned(id)
+    local raw = redis.call('HGET', owners_key, id)
+    if not raw then
+        return
+    end
+    local decoded, assignment = pcall(cjson.decode, raw)
+    if decoded and type(assignment) == 'table' and
+        type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
+        redis.call('ZREM', assignment.workerIndexKey, id)
+    end
+end
+
 local function remove_proxy(id)
     local raw = redis.call('HGET', records_key, id)
     local record = nil
@@ -97,6 +109,7 @@ local function remove_proxy(id)
     redis.call('HDEL', idkeys_key, id)
     redis.call('HDEL', records_key, id)
     redis.call('ZREM', expiry_key, id)
+    remove_worker_owned(id)
     redis.call('HDEL', owners_key, id)
     redis.call('ZREM', owner_expiry_key, id)
 end

internal/adapters/redisactivity/scripts/ownership.lua

@@ -20,6 +20,7 @@ local epoch = tonumber(ARGV[7])
 local value = tonumber(ARGV[8])
 local active = tonumber(ARGV[9])
 local reserved = tonumber(ARGV[10])
+local worker_index_key = ARGV[11]
 local mutating = operation ~= 'get'
 
 local function finish(reply)
@@ -185,6 +186,9 @@ local function clear_owner(id, assignment, at_ms, restore)
     end
     redis.call('HDEL', owners_key, id)
     redis.call('ZREM', owner_expiry_key, id)
+    if assignment and type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
+        redis.call('ZREM', assignment.workerIndexKey, id)
+    end
 end
 
 if operation == 'assign' then
@@ -209,6 +213,10 @@ if operation == 'assign' then
         not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
         return finish({status = 'unavailable'})
     end
+    if type(worker_index_key) ~= 'string' or worker_index_key == '' or
+        not string.find(worker_index_key, '{activity}', 1, true) then
+        return finish({status = 'invalid'})
+    end
     local expires_at_ms = now_ms + value
     if tonumber(record.usableUntilMs) < expires_at_ms then
         expires_at_ms = tonumber(record.usableUntilMs)
@@ -219,6 +227,7 @@ if operation == 'assign' then
         version = 1,
         proxyId = proxy_id,
         workerId = worker_id,
+        workerIndexKey = worker_index_key,
         epoch = next_epoch,
         assignmentVersion = 1,
         expiresAtMs = expires_at_ms,
@@ -234,6 +243,8 @@ if operation == 'assign' then
     remove_available(proxy_id, record)
     redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
     touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
+    redis.call('ZADD', worker_index_key, expires_at_ms, proxy_id)
+    touch(worker_index_key, tonumber(record.expiresAtMs))
     return finish({status = 'ok', record = encoded})
 end
 
@@ -264,6 +275,10 @@ if operation == 'renew' then
     redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
     touch(owners_key, tonumber(record.expiresAtMs))
     touch(owner_expiry_key, tonumber(record.expiresAtMs))
+    if not current.draining and type(current.workerIndexKey) == 'string' and current.workerIndexKey ~= '' then
+        redis.call('ZADD', current.workerIndexKey, expires_at_ms, proxy_id)
+        touch(current.workerIndexKey, tonumber(record.expiresAtMs))
+    end
     return finish({status = 'ok', record = encoded})
 end
 
@@ -279,6 +294,9 @@ if operation == 'begin_drain' then
         redis.call('HSET', owners_key, proxy_id, encoded)
         local record = decode_table(redis.call('HGET', records_key, proxy_id))
         remove_owned(proxy_id, record)
+        if type(current.workerIndexKey) == 'string' and current.workerIndexKey ~= '' then
+            redis.call('ZREM', current.workerIndexKey, proxy_id)
+        end
         return finish({status = 'ok', record = encoded})
     end
     return finish({status = 'ok', record = cjson.encode(current)})

internal/adapters/redisactivity/scripts/sweep.lua

@@ -79,6 +79,18 @@ local function remove_owned(proxy_id, record)
     end
 end
 
+local function remove_worker_owned(proxy_id)
+    local raw = redis.call('HGET', owners_key, proxy_id)
+    if not raw then
+        return
+    end
+    local decoded, assignment = pcall(cjson.decode, raw)
+    if decoded and type(assignment) == 'table' and
+        type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
+        redis.call('ZREM', assignment.workerIndexKey, proxy_id)
+    end
+end
+
 local function remove_proxy(proxy_id)
     local raw = redis.call('HGET', records_key, proxy_id)
     local record = nil
@@ -103,6 +115,7 @@ local function remove_proxy(proxy_id)
     redis.call('HDEL', idkeys_key, proxy_id)
     redis.call('HDEL', records_key, proxy_id)
     redis.call('ZREM', expiry_key, proxy_id)
+    remove_worker_owned(proxy_id)
     redis.call('HDEL', owners_key, proxy_id)
     redis.call('ZREM', owner_expiry_key, proxy_id)
 end

internal/adapters/redisactivity/scripts/upsert.lua

@@ -76,6 +76,18 @@ local function remove_owned(proxy_id, record)
     end
 end
 
+local function remove_worker_owned(proxy_id)
+    local raw = redis.call('HGET', owners_key, proxy_id)
+    if not raw then
+        return
+    end
+    local decoded, assignment = pcall(cjson.decode, raw)
+    if decoded and type(assignment) == 'table' and
+        type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
+        redis.call('ZREM', assignment.workerIndexKey, proxy_id)
+    end
+end
+
 local function remove_proxy(proxy_id)
     local raw = redis.call('HGET', records_key, proxy_id)
     local record = nil
@@ -97,6 +109,7 @@ local function remove_proxy(proxy_id)
     redis.call('HDEL', idkeys_key, proxy_id)
     redis.call('HDEL', records_key, proxy_id)
     redis.call('ZREM', expiry_key, proxy_id)
+    remove_worker_owned(proxy_id)
     redis.call('HDEL', owners_key, proxy_id)
     redis.call('ZREM', owner_expiry_key, proxy_id)
 end