handle SIGINT in StatusTracker properly

2 files changed, 42 insertions, 8 deletions

diff --git a/.travis/track_status.sh b/.travis/track_status.sh
index 3702676..0e3d2d8 100755
--- a/.travis/track_status.sh
+++ b/.travis/track_status.sh
@@ -2,18 +2,35 @@
 
 set -o errexit -o nounset -o pipefail
 
-main() {
+track_status() {
     local log_path
     log_path="$( mktemp )"
+    echo "Log file path: $log_path"
+
+    local rm_log_path
+    rm_log_path="$( printf -- 'rm -f -- %q' "$log_path" )"
 
-    nohup python3 -m bin.track_status egor.tensin > "$log_path" 2>&1 &
+    trap "$rm_log_path" RETURN
+
+    echo 'Running track_status.py...'
+    python3 -m bin.track_status egor.tensin --log "$log_path" &
     local pid="$!"
+    echo "Its PID is $pid"
+
+    local timeout=15
+    echo "Sleeping for $timeout seconds..."
+    sleep "$timeout"
 
-    sleep 15
+    echo 'Terminating track_status.py...'
     kill -SIGINT "$pid"
+    echo 'Waiting for track_status.py to terminate...'
     wait "$pid"
 
     cat "$log_path"
 }
 
-main
+main() {
+    track_status
+}
+
+main "$@"
diff --git a/vk/tracking/status_tracker.py b/vk/tracking/status_tracker.py
index 8b90d1c..30e0f97 100644
--- a/vk/tracking/status_tracker.py
+++ b/vk/tracking/status_tracker.py
@@ -4,7 +4,9 @@
 # Distributed under the MIT License.
 
 from collections.abc import Callable
+import contextlib
 import time
+import signal
 
 import vk.error
 from vk.user import UserField
@@ -98,8 +100,23 @@ class StatusTracker:
             for user in self._filter_status_updates(users, updated_users):
                 self._notify_status_update(user)
 
-    def loop(self, uids):
+    @staticmethod
+    @contextlib.contextmanager
+    def _handle_sigint():
+        # Python doesn't raise KeyboardInterrupt in case a real SIGINT is sent
+        # from outside, surprisingly.
+        def _raise_keyboard_interrupt(signum, frame):
+            raise KeyboardInterrupt()
+        old_handler = signal.getsignal(signal.SIGINT)
+        signal.signal(signal.SIGINT, _raise_keyboard_interrupt)
         try:
-            self._do_loop(uids)
-        except KeyboardInterrupt:
-            pass
+            yield
+        finally:
+            signal.signal(signal.SIGINT, old_handler)
+
+    def loop(self, uids):
+        with self._handle_sigint():
+            try:
+                self._do_loop(uids)
+            except KeyboardInterrupt:
+                pass