Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 7 additions & 1 deletion concore.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,13 +112,19 @@ def _write_kill_script():
try:
reg_name = os.path.basename(_PID_REGISTRY_FILE)
bat_name = os.path.basename(_KILL_SCRIPT_FILE)
# nodes run as "python <script>.py", so check for the script name
script_name = os.path.basename(sys.argv[0]) if sys.argv else ""
marker = script_name if script_name.endswith(".py") else "concore"
if not marker.isascii():
marker = "concore"
marker = marker.replace("%", "%%")
script = "@echo off\r\n"
script += 'if not exist "%~dp0' + reg_name + '" (\r\n'
script += " echo No PID registry found. Nothing to kill.\r\n"
script += " exit /b 0\r\n"
script += ")\r\n"
script += 'for /f "usebackq tokens=*" %%p in ("%~dp0' + reg_name + '") do (\r\n'
script += ' powershell -NoProfile -Command "(Get-CimInstance Win32_Process -Filter \'ProcessId=%%p\').CommandLine" 2>nul | findstr /i "concore" >nul\r\n'
script += ' powershell -NoProfile -Command "(Get-CimInstance Win32_Process -Filter \'ProcessId=%%p\').CommandLine" 2>nul | findstr /i /l /c:"' + marker + '" >nul\r\n'
script += " if not errorlevel 1 (\r\n"
script += " echo Killing concore process %%p\r\n"
script += " taskkill /F /PID %%p >nul 2>&1\r\n"
Expand Down
34 changes: 25 additions & 9 deletions tests/test_concore.py
Original file line number Diff line number Diff line change
Expand Up @@ -842,9 +842,10 @@ def test_cleanup_pid_handles_missing_registry(self):
assert not os.path.exists(_PID_REGISTRY_FILE)
_cleanup_pid() # Should not raise

def test_write_kill_script_generates_bat_file(self):
def test_write_kill_script_generates_bat_file(self, monkeypatch):
from concore import _write_kill_script, _KILL_SCRIPT_FILE, _PID_REGISTRY_FILE

monkeypatch.setattr(sys, "argv", ["controller.py"])
_write_kill_script()
assert os.path.exists(_KILL_SCRIPT_FILE)
with open(_KILL_SCRIPT_FILE) as f:
Expand All @@ -853,27 +854,42 @@ def test_write_kill_script_generates_bat_file(self):
assert "wmic" not in content
assert "Get-CimInstance" in content
assert "taskkill" in content
assert "concore" in content.lower()
assert 'findstr /i /l /c:"controller.py"' in content

@pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
def test_kill_script_kills_concore_process(self):
def test_kill_script_kills_concore_process(self, monkeypatch):
import subprocess
import time
from concore import _write_kill_script, _KILL_SCRIPT_FILE, _PID_REGISTRY_FILE

proc = subprocess.Popen(
[sys.executable, "-c", "import time; time.sleep(60)", "concore_node"]
)
# a node runs as "python controller.py" and registers its own PID
with open("controller.py", "w") as f:
f.write(
"import os, time\n"
"with open('concorekill_pids.txt', 'a') as f:\n"
" f.write(str(os.getpid()) + '\\n')\n"
"time.sleep(60)\n"
)
monkeypatch.setattr(sys, "argv", ["controller.py"])
proc = subprocess.Popen([sys.executable, "controller.py"])
try:
with open(_PID_REGISTRY_FILE, "w") as f:
f.write(str(proc.pid) + "\n")
node_pid = ""
for _ in range(100):
if os.path.exists(_PID_REGISTRY_FILE):
with open(_PID_REGISTRY_FILE) as f:
node_pid = f.read().strip()
if node_pid:
break
time.sleep(0.1)
assert node_pid
_write_kill_script()
result = subprocess.run(
["cmd", "/c", _KILL_SCRIPT_FILE],
capture_output=True,
text=True,
timeout=60,
)
assert "Killing concore process " + str(proc.pid) in result.stdout
assert "Killing concore process " + node_pid in result.stdout
proc.wait(timeout=10)
assert not os.path.exists(_PID_REGISTRY_FILE)
assert not os.path.exists(_KILL_SCRIPT_FILE)
Expand Down
Loading