fix(fabric): harden port lock probes and state queue

Author: littlebullGit

src/lightning/fabric/utilities/file_lock.py

@@ -20,7 +20,8 @@
 from abc import ABC, abstractmethod
 from contextlib import suppress
 from pathlib import Path
-from typing import Optional
+from types import TracebackType
+from typing import Literal, Optional
 
 log = logging.getLogger(__name__)
 
@@ -75,7 +76,12 @@ def __enter__(self) -> "FileLock":
             raise TimeoutError(f"Failed to acquire lock on {self._lock_file} within timeout")
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Literal[False]:
         """Exit context manager - release lock."""
         self.release()
         return False  # Don't suppress exceptions

src/lightning/fabric/utilities/port_manager.py

@@ -24,7 +24,8 @@
 from collections.abc import Iterator
 from contextlib import contextmanager, suppress
 from pathlib import Path
-from typing import Optional
+from types import TracebackType
+from typing import Literal, Optional
 
 from lightning.fabric.utilities.file_lock import create_file_lock
 from lightning.fabric.utilities.port_state import PortState
@@ -76,11 +77,16 @@ def _get_lock_dir() -> Path:
         except OSError as e:
             log.debug(f"Port manager probe file removal failed with {e}; scheduling cleanup")
 
-    atexit.register(lambda p=test_file: p.unlink(missing_ok=True))
+    atexit.register(_cleanup_probe_file, test_file)
 
     return lock_path
 
 
+def _cleanup_probe_file(path: Path) -> None:
+    """Best-effort removal of a temporary probe file at exit."""
+    path.unlink(missing_ok=True)
+
+
 def _get_lock_file() -> Path:
     """Get path to the port manager lock file.
 
@@ -218,38 +224,13 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
         with self._lock:  # Thread-safety
             try:
                 with self._file_lock:  # Process-safety
-                    # Read current state from file
                     state = self._read_state()
-
-                    # Try preferred port if specified
-                    if preferred_port is not None and self._is_port_available(preferred_port, state):
-                        port = preferred_port
-                    else:
-                        # Find a free port
-                        port = None
-                        for _ in range(max_attempts):
-                            candidate = self._find_free_port()
-                            if self._is_port_available(candidate, state):
-                                port = candidate
-                                break
-
-                        if port is None:
-                            # Provide detailed diagnostics
-                            allocated_count = len(state.allocated_ports)
-                            queue_count = len(state.recently_released)
-                            raise RuntimeError(
-                                f"Failed to allocate a free port after {max_attempts} attempts. "
-                                f"Diagnostics: allocated={allocated_count}, recently_released={queue_count}"
-                            )
-
-                    # Allocate in shared state
+                    port = self._select_port(state, preferred_port, max_attempts)
                     state.allocate_port(port, pid=os.getpid())
                     self._write_state(state)
 
-                    # Update in-memory cache
                     self._allocated_ports.add(port)
 
-                    # Log diagnostics if queue utilization is high
                     queue_count = len(state.recently_released)
                     if queue_count > 800:  # >78% of typical 1024 capacity
                         log.warning(
@@ -261,7 +242,6 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
                     return port
 
             except TimeoutError as e:
-                # File lock timeout - fail fast to prevent state divergence
                 log.error(
                     "Failed to acquire file lock for port allocation. "
                     "Remediation: (1) Retry the operation after a short delay, "
@@ -274,6 +254,30 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
                     "Check if another process is holding the lock or if the lock file is inaccessible."
                 ) from e
 
+        raise RuntimeError("Unexpected error allocating port")
+
+    def _select_port(
+        self,
+        state: PortState,
+        preferred_port: Optional[int],
+        max_attempts: int,
+    ) -> int:
+        """Choose an available port based on preference and state."""
+        if preferred_port is not None and self._is_port_available(preferred_port, state):
+            return preferred_port
+
+        for _ in range(max_attempts):
+            candidate = self._find_free_port()
+            if self._is_port_available(candidate, state):
+                return candidate
+
+        allocated_count = len(state.allocated_ports)
+        queue_count = len(state.recently_released)
+        raise RuntimeError(
+            f"Failed to allocate a free port after {max_attempts} attempts. "
+            f"Diagnostics: allocated={allocated_count}, recently_released={queue_count}"
+        )
+
     def _is_port_available(self, port: int, state: PortState) -> bool:
         """Check if a port is available for allocation.
 
@@ -480,7 +484,12 @@ def __enter__(self) -> "PortManager":
         """
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Literal[False]:
         """Exit context manager - cleanup ports from this process."""
         self.release_all()
         return False  # Don't suppress exceptions

src/lightning/fabric/utilities/port_state.py

@@ -27,6 +27,9 @@
 # Maximum age for recently released entries (2 hours)
 _RECENTLY_RELEASED_MAX_AGE_SECONDS = 7200
 
+# Maximum number of recently released entries to retain
+_RECENTLY_RELEASED_MAX_LEN = 1024
+
 
 @dataclass
 class PortAllocation:
@@ -179,6 +182,7 @@ def release_port(self, port: int) -> None:
                 pid=allocation.pid,
             )
             self.recently_released.append(entry)
+            self._trim_recently_released()
             # Remove from allocated
             del self.allocated_ports[port_str]
 
@@ -239,6 +243,9 @@ def cleanup_stale_entries(self) -> int:
         # Clean up stale recently released entries
         original_count = len(self.recently_released)
         self.recently_released = [entry for entry in self.recently_released if not entry.is_stale(current_time)]
+        if len(self.recently_released) > _RECENTLY_RELEASED_MAX_LEN:
+            # Keep only the most recent entries if stale cleanup still exceeds max length
+            self.recently_released = self.recently_released[-_RECENTLY_RELEASED_MAX_LEN:]
         stale_count += original_count - len(self.recently_released)
 
         return stale_count
@@ -288,12 +295,22 @@ def from_dict(cls, data: dict[str, Any]) -> "PortState":
             RecentlyReleasedEntry.from_dict(entry_data) for entry_data in data.get("recently_released", [])
         ]
 
-        return cls(
+        state = cls(
             version=data.get("version", "1.0"),
             allocated_ports=allocated_ports,
             recently_released=recently_released,
         )
 
+        state._trim_recently_released()
+        return state
+
+    def _trim_recently_released(self) -> None:
+        """Ensure recently released queue stays within configured bound."""
+        if len(self.recently_released) > _RECENTLY_RELEASED_MAX_LEN:
+            excess = len(self.recently_released) - _RECENTLY_RELEASED_MAX_LEN
+            # Remove the oldest entries (front of the list)
+            self.recently_released = self.recently_released[excess:]
+
 
 def _is_pid_alive(pid: int) -> bool:
     """Check if a process with given PID is still running.