Coverage for src/infra/tools/locking.py: 63%

1"""Centralized file locking for multi-agent coordination.

2

3 Consolidates locking behavior from shell scripts.

4"""

5

6 import hashlib

7 import logging

8 import os

9 import sys

10 from collections.abc import Callable

11 from dataclasses import dataclass

12 from pathlib import Path

13

14 from .env import get_lock_dir

15

16 logger = logging.getLogger(__name__)

17

18 __all__ = [

19 "LockManager",

20 "canonicalize_path",

21 "get_all_locks",

22 "get_lock_dir",

23 "lock_path",

24 "parse_lock_file",

25]

26

27

28 def canonicalize_path(filepath: str, repo_namespace: str | None = None) -> str:

29 """Canonicalize a file path for consistent lock key generation.

30

31 Public wrapper for _canonicalize_path. See _canonicalize_path for details.

32

33 Args:

34 filepath: The file path to canonicalize.

35 repo_namespace: Optional repo root path for resolving relative paths.

36

37 Returns:

38 A canonicalized absolute path string, or the literal key as-is.

39 """

40 return _canonicalize_path(filepath, repo_namespace)

41

42

43 def _get_lock_dir() -> Path:

44 """Get the lock directory using the accessor from env module.

45

46 This allows tests to either:

47 1. Patch os.environ["MALA_LOCK_DIR"] before calling

48 2. Patch src.infra.tools.env.get_lock_dir if more control needed

49 """

50 return get_lock_dir()

51

52

53 def _is_literal_key(filepath: str) -> bool:

54 """Check if a filepath is a literal key (not a real path).

55

56 Literal keys are special identifiers like __test_mutex__ that should

57 not be normalized as file paths. They are used for global locks that

58 are not tied to specific files.

59 """

60 return filepath.startswith("__") and filepath.endswith("__")

61

62

63 def _resolve_with_parents(path: Path) -> Path:

64 """Resolve a path by resolving existing parent directories.

65

66 For non-existent paths, walks up to find the first existing ancestor,

67 resolves its symlinks, then appends the remaining path components.

68 This ensures consistent lock keys for paths through symlinked directories.

69

70 Args:

71 path: The path to resolve (should be absolute).

72

73 Returns:

74 The resolved path with parent symlinks resolved.

75 """

76 if path.exists():

77 return path.resolve()

78

79 # Walk up to find an existing ancestor

80 # Collect the parts that don't exist yet

81 missing_parts: list[str] = []

82 current = path

83

84 max_iterations = 100

85 iterations = 0

86 while not current.exists() and iterations < max_iterations:

87 iterations += 1

88 missing_parts.append(current.name)

89 parent = current.parent

90 if parent == current:

91 # Reached root without finding existing path

92 break

93 current = parent

94

95 if iterations >= max_iterations:

96 import logging

97

98 logging.warning(

99 f"_resolve_with_parents: max iterations reached for path {path}, "

100 "using unresolved path which may cause inconsistent lock keys"

101 )

102 return path

103

104 # Resolve the existing ancestor (resolves symlinks)

105 resolved_base = current.resolve()

106

107 # Append the missing parts back

108 for part in reversed(missing_parts):

109 resolved_base = resolved_base / part

110

111 return resolved_base

112

113

114 def _canonicalize_path(filepath: str, repo_namespace: str | None = None) -> str:

115 """Canonicalize a file path for consistent lock key generation.

116

117 Normalizes paths by:

118 - Resolving symlinks (including parent directory symlinks for non-existent paths)

119 - Making paths absolute

120 - Normalizing . and .. segments

121

122 Literal keys (like __test_mutex__) are returned as-is without normalization.

123

124 This matches the shell script behavior (realpath -m), which always produces

125 absolute paths. The repo_namespace is used by _lock_key to build the final

126 key as "namespace:absolute_path".

127

128 Args:

129 filepath: The file path to canonicalize.

130 repo_namespace: Optional repo root path for resolving relative paths.

131 When provided and filepath is relative, the path is resolved

132 relative to the namespace directory (mimicking cwd=repo behavior).

133

134 Returns:

135 A canonicalized absolute path string, or the literal key as-is.

136 """

