From 4a229e4070f35f0f2109d83aeaa0953cf294d8d6 Mon Sep 17 00:00:00 2001
From: slicklash <slicklash@gmail.com>
Date: Mon, 15 Jul 2024 14:15:25 +0300
Subject: [PATCH] use setup_singnals

---
 gprofiler/main.py           | 30 ++----------------------------
 gprofiler/utils/__init__.py | 34 +++++++++++++++++++++++++---------
 2 files changed, 27 insertions(+), 37 deletions(-)

diff --git a/gprofiler/main.py b/gprofiler/main.py
index 9e845e891..597f85820 100644
--- a/gprofiler/main.py
+++ b/gprofiler/main.py
@@ -20,13 +20,12 @@
 import logging.handlers
 import os
 import shutil
-import signal
 import sys
 import time
 import traceback
 from pathlib import Path
 from threading import Event
-from types import FrameType, TracebackType
+from types import TracebackType
 from typing import Iterable, List, Optional, Type, cast
 
 import configargparse
@@ -74,6 +73,7 @@
     reset_umask,
     resource_path,
     run_process,
+    setup_signals,
 )
 from gprofiler.utils.fs import escape_filename, mkdir_owned_root
 from gprofiler.utils.proxy import get_https_proxy
@@ -98,21 +98,6 @@
 
 UPLOAD_FILE_SUBCOMMAND = "upload-file"
 
-# 1 KeyboardInterrupt raised per this many seconds, no matter how many SIGINTs we get.
-SIGINT_RATELIMIT = 0.5
-
-last_signal_ts: Optional[float] = None
-
-
-def sigint_handler(sig: int, frame: Optional[FrameType]) -> None:
-    global last_signal_ts
-    ts = time.monotonic()
-    # no need for atomicity here: we can't get another SIGINT before this one returns.
-    # https://www.gnu.org/software/libc/manual/html_node/Signals-in-Handler.html#Signals-in-Handler
-    if last_signal_ts is None or ts > last_signal_ts + SIGINT_RATELIMIT:
-        last_signal_ts = ts
-        raise KeyboardInterrupt
-
 
 class GProfiler:
     def __init__(
@@ -939,17 +924,6 @@ def verify_preconditions(args: configargparse.Namespace, processes_to_profile: O
             sys.exit(1)
 
 
-def setup_signals() -> None:
-    # When we run under staticx & PyInstaller, both of them forward (some of the) signals to gProfiler.
-    # We catch SIGINTs and ratelimit them, to avoid being interrupted again during the handling of the
-    # first INT.
-    # See my commit message for more information.
-    signal.signal(signal.SIGINT, sigint_handler)
-    # handle SIGTERM in the same manner - gracefully stop gProfiler.
-    # SIGTERM is also forwarded by staticx & PyInstaller, so we need to ratelimit it.
-    signal.signal(signal.SIGTERM, sigint_handler)
-
-
 def log_system_info() -> None:
     system_info = get_static_system_info()
     logger.info(f"gProfiler Python version: {system_info.python_version}")
diff --git a/gprofiler/utils/__init__.py b/gprofiler/utils/__init__.py
index 98143c5ba..856841e13 100644
--- a/gprofiler/utils/__init__.py
+++ b/gprofiler/utils/__init__.py
@@ -34,6 +34,7 @@
 from subprocess import CompletedProcess, Popen, TimeoutExpired
 from tempfile import TemporaryDirectory
 from threading import Event
+from types import FrameType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union, cast
 
 import importlib_resources
@@ -71,8 +72,10 @@
 
 gprofiler_mutex: Optional[socket.socket] = None
 
-STATUS_KILL = 137
-STATUS_INTERRUPT = 130
+# 1 KeyboardInterrupt raised per this many seconds, no matter how many SIGINTs we get.
+SIGINT_RATELIMIT = 0.5
+
+_last_signal_ts: Optional[float] = None
 _processes: List[Popen] = []
 
 
@@ -522,10 +525,23 @@ def _exit_handler() -> None:
         process.kill()
 
 
-def _kill_handler(*args: Any) -> None:
-    sys.exit(STATUS_KILL if args[0] == signal.SIGTERM else STATUS_INTERRUPT)
-
-
-atexit.register(_exit_handler)
-signal.signal(signal.SIGINT, _kill_handler)
-signal.signal(signal.SIGTERM, _kill_handler)
+def _sigint_handler(sig: int, frame: Optional[FrameType]) -> None:
+    global _last_signal_ts
+    ts = time.monotonic()
+    # no need for atomicity here: we can't get another SIGINT before this one returns.
+    # https://www.gnu.org/software/libc/manual/html_node/Signals-in-Handler.html#Signals-in-Handler
+    if _last_signal_ts is None or ts > _last_signal_ts + SIGINT_RATELIMIT:
+        _last_signal_ts = ts
+        raise KeyboardInterrupt
+
+
+def setup_signals() -> None:
+    atexit.register(_exit_handler)
+    # When we run under staticx & PyInstaller, both of them forward (some of the) signals to gProfiler.
+    # We catch SIGINTs and ratelimit them, to avoid being interrupted again during the handling of the
+    # first INT.
+    # See my commit message for more information.
+    signal.signal(signal.SIGINT, _sigint_handler)
+    # handle SIGTERM in the same manner - gracefully stop gProfiler.
+    # SIGTERM is also forwarded by staticx & PyInstaller, so we need to ratelimit it.
+    signal.signal(signal.SIGTERM, _sigint_handler)