Coverage for agentos/core/distributed_lock.py: 0%

1"""AgentOS Distributed Lock — production-grade distributed mutex. 

2 

3Backends: 

4- InMemoryLock: single-process (test/local dev) 

5- PostgresLock: advisory lock via pg_advisory_lock 

6- RedisLock: Redlock-inspired with TTL + renew 

7 

8Design: ~340 lines, async-first, context-manager compatible. 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14import logging 

15import time 

16import uuid 

17from abc import ABC, abstractmethod 

18from collections.abc import AsyncIterator 

19from contextlib import asynccontextmanager 

20from dataclasses import dataclass, field 

21from enum import StrEnum 

22from typing import Any 

23 

24logger = logging.getLogger(__name__) 

25 

26 

27# ============================================================================ 

28# Data types 

29# ============================================================================ 

30 

31 

32class LockBackend(StrEnum): 

33 IN_MEMORY = "in_memory" 

34 POSTGRES = "postgres" 

35 REDIS = "redis" 

36 

37 

38@dataclass 

39class LockConfig: 

40 """Configuration for distributed lock acquisition.""" 

41 

42 ttl: float = 30.0 # Seconds until lock auto-expires 

43 retry_interval: float = 0.1 # Polling interval when waiting 

44 acquire_timeout: float = 10.0 # Max time to wait for lock 

45 renew_interval: float = 0.0 # Auto-renew interval (0 = disabled) 

46 

47 

48@dataclass 

49class LockToken: 

50 """Token representing a held lock — required to release.""" 

51 

52 key: str 

53 owner_id: str 

54 acquired_at: float 

55 ttl: float 

56 _backend: Any = field(repr=False) # Backend reference for release 

57 

58 

59class LockAcquireError(Exception): 

60 """Failed to acquire lock within timeout.""" 

61 

62 

63class LockNotHeldError(Exception): 

64 """Attempted to release a lock not held by this owner.""" 

65 

66 

67# ============================================================================ 

68# Abstract backend 

69# ============================================================================ 

70 

71 

72class AbstractLockBackend(ABC): 

73 """Interface all lock backends must implement.""" 

74 

75 @abstractmethod 

76 async def acquire(self, key: str, owner_id: str, ttl: float) -> bool: 

77 """Try to acquire lock. Returns True on success.""" 

78 

79 @abstractmethod 

80 async def release(self, key: str, owner_id: str) -> bool: 

81 """Release lock. Returns True if this owner held it.""" 

82 

83 @abstractmethod 

84 async def extend(self, key: str, owner_id: str, ttl: float) -> bool: 

85 """Extend TTL. Returns True if this owner still holds it.""" 

86 

87 @abstractmethod 

88 async def is_held(self, key: str, owner_id: str) -> bool: 

89 """Check if this owner holds the lock.""" 

90 

91 

92# ============================================================================ 

93# In-Memory backend 

94# ============================================================================ 

95 

96 

97class InMemoryLockBackend(AbstractLockBackend): 

98 """Single-process in-memory lock — for testing and single-worker scenarios.""" 

99 

100 def __init__(self): 

101 self._locks: dict[str, tuple[str, float]] = {} # key → (owner_id, expiry) 

102 self._lock = asyncio.Lock() 

103 

104 async def acquire(self, key: str, owner_id: str, ttl: float) -> bool: 

105 async with self._lock: 

106 now = time.monotonic() 

107 if key in self._locks: 

108 owner, expiry = self._locks[key] 

109 if expiry > now and owner != owner_id: 

110 return False 

111 self._locks[key] = (owner_id, now + ttl) 

112 return True 

113 

114 async def release(self, key: str, owner_id: str) -> bool: 

115 async with self._lock: 

116 if key not in self._locks: 

117 return False 

118 owner, expiry = self._locks[key] 

119 if owner != owner_id: 

120 return False 

121 del self._locks[key] 

122 return True 

123 

124 async def extend(self, key: str, owner_id: str, ttl: float) -> bool: 

125 async with self._lock: 

126 if key not in self._locks: 

127 return False 

128 owner, expiry = self._locks[key] 

129 if owner != owner_id: 

130 return False 

131 if expiry < time.monotonic(): 

132 del self._locks[key] 

133 return False 

134 self._locks[key] = (owner_id, time.monotonic() + ttl) 

135 return True 

136 

137 async def is_held(self, key: str, owner_id: str) -> bool: 

138 async with self._lock: 

139 if key not in self._locks: 

140 return False 

141 owner, expiry = self._locks[key] 

142 return owner == owner_id and expiry > time.monotonic() 

143 

144 

145# ============================================================================ 

146# Postgres advisory lock backend 

147# ============================================================================ 

148 

149ADVISORY_LOCK_SQL = """ 