137 # Skip normalization for literal keys (non-path identifiers like __test_mutex__)

138 if _is_literal_key(filepath):

139 return filepath

140

141 path = Path(filepath)

142

143 # When we have a namespace and a relative path, resolve relative to the namespace

144 # This mimics shell script behavior when cwd is the repo directory

145 if repo_namespace and not path.is_absolute():

146 namespace_path = Path(repo_namespace).resolve()

147 candidate = namespace_path / path

148

149 if candidate.exists():

150 # Path exists - resolve symlinks

151 return str(candidate.resolve())

152 else:

153 # Normalize and resolve parent symlinks for non-existent paths

154 normalized = Path(os.path.normpath(candidate))

155 return str(_resolve_with_parents(normalized))

156

157 # Absolute path or no namespace - resolve to absolute

158 if path.exists():

159 return str(path.resolve()) # Resolves symlinks

160 else:

161 if path.is_absolute():

162 resolved = path

163 else:

164 resolved = Path.cwd() / path

165 # Normalize . and .. segments, then resolve parent symlinks

166 normalized = Path(os.path.normpath(resolved))

167 return str(_resolve_with_parents(normalized))

168

169

170 def _lock_key(filepath: str, repo_namespace: str | None = None) -> str:

171 """Build a canonical key for the lock.

172

173 Args:

174 filepath: The file path to lock.

175 repo_namespace: Optional repo namespace for cross-repo disambiguation.

176

177 Returns:

178 The canonical key string.

179 """

180 # Treat empty namespace as None

181 if repo_namespace == "":

182 repo_namespace = None

183

184 canonical_path = _canonicalize_path(filepath, repo_namespace)

185

186 if repo_namespace:

187 # Use namespace as-is to match shell script behavior

188 # Shell scripts pass REPO_NAMESPACE directly without normalizing

189 return f"{repo_namespace}:{canonical_path}"

190 return canonical_path

191

192

193 def lock_path(filepath: str, repo_namespace: str | None = None) -> Path:

194 """Convert a file path to its lock file path.

195

196 Uses SHA-256 hash of the canonical key to avoid collisions

197 (e.g., 'a/b' vs 'a_b' which would both become 'a_b.lock' with simple replacement).

198

199 Args:

200 filepath: The file path to lock.

201 repo_namespace: Optional repo namespace for cross-repo disambiguation.

202

203 Returns:

204 Path to the lock file.

205 """

206 key = _lock_key(filepath, repo_namespace)

207 key_hash = hashlib.sha256(key.encode()).hexdigest()[:16]

208 return _get_lock_dir() / f"{key_hash}.lock"

209

210

211 def release_all_locks() -> None:

212 """Release all locks in the lock directory."""

213 lock_dir = _get_lock_dir()

214 if lock_dir.exists():

215 for lock in lock_dir.glob("*.lock"):

216 # Also remove companion .meta file

217 lock.with_suffix(".meta").unlink(missing_ok=True)

218 lock.unlink(missing_ok=True)

219

220

221 def release_run_locks(agent_ids: list[str]) -> int:

222 """Release locks owned by the specified agent IDs.

223

224 Used by orchestrator shutdown to only clean up locks from this run,

225 leaving locks from other concurrent runs intact.

226

227 Args:

228 agent_ids: List of agent IDs whose locks should be released.

229

230 Returns:

231 Number of locks released.

232 """

233 lock_dir = _get_lock_dir()

234 if not lock_dir.exists() or not agent_ids:

235 return 0

236

237 agent_set = set(agent_ids)

238 released = 0

239 for lock in lock_dir.glob("*.lock"):

240 try:

241 if lock.is_file() and lock.read_text().strip() in agent_set:

242 # Also remove companion .meta file

243 lock.with_suffix(".meta").unlink(missing_ok=True)

244 lock.unlink()

245 released += 1

246 except OSError:

