diff --git a/internal/adapters/redisactivity/codec.go b/internal/adapters/redisactivity/codec.go index 3ddbdf8..0ba5e4c 100644 --- a/internal/adapters/redisactivity/codec.go +++ b/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 } diff --git a/internal/adapters/redisactivity/keys.go b/internal/adapters/redisactivity/keys.go index c514bad..5b7d976 100644 --- a/internal/adapters/redisactivity/keys.go +++ b/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) } diff --git a/internal/adapters/redisactivity/ownership.go b/internal/adapters/redisactivity/ownership.go index 9cb350b..3bb6389 100644 --- a/internal/adapters/redisactivity/ownership.go +++ b/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 } diff --git a/internal/adapters/redisactivity/ownership_integration_test.go b/internal/adapters/redisactivity/ownership_integration_test.go index 0b82e35..35cc0f0 100644 --- a/internal/adapters/redisactivity/ownership_integration_test.go +++ b/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 { diff --git a/internal/adapters/redisactivity/scripts/extract.lua b/internal/adapters/redisactivity/scripts/extract.lua index 44c15e0..b4b4423 100644 --- a/internal/adapters/redisactivity/scripts/extract.lua +++ b/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 diff --git a/internal/adapters/redisactivity/scripts/health.lua b/internal/adapters/redisactivity/scripts/health.lua index 30bfe1a..f25e727 100644 --- a/internal/adapters/redisactivity/scripts/health.lua +++ b/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 diff --git a/internal/adapters/redisactivity/scripts/ownership.lua b/internal/adapters/redisactivity/scripts/ownership.lua index 2bd33ab..f5a279d 100644 --- a/internal/adapters/redisactivity/scripts/ownership.lua +++ b/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)}) diff --git a/internal/adapters/redisactivity/scripts/sweep.lua b/internal/adapters/redisactivity/scripts/sweep.lua index fabb9ed..8a23a0c 100644 --- a/internal/adapters/redisactivity/scripts/sweep.lua +++ b/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 diff --git a/internal/adapters/redisactivity/scripts/upsert.lua b/internal/adapters/redisactivity/scripts/upsert.lua index aa2d123..dd62a3c 100644 --- a/internal/adapters/redisactivity/scripts/upsert.lua +++ b/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