150SELECT pg_try_advisory_lock(%s) AS acquired; 

151""" 

152 

153ADVISORY_UNLOCK_SQL = """ 

154SELECT pg_advisory_unlock(%s) AS released; 

155""" 

156 

157 

158# Hash key to an int64 for advisory lock 

159def _key_to_int64(key: str) -> int: 

160 import hashlib 

161 

162 return int(hashlib.sha256(key.encode()).hexdigest()[:16], 16) % (2**63 - 1) 

163 

164 

165class PostgresLockBackend(AbstractLockBackend): 

166 """PostgreSQL advisory-lock based backend. 

167 

168 Uses pg_try_advisory_lock for non-blocking acquire. 

169 Client must provide a pool (asyncpg or similar) via _pool attribute. 

170 """ 

171 

172 def __init__(self, pool: Any = None): 

173 self._pool = pool 

174 self._acquired: set[tuple[str, str]] = set() # (key, owner_id) tracking 

175 

176 async def _get_conn(self): 

177 if self._pool is None: 

178 raise RuntimeError("PostgresLockBackend requires a pool") 

179 return await self._pool.acquire() 

180 

181 async def acquire(self, key: str, owner_id: str, ttl: float) -> bool: 

182 conn = await self._get_conn() 

183 try: 

184 lock_id = _key_to_int64(key) 

185 result = await conn.fetchval("SELECT pg_try_advisory_lock($1) AS acquired;", lock_id) 

186 if result: 

187 self._acquired.add((key, owner_id)) 

188 return True 

189 return False 

190 finally: 

191 await self._pool.release(conn) 

192 

193 async def release(self, key: str, owner_id: str) -> bool: 

194 if (key, owner_id) not in self._acquired: 

195 return False 

196 conn = await self._get_conn() 

197 try: 

198 lock_id = _key_to_int64(key) 

199 result = await conn.fetchval("SELECT pg_advisory_unlock($1) AS released;", lock_id) 

200 if result: 

201 self._acquired.discard((key, owner_id)) 

202 return bool(result) 

203 finally: 

204 await self._pool.release(conn) 

205 

206 async def extend(self, key: str, owner_id: str, ttl: float) -> bool: 

207 # Advisory locks don't expire — always held until released 

208 return (key, owner_id) in self._acquired 

209 

210 async def is_held(self, key: str, owner_id: str) -> bool: 

211 return (key, owner_id) in self._acquired 

212 

213 

214# ============================================================================ 

215# Redis lock backend (Redlock-inspired) 

216# ============================================================================ 

217 

218SET_IF_NOT_EXISTS = """ 

219local v = redis.call('GET', KEYS[1]) 

220if v == false then 

221 redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2]) 

222 return 1 

223end 

224return 0 

225""" 

226 

227RELEASE_SCRIPT = """ 

228if redis.call('GET', KEYS[1]) == ARGV[1] then 

229 return redis.call('DEL', KEYS[1]) 

230end 

231return 0 

232""" 

233 

234EXTEND_SCRIPT = """ 

235if redis.call('GET', KEYS[1]) == ARGV[1] then 

236 return redis.call('PEXPIRE', KEYS[1], ARGV[2]) 

237end 

238return 0 

239""" 

240 

241 

242class RedisLockBackend(AbstractLockBackend): 

243 """Redis-based distributed lock using Lua scripts for atomicity. 

244 

245 Uses SET NX PX for acquire, Lua-scripted DEL for safe release. 