247 pass

248

249 return released

250

251

252 def try_lock(filepath: str, agent_id: str, repo_namespace: str | None = None) -> bool:

253 """Try to acquire a lock on a file.

254

255 Args:

256 filepath: The file path to lock.

257 agent_id: The agent ID to record in the lock.

258 repo_namespace: Optional repo namespace for cross-repo disambiguation.

259

260 Returns:

261 True if lock was acquired, False if already locked.

262 """

263 lp = lock_path(filepath, repo_namespace)

264 lock_dir = _get_lock_dir()

265 lock_dir.mkdir(parents=True, exist_ok=True)

266

267 # Fast-path if already locked

268 if lp.exists():

269 # Get holder for contention logging

270 try:

271 holder = lp.read_text().strip()

272 logger.debug(

273 "Lock contention: path=%s holder=%s requester=%s",

274 filepath,

275 holder,

276 agent_id,

277 )

278 except OSError:

279 pass

280 return False

281

282 # Atomic lock creation using temp file + rename

283 import tempfile

284

285 try:

286 fd, tmp_path = tempfile.mkstemp(

287 prefix=f".locktmp.{agent_id}.", dir=lock_dir, text=True

288 )

289 # Lock file contains only agent_id (simple, atomic reads)

290 # Filepath is stored in companion .meta file for diagnostics

291 canonical = _canonicalize_path(filepath, repo_namespace)

292 os.write(fd, f"{agent_id}\n".encode())

293 os.close(fd)

294

295 # Atomic hardlink attempt

296 try:

297 os.link(tmp_path, lp)

298 os.unlink(tmp_path)

299 # Write meta file after successful lock acquisition

300 # If meta write fails, we still own the lock - proceed anyway

301 try:

302 meta_path = lp.with_suffix(".meta")

303 meta_path.write_text(f"{canonical}\n")

304 except OSError:

305 pass # Lock acquired; meta is optional

306 logger.debug("Lock acquired: path=%s agent_id=%s", filepath, agent_id)

307 return True

308 except OSError:

309 os.unlink(tmp_path)

310 return False

311 except OSError:

312 return False

313

314

315 def wait_for_lock(

316 filepath: str,

317 agent_id: str,

318 repo_namespace: str | None = None,

319 timeout_seconds: float = 30.0,

320 poll_interval_ms: int = 100,

321) -> bool:

322 """Wait for and acquire a lock on a file.

323

324 Polls until the lock becomes available or timeout is reached.

325

326 Args:

327 filepath: The file path to lock.

328 agent_id: The agent ID to record in the lock.

329 repo_namespace: Optional repo namespace for cross-repo disambiguation.

330 timeout_seconds: Maximum time to wait for the lock (default 30).

331 poll_interval_ms: Polling interval in milliseconds (default 100).

332

333 Returns:

334 True if lock was acquired, False if timeout.

335 """

336 import time

337

338 deadline = time.monotonic() + timeout_seconds

339 poll_interval_sec = poll_interval_ms / 1000.0

340

341 while True:

342 if try_lock(filepath, agent_id, repo_namespace):

343 return True

344

345 if time.monotonic() >= deadline:

346 logger.warning(

347 "Lock timeout: path=%s agent_id=%s after=%.1fs",

348 filepath,

349 agent_id,

350 timeout_seconds,

351 )

352 return False

353

354 time.sleep(poll_interval_sec)

355

356

357 def is_locked(filepath: str, repo_namespace: str | None = None) -> bool:

358 """Check if a file is currently locked.

359

360 Args:

361 filepath: The file path to check.

362 repo_namespace: Optional repo namespace for cross-repo disambiguation.

363

364 Returns:

365 True if the file is locked, False otherwise.

366 """

367 return lock_path(filepath, repo_namespace).exists()

368

369

370 def release_lock(

371 filepath: str, agent_id: str, repo_namespace: str | None = None

372) -> bool:

