Use pidfile to kill duplicate Astrality processes

CHANGELOG.rst

@@ -156,6 +156,8 @@ Changed
   ``$XDG_DATA_HOME/astrality/compilations`` to be used as the compilation
   target. This behaves better than temporary files when programs expect
   files to still be present after Astrality restarts.
+- Astrality is now more conservative when killing duplicate Astrality processes
+  by using a *pidfile* instead of ``pgrep -f astrality``.
 
 
 Fixed
@@ -170,3 +172,5 @@ Fixed
   paths.
 
 - Module option ``requires_timeout`` is now respected.
+- Astrality no longer kills processes containing "astrality" in their command
+  line invocation.

astrality/astrality.py

@@ -5,13 +5,16 @@
 import logging
 import os
 import signal
-import subprocess
 import sys
 import time
-from typing import Set, List
+from typing import List
 
+import psutil
+
+from astrality import utils
 from astrality.config import user_configuration
 from astrality.module import ModuleManager
+from astrality.xdg import XDG
 
 
 logger = logging.getLogger(__name__)
@@ -38,7 +41,7 @@ def main(
     # Set the logging level to the configured setting
     logging.basicConfig(level=logging_level)
 
-    if not modules and not dry_run:
+    if not modules and not dry_run and not test:
         # Quit old astrality instances
         kill_old_astrality_processes()
 
@@ -145,39 +148,50 @@ def exit_handler(signal=None, frame=None) -> None:
         exit_handler()
 
 
-def other_astrality_pids() -> Set[int]:
-    """Return the process ids (PIDs) of any other Astrality instances."""
-    # Get all processes instanciated from this file
-    result = subprocess.Popen(
-        ['pgrep', '-f', 'astrality'],
-        stdout=subprocess.PIPE,
-        universal_newlines=True,
+def kill_old_astrality_processes() -> None:
+    """
+    Kill any previous Astrality process instance.
+
+    This process kills the last process which invoked this function.
+    If the process is no longer running, it is owned by another user, or has
+    a new create_time, it will *not* be killed.
+    """
+    # The current process
+    new_process = psutil.Process()
+
+    # Fetch info of possible previous process instance
+    pidfile = XDG().data('astrality.pid')
+    old_process_info = utils.load_yaml(path=pidfile)
+    utils.dump_yaml(
+        data=new_process.as_dict(attrs=['pid', 'create_time', 'username']),
+        path=pidfile,
     )
-    pids = set(int(pid.strip()) for pid in result.stdout)
 
-    # Return all the PIDs except for the PID of this process
-    this_process_pid = os.getpid()
-    return pids - set((this_process_pid,))
+    if not old_process_info or not psutil.pid_exists(old_process_info['pid']):
+        return
+
+    try:
+        old_process = psutil.Process(pid=old_process_info['pid'])
+    except BaseException:
+        return
 
+    if not old_process.as_dict(
+        attrs=['pid', 'create_time', 'username'],
+    ) == old_process_info:
+        return
 
-def kill_old_astrality_processes() -> None:
-    """Kill all other instances of this script, to prevent duplicates."""
-    pids = other_astrality_pids()
-    failed_exits = 0
-    for pid in pids:
-        try:
-            logger.info(f'Killing duplicate Astrality process with pid {pid}.')
-            os.kill(pid, signal.SIGTERM)
-        except OSError:
-            logger.error(
-                f'Could not kill old instance of astrality with pid {pid}.',
-            )
-            logger.error('Continuing anyway...')
-            failed_exits += 1
-
-    while len(other_astrality_pids()) > failed_exits:
-        # Wait for all the processes to exit properly
-        time.sleep(0.2)
+    try:
+        logger.info(
+            'Killing duplicate Astrality process with pid: '
+            f'{old_process.pid}.',
+        )
+        old_process.terminate()
+        old_process.wait()
+    except BaseException:
+        logger.error(
+            f'Could not kill old instance of astrality with pid: '
+            f'{old_process.pid}. Continuing anyway...',
+        )
 
 
 if __name__ == '__main__':

astrality/tests/test_astrality.py

@@ -1,12 +1,16 @@
 import logging
 import os
+import psutil
 import signal
 import subprocess
 import time
 
 import pytest
 
-from astrality.astrality import main
+from astrality.astrality import main, kill_old_astrality_processes
+from astrality import utils
+from astrality.tests.utils import Retry
+from astrality.xdg import XDG
 
 
 @pytest.mark.slow
@@ -57,3 +61,70 @@ def test_enabling_specific_module_from_command_line(
         logging.INFO,
         'Greetings from Dhaka!',
     ) in caplog.record_tuples
+
+
+class TestKillOldAstralityProcesses:
+    """Tests for astrality.astrality.kill_old_astrality_processes."""
+
+    def test_killing_old_running_process(self):
+        """The same running process should be killed."""
+        perpetual_process = psutil.Popen([
+            'python',
+            '-c',
+            '"from time import sleep; sleep(9999999999999)"',
+        ])
+        pidfile = XDG().data('astrality.pid')
+        utils.dump_yaml(
+            data=perpetual_process.as_dict(
+                attrs=['pid', 'create_time', 'username'],
+            ),
+            path=pidfile,
+        )
+        kill_old_astrality_processes()
+        assert Retry()(lambda: not perpetual_process.is_running())
+
+    def test_not_killing_new_procces_with_same_pid(self):
+        """The process should not be killed when it is not the original saved"""
+        perpetual_process = psutil.Popen([
+            'python',
+            '-c',
+            '"from time import sleep; sleep(9999999999999)"',
+        ])
+
+        process_data = perpetual_process.as_dict(
+            attrs=['pid', 'create_time', 'username'],
+        )
+        process_data['create_time'] += 1
+
+        utils.dump_yaml(
+            data=process_data,
+            path=XDG().data('astrality.pid'),
+        )
+        kill_old_astrality_processes()
+        assert Retry()(lambda: perpetual_process.is_running())
+        perpetual_process.kill()
+
+    def test_trying_to_kill_process_no_longer_running(self):
+        """No longer running processes should be handled gracefully."""
+        finished_process = psutil.Popen(['echo', 'Done!'])
+        process_data = finished_process.as_dict(
+            attrs=['pid', 'create_time', 'username'],
+        )
+        finished_process.wait()
+
+        utils.dump_yaml(
+            data=process_data,
+            path=XDG().data('astrality.pid'),
+        )
+        kill_old_astrality_processes()
+
+    def test_killing_processes_when_no_previous_command_has_been_run(self):
+        """The first ever invocation of the function should be handled."""
+        pidfile = XDG().data('astrality.pid')
+        pidfile.unlink()
+        kill_old_astrality_processes()
+        assert utils.load_yaml(
+            path=pidfile,
+        ) == psutil.Process().as_dict(
+            attrs=['pid', 'create_time', 'username'],
+        )