tumblr_backup: Worker early-stop and graceful SIGTERM/SIGHUP

Included revisions:
- Be more paranoid to avoid swallowing exceptions

tumblr_backup.py

@@ -13,6 +13,7 @@
 import os
 import re
 import shutil
+import signal
 import sys
 import threading
 import time
@@ -22,8 +23,9 @@
 from os.path import join, split, splitext
 from xml.sax.saxutils import escape
 
-from util import (AsyncCallable, ConnectionFile, LockedQueue, MultiCondition, PY3, is_dns_working,
-                  make_requests_session, no_internet, nullcontext, path_is_on_vfat, to_bytes, to_unicode)
+from util import (AsyncCallable, ConnectionFile, LockedQueue, MultiCondition, PY3, disable_unraisable_hook,
+                  is_dns_working, make_requests_session, no_internet, nullcontext, path_is_on_vfat, to_bytes,
+                  to_unicode)
 from wget import HTTPError, HTTP_RETRY, HTTP_TIMEOUT, WGError, WgetRetrieveWrapper, setup_wget, urlopen
 
 try:
@@ -1447,9 +1449,18 @@ def cancel(self):
             self.abort_flag = True
             self.quit.notify_all()
 
-        for i, t in enumerate(self.threads, start=1):
-            log.status('Stopping threads {}{}\r'.format(' ' * i, '.' * (len(self.threads) - i)))
-            t.join()
+        duh = disable_unraisable_hook() if hasattr(signal, 'pthread_kill') else nullcontext()  # type: Any
+        with duh:
+            # The SIGTERM handler raises SystemExit to gracefully stop the worker
+            # threads. Otherwise, we must wait for them to finish their posts.
+            if hasattr(signal, 'pthread_kill'):
+                for t in self.threads:
+                    if t.ident is not None:
+                        signal.pthread_kill(t.ident, signal.SIGTERM)  # type: ignore[attr-defined]
+
+            for i, t in enumerate(self.threads, start=1):
+                log.status('Stopping threads {}{}\r'.format(' ' * i, '.' * (len(self.threads) - i)))
+                t.join()
 
         with main_thread_lock:
             self.queue.queue.clear()
@@ -1491,6 +1502,19 @@ def wait_for_work():
     else:
         multiprocessing.set_start_method('spawn')  # Slow but safe
 
+    # Raises SystemExit to terminate gracefully
+    def handle_term_signal(signum, frame):
+        if sys.exc_info() != (None, None, None):
+            return  # Not a good time to exit
+        if not PY3:
+            pass  # No is_finalizing
+        elif sys.is_finalizing():
+            return  # Not a good time to exit
+        sys.exit(1)
+    signal.signal(signal.SIGTERM, handle_term_signal)
+    if hasattr(signal, 'SIGHUP'):
+        signal.signal(signal.SIGHUP, handle_term_signal)
+
     no_internet.setup(main_thread_lock)
 
     import argparse

util.py

@@ -3,6 +3,7 @@
 from __future__ import absolute_import, division, print_function, with_statement
 
 import collections
+import contextlib
 import io
 import os
 import socket
@@ -140,6 +141,19 @@ def __exit__(self, *excinfo):
         pass
 
 
+@contextlib.contextmanager
+def disable_unraisable_hook():
+    if sys.version_info >= (3, 8):
+        def hook(hookargs): pass
+        old_hook, sys.unraisablehook = sys.unraisablehook, hook  # pytype: disable=module-attr
+        try:
+            yield None
+        finally:
+            sys.unraisablehook = old_hook  # pytype: disable=module-attr
+    else:
+        yield None
+
+
 KNOWN_GOOD_NAMESERVER = '8.8.8.8'
 # DNS query for 'A' record of 'google.com'.
 # Generated using python -c "import dnslib; print(bytes(dnslib.DNSRecord.question('google.com').pack()))"