tests/topotests/munet/cleanup.py

   1 # -*- coding: utf-8 eval: (blacken-mode 1) -*-
   2 # SPDX-License-Identifier: GPL-2.0-or-later
   3 #
   4 # September 30 2021, Christian Hopps <chopps@labn.net>
   5 #
   6 # Copyright 2021, LabN Consulting, L.L.C.
   7 #
   8 """Provides functionality to cleanup processes on posix systems."""
   9 import glob
  10 import logging
  11 import os
  12 import signal
  13
  14
  15 def get_pids_with_env(has_var, has_val=None):
  16     result = {}
  17     for pidenv in glob.iglob("/proc/*/environ"):
  18         pid = pidenv.split("/")[2]
  19         try:
  20             with open(pidenv, "rb") as rfb:
  21                 envlist = [
  22                     x.decode("utf-8").split("=", 1) for x in rfb.read().split(b"\0")
  23                 ]
  24                 envlist = [[x[0], ""] if len(x) == 1 else x for x in envlist]
  25                 envdict = dict(envlist)
  26                 if has_var not in envdict:
  27                     continue
  28                 if has_val is None:
  29                     result[pid] = envdict
  30                 elif envdict[has_var] == str(has_val):
  31                     result[pid] = envdict
  32         except Exception:
  33             # E.g., process exited and files are gone
  34             pass
  35     return result
  36
  37
  38 def _kill_piddict(pids_by_upid, sig):
  39     ourpid = str(os.getpid())
  40     for upid, pids in pids_by_upid:
  41         logging.info("Sending %s to (%s) of munet pid %s", sig, ", ".join(pids), upid)
  42         for pid in pids:
  43             try:
  44                 if pid != ourpid:
  45                     cmdline = open(f"/proc/{pid}/cmdline", "r", encoding="ascii").read()
  46                     cmdline = cmdline.replace("\x00", " ")
  47                     logging.info("killing proc %s (%s)", pid, cmdline)
  48                     os.kill(int(pid), sig)
  49             except Exception:
  50                 pass
  51
  52
  53 def _get_our_pids():
  54     ourpid = str(os.getpid())
  55     piddict = get_pids_with_env("MUNET_PID", ourpid)
  56     pids = [x for x in piddict if x != ourpid]
  57     if pids:
  58         return {ourpid: pids}
  59     return {}
  60
  61
  62 def _get_other_pids():
  63     piddict = get_pids_with_env("MUNET_PID")
  64     unet_pids = {d["MUNET_PID"] for d in piddict.values()}
  65     pids_by_upid = {p: set() for p in unet_pids}
  66     for pid, envdict in piddict.items():
  67         unet_pid = envdict["MUNET_PID"]
  68         pids_by_upid[unet_pid].add(pid)
  69     # Filter out any child pid sets whos munet pid is still running
  70     return {x: y for x, y in pids_by_upid.items() if x not in y}
  71
  72
  73 def _get_pids_by_upid(ours):
  74     if ours:
  75         return _get_our_pids()
  76     return _get_other_pids()
  77
  78
  79 def _cleanup_pids(ours):
  80     pids_by_upid = _get_pids_by_upid(ours).items()
  81     if not pids_by_upid:
  82         return
  83
  84     t = "current" if ours else "previous"
  85     logging.info("Reaping %s munet processes", t)
  86
  87     # _kill_piddict(pids_by_upid, signal.SIGTERM)
  88
  89     # # Give them 5 second to exit cleanly
  90     # logging.info("Waiting up to 5s to allow for clean exit of abandon'd pids")
  91     # for _ in range(0, 5):
  92     #     pids_by_upid = _get_pids_by_upid(ours).items()
  93     #     if not pids_by_upid:
  94     #         return
  95     #     time.sleep(1)
  96
  97     pids_by_upid = _get_pids_by_upid(ours).items()
  98     _kill_piddict(pids_by_upid, signal.SIGKILL)
  99
 100
 101 def cleanup_current():
 102     """Attempt to cleanup preview runs.
 103
 104     Currently this only scans for old processes.
 105     """
 106     _cleanup_pids(True)
 107
 108
 109 def cleanup_previous():
 110     """Attempt to cleanup preview runs.
 111
 112     Currently this only scans for old processes.
 113     """
 114     _cleanup_pids(False)