Clean up popen objects when calling start_process directly

Signed-off-by: Min Lim <min.yeol.lim@intel.com>

Author: mlim19

gprofiler/profilers/python_ebpf.py

@@ -34,7 +34,6 @@
 from gprofiler.profilers import python
 from gprofiler.profilers.profiler_base import ProfilerBase
 from gprofiler.utils import (
-    poll_process,
     random_prefix,
     reap_process,
     resource_path,
@@ -187,14 +186,14 @@ def test(self) -> None:
         # pyperf sometimes has a lot of output to stdout and stderr, which makes the process halt until read.
         process = start_process(cmd, tmpdir=self._pyperf_staticx_tmpdir, pipesize=1024 * 1024)
         try:
-            poll_process(process, self._POLL_TIMEOUT, self._profiler_state.stop_event)
-        except TimeoutError:
+            wait_event(self._POLL_TIMEOUT, self._profiler_state.stop_event, lambda: process.poll() is not None)
+        except (TimeoutError, StopEventSetException):
             process.kill()
             raise
         else:
-            cleanup_process_reference(process)
             self._check_output(process, self.output_path)
         finally:
+            cleanup_process_reference(process)
             self._staticx_cleanup()
 
     def start(self) -> None:

gprofiler/utils/__init__.py

@@ -33,7 +33,7 @@
 from pathlib import Path
 from subprocess import CompletedProcess, Popen, TimeoutExpired
 from tempfile import TemporaryDirectory
-from threading import Event
+from threading import Event, Thread
 from types import FrameType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
@@ -93,6 +93,7 @@ def start_process(
     cmd: Union[str, List[str]],
     via_staticx: bool = False,
     tmpdir: Optional[Path] = None,
+    direct_call: bool = True,
     **kwargs: Any,
 ) -> Popen:
     if isinstance(cmd, str):
@@ -135,7 +136,16 @@ def start_process(
         env=env,
         **kwargs,
     )
+
+    def process_exit_handler() -> None:
+        process.communicate()
+        # Remove from _processes and do any cleanup
+        cleanup_process_reference(process)
+        logger.critical(f"Process {process.pid} exited, total processes tracked {len(_processes)}")
+
     _processes.append(process)
+    if direct_call:
+        Thread(target=process_exit_handler, daemon=True).start()
     return process
 
 
@@ -152,14 +162,6 @@ def wait_event(timeout: float, stop_event: Event, condition: Callable[[], bool],
             raise TimeoutError()
 
 
-def poll_process(process: Popen, timeout: float, stop_event: Event) -> None:
-    try:
-        wait_event(timeout, stop_event, lambda: process.poll() is not None)
-    except StopEventSetException:
-        process.kill()
-        raise
-
-
 def remove_files_by_prefix(prefix: str) -> None:
     for f in glob.glob(f"{prefix}*"):
         os.unlink(f)
@@ -228,7 +230,7 @@ def run_process(
     stderr: bytes
 
     reraise_exc: Optional[BaseException] = None
-    with start_process(cmd, via_staticx, **kwargs) as process:
+    with start_process(cmd, via_staticx, direct_call=False, **kwargs) as process:
         assert isinstance(process.args, str) or (
             isinstance(process.args, list) and all(isinstance(s, str) for s in process.args)
         ), process.args  # mypy
@@ -536,6 +538,8 @@ def cleanup_process_reference(process: Popen) -> None:
 def _exit_handler() -> None:
     for process in _processes:
         process.kill()
+        # remove process in _processes
+        cleanup_process_reference(process)
 
 
 def _sigint_handler(sig: int, frame: Optional[FrameType]) -> None: