use process_name_prefix to kill ustreamer

2 files changed, 17 insertions, 23 deletions

diff --git a/kvmd/apps/cleanup/__init__.py b/kvmd/apps/cleanup/__init__.py
index 3314066c..57823ff1 100644
--- a/kvmd/apps/cleanup/__init__.py
+++ b/kvmd/apps/cleanup/__init__.py
@@ -72,22 +72,23 @@ def _clear_gpio(config: Section) -> None:
 def _kill_streamer(config: Section) -> None:
     logger = get_logger(0)
 
-    streamer = os.path.basename(config.streamer.cmd[0])
+    if config.streamer.process_name_prefix:
+        prefix = config.streamer.process_name_prefix + ":"
+        logger.info("Trying to find and kill the streamer %r ...", prefix + " <app>")
 
-    logger.info("Trying to find and kill %r ...", streamer)
-    for proc in psutil.process_iter():
-        attrs = proc.as_dict(attrs=["name"])
-        if os.path.basename(attrs.get("name", "")) == streamer:
-            try:
-                proc.send_signal(signal.SIGTERM)
-            except Exception:
-                logger.exception("Can't send SIGTERM to streamer with pid=%d", proc.pid)
-            time.sleep(3)
-            if proc.is_running():
+        for proc in psutil.process_iter():
+            attrs = proc.as_dict(attrs=["name"])
+            if attrs.get("name", "").startswith(prefix):
                 try:
-                    proc.send_signal(signal.SIGKILL)
+                    proc.send_signal(signal.SIGTERM)
                 except Exception:
-                    logger.exception("Can't send SIGKILL to streamer with pid=%d", proc.pid)
+                    logger.exception("Can't send SIGTERM to streamer with pid=%d", proc.pid)
+                time.sleep(3)
+                if proc.is_running():
+                    try:
+                        proc.send_signal(signal.SIGKILL)
+                    except Exception:
+                        logger.exception("Can't send SIGKILL to streamer with pid=%d", proc.pid)
 
 
 def _remove_sockets(config: Section) -> None:
diff --git a/testenv/tests/apps/cleanup/test_main.py b/testenv/tests/apps/cleanup/test_main.py
index 3e2e4c72..f9b74a06 100644
--- a/testenv/tests/apps/cleanup/test_main.py
+++ b/testenv/tests/apps/cleanup/test_main.py
@@ -21,7 +21,6 @@
 
 
 import os
-import secrets
 import multiprocessing
 import multiprocessing.queues
 import time
@@ -35,16 +34,13 @@ from kvmd.apps.cleanup import main
 def test_ok(tmpdir) -> None:  # type: ignore
     queue: multiprocessing.queues.Queue = multiprocessing.Queue()
 
-    ustreamer_tmp_path = os.path.abspath(str(tmpdir.join("ustr-" + secrets.token_hex(3))))
-    os.symlink("/usr/bin/ustreamer", ustreamer_tmp_path)
-
     ustreamer_sock_path = os.path.abspath(str(tmpdir.join("ustreamer-fake.sock")))
     open(ustreamer_sock_path, "w").close()
     kvmd_sock_path = os.path.abspath(str(tmpdir.join("kvmd-fake.sock")))
     open(kvmd_sock_path, "w").close()
 
     def ustreamer_fake() -> None:
-        setproctitle.setproctitle(os.path.basename(ustreamer_tmp_path))
+        setproctitle.setproctitle("kvmd/streamer: /usr/bin/ustreamer")
         queue.put(True)
         while True:
             time.sleep(1)
@@ -57,11 +53,8 @@ def test_ok(tmpdir) -> None:  # type: ignore
     main([
         "kvmd-cleanup",
         "--set-options",
-        "kvmd/server/port=0",
-        "kvmd/server/unix=" + kvmd_sock_path,
-        "kvmd/streamer/port=0",
-        "kvmd/streamer/unix=" + ustreamer_sock_path,
-        "kvmd/streamer/cmd=" + ustreamer_tmp_path,
+        f"kvmd/server/unix={kvmd_sock_path}",
+        f"kvmd/streamer/unix={ustreamer_sock_path}",
     ])
 
     assert not os.path.exists(ustreamer_sock_path)