From 479c09abefb55db4a9f98111188a1d9ba530f7a3 Mon Sep 17 00:00:00 2001
From: "Caleb St. John" <30729806+yocalebo@users.noreply.github.com>
Date: Tue, 17 Dec 2024 14:45:08 -0500
Subject: [PATCH] NAS-133125 / 25.04 / remove psutil from webshell_app.py
 (#15224)

* remove psutil from webshell_app.py

* address review

* address review
---
 .../middlewared/apps/webshell_app.py          | 28 ++---------------
 src/middlewared/middlewared/utils/os.py       | 31 ++++++++++++++-----
 2 files changed, 27 insertions(+), 32 deletions(-)

diff --git a/src/middlewared/middlewared/apps/webshell_app.py b/src/middlewared/middlewared/apps/webshell_app.py
index 679b8133c038..8663db1d0131 100644
--- a/src/middlewared/middlewared/apps/webshell_app.py
+++ b/src/middlewared/middlewared/apps/webshell_app.py
@@ -5,16 +5,12 @@
 import json
 import os
 import queue
-import signal
 import struct
 import termios
 import threading
 import time
 import uuid
 
-
-import psutil
-
 from middlewared.api.base.server.ws_handler.base import BaseWebSocketHandler
 from middlewared.service_exception import (
     CallError,
@@ -22,7 +18,7 @@
     InstanceNotFound,
     MatchNotFound,
 )
-from middlewared.utils.os import close_fds
+from middlewared.utils.os import close_fds, terminate_pid
 
 __all__ = ("ShellApplication",)
 
@@ -200,7 +196,7 @@ def abort(self):
         asyncio.run_coroutine_threadsafe(self.ws.close(), self.loop)
 
         with contextlib.suppress(ProcessLookupError):
-            os.kill(self.shell_pid, signal.SIGTERM)
+            terminate_pid(self.shell_pid, timeout=5, use_pgid=True)
 
         self.die()
 
@@ -363,22 +359,4 @@ async def run(self, ws, origin, conndata):
         return ws
 
     async def worker_kill(self, t_worker):
-        def worker_kill_impl():
-            # If connection has been closed lets make sure shell is killed
-            if t_worker.shell_pid:
-                with contextlib.suppress(psutil.NoSuchProcess):
-                    shell = psutil.Process(t_worker.shell_pid)
-                    to_terminate = [shell] + shell.children(recursive=True)
-
-                    for p in to_terminate:
-                        with contextlib.suppress(psutil.NoSuchProcess):
-                            p.terminate()
-                    gone, alive = psutil.wait_procs(to_terminate, timeout=2)
-
-                    for p in alive:
-                        with contextlib.suppress(psutil.NoSuchProcess):
-                            p.kill()
-
-            t_worker.join()
-
-        await self.middleware.run_in_thread(worker_kill_impl)
+        await self.middleware.run_in_thread(t_worker.abort)
diff --git a/src/middlewared/middlewared/utils/os.py b/src/middlewared/middlewared/utils/os.py
index 665069778fae..242e6cbb84a1 100644
--- a/src/middlewared/middlewared/utils/os.py
+++ b/src/middlewared/middlewared/utils/os.py
@@ -1,7 +1,7 @@
 from collections.abc import Generator
 from dataclasses import dataclass
 from functools import cached_property
-from os import closerange, kill, scandir
+from os import closerange, getpgid, kill, killpg, scandir
 from resource import getrlimit, RLIMIT_NOFILE, RLIM_INFINITY
 from signal import SIGKILL, SIGTERM
 from time import sleep, time
@@ -42,12 +42,29 @@ def close_fds(low_fd, max_fd=None):
     closerange(low_fd, max_fd)
 
 
-def terminate_pid(pid: int, timeout: int = 10) -> bool:
-    # Send SIGTERM to request the process to terminate
-    kill(pid, SIGTERM)
+def terminate_pid(pid: int, timeout: int = 10, use_pgid: bool = False) -> bool:
+    """
+    Send SIGTERM to `pid` and wait `timeout` seconds for the process to terminate.
+    If the process is still alive after `timeout`, proceed to send SIGKILL to the
+    process.
+
+    `pid`: integer represents the PID to terminate.
+    `timeout`: integer represents the total time (in seconds) to wait
+        before sending SIGKILL to `pid` if the process doesn't honor
+        SIGTERM or takes longer than `timeout` seconds to terminate.
+    `use_pgid`: boolean, if true will lookup the process group id of
+        `pid` and apply the same logic.
+    """
+    pid_or_pgid, method = pid, kill
+    if use_pgid:
+        method = killpg
+        pid_or_pgid = getpgid(pid)
+
+    # Send SIGTERM to request the process (group) to terminate
+    method(pid_or_pgid, SIGTERM)
 
     try:
-        kill(pid, ALIVE_SIGNAL)
+        method(pid_or_pgid, ALIVE_SIGNAL)
     except ProcessLookupError:
         # SIGTERM was honored
         return True
@@ -56,7 +73,7 @@ def terminate_pid(pid: int, timeout: int = 10) -> bool:
     start_time = time()
     while True:
         try:
-            kill(pid, ALIVE_SIGNAL)
+            method(pid_or_pgid, ALIVE_SIGNAL)
         except ProcessLookupError:
             # SIGTERM was honored (eventually)
             return True
@@ -70,7 +87,7 @@ def terminate_pid(pid: int, timeout: int = 10) -> bool:
 
     try:
         # Send SIGKILL to forcefully terminate the process
-        kill(pid, SIGKILL)
+        method(pid_or_pgid, SIGKILL)
         return False
     except ProcessLookupError:
         # Process may have terminated between checks