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 contextlib import asynccontextmanager 

19from dataclasses import dataclass, field 

20from enum import Enum 

21from typing import Any, AsyncIterator, Optional 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26# ============================================================================ 

27# Data types 

28# ============================================================================ 

29 

30class LockBackend(str, Enum): 

31 IN_MEMORY = "in_memory" 

32 POSTGRES = "postgres" 

33 REDIS = "redis" 

34 

35 

36@dataclass 

37class LockConfig: 

38 """Configuration for distributed lock acquisition.""" 

39 

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

41 retry_interval: float = 0.1 # Polling interval when waiting 

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

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

44 

45 

46@dataclass 

47class LockToken: 

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

49 

50 key: str 

51 owner_id: str 

52 acquired_at: float 

53 ttl: float 

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

55 

56 

57class LockAcquireError(Exception): 

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

59 

60 

61class LockNotHeldError(Exception): 

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

63 

64 

65# ============================================================================ 

66# Abstract backend 

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

68 

69class AbstractLockBackend(ABC): 

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

71 

72 @abstractmethod 

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

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

75 

76 @abstractmethod 

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

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

79 

80 @abstractmethod 

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

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

83 

84 @abstractmethod 

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

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

87 

88 

89# ============================================================================ 

90# In-Memory backend 

91# ============================================================================ 

92 

93class InMemoryLockBackend(AbstractLockBackend): 

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

95 

96 def __init__(self): 

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

98 self._lock = asyncio.Lock() 

99 

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

101 async with self._lock: 

102 now = time.monotonic() 

103 if key in self._locks: 

104 owner, expiry = self._locks[key] 

105 if expiry > now and owner != owner_id: 

106 return False 

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

108 return True 

109 

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

111 async with self._lock: 

112 if key not in self._locks: 

113 return False 

114 owner, expiry = self._locks[key] 

115 if owner != owner_id: 

116 return False 

117 del self._locks[key] 

118 return True 

119 

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

121 async with self._lock: 

122 if key not in self._locks: 

123 return False 

124 owner, expiry = self._locks[key] 

125 if owner != owner_id: 

126 return False 

127 if expiry < time.monotonic(): 

128 del self._locks[key] 

129 return False 

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

131 return True 

132 

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

134 async with self._lock: 

135 if key not in self._locks: 

136 return False 

137 owner, expiry = self._locks[key] 

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

139 

140 

141# ============================================================================ 

142# Postgres advisory lock backend 

143# ============================================================================ 

144 

145ADVISORY_LOCK_SQL = """ 

146SELECT pg_try_advisory_lock(%s) AS acquired; 

147""" 

148 

149ADVISORY_UNLOCK_SQL = """ 

150SELECT pg_advisory_unlock(%s) AS released; 

151""" 

152 

153# Hash key to an int64 for advisory lock 

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

155 import hashlib 

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

157 

158 

159class PostgresLockBackend(AbstractLockBackend): 

160 """PostgreSQL advisory-lock based backend. 

161 

162 Uses pg_try_advisory_lock for non-blocking acquire. 

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

164 """ 

165 

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

167 self._pool = pool 

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

169 

170 async def _get_conn(self): 

171 if self._pool is None: 

172 raise RuntimeError("PostgresLockBackend requires a pool") 

173 return await self._pool.acquire() 

174 

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

176 conn = await self._get_conn() 

177 try: 

178 lock_id = _key_to_int64(key) 

179 result = await conn.fetchval( 

180 "SELECT pg_try_advisory_lock($1) AS acquired;", lock_id 

181 ) 

182 if result: 

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

184 return True 

185 return False 

186 finally: 

187 await self._pool.release(conn) 

188 

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

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

191 return False 

192 conn = await self._get_conn() 

193 try: 

194 lock_id = _key_to_int64(key) 

195 result = await conn.fetchval( 

196 "SELECT pg_advisory_unlock($1) AS released;", lock_id 

197 ) 

198 if result: 

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

200 return bool(result) 

201 finally: 

202 await self._pool.release(conn) 

203 

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

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

206 return (key, owner_id) in self._acquired 

207 

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

209 return (key, owner_id) in self._acquired 

210 

211 

212# ============================================================================ 

213# Redis lock backend (Redlock-inspired) 

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

215 

216SET_IF_NOT_EXISTS = """ 

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

218if v == false then 

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

220 return 1 

221end 

222return 0 

223""" 

224 

225RELEASE_SCRIPT = """ 

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

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

228end 

229return 0 

230""" 

231 

232EXTEND_SCRIPT = """ 

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

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

235end 

236return 0 

237""" 

238 

239 

240class RedisLockBackend(AbstractLockBackend): 

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

242 

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

244 """ 

245 

246 def __init__(self, redis_client: Any): 

247 self._redis = redis_client 

248 

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

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

251 return bool(result) 

252 

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

254 script = self._redis.register_script(RELEASE_SCRIPT) if hasattr(self._redis, 'register_script') else None 

255 if script: 

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

257 return int(result) == 1 

258 # Fallback for sync redis clients 

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

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

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

262 return False 

263 

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

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

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

267 if owner == owner_id: 

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

269 return False 

270 

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

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

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

274 return owner == owner_id 

275 

276 

277# ============================================================================ 

278# High-level Lock Manager 

279# ============================================================================ 

280 

281class DistributedLock: 

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

283 

284 Usage: 

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

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

287 try: 

288 # critical section 

289 finally: 

290 await lock.release(token) 

291 

292 # Context manager 

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

294 # critical section 

295 """ 

296 

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

298 self._backend = backend 

299 self._config = config 

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

301 

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

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

304 owner_id = uuid.uuid4().hex 

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

306 

307 while True: 

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

309 token = LockToken( 

310 key=key, 

311 owner_id=owner_id, 

312 acquired_at=time.monotonic(), 

313 ttl=self._config.ttl, 

314 _backend=self, 

315 ) 

316 if self._config.renew_interval > 0: 

317 self._start_renew(key, owner_id) 

318 return token 

319 

320 if time.monotonic() >= deadline: 

321 raise LockAcquireError( 

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

323 ) 

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

325 

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

327 """Release a held lock.""" 

328 if token.key in self._renew_tasks: 

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

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

331 

332 async def extend(self, token: LockToken, ttl: Optional[float] = None) -> bool: 

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

334 return await self._backend.extend( 

335 token.key, token.owner_id, ttl or self._config.ttl 

336 ) 

337 

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

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

340 

341 async def _renew(): 

342 while True: 

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

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

345 if not ok: 

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

347 break 

348 

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

350 

351 @asynccontextmanager 

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

353 token = await self.acquire(key) 

354 try: 

355 yield token 

356 finally: 

357 await self.release(token) 

358 

359 

360# ============================================================================ 

361# Factory 

362# ============================================================================ 

363 

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

365 """Factory for creating lock backends.""" 

366 if backend == LockBackend.IN_MEMORY: 

367 return InMemoryLockBackend() 

368 if backend == LockBackend.REDIS: 

369 redis_client = kwargs.get("redis_client") 

370 if redis_client is None: 

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

372 return RedisLockBackend(redis_client) 

373 if backend == LockBackend.POSTGRES: 

374 pool = kwargs.get("pool") 

375 return PostgresLockBackend(pool=pool) 

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