246 """ 

247 

248 def __init__(self, redis_client: Any): 

249 self._redis = redis_client 

250 

251 async def acquire(self, key: str, owner_id: str, ttl: float) -> bool: 

252 result = await self._redis.set(key, owner_id, nx=True, px=int(ttl * 1000)) 

253 return bool(result) 

254 

255 async def release(self, key: str, owner_id: str) -> bool: 

256 script = ( 

257 self._redis.register_script(RELEASE_SCRIPT) 

258 if hasattr(self._redis, "register_script") 

259 else None 

260 ) 

261 if script: 

262 result = await script(keys=[key], args=[owner_id]) 

263 return int(result) == 1 

264 # Fallback for sync redis clients 

265 current = await self._redis.get(key) 

266 if current and current.decode() if isinstance(current, bytes) else current == owner_id: 

267 return bool(await self._redis.delete(key)) 

268 return False 

269 

270 async def extend(self, key: str, owner_id: str, ttl: float) -> bool: 

271 current = await self._redis.get(key) 

272 owner = current.decode() if isinstance(current, bytes) else current 

273 if owner == owner_id: 

274 return bool(await self._redis.pexpire(key, int(ttl * 1000))) 

275 return False 

276 

277 async def is_held(self, key: str, owner_id: str) -> bool: 

278 current = await self._redis.get(key) 

279 owner = current.decode() if isinstance(current, bytes) else current 

280 return owner == owner_id 

281 

282 

283# ============================================================================ 

284# High-level Lock Manager 

285# ============================================================================ 

286 

287 

288class DistributedLock: 

289 """High-level distributed lock with auto-renew and context manager support. 

290 

291 Usage: 

292 lock = DistributedLock(backend, LockConfig(ttl=30)) 

293 token = await lock.acquire("job:123") 

294 try: 

295 # critical section 

296 finally: 

297 await lock.release(token) 

298 

299 # Context manager 

300 async with lock("job:123"): 

301 # critical section 

302 """ 

303 

304 def __init__(self, backend: AbstractLockBackend, config: LockConfig = LockConfig()): 

305 self._backend = backend 

306 self._config = config 

307 self._renew_tasks: dict[str, asyncio.Task] = {} 

308 

309 async def acquire(self, key: str) -> LockToken: 

310 """Acquire lock, blocking up to acquire_timeout.""" 

311 owner_id = uuid.uuid4().hex 

312 deadline = time.monotonic() + self._config.acquire_timeout 

313 

314 while True: 

315 if await self._backend.acquire(key, owner_id, self._config.ttl): 

316 token = LockToken( 

317 key=key, 

318 owner_id=owner_id, 

319 acquired_at=time.monotonic(), 

320 ttl=self._config.ttl, 

321 _backend=self, 

322 ) 

323 if self._config.renew_interval > 0: 

324 self._start_renew(key, owner_id) 

325 return token 

326 

327 if time.monotonic() >= deadline: 

328 raise LockAcquireError( 

329 f"Failed to acquire lock '{key}' within {self._config.acquire_timeout}s" 

330 ) 

331 await asyncio.sleep(self._config.retry_interval) 

332 

333 async def release(self, token: LockToken) -> bool: 

334 """Release a held lock.""" 

335 if token.key in self._renew_tasks: 

336 self._renew_tasks.pop(token.key).cancel() 

337 return await self._backend.release(token.key, token.owner_id) 

338 

339 async def extend(self, token: LockToken, ttl: float | None = None) -> bool: 

340 """Extend the TTL of a held lock.""" 

341 return await self._backend.extend(token.key, token.owner_id, ttl or self._config.ttl) 

342 

343 def _start_renew(self, key: str, owner_id: str): 

344 """Start auto-renew background task.""" 

345 

346 async def _renew(): 

347 while True: 

348 await asyncio.sleep(self._config.renew_interval) 

349 ok = await self._backend.extend(key, owner_id, self._config.ttl) 

350 if not ok: 

351 logger.warning("Lock renew failed for key=%s — lost ownership", key) 

352 break 

353 

354 self._renew_tasks[key] = asyncio.ensure_future(_renew()) 

355 

356 @asynccontextmanager 

357 async def __call__(self, key: str) -> AsyncIterator[LockToken]: 

358 token = await self.acquire(key) 

359 try: 

360 yield token 

361 finally: 

362 await self.release(token) 

363 

364 

365# ============================================================================ 

366# Factory 

367# ============================================================================ 

368 

369 

370def create_lock_backend(backend: LockBackend, **kwargs: Any) -> AbstractLockBackend: 

371 """Factory for creating lock backends.""" 

372 if backend == LockBackend.IN_MEMORY: 

373 return InMemoryLockBackend() 

374 if backend == LockBackend.REDIS: 

375 redis_client = kwargs.get("redis_client") 

376 if redis_client is None: 

377 raise ValueError("RedisLockBackend requires 'redis_client'") 

378 return RedisLockBackend(redis_client) 

379 if backend == LockBackend.POSTGRES: 

380 pool = kwargs.get("pool") 

381 return PostgresLockBackend(pool=pool) 

382 raise ValueError(f"Unknown lock backend: {backend}")