From 9ecbdea371d290800860300a2cf0bf719a661e4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Doma=C5=84ski?= <jandom@gmail.com>
Date: Wed, 15 Nov 2023 07:16:28 +0000
Subject: [PATCH] fix #252: handle keyboard interrupts (#270)

* fix #252: Handle KeyboardInterrupts when using MDRunner.run
* adding tests
* update CHANGES
---
 CHANGES           |  1 +
 gromacs/run.py    | 30 +++++++++++++++++++++++++-----
 tests/test_run.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/CHANGES b/CHANGES
index a533a43e..f6602ad8 100644
--- a/CHANGES
+++ b/CHANGES
@@ -11,6 +11,7 @@ orbeckst, jandom, njzjz
   file (#261)
 * fixed Python 3.12: No module named ('pkg_resources' #263)
 * fixed AttributeDict does not support hasattr (#214)
+* fixed handle KeyboardInterrupts when using MDRunner.run() (#252)
 
 
 2023-09-16      0.8.5
diff --git a/gromacs/run.py b/gromacs/run.py
index af7094c2..3acf9313 100644
--- a/gromacs/run.py
+++ b/gromacs/run.py
@@ -102,6 +102,7 @@ class MDrunnerMPI(gromacs.run.MDrunner):
 import subprocess
 import os.path
 import errno
+import signal
 
 # logging
 import logging
@@ -230,6 +231,20 @@ def __init__(self, dirname=os.path.curdir, **kwargs):
         self.logname = os.path.realpath(
             os.path.join(self.dirname, self.filename(logname, ext="log"))
         )
+        self.process = None
+        self.signal_handled = False
+        signal.signal(signal.SIGINT, self.signal_handler)
+
+    def signal_handler(self, signum, frame):
+        """Custom signal handler for SIGINT."""
+        if self.process is not None:
+            try:
+                self.process.terminate()  # Attempt to terminate the subprocess
+                self.process.wait()  # Wait for the subprocess to terminate
+                self.signal_handled = True
+            except Exception as e:
+                logger.error(f"Error terminating subprocess: {e}")
+        raise KeyboardInterrupt  # Re-raise the KeyboardInterrupt to exit the main script
 
     def commandline(self, **mpiargs):
         """Returns simple command line to invoke mdrun.
@@ -308,17 +323,22 @@ def run(self, pre=None, post=None, mdrunargs=None, **mpiargs):
             try:
                 self.prehook(**pre)
                 logger.info(" ".join(cmd))
-                rc = subprocess.call(cmd)
+                self.process = subprocess.Popen(cmd)  # Use Popen instead of call
+                returncode = self.process.wait()  # Wait for the process to complete
+            except KeyboardInterrupt:
+                # Handle the keyboard interrupt gracefully
+                logger.info("Keyboard Interrupt received, terminating the subprocess.")
+                raise
             except:
                 logger.exception("Failed MD run for unknown reasons.")
                 raise
             finally:
                 self.posthook(**post)
-        if rc == 0:
-            logger.info("MDrun completed ok, returncode = {0:d}".format(rc))
+        if returncode == 0:
+            logger.info("MDrun completed ok, returncode = {0:d}".format(returncode))
         else:
-            logger.critical("Failure in MDrun, returncode = {0:d}".format(rc))
-        return rc
+            logger.critical("Failure in MDrun, returncode = {0:d}".format(returncode))
+        return returncode
 
     def run_check(self, **kwargs):
         """Run :program:`mdrun` and check if run completed when it finishes.
diff --git a/tests/test_run.py b/tests/test_run.py
index d4693fc3..12dd64ca 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -5,7 +5,11 @@
 
 from __future__ import division, absolute_import, print_function
 
+import time
 import pytest
+import os
+import signal
+import threading
 
 from .datafiles import datafile
 
@@ -54,6 +58,49 @@ def test_MDRunner():
     assert rc == 0, "mdrun failed to run through MDrunner"
 
 
+def test_MDRunner_keyboard_interrupt(monkeypatch):
+    """Test that a keyboard interrupt is handled correctly."""
+
+    # Create a mock for subprocess.Popen
+    class MockPopen:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def wait(self):
+            time.sleep(2)  # Simulate a long-running process
+
+        def terminate(self):
+            pass
+
+    # Use monkeypatch to replace subprocess.Popen with MockPopen
+    monkeypatch.setattr("gromacs.run.subprocess.Popen", MockPopen)
+
+    # Initialize MDrunner
+    mdrunner = gromacs.run.MDrunner()
+
+    # Function to send SIGINT after a delay
+    def send_interrupt():
+        time.sleep(1)  # Short delay before sending SIGINT
+        os.kill(os.getpid(), signal.SIGINT)
+
+    # Start a thread to send the SIGINT
+    interrupt_thread = threading.Thread(target=send_interrupt)
+    interrupt_thread.start()
+
+    # Run the MDrunner in the main thread to handle the signal
+    try:
+        mdrunner.run()
+    except KeyboardInterrupt:
+        # Handle the KeyboardInterrupt within the test
+        pass
+
+    # Ensure the signal handler was called
+    assert mdrunner.signal_handled, "The signal handler was not called as expected"
+
+    # Ensure the interrupt thread has finished
+    interrupt_thread.join()
+
+
 class Test_find_gromacs_command(object):
     # Gromacs 4 or Gromacs 5 (in this order)
     commands = ["grompp", "gmx grompp"]