Skip to content

perf: only probe node processes when looking for stale managers #57

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
donbowman wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: only probe node processes when looking for stale managers #57

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 30 additions & 16 deletions src/py_pglite/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ def __init__(self, config: PGliteConfig | None = None):
self.work_dir: Path | None = None
self._original_cwd: str | None = None
self._shared_engine: Any | None = None
self._auto_work_dir: bool = False

# Set up logging
self.logger = logging.getLogger(__name__)
Expand All @@ -63,6 +64,7 @@ def _setup_work_dir(self) -> Path:
work_dir.mkdir(parents=True, exist_ok=True)
else:
work_dir = Path(tempfile.mkdtemp(prefix="py-pglite-"))
self._auto_work_dir = True

# Create package.json if it doesn't exist
package_json = work_dir / "package.json"
Expand Down Expand Up @@ -277,6 +279,11 @@ def _cleanup_socket(self) -> None:

def _kill_existing_processes(self) -> None:
"""Kill any existing PGlite processes that might conflict with this socket."""
# An auto-created temporary work directory is brand new, so no
# process can be using it; skip the process scan entirely.
if self._auto_work_dir:
return

try:
# Fix for issue #31: Compare work directory, not socket directory
# Socket and work directories are different by design for isolation
Expand All @@ -288,23 +295,30 @@ def _kill_existing_processes(self) -> None:
my_target_dir = str(Path(self.config.socket_path).parent)
comparison_type = "socket directory"

for proc in psutil.process_iter(["pid", "name", "cmdline", "cwd"]):
if proc.info["cmdline"] and any(
"pglite_manager.js" in cmd for cmd in proc.info["cmdline"]
):
# Use exact directory match to avoid killing processes in similar paths
try:
proc_cwd = proc.info.get("cwd", "")
if proc_cwd == my_target_dir:
pid = proc.info["pid"]
self.logger.info(
f"Killing existing PGlite process: {pid} (matching {comparison_type})"
)
proc.kill()
proc.wait(timeout=5)
except (psutil.NoSuchProcess, psutil.AccessDenied):
# Process already gone or can't access it
# Only processes named like node can be running the manager
# script, so fetch the more expensive cmdline/cwd lazily for
# those instead of for every process on the system.
for proc in psutil.process_iter(["pid", "name"]):
name = (proc.info.get("name") or "").lower()
if "node" not in name:
continue
try:
cmdline = proc.cmdline()
if not any("pglite_manager.js" in cmd for cmd in cmdline):
continue
# Use exact directory match to avoid killing processes
# in similar paths.
if proc.cwd() != my_target_dir:
continue
pid = proc.info["pid"]
self.logger.info(
f"Killing existing PGlite process: {pid} (matching {comparison_type})"
)
proc.kill()
proc.wait(timeout=5)
except (psutil.NoSuchProcess, psutil.AccessDenied):
# Process already gone or can't access it
continue
except Exception as e:
self.logger.warning(f"Error killing existing PGlite processes: {e}")

Expand Down
67 changes: 49 additions & 18 deletions tests/test_manager_comprehensive.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,20 +251,12 @@ def test_kill_existing_processes_success(self):
manager.work_dir = Path("/tmp/pglite-work-dir")

mock_proc1 = Mock()
mock_proc1.info = {
"pid": 1234,
"name": "node",
"cmdline": ["node", "pglite_manager.js"],
"cwd": "/tmp/pglite-work-dir", # Same directory as work_dir
}
mock_proc1.info = {"pid": 1234, "name": "node"}
mock_proc1.cmdline.return_value = ["node", "pglite_manager.js"]
mock_proc1.cwd.return_value = "/tmp/pglite-work-dir"

mock_proc2 = Mock()
mock_proc2.info = {
"pid": 5678,
"name": "python",
"cmdline": ["python", "test.py"],
"cwd": "/home/user",
}
mock_proc2.info = {"pid": 5678, "name": "python"}

with patch("psutil.process_iter", return_value=[mock_proc1, mock_proc2]):
manager._kill_existing_processes()
Expand All @@ -273,6 +265,8 @@ def test_kill_existing_processes_success(self):
mock_proc1.kill.assert_called_once()
mock_proc1.wait.assert_called_once_with(timeout=5)
mock_proc2.kill.assert_not_called()
# Non-node processes are not probed at all.
mock_proc2.cmdline.assert_not_called()

def test_kill_existing_processes_no_work_dir(self):
"""Test killing existing processes when work_dir is None."""
Expand All @@ -281,19 +275,56 @@ def test_kill_existing_processes_no_work_dir(self):
# work_dir is None initially

mock_proc1 = Mock()
mock_proc1.info = {
"pid": 1234,
"name": "node",
"cmdline": ["node", "pglite_manager.js"],
"cwd": "/tmp/some-dir",
}
mock_proc1.info = {"pid": 1234, "name": "node"}
mock_proc1.cmdline.return_value = ["node", "pglite_manager.js"]
mock_proc1.cwd.return_value = "/tmp/some-dir"

with patch("psutil.process_iter", return_value=[mock_proc1]):
manager._kill_existing_processes()

# Should not kill any processes when work_dir is None
mock_proc1.kill.assert_not_called()

def test_kill_existing_processes_skips_auto_work_dir(self):
"""An auto-created work directory cannot have a stale process."""
manager = PGliteManager()
manager._auto_work_dir = True

with patch("psutil.process_iter") as mock_process_iter:
manager._kill_existing_processes()

mock_process_iter.assert_not_called()

def test_kill_existing_processes_ignores_unrelated_node_processes(self):
"""Node processes not running the manager script are left alone."""
config = PGliteConfig(socket_path="/tmp/pglite-socket/socket")
manager = PGliteManager(config)
manager.work_dir = Path("/tmp/pglite-work-dir")

mock_proc = Mock()
mock_proc.info = {"pid": 99, "name": "node"}
mock_proc.cmdline.return_value = ["node", "some-other-server.js"]
mock_proc.cwd.return_value = "/tmp/pglite-work-dir"

with patch("psutil.process_iter", return_value=[mock_proc]):
manager._kill_existing_processes()

mock_proc.kill.assert_not_called()

def test_kill_existing_processes_handles_missing_name(self):
"""Processes whose name cannot be read are skipped safely."""
config = PGliteConfig(socket_path="/tmp/pglite-socket/socket")
manager = PGliteManager(config)
manager.work_dir = Path("/tmp/pglite-work-dir")

mock_proc = Mock()
mock_proc.info = {"pid": 7, "name": None}

with patch("psutil.process_iter", return_value=[mock_proc]):
manager._kill_existing_processes()

mock_proc.kill.assert_not_called()

def test_kill_existing_processes_exception_handling(self):
"""Test exception handling in process killing."""
manager = PGliteManager()
Expand Down