373 """Release a lock on a file.

374

375 Only releases the lock if it is held by the specified agent_id.

376 This prevents accidental or malicious release of locks held by

377 other agents.

378

379 Args:

380 filepath: Path to the file to unlock.

381 agent_id: Identifier of the agent releasing the lock.

382 repo_namespace: Optional repo namespace for cross-repo disambiguation.

383

384 Returns:

385 True if lock was released, False if lock was not held by agent_id.

386 """

387 holder = get_lock_holder(filepath, repo_namespace)

388 if holder != agent_id:

389 return False

390 lp = lock_path(filepath, repo_namespace)

391 # Also remove companion .meta file

392 lp.with_suffix(".meta").unlink(missing_ok=True)

393 lp.unlink(missing_ok=True)

394 logger.debug("Lock released: path=%s agent_id=%s", filepath, agent_id)

395 return True

396

397

398 def get_lock_holder(filepath: str, repo_namespace: str | None = None) -> str | None:

399 """Get the agent ID holding a lock, or None if not locked.

400

401 Args:

402 filepath: The file path to check.

403 repo_namespace: Optional repo namespace for cross-repo disambiguation.

404

405 Returns:

406 The agent ID of the lock holder, or None if not locked.

407 """

408 lp = lock_path(filepath, repo_namespace)

409 if lp.exists():

410 try:

411 return lp.read_text().strip()

412 except OSError:

413 return None

414 return None

415

416

417 def parse_lock_file(lock_file: Path) -> tuple[str, str | None] | None:

418 """Parse a lock file to get agent_id and original filepath.

419

420 Args:

421 lock_file: Path to the lock file (.lock file).

422

423 Returns:

424 Tuple of (agent_id, filepath) or None if file cannot be read.

425 filepath may be None for legacy lock files without a .meta file.

426 """

427 try:

428 agent_id = lock_file.read_text().strip()

429 if not agent_id:

430 return None

431 # Read filepath from companion .meta file

432 meta_file = lock_file.with_suffix(".meta")

433 filepath = meta_file.read_text().strip() if meta_file.exists() else None

434 return (agent_id, filepath)

435 except OSError:

436 return None

437

438

439 def get_all_locks() -> dict[str, list[str]]:

440 """Get all active locks grouped by agent ID.

441

442 Returns:

443 Dictionary mapping agent_id -> list of locked filepaths.

444 Filepaths may be the hash stem for legacy locks without filepath info.

445 """

446 lock_dir = _get_lock_dir()

447 if not lock_dir.exists():

448 return {}

449

450 locks_by_agent: dict[str, list[str]] = {}

451 for lock in lock_dir.glob("*.lock"):

452 parsed = parse_lock_file(lock)

453 if parsed:

454 agent_id, filepath = parsed

455 if agent_id not in locks_by_agent:

456 locks_by_agent[agent_id] = []

457 # Use filepath if available, else fall back to hash stem

458 locks_by_agent[agent_id].append(filepath or lock.stem)

459

460 # Clean up orphaned .meta files (whose .lock was deleted externally)

461 for meta in lock_dir.glob("*.meta"):

462 if not meta.with_suffix(".lock").exists():

463 try:

464 meta.unlink()

465 except OSError:

466 pass

467

468 return locks_by_agent

469

470

471 def cleanup_agent_locks(agent_id: str) -> int:

472 """Remove locks held by a specific agent (crash/timeout cleanup).

473

474 Args:

475 agent_id: The agent ID whose locks should be cleaned up.

476

477 Returns:

478 Number of locks cleaned up.

479 """

480 if not _get_lock_dir().exists():

481 return 0

482

483 cleaned = 0

484 for lock in _get_lock_dir().glob("*.lock"):

485 try:

486 if lock.is_file() and lock.read_text().strip() == agent_id:

487 # Also remove companion .meta file

488 lock.with_suffix(".meta").unlink(missing_ok=True)

489 lock.unlink()

490 cleaned += 1

491 except OSError:

492 pass

493

494 logger.info("Agent locks cleaned: agent_id=%s count=%d", agent_id, cleaned)

495 return cleaned

496

497

