84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
 
|
"""Redis lease primitives for account ownership (refactor §2).
|
||
|
|
|
||
|
|
A sync worker must hold the lease ``sync:lease:{account_id}`` before it runs an
|
||
|
|
account's sync loops, and no two workers may hold the same lease. Claims are
|
||
|
|
atomic (``SET NX EX``); renewal and release are owner-checked Lua so a worker can
|
||
|
|
never renew or delete a lease that has been taken over by another worker.
|
||
|
|
|
||
|
|
The set ``sync:live`` is the registry of accounts that *should* be synced; the
|
||
|
|
claimer only ever claims ids drawn from it.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
import redis.asyncio as aioredis
|
||
|
|
|
||
|
|
from app.redis_client import aw
|
||
|
|
|
||
|
|
LIVE_SET = "sync:live"
|
||
|
|
_LEASE_PREFIX = "sync:lease:"
|
||
|
|
WORKERS_PREFIX = "sync:workers:"
|
||
|
|
|
||
|
|
|
||
|
|
def lease_key(account_id: uuid.UUID) -> str:
|
||
|
|
return f"{_LEASE_PREFIX}{account_id}"
|
||
|
|
|
||
|
|
|
||
|
|
def worker_key(worker_id: str) -> str:
|
||
|
|
return f"{WORKERS_PREFIX}{worker_id}"
|
||
|
|
|
||
|
|
|
||
|
|
# Renew only if we still own the lease (guards against a stale worker extending a
|
||
|
|
# lease another worker has already claimed after expiry). Returns 1 on success.
|
||
|
|
_RENEW_LUA = """
|
||
|
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||
|
|
else
|
||
|
|
return 0
|
||
|
|
end
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Delete only if we still own the lease. Returns 1 if we released it, else 0.
|
||
|
|
_RELEASE_LUA = """
|
||
|
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call('DEL', KEYS[1])
|
||
|
|
else
|
||
|
|
return 0
|
||
|
|
end
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
class LeaseScripts:
|
||
|
|
"""Server-registered Lua for owner-checked renew/release."""
|
||
|
|
|
||
|
|
def __init__(self, redis: aioredis.Redis) -> None:
|
||
|
|
self._renew = redis.register_script(_RENEW_LUA)
|
||
|
|
self._release = redis.register_script(_RELEASE_LUA)
|
||
|
|
|
||
|
|
async def renew(self, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool:
|
||
|
|
result = await self._renew(keys=[lease_key(account_id)], args=[worker_id, ttl])
|
||
|
|
return bool(result)
|
||
|
|
|
||
|
|
async def release(self, account_id: uuid.UUID, worker_id: str) -> bool:
|
||
|
|
result = await self._release(keys=[lease_key(account_id)], args=[worker_id])
|
||
|
|
return bool(result)
|
||
|
|
|
||
|
|
|
||
|
|
async def claim(redis: aioredis.Redis, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool:
|
||
|
|
"""Atomically claim the lease for ``account_id``; True iff we now own it."""
|
||
|
|
ok = await redis.set(lease_key(account_id), worker_id, nx=True, ex=ttl)
|
||
|
|
return bool(ok)
|
||
|
|
|
||
|
|
|
||
|
|
async def mark_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None:
|
||
|
|
"""Register an account as one that should be synced (idempotent)."""
|
||
|
|
await aw(redis.sadd(LIVE_SET, str(account_id)))
|
||
|
|
|
||
|
|
|
||
|
|
async def mark_not_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None:
|
||
|
|
"""Deregister an account and drop any lease so its owner stops it (fencing)."""
|
||
|
|
await aw(redis.srem(LIVE_SET, str(account_id)))
|
||
|
|
await aw(redis.delete(lease_key(account_id)))
|