diff --git a/concore.java b/concore.java index 94f92cd..133c5ad 100644 --- a/concore.java +++ b/concore.java @@ -39,75 +39,85 @@ private static boolean isWindows() { private static void registerPid() { try { - String pid = String.valueOf(ProcessHandle.current().pid()); + ProcessHandle current = ProcessHandle.current(); + long pid = current.pid(); + long startTimeMs = current.info().startInstant() + .orElseThrow(() -> new IllegalStateException( + "Process start time is unavailable")) + .toEpochMilli(); + String record = pid + "|" + startTimeMs + System.lineSeparator(); try (FileChannel channel = FileChannel.open(PID_REGISTRY_FILE, - StandardOpenOption.CREATE, StandardOpenOption.WRITE, StandardOpenOption.APPEND)) { + StandardOpenOption.CREATE, StandardOpenOption.WRITE, + StandardOpenOption.APPEND)) { try (FileLock lock = channel.lock()) { - channel.write(ByteBuffer.wrap((pid + System.lineSeparator()).getBytes(java.nio.charset.StandardCharsets.UTF_8))); + ByteBuffer buffer = ByteBuffer.wrap( + record.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + while (buffer.hasRemaining()) { + channel.write(buffer); + } } } - } catch (IOException e) { + } catch (IOException | RuntimeException e) { } } - private static void cleanupPid() { String pid = String.valueOf(ProcessHandle.current().pid()); if (!Files.exists(PID_REGISTRY_FILE)) return; - List remaining = new ArrayList<>(); try (FileChannel channel = FileChannel.open(PID_REGISTRY_FILE, StandardOpenOption.READ, StandardOpenOption.WRITE)) { try (FileLock lock = channel.lock()) { ByteBuffer buf = ByteBuffer.allocate((int) channel.size()); - channel.read(buf); + while (buf.hasRemaining() && channel.read(buf) != -1) { + } buf.flip(); - String content = java.nio.charset.StandardCharsets.UTF_8.decode(buf).toString(); + String content = java.nio.charset.StandardCharsets.UTF_8 + .decode(buf).toString(); + List remaining = new ArrayList<>(); for (String line : content.split("\\R")) { String trimmed = line.trim(); - if (!trimmed.isEmpty() && !trimmed.equals(pid)) { - remaining.add(trimmed); + if (!trimmed.isEmpty()) { + String entryPid = trimmed.split("\\|", 2)[0].trim(); + if (!entryPid.equals(pid)) { + remaining.add(trimmed); + } } } + StringBuilder sb = new StringBuilder(); + for (String entry : remaining) { + sb.append(entry).append(System.lineSeparator()); + } channel.truncate(0); channel.position(0); - if (!remaining.isEmpty()) { - StringBuilder sb = new StringBuilder(); - for (String r : remaining) sb.append(r).append(System.lineSeparator()); - channel.write(ByteBuffer.wrap(sb.toString().getBytes(java.nio.charset.StandardCharsets.UTF_8))); + ByteBuffer output = ByteBuffer.wrap( + sb.toString().getBytes( + java.nio.charset.StandardCharsets.UTF_8)); + while (output.hasRemaining()) { + channel.write(output); } } } catch (IOException e) { } - if (remaining.isEmpty()) { - try { Files.deleteIfExists(PID_REGISTRY_FILE); } catch (IOException e) {} - try { Files.deleteIfExists(KILL_SCRIPT_FILE); } catch (IOException e) {} - } } - private static void writeKillScript() { try { - String regName = PID_REGISTRY_FILE.getFileName().toString(); - String batName = KILL_SCRIPT_FILE.getFileName().toString(); - String script = "@echo off\r\n"; - script += "if not exist \"%~dp0" + regName + "\" (\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" + regName + "\") do (\r\n"; - script += " powershell -NoProfile -Command \"(Get-CimInstance Win32_Process -Filter 'ProcessId=%%p').CommandLine\" 2>nul | findstr /i \"concore\" >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"; - script += " ) else (\r\n"; - script += " echo Skipping PID %%p - not a concore process or not running\r\n"; - script += " )\r\n"; - script += ")\r\n"; - script += "del /q \"%~dp0" + regName + "\" 2>nul\r\n"; - script += "del /q \"%~dp0" + batName + "\" 2>nul\r\n"; - Files.write(KILL_SCRIPT_FILE, script.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + String helperName = "concorekill.ps1"; + String script = + "@echo off\r\n" + + "setlocal\r\n" + + "if not exist \"%~dp0" + helperName + "\" (\r\n" + + " echo Missing process cleanup helper: \"%~dp0" + + helperName + "\"\r\n" + + " exit /b 1\r\n" + + ")\r\n" + + "powershell.exe -NoLogo -NoProfile -NonInteractive " + + "-ExecutionPolicy Bypass -File \"%~dp0" + + helperName + "\"\r\n" + + "exit /b %ERRORLEVEL%\r\n"; + Files.write(KILL_SCRIPT_FILE, + script.getBytes(java.nio.charset.StandardCharsets.UTF_8)); } catch (IOException e) { } } - /** * Sets maxtime from concore.maxtime file, or uses defaultValue if file not found. * Catches both IOException and RuntimeException to match Python safe_literal_eval. diff --git a/concore.py b/concore.py index 01c4e0f..fe9471f 100644 --- a/concore.py +++ b/concore.py @@ -31,103 +31,83 @@ _PID_REGISTRY_FILE = os.path.join(_BASE_DIR, "concorekill_pids.txt") _KILL_SCRIPT_FILE = os.path.join(_BASE_DIR, "concorekill.bat") +def _current_pid_record(): + """Return current PID and process creation time in Unix milliseconds.""" + import psutil + process = psutil.Process(os.getpid()) + start_time_ms = int(process.create_time() * 1000) + return f"{process.pid}|{start_time_ms}" def _register_pid(): - """Append current PID to the shared registry file. Uses file locking on Windows.""" + """Append current process identity to the shared registry.""" + import psutil try: + record = _current_pid_record() with open(_PID_REGISTRY_FILE, "a") as f: if hasattr(sys, 'getwindowsversion'): import msvcrt try: + f.seek(0) msvcrt.locking(f.fileno(), msvcrt.LK_LOCK, _LOCK_LEN) - f.write(str(os.getpid()) + "\n") + f.write(record + "\n") + f.flush() finally: try: f.seek(0) msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, _LOCK_LEN) - except OSError: + except (OSError, ValueError): pass else: - f.write(str(os.getpid()) + "\n") - except OSError: + f.write(record + "\n") + except (OSError, psutil.Error): pass - def _cleanup_pid(): - """Remove current PID from registry on exit. Uses file locking on Windows.""" + """Remove current process entries while preserving the shared files.""" pid = str(os.getpid()) try: if not os.path.exists(_PID_REGISTRY_FILE): return with open(_PID_REGISTRY_FILE, "r+") as f: - if hasattr(sys, 'getwindowsversion'): + is_windows = hasattr(sys, "getwindowsversion") + if is_windows: import msvcrt - try: - msvcrt.locking(f.fileno(), msvcrt.LK_LOCK, _LOCK_LEN) - pids = [line.strip() for line in f if line.strip()] - remaining = [p for p in pids if p != pid] - if remaining: - f.seek(0) - f.truncate() - for p in remaining: - f.write(p + "\n") - else: - f.close() - try: - os.remove(_PID_REGISTRY_FILE) - except OSError: - pass - try: - os.remove(_KILL_SCRIPT_FILE) - except OSError: - pass - return - finally: + f.seek(0) + msvcrt.locking(f.fileno(), msvcrt.LK_LOCK, _LOCK_LEN) + try: + f.seek(0) + entries = [line.strip() for line in f if line.strip()] + remaining = [ + entry for entry in entries + if entry.split("|", 1)[0].strip() != pid + ] + f.seek(0) + f.truncate() + for entry in remaining: + f.write(entry + "\n") + f.flush() + finally: + if is_windows: try: f.seek(0) msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, _LOCK_LEN) except (OSError, ValueError): pass - else: - pids = [line.strip() for line in f if line.strip()] - remaining = [p for p in pids if p != pid] - if remaining: - f.seek(0) - f.truncate() - for p in remaining: - f.write(p + "\n") - else: - f.close() - try: - os.remove(_PID_REGISTRY_FILE) - except OSError: - pass - try: - os.remove(_KILL_SCRIPT_FILE) - except OSError: - pass except OSError: pass def _write_kill_script(): - """Generate concorekill.bat that validates each PID before killing.""" + """Generate a batch launcher for the shared PowerShell cleanup helper.""" try: - reg_name = os.path.basename(_PID_REGISTRY_FILE) - bat_name = os.path.basename(_KILL_SCRIPT_FILE) - 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 += " if not errorlevel 1 (\r\n" - script += " echo Killing concore process %%p\r\n" - script += " taskkill /F /PID %%p >nul 2>&1\r\n" - script += " ) else (\r\n" - script += " echo Skipping PID %%p - not a concore process or not running\r\n" - script += " )\r\n" - script += ")\r\n" - script += 'del /q "%~dp0' + reg_name + '" 2>nul\r\n' - script += 'del /q "%~dp0' + bat_name + '" 2>nul\r\n' + helper_name = "concorekill.ps1" + script = ( + "@echo off\r\n" + "setlocal\r\n" + 'if not exist "%~dp0' + helper_name + '" (\r\n' + ' echo Missing process cleanup helper: "%~dp0' + helper_name + '"\r\n' + " exit /b 1\r\n" + ")\r\n" + 'powershell.exe -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Bypass -File "%~dp0' + helper_name + '"\r\n' + "exit /b %ERRORLEVEL%\r\n" + ) with open(_KILL_SCRIPT_FILE, "w", newline="") as f: f.write(script) except OSError: diff --git a/concorekill.ps1 b/concorekill.ps1 new file mode 100644 index 0000000..2e5276d --- /dev/null +++ b/concorekill.ps1 @@ -0,0 +1,244 @@ +# Windows process cleanup helper shared by the Python and Java runtimes. +# Registry format: PID|START_TIME_UNIX_MS. PID-only entries use legacy validation. + +$ErrorActionPreference = 'Stop' +$registryPath = Join-Path $PSScriptRoot 'concorekill_pids.txt' + +if (-not (Test-Path -LiteralPath $registryPath)) { + Write-Output 'No PID registry found. Nothing to kill.' + exit 0 +} + +if (-not ('ConcoreNativeMethods' -as [type])) { + Add-Type -TypeDefinition @' +using System; +using System.Runtime.InteropServices; + +public static class ConcoreNativeMethods +{ + // PROCESS_TERMINATE | PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE + [DllImport("kernel32.dll", SetLastError = true)] + public static extern IntPtr OpenProcess(uint desiredAccess, bool inheritHandle, uint processId); + + [DllImport("kernel32.dll", SetLastError = true)] + public static extern bool GetProcessTimes( + IntPtr processHandle, + out long creationTime, + out long exitTime, + out long kernelTime, + out long userTime); + + [DllImport("kernel32.dll", SetLastError = true)] + public static extern bool TerminateProcess(IntPtr processHandle, uint exitCode); + + [DllImport("kernel32.dll", SetLastError = true)] + public static extern uint WaitForSingleObject(IntPtr handle, uint milliseconds); + + [DllImport("kernel32.dll", SetLastError = true)] + public static extern bool CloseHandle(IntPtr handle); +} +'@ +} + +$remaining = New-Object 'System.Collections.Generic.List[string]' +$failed = $false +$stream = $null +$locked = $false + +try { + $stream = [System.IO.File]::Open( + $registryPath, + [System.IO.FileMode]::Open, + [System.IO.FileAccess]::ReadWrite, + [System.IO.FileShare]::ReadWrite + ) + + # Coordinate with Python msvcrt.locking and Java FileChannel.lock(). + $lockDeadline = [DateTime]::UtcNow.AddSeconds(30) + while (-not $locked) { + try { + $stream.Lock(0, 2147483647) + $locked = $true + } + catch [System.IO.IOException] { + if ([DateTime]::UtcNow -ge $lockDeadline) { + throw + } + Start-Sleep -Milliseconds 100 + } + } + + $stream.Position = 0 + $reader = [System.IO.StreamReader]::new( + $stream, [System.Text.Encoding]::UTF8, $true, 1024, $true + ) + $lines = New-Object 'System.Collections.Generic.List[string]' + while ($null -ne ($line = $reader.ReadLine())) { + $lines.Add($line) + } + $reader.Dispose() + + $desiredAccess = [uint32]0x00101001 + $unixEpochFileTime = 116444736000000000L + + foreach ($line in $lines) { + if ([string]::IsNullOrWhiteSpace($line)) { + continue + } + + $fields = $line.Split([char]'|') + $targetPid = [uint32]0 + + if ($fields.Length -gt 2 -or + -not [uint32]::TryParse($fields[0], [ref]$targetPid) -or + $targetPid -eq 0) { + Write-Output ('Preserving invalid registry entry: ' + $line) + $remaining.Add($line) + $failed = $true + continue + } + + $handle = [ConcoreNativeMethods]::OpenProcess( + $desiredAccess, $false, $targetPid + ) + + if ($handle -eq [IntPtr]::Zero) { + $errorCode = [Runtime.InteropServices.Marshal]::GetLastWin32Error() + if ($errorCode -eq 87 -or $errorCode -eq 1168) { + Write-Output ('Skipping exited PID ' + $targetPid) + } + else { + Write-Output ('Unable to open PID ' + $targetPid + ': Windows error ' + $errorCode) + $remaining.Add($line) + $failed = $true + } + continue + } + + try { + $creationTime = [long]0 + $exitTime = [long]0 + $kernelTime = [long]0 + $userTime = [long]0 + + if (-not [ConcoreNativeMethods]::GetProcessTimes( + $handle, + [ref]$creationTime, + [ref]$exitTime, + [ref]$kernelTime, + [ref]$userTime + )) { + $errorCode = [Runtime.InteropServices.Marshal]::GetLastWin32Error() + throw "GetProcessTimes failed: Windows error $errorCode" + } + + $actualStartMs = [long][Math]::Floor( + ($creationTime - $unixEpochFileTime) / 10000.0 + ) + + if ($fields.Length -eq 2 -and -not [string]::IsNullOrWhiteSpace($fields[1])) { + $expectedStartMs = [long]0 + $parsed = [long]::TryParse( + $fields[1], + [Globalization.NumberStyles]::Integer, + [Globalization.CultureInfo]::InvariantCulture, + [ref]$expectedStartMs + ) + if (-not $parsed) { + Write-Output ('Preserving invalid timestamp for PID ' + $targetPid) + $remaining.Add($line) + $failed = $true + continue + } + } + elseif ($fields.Length -eq 1) { + # Backward compatibility for old PID-only entries. Compare the + # CIM creation time with the already-open native process handle. + $candidate = Get-CimInstance -ClassName Win32_Process ` + -Filter ('ProcessId=' + $targetPid) -ErrorAction Stop + + if (-not $candidate -or $candidate.CommandLine -notmatch '(?i)concore') { + Write-Output ('Skipping unidentified legacy PID ' + $targetPid) + continue + } + + $expectedStartMs = ([DateTimeOffset]::new( + $candidate.CreationDate.ToUniversalTime() + )).ToUnixTimeMilliseconds() + } + else { + Write-Output ('Preserving invalid registry entry: ' + $line) + $remaining.Add($line) + $failed = $true + continue + } + + if ($actualStartMs -ne $expectedStartMs) { + Write-Output ('Skipping stale PID ' + $targetPid) + continue + } + + # Identity validation and termination operate on this same handle. + if (-not [ConcoreNativeMethods]::TerminateProcess($handle, 1)) { + $errorCode = [Runtime.InteropServices.Marshal]::GetLastWin32Error() + throw "TerminateProcess failed: Windows error $errorCode" + } + + $waitResult = [ConcoreNativeMethods]::WaitForSingleObject($handle, 5000) + if ($waitResult -ne 0) { + throw "Process $targetPid did not signal termination within five seconds" + } + + Write-Output ('Killing concore process ' + $targetPid) + } + catch { + Write-Output ('Unable to stop PID ' + $targetPid + ': ' + $_.Exception.Message) + $remaining.Add($line) + $failed = $true + } + finally { + [void][ConcoreNativeMethods]::CloseHandle($handle) + } + } + + # Keep the registry file itself. Deleting it after unlocking could race + # with a newly starting node appending a fresh registration. + $remainingText = '' + if ($remaining.Count -gt 0) { + $remainingText = [string]::Join( + [Environment]::NewLine, + $remaining.ToArray() + ) + [Environment]::NewLine + } + + $bytes = [System.Text.UTF8Encoding]::new($false).GetBytes($remainingText) + $stream.SetLength(0) + $stream.Position = 0 + if ($bytes.Length -gt 0) { + $stream.Write($bytes, 0, $bytes.Length) + } + $stream.Flush() +} +catch { + Write-Output ('Kill script failed: ' + $_.Exception.Message) + $failed = $true +} +finally { + if ($stream) { + if ($locked) { + try { + $stream.Unlock(0, 2147483647) + } + catch { + Write-Output ('Registry unlock failed: ' + $_.Exception.Message) + $failed = $true + } + } + $stream.Dispose() + } +} + +if ($failed) { + exit 1 +} +exit 0 diff --git a/mkconcore.py b/mkconcore.py index a1f44ef..e294666 100644 --- a/mkconcore.py +++ b/mkconcore.py @@ -269,7 +269,7 @@ def _resolve_concore_path(): concoretype = sys.argv[4] if not (concoretype in ["posix","windows","docker","macos","ubuntu"]): logging.error(" type must be posix (macos or ubuntu), windows, or docker") - sys.exit(1) + sys.exit(1) ubuntu = False #6/24/21 if concoretype == "ubuntu": concoretype = "posix" @@ -632,7 +632,7 @@ def cleanup_script_files(): fsource = open(CONCOREPATH+"/concore.py") except (FileNotFoundError, IOError): print(CONCOREPATH+" is not correct path to concore (missing python files)") - sys.exit(1) + sys.exit(1) with open(outdir+"/src/concore.py","w") as fcopy: fcopy.write(fsource.read()) fsource.close() @@ -662,7 +662,7 @@ def cleanup_script_files(): fsource = open(CONCOREPATH+"/concore.hpp") except (FileNotFoundError, IOError): print(CONCOREPATH+" is not correct path to concore (missing C++ files)") - sys.exit(1) + sys.exit(1) with open(outdir+"/src/concore.hpp","w") as fcopy: fcopy.write(fsource.read()) fsource.close() @@ -672,7 +672,7 @@ def cleanup_script_files(): fcopy.write(fbase.read()) except (FileNotFoundError, IOError): print(CONCOREPATH+" is not correct path to concore (missing concore_base.hpp)") - sys.exit(1) + sys.exit(1) if 'v' in required_langs: try: @@ -682,7 +682,7 @@ def cleanup_script_files(): fsource = open(CONCOREPATH+"/concore.v") except (FileNotFoundError, IOError): print(CONCOREPATH+" is not correct path to concore (missing Verilog files)") - sys.exit(1) + sys.exit(1) with open(outdir+"/src/concore.v","w") as fcopy: fcopy.write(fsource.read()) fsource.close() @@ -702,6 +702,14 @@ def cleanup_script_files(): fcopy.write(fcore.read()) fcore.close() +if concoretype == "windows" and required_langs.intersection({"py", "java"}): + helper_source = os.path.join(CONCOREPATH, "concorekill.ps1") + if not os.path.isfile(helper_source): + logging.error("Missing Windows process cleanup helper: " + helper_source) + sys.exit(1) + os.makedirs(os.path.join(outdir, "src"), exist_ok=True) + shutil.copy2(helper_source, os.path.join(outdir, "src", "concorekill.ps1")) + if 'm' in required_langs: try: fsource = open(CONCOREPATH+"/concore_default_maxtime.m") @@ -745,7 +753,7 @@ def cleanup_script_files(): os.chmod(outdir+"/src/mkcompile",stat.S_IRWXU) except Exception as e: print(CONCOREPATH+" is not correct path to concore (missing MATLAB files):", e) - sys.exit(1) + sys.exit(1) # --- Generate iport and oport mappings --- logging.info("Generating iport/oport mappings...") @@ -1128,6 +1136,8 @@ def cleanup_script_files(): fbuild.write("mkdir -p ./"+containername+"/"+source_subdir+"\n") if concoretype == "windows": fbuild.write("copy .\\src\\"+sourcecode+" .\\"+containername+"\\"+sourcecode+"\n") + if langext in ("py", "java"): + fbuild.write("copy .\\src\\concorekill.ps1 .\\" + containername + "\\concorekill.ps1\n") if langext == "py": fbuild.write("copy .\\src\\concore.py .\\" + containername + "\\concore.py\n") fbuild.write("copy .\\src\\concore_base.py .\\" + containername + "\\concore_base.py\n") @@ -1376,7 +1386,9 @@ def cleanup_script_files(): if concoretype=="windows": q_container = f'"{containername}"' fstop.write('cmd /C '+q_container+"\\concorekill\n") - fstop.write('del '+q_container+"\\concorekill.bat\n") + fstop.write('if not errorlevel 1 (\n') + fstop.write(' del '+q_container+"\\concorekill.bat\n") + fstop.write(')\n') else: safe_pidfile = shlex.quote(f"{containername}/concorepid") fstop.write('kill -9 `cat '+safe_pidfile+'`\n') diff --git a/tests/test_concore.py b/tests/test_concore.py index 943075d..4d6b364 100644 --- a/tests/test_concore.py +++ b/tests/test_concore.py @@ -770,7 +770,17 @@ class TestPidRegistry: @pytest.fixture(autouse=True) def use_temp_dir(self, temp_dir, monkeypatch): + import shutil self.temp_dir = temp_dir + helper_source = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "concorekill.ps1", + ) + assert os.path.isfile(helper_source), "concorekill.ps1 is missing" + shutil.copy2( + helper_source, + os.path.join(temp_dir, "concorekill.ps1"), + ) monkeypatch.chdir(temp_dir) import concore @@ -793,7 +803,7 @@ def test_register_pid_creates_registry_file(self): assert os.path.exists(_PID_REGISTRY_FILE) with open(_PID_REGISTRY_FILE) as f: pids = [line.strip() for line in f if line.strip()] - assert str(os.getpid()) in pids + assert any(entry.split("|", 1)[0] == str(os.getpid()) for entry in pids) def test_register_pid_appends_not_overwrites(self): from concore import _register_pid, _PID_REGISTRY_FILE @@ -806,7 +816,7 @@ def test_register_pid_appends_not_overwrites(self): pids = [line.strip() for line in f if line.strip()] assert "11111" in pids assert "22222" in pids - assert str(os.getpid()) in pids + assert any(entry.split("|", 1)[0] == str(os.getpid()) for entry in pids) assert len(pids) == 3 def test_cleanup_pid_removes_current_pid(self): @@ -824,17 +834,19 @@ def test_cleanup_pid_removes_current_pid(self): assert "99999" in pids assert "88888" in pids - def test_cleanup_pid_deletes_files_when_last_pid(self): + def test_cleanup_pid_preserves_files_when_last_pid(self): from concore import _cleanup_pid, _PID_REGISTRY_FILE, _KILL_SCRIPT_FILE - current_pid = str(os.getpid()) with open(_PID_REGISTRY_FILE, "w") as f: f.write(current_pid + "\n") with open(_KILL_SCRIPT_FILE, "w") as f: f.write("@echo off\n") _cleanup_pid() - assert not os.path.exists(_PID_REGISTRY_FILE) - assert not os.path.exists(_KILL_SCRIPT_FILE) + assert os.path.isfile(_PID_REGISTRY_FILE) + with open(_PID_REGISTRY_FILE) as f: + assert f.read().strip() == "" + assert os.path.isfile(_KILL_SCRIPT_FILE) + def test_cleanup_pid_handles_missing_registry(self): from concore import _cleanup_pid, _PID_REGISTRY_FILE @@ -849,10 +861,10 @@ def test_write_kill_script_generates_bat_file(self): assert os.path.exists(_KILL_SCRIPT_FILE) with open(_KILL_SCRIPT_FILE) as f: content = f.read() - assert os.path.basename(_PID_REGISTRY_FILE) in content - assert "wmic" not in content - assert "Get-CimInstance" in content - assert "taskkill" in content + assert "concorekill.ps1" in content + assert "powershell.exe" in content.lower() + assert "Get-CimInstance" not in content + assert "taskkill" not in content.lower() assert "concore" in content.lower() @pytest.mark.skipif(sys.platform != "win32", reason="Windows only") @@ -875,12 +887,95 @@ def test_kill_script_kills_concore_process(self): ) assert "Killing concore process " + str(proc.pid) in result.stdout proc.wait(timeout=10) - assert not os.path.exists(_PID_REGISTRY_FILE) - assert not os.path.exists(_KILL_SCRIPT_FILE) + assert os.path.isfile(_PID_REGISTRY_FILE) + with open(_PID_REGISTRY_FILE) as f: + assert f.read().strip() == "" + assert os.path.isfile(_KILL_SCRIPT_FILE) finally: if proc.poll() is None: proc.kill() + + @pytest.mark.skipif(sys.platform != "win32", reason="Windows only") + def test_kill_script_kills_registered_controller(self): + import subprocess + import time + from concore import _write_kill_script, _KILL_SCRIPT_FILE, _PID_REGISTRY_FILE + base_python = os.path.join(sys.base_prefix, "python.exe") + controller_file = os.path.join(self.temp_dir, "controller.py") + assert "concore" not in base_python.lower() + assert "concore" not in controller_file.lower() + with open(controller_file, "w") as f: + f.write("import concore\n") + f.write("import time\n") + f.write("time.sleep(60)\n") + env = os.environ.copy() + project_root = os.path.dirname( + os.path.dirname(os.path.abspath(__file__)) + ) + venv_site_packages = os.path.join( + sys.prefix, "Lib", "site-packages" + ) + pythonpath_entries = [ + project_root, + venv_site_packages + ] + existing_path = env.get("PYTHONPATH") + if existing_path: + pythonpath_entries.append(existing_path) + env["PYTHONPATH"] = os.pathsep.join(pythonpath_entries) + proc = subprocess.Popen( + [base_python, controller_file], + cwd=self.temp_dir, + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + controller_entry = None + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + if proc.poll() is not None: + raise AssertionError( + f"Controller exited early: {proc.returncode}" + ) + try: + with open(_PID_REGISTRY_FILE) as f: + pids = [ + line.strip() + for line in f + if line.strip() + ] + except FileNotFoundError: + pids = [] + for entry in pids: + if entry.split('|', 1)[0] == str(proc.pid): + controller_entry = entry + break + if controller_entry is not None: + break + time.sleep(0.1) + else: + raise AssertionError( + "Controller did not register its PID" + ) + with open(_PID_REGISTRY_FILE, "w") as f: + f.write(controller_entry + "\n") + _write_kill_script() + result = subprocess.run( + ["cmd", "/c", _KILL_SCRIPT_FILE], + capture_output=True, + text=True, + timeout=60, + ) + assert ( + f"Killing concore process {proc.pid}" in result.stdout + ), f"Unexpected output: {result.stdout!r}" + proc.wait(timeout=10) + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=10) def test_multi_node_registration(self): from concore import _register_pid, _PID_REGISTRY_FILE @@ -893,7 +988,7 @@ def test_multi_node_registration(self): pids = [line.strip() for line in f if line.strip()] for pid in fake_pids: assert pid in pids - assert str(os.getpid()) in pids + assert any(entry.split("|", 1)[0] == str(os.getpid()) for entry in pids) assert len(pids) == 4 def test_cleanup_preserves_other_pids(self): @@ -927,5 +1022,5 @@ def test_import_registers_pid_on_windows(self): assert os.path.exists(concore._PID_REGISTRY_FILE) with open(concore._PID_REGISTRY_FILE) as f: pids = [line.strip() for line in f if line.strip()] - assert str(os.getpid()) in pids + assert any(entry.split("|", 1)[0] == str(os.getpid()) for entry in pids) importlib.reload(concore)