498 class LockManager:

499 """Implementation of LockManagerPort using standalone locking functions.

500

501 This class provides an object-oriented wrapper around the standalone locking

502 functions, enabling dependency injection into domain modules.

503 """

504

505 def lock_path(self, filepath: str, repo_namespace: str | None = None) -> Path:

506 """Get the lock file path for a given filepath."""

507 return lock_path(filepath, repo_namespace)

508

509 def try_lock(

510 self, filepath: str, agent_id: str, repo_namespace: str | None = None

511 ) -> bool:

512 """Try to acquire a lock without blocking."""

513 return try_lock(filepath, agent_id, repo_namespace)

514

515 def wait_for_lock(

516 self,

517 filepath: str,

518 agent_id: str,

519 repo_namespace: str | None = None,

520 timeout_seconds: float = 30.0,

521 poll_interval_ms: int = 100,

522 ) -> bool:

523 """Wait for and acquire a lock on a file."""

524 return wait_for_lock(

525 filepath, agent_id, repo_namespace, timeout_seconds, poll_interval_ms

526 )

527

528 def release_lock(

529 self, filepath: str, agent_id: str, repo_namespace: str | None = None

530 ) -> bool:

531 """Release a lock on a file."""

532 return release_lock(filepath, agent_id, repo_namespace)

533

534

535# ---------------------------------------------------------------------------

536# CLI Command Dispatch

537# ---------------------------------------------------------------------------

538

539

540@dataclass(frozen=True)

541 class CliContext:

542 """Parsed CLI context for command dispatch."""

543

544 command: str

545 lock_dir: str

546 agent_id: str

547 repo_namespace: str | None

548 filepath: str | None

549 timeout: float

550 poll_ms: int

551

552

553 def _cmd_try(ctx: CliContext) -> int:

554 """Handle 'try' command: attempt to acquire lock."""

555 if try_lock(ctx.filepath, ctx.agent_id, ctx.repo_namespace): # type: ignore[arg-type]

556 return 0

557 return 1

558

559

560 def _cmd_wait(ctx: CliContext) -> int:

561 """Handle 'wait' command: wait for lock with timeout."""

562 if wait_for_lock(

563 ctx.filepath, # type: ignore[arg-type]

564 ctx.agent_id,

565 ctx.repo_namespace,

566 ctx.timeout,

567 ctx.poll_ms,

568 ):

569 return 0

570 return 1

571

572

573 def _cmd_check(ctx: CliContext) -> int:

574 """Handle 'check' command: check if we hold the lock."""

575 holder = get_lock_holder(ctx.filepath, ctx.repo_namespace) # type: ignore[arg-type]

576 if holder == ctx.agent_id:

577 return 0

578 return 1

579

580

581 def _cmd_holder(ctx: CliContext) -> int:

582 """Handle 'holder' command: print lock holder."""

583 holder = get_lock_holder(ctx.filepath, ctx.repo_namespace) # type: ignore[arg-type]

584 if holder:

585 print(holder)

586 return 0

587

588

589 def _cmd_release(ctx: CliContext) -> int:

590 """Handle 'release' command: release lock if we hold it."""

591 holder = get_lock_holder(ctx.filepath, ctx.repo_namespace) # type: ignore[arg-type]

592 if holder == ctx.agent_id:

593 lp = lock_path(ctx.filepath, ctx.repo_namespace) # type: ignore[arg-type]

594 lp.with_suffix(".meta").unlink(missing_ok=True)

595 lp.unlink(missing_ok=True)

596 return 0

597

598

599 def _cmd_release_all(ctx: CliContext) -> int:

600 """Handle 'release-all' command: release all locks for agent."""

601 cleanup_agent_locks(ctx.agent_id)

602 return 0

603

604

605 COMMANDS: dict[str, tuple[Callable[[CliContext], int], bool, bool]] = {

606 # command: (handler, requires_filepath, requires_agent_id)

607 "try": (_cmd_try, True, True),

608 "wait": (_cmd_wait, True, True),

609 "check": (_cmd_check, True, True),

610 "holder": (_cmd_holder, True, False),

611 "release": (_cmd_release, True, True),

612 "release-all": (_cmd_release_all, False, True),

613}

