caterpillar, download: improve SIGINT and SIGTERM handling

Author: zmwangx

src/caterpillar/caterpillar.py

@@ -7,7 +7,7 @@
 import sys
 import time
 import urllib.parse
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import peewee
 
@@ -255,6 +255,7 @@ def rmdir_p(path: pathlib.Path, *, root: pathlib.Path = None) -> None:
 def process_entry(
     m3u8_url: str,
     output: pathlib.Path,
+    *,
     force: bool = False,
     exist_ok: bool = False,
     workdir: pathlib.Path = None,
@@ -376,6 +377,65 @@ def process_entry(
     return 0
 
 
+def process_batch(
+    manifest: pathlib.Path,
+    remove_manifest_on_success: bool = False,
+    debug: bool = False,
+    **processing_kwargs: Any,
+) -> int:
+    target_dir = manifest.parent
+    try:
+        entries = []
+        with manifest.open(encoding="utf-8") as fp:
+            manifest_content = fp.read()
+    except OSError:
+        logger.critical("cannot open batch mode manifest", exc_info=debug)
+        if debug:
+            raise
+        return 1
+    except UnicodeDecodeError:
+        logger.critical(
+            "cannot decode batch mode manifest as utf-8; "
+            "see https://git.io/caterpillar-encoding",
+            exc_info=debug,
+        )
+        if debug:
+            raise
+        return 1
+
+    for line in manifest_content.splitlines():
+        try:
+            m3u8_url, filename = line.strip().split("\t")
+            output = target_dir.joinpath(filename)
+            entries.append((m3u8_url, output))
+        except Exception:
+            logger.critical(
+                "malformed line in batch mode manifest: %s", line, exc_info=debug
+            )
+            if debug:
+                raise
+            return 1
+
+    retvals = []
+    count = len(entries)
+    for i, (m3u8_url, output) in enumerate(entries):
+        sys.stderr.write(
+            f'[{i + 1}/{count}] Downloading {m3u8_url} into "{output}"...\n'
+        )
+        retvals.append(process_entry(m3u8_url, output, **processing_kwargs))
+        sys.stderr.write("\n")
+    retval = int(any(retvals))
+    if retval == 0 and remove_manifest_on_success:
+        try:
+            manifest.unlink()
+        except OSError:
+            logger.error("cannot remove batch mode manifest")
+            if debug:
+                raise
+            return 1
+    return retval
+
+
 def main() -> int:
     user_config_options = [] if USER_CONFIG_DISABLED else load_user_config()
 
@@ -545,63 +605,19 @@ def main() -> int:
         logger.critical("ffmpeg not found")
         return 1
 
-    if not args.batch:
-        return process_entry(args.m3u8_url, args.output, **kwargs)
-    else:
-        manifest = pathlib.Path(args.m3u8_url).resolve()
-        target_dir = manifest.parent
-        try:
-            entries = []
-            with manifest.open(encoding="utf-8") as fp:
-                manifest_content = fp.read()
-        except OSError:
-            logger.critical("cannot open batch mode manifest", exc_info=args.debug)
-            if args.debug:
-                raise
-            return 1
-        except UnicodeDecodeError:
-            logger.critical(
-                "cannot decode batch mode manifest as utf-8; "
-                "see https://git.io/caterpillar-encoding",
-                exc_info=args.debug,
-            )
-            if args.debug:
-                raise
-            return 1
-
-        for line in manifest_content.splitlines():
-            try:
-                m3u8_url, filename = line.strip().split("\t")
-                output = target_dir.joinpath(filename)
-                entries.append((m3u8_url, output))
-            except Exception:
-                logger.critical(
-                    "malformed line in batch mode manifest: %s",
-                    line,
-                    exc_info=args.debug,
-                )
-                if args.debug:
-                    raise
-                return 1
-
-        retvals = []
-        count = len(entries)
-        for i, (m3u8_url, output) in enumerate(entries):
-            sys.stderr.write(
-                f'[{i + 1}/{count}] Downloading {m3u8_url} into "{output}"...\n'
+    try:
+        if not args.batch:
+            return process_entry(args.m3u8_url, args.output, **kwargs)
+        else:
+            manifest = pathlib.Path(args.m3u8_url).resolve()
+            return process_batch(
+                manifest,
+                remove_manifest_on_success=args.remove_manifest_on_success,
+                debug=args.debug,
+                **kwargs,
             )
-            retvals.append(process_entry(m3u8_url, output, **kwargs))
-            sys.stderr.write("\n")
-        retval = int(any(retvals))
-        if retval == 0 and args.remove_manifest_on_success:
-            try:
-                manifest.unlink()
-            except OSError:
-                logger.error("cannot remove batch mode manifest")
-                if args.debug:
-                    raise
-                return 1
-        return retval
+    except KeyboardInterrupt:
+        return 1
 
 
 if __name__ == "__main__":

src/caterpillar/download.py

@@ -128,14 +128,20 @@ def download_segment(
 
 # download_segment wrapper that takes all arguments as a single tuple,
 # so that we can use it with multiprocessing.pool.Pool.map and company.
+# It also gracefully consumes KeyboardInterrupt.
 def _download_segment_mappable(args: Tuple[str, int, pathlib.Path]) -> bool:
-    return download_segment(*args)
+    try:
+        return download_segment(*args)
+    except KeyboardInterrupt:
+        url, *_ = args
+        logger.debug(f"download of {url} has been interrupted")
+        return False
 
 
-def _init_worker():
-    # Ignore SIGINT in worker processes to disable traceback from
-    # each worker on keyboard interrupt.
-    signal.signal(signal.SIGINT, signal.SIG_IGN)
+def _raise_keyboard_interrupt(signum, _):
+    pid = os.getpid()
+    logger.debug(f"pid {pid} received signal {signum}; transforming into SIGINT")
+    raise KeyboardInterrupt
 
 
 # Download all segments in remote_m3u8_file (downloaded from
@@ -179,33 +185,48 @@ def download_m3u8_segments(
         logger.error(f"{remote_m3u8_file}: empty playlist")
         return False
     jobs = min(jobs, total)
-    with multiprocessing.Pool(jobs, _init_worker) as pool:
-        num_success = 0
-        num_failure = 0
-        logger.info(f"downloading {total} segments with {jobs} workers...")
-        progress_bar_generator = (
-            click.progressbar if should_log_warning() else stub_context_manager
-        )
-        progress_bar_props = dict(
-            width=0,  # Full width
-            bar_template="[%(bar)s] %(info)s",
-            show_pos=True,
-            length=total,
-        )
-        with progress_bar_generator(**progress_bar_props) as bar:  # type: ignore
-            for success in pool.imap_unordered(
-                _download_segment_mappable, download_args
-            ):
-                if success:
-                    num_success += 1
-                else:
-                    num_failure += 1
-                logger.debug(f"progress: {num_success}/{num_failure}/{total}")
-                bar.update(1)
-
-        if num_failure > 0:
-            logger.error(f"failed to download {num_failure} segments")
-            return False
-        else:
-            logger.info(f"finished downloading all {total} segments")
-            return True
+    with multiprocessing.Pool(jobs) as pool:
+        # For the duration of the worker pool, map SIGTERM to SIGINT on
+        # the main process. We only do this after the fork, and restore
+        # the original SIGTERM handler (usually SIG_DFL) at the end of
+        # the pool, because using _raise_keyboard_interrupt as the
+        # SIGTERM handler on workers could somehow lead to dead locks.
+        old_sigterm_handler = signal.signal(signal.SIGTERM, _raise_keyboard_interrupt)
+        try:
+            num_success = 0
+            num_failure = 0
+            logger.info(f"downloading {total} segments with {jobs} workers...")
+            progress_bar_generator = (
+                click.progressbar if should_log_warning() else stub_context_manager
+            )
+            progress_bar_props = dict(
+                width=0,  # Full width
+                bar_template="[%(bar)s] %(info)s",
+                show_pos=True,
+                length=total,
+            )
+            with progress_bar_generator(**progress_bar_props) as bar:  # type: ignore
+                for success in pool.imap_unordered(
+                    _download_segment_mappable, download_args
+                ):
+                    if success:
+                        num_success += 1
+                    else:
+                        num_failure += 1
+                    logger.debug(f"progress: {num_success}/{num_failure}/{total}")
+                    bar.update(1)
+
+            if num_failure > 0:
+                logger.error(f"failed to download {num_failure} segments")
+                return False
+            else:
+                logger.info(f"finished downloading all {total} segments")
+                return True
+        except KeyboardInterrupt:
+            pool.terminate()
+            pool.join()
+            logger.critical("interrupted")
+            # Bubble KeyboardInterrupt to stop retries.
+            raise KeyboardInterrupt
+        finally:
+            signal.signal(signal.SIGTERM, old_sigterm_handler)