1 files changed, 114 insertions, 0 deletions
diff --git a/tests/topotests/munet/cleanup.py b/tests/topotests/munet/cleanup.py
new file mode 100644
index 0000000..c641cda
--- /dev/null
+++ b/tests/topotests/munet/cleanup.py
@@ -0,0 +1,114 @@
+# -*- coding: utf-8 eval: (blacken-mode 1) -*-
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# September 30 2021, Christian Hopps <chopps@labn.net>
+#
+# Copyright 2021, LabN Consulting, L.L.C.
+#
+"""Provides functionality to cleanup processes on posix systems."""
+import glob
+import logging
+import os
+import signal
+
+
+def get_pids_with_env(has_var, has_val=None):
+    result = {}
+    for pidenv in glob.iglob("/proc/*/environ"):
+        pid = pidenv.split("/")[2]
+        try:
+            with open(pidenv, "rb") as rfb:
+                envlist = [
+                    x.decode("utf-8").split("=", 1) for x in rfb.read().split(b"\0")
+                ]
+                envlist = [[x[0], ""] if len(x) == 1 else x for x in envlist]
+                envdict = dict(envlist)
+                if has_var not in envdict:
+                    continue
+                if has_val is None:
+                    result[pid] = envdict
+                elif envdict[has_var] == str(has_val):
+                    result[pid] = envdict
+        except Exception:
+            # E.g., process exited and files are gone
+            pass
+    return result
+
+
+def _kill_piddict(pids_by_upid, sig):
+    ourpid = str(os.getpid())
+    for upid, pids in pids_by_upid:
+        logging.info("Sending %s to (%s) of munet pid %s", sig, ", ".join(pids), upid)
+        for pid in pids:
+            try:
+                if pid != ourpid:
+                    cmdline = open(f"/proc/{pid}/cmdline", "r", encoding="ascii").read()
+                    cmdline = cmdline.replace("\x00", " ")
+                    logging.info("killing proc %s (%s)", pid, cmdline)
+                    os.kill(int(pid), sig)
+            except Exception:
+                pass
+
+
+def _get_our_pids():
+    ourpid = str(os.getpid())
+    piddict = get_pids_with_env("MUNET_PID", ourpid)
+    pids = [x for x in piddict if x != ourpid]
+    if pids:
+        return {ourpid: pids}
+    return {}
+
+
+def _get_other_pids():
+    piddict = get_pids_with_env("MUNET_PID")
+    unet_pids = {d["MUNET_PID"] for d in piddict.values()}
+    pids_by_upid = {p: set() for p in unet_pids}
+    for pid, envdict in piddict.items():
+        unet_pid = envdict["MUNET_PID"]
+        pids_by_upid[unet_pid].add(pid)
+    # Filter out any child pid sets whos munet pid is still running
+    return {x: y for x, y in pids_by_upid.items() if x not in y}
+
+
+def _get_pids_by_upid(ours):
+    if ours:
+        return _get_our_pids()
+    return _get_other_pids()
+
+
+def _cleanup_pids(ours):
+    pids_by_upid = _get_pids_by_upid(ours).items()
+    if not pids_by_upid:
+        return
+
+    t = "current" if ours else "previous"
+    logging.info("Reaping %s munet processes", t)
+
+    # _kill_piddict(pids_by_upid, signal.SIGTERM)
+
+    # # Give them 5 second to exit cleanly
+    # logging.info("Waiting up to 5s to allow for clean exit of abandon'd pids")
+    # for _ in range(0, 5):
+    #     pids_by_upid = _get_pids_by_upid(ours).items()
+    #     if not pids_by_upid:
+    #         return
+    #     time.sleep(1)
+
+    pids_by_upid = _get_pids_by_upid(ours).items()
+    _kill_piddict(pids_by_upid, signal.SIGKILL)
+
+
+def cleanup_current():
+    """Attempt to cleanup preview runs.
+
+    Currently this only scans for old processes.
+    """
+    _cleanup_pids(True)
+
+
+def cleanup_previous():
+    """Attempt to cleanup preview runs.
+
+    Currently this only scans for old processes.
+    """
+    _cleanup_pids(False)