614

615

616 def _cli_main() -> int:

617 """CLI entry point for shell script delegation.

618

619 Usage:

620 python -m src.infra.tools.locking try <filepath>

621 python -m src.infra.tools.locking wait <filepath> [timeout_seconds] [poll_interval_ms]

622 python -m src.infra.tools.locking check <filepath>

623 python -m src.infra.tools.locking holder <filepath>

624 python -m src.infra.tools.locking release <filepath>

625 python -m src.infra.tools.locking release-all

626

627 Environment variables:

628 LOCK_DIR: Directory for lock files (required for most commands)

629 AGENT_ID: Agent identifier (required for most commands)

630 REPO_NAMESPACE: Optional repo namespace for cross-repo disambiguation

631

632 Exit codes:

633 0: Success (lock acquired, held, released, etc.)

634 1: Failure (lock blocked, timeout, not held, etc.)

635 2: Usage error (missing env vars, invalid arguments)

636 """

637 if len(sys.argv) < 2:

638 print(

639 "Usage: python -m src.infra.tools.locking <command> [args...]",

640 file=sys.stderr,

641 )

642 print(

643 "Commands: try, wait, check, holder, release, release-all", file=sys.stderr

644 )

645 return 2

646

647 command = sys.argv[1]

648

649 # Validate command exists

650 if command not in COMMANDS:

651 print(f"Unknown command: {command}", file=sys.stderr)

652 print(

653 "Commands: try, wait, check, holder, release, release-all", file=sys.stderr

654 )

655 return 2

656

657 handler, requires_filepath, requires_agent_id = COMMANDS[command]

658

659 # Parse environment

660 lock_dir = os.environ.get("LOCK_DIR")

661 agent_id = os.environ.get("AGENT_ID")

662 repo_namespace = os.environ.get("REPO_NAMESPACE") or None

663

664 # Validate LOCK_DIR (required for all commands)

665 if not lock_dir:

666 print("Error: LOCK_DIR must be set", file=sys.stderr)

667 return 2

668

669 # Validate AGENT_ID (required for most commands)

670 if requires_agent_id and not agent_id:

671 print("Error: AGENT_ID must be set", file=sys.stderr)

672 return 2

673

674 # Parse filepath argument

675 filepath: str | None = None

676 if requires_filepath:

677 if len(sys.argv) < 3:

678 print(

679 f"Usage: python -m src.infra.tools.locking {command} <filepath>",

680 file=sys.stderr,

681 )

682 return 2

683 filepath = sys.argv[2]

684 # Enforce exact arg count for commands without optional arguments

685 if command != "wait" and len(sys.argv) > 3:

686 print(

687 f"Usage: python -m src.infra.tools.locking {command} <filepath>",

688 file=sys.stderr,

689 )

690 return 2

691 elif command == "release-all" and len(sys.argv) != 2:

692 print("Usage: python -m src.infra.tools.locking release-all", file=sys.stderr)

693 return 2

694

695 # Parse wait-specific arguments

696 timeout = 30.0

697 poll_ms = 100

698 if command == "wait":

699 timeout = float(sys.argv[3]) if len(sys.argv) > 3 else 30.0

700 poll_ms = int(sys.argv[4]) if len(sys.argv) > 4 else 100

701

702 # Set MALA_LOCK_DIR so our functions use the shell script's LOCK_DIR

703 os.environ["MALA_LOCK_DIR"] = lock_dir

704

705 # Build context and dispatch

706 ctx = CliContext(

707 command=command,

708 lock_dir=lock_dir,

709 agent_id=agent_id or "",

710 repo_namespace=repo_namespace,

711 filepath=filepath,

712 timeout=timeout,

713 poll_ms=poll_ms,

714 )

715 return handler(ctx)

716

717

718 if __name__ == "__main__":

719 sys.exit(_cli_main())