diff --git a/app/tests/test_ytdl_utils.py b/app/tests/test_ytdl_utils.py index 72c716e..e3315a8 100644 --- a/app/tests/test_ytdl_utils.py +++ b/app/tests/test_ytdl_utils.py @@ -49,6 +49,7 @@ sys.modules.setdefault("yt_dlp.postprocessor", fake_postprocessor) sys.modules.setdefault("yt_dlp.postprocessor.common", fake_postprocessor_common) sys.modules.setdefault("yt_dlp.utils", fake_utils) +import ytdl from ytdl import ( Download, DownloadInfo, @@ -405,11 +406,16 @@ class ProgressThrottleTests(unittest.TestCase): class CancelProcessGroupTests(unittest.TestCase): - def test_cancel_kills_group_when_child_is_group_leader(self): + # cancel() now sends SIGINT first (so yt-dlp/ffmpeg can finalize the + # partial file) and schedules a SIGKILL escalation via the event loop + # after ytdl._CANCEL_GRACE_SECONDS, instead of SIGKILLing immediately. + + def test_cancel_sends_sigint_to_group_and_schedules_sigkill_escalation(self): # Child successfully ran os.setpgrp(): its pgid equals its own pid. dl = _make_test_download() dl.proc = types.SimpleNamespace(pid=4321) dl.status_queue = types.SimpleNamespace(put=lambda _item: None) + dl.loop = MagicMock() with patch.object(Download, "running", return_value=True), \ patch("ytdl.os.getpgid", return_value=4321) as mock_getpgid, \ @@ -417,38 +423,109 @@ class CancelProcessGroupTests(unittest.TestCase): dl.cancel() mock_getpgid.assert_called_once_with(4321) - mock_killpg.assert_called_once_with(4321, signal.SIGKILL) + mock_killpg.assert_called_once_with(4321, signal.SIGINT) + dl.loop.call_later.assert_called_once_with(ytdl._CANCEL_GRACE_SECONDS, dl._kill_if_alive) self.assertTrue(dl.canceled) - def test_cancel_does_not_killpg_parent_group_kills_child_only(self): + def test_cancel_does_not_killpg_parent_group_signals_child_only(self): # Child has NOT become its own group leader yet (pgid != pid, e.g. it is # still in the server's process group). killpg must NOT be called — that - # would SIGKILL the whole server — and we fall back to proc.kill(). + # would signal the whole server — and we fall back to os.kill(pid, SIGINT). dl = _make_test_download() - dl.proc = types.SimpleNamespace(pid=4321, kill=MagicMock()) + dl.proc = types.SimpleNamespace(pid=4321) dl.status_queue = types.SimpleNamespace(put=lambda _item: None) + dl.loop = MagicMock() with patch.object(Download, "running", return_value=True), \ patch("ytdl.os.getpgid", return_value=999), \ - patch("ytdl.os.killpg") as mock_killpg: + patch("ytdl.os.killpg") as mock_killpg, \ + patch("ytdl.os.kill") as mock_kill: dl.cancel() mock_killpg.assert_not_called() - dl.proc.kill.assert_called_once() + mock_kill.assert_called_once_with(4321, signal.SIGINT) + dl.loop.call_later.assert_called_once_with(ytdl._CANCEL_GRACE_SECONDS, dl._kill_if_alive) self.assertTrue(dl.canceled) - def test_cancel_falls_back_to_proc_kill_when_getpgid_unavailable(self): + def test_cancel_falls_back_to_pid_signal_when_getpgid_unavailable(self): + dl = _make_test_download() + dl.proc = types.SimpleNamespace(pid=4321) + dl.status_queue = types.SimpleNamespace(put=lambda _item: None) + dl.loop = MagicMock() + + with patch.object(Download, "running", return_value=True), \ + patch("ytdl.os.getpgid", side_effect=OSError("no such process")), \ + patch("ytdl.os.kill") as mock_kill: + dl.cancel() + + mock_kill.assert_called_once_with(4321, signal.SIGINT) + dl.loop.call_later.assert_called_once_with(ytdl._CANCEL_GRACE_SECONDS, dl._kill_if_alive) + self.assertTrue(dl.canceled) + + def test_cancel_kills_immediately_when_signal_delivery_fails(self): + # Neither killpg nor os.kill succeed (process already gone): cancel() + # must fall through to an immediate SIGKILL attempt instead of + # scheduling a pointless escalation. dl = _make_test_download() dl.proc = types.SimpleNamespace(pid=4321, kill=MagicMock()) dl.status_queue = types.SimpleNamespace(put=lambda _item: None) + dl.loop = MagicMock() with patch.object(Download, "running", return_value=True), \ - patch("ytdl.os.getpgid", side_effect=OSError("no such process")): + patch("ytdl.os.getpgid", side_effect=OSError("no such process")), \ + patch("ytdl.os.kill", side_effect=ProcessLookupError()): dl.cancel() + dl.loop.call_later.assert_not_called() dl.proc.kill.assert_called_once() self.assertTrue(dl.canceled) + def test_cancel_schedules_escalation_even_without_running_loop(self): + # dl.loop is None (e.g. cancel() called before start()'s run_in_executor + # set it up): _kill_if_alive() must run synchronously instead of being + # scheduled, since there's no loop to schedule it on. + dl = _make_test_download() + dl.proc = types.SimpleNamespace(pid=4321) + dl.status_queue = types.SimpleNamespace(put=lambda _item: None) + self.assertIsNone(dl.loop) + + with patch.object(Download, "running", side_effect=[True, True]), \ + patch("ytdl.os.getpgid", return_value=4321), \ + patch("ytdl.os.killpg") as mock_killpg: + dl.cancel() + + # First SIGINT, then _kill_if_alive() ran inline and sent SIGKILL. + mock_killpg.assert_has_calls([ + unittest.mock.call(4321, signal.SIGINT), + unittest.mock.call(4321, signal.SIGKILL), + ]) + self.assertTrue(dl.canceled) + + +class KillIfAliveTests(unittest.TestCase): + def test_kill_if_alive_sigkills_running_process(self): + dl = _make_test_download() + dl.proc = types.SimpleNamespace(pid=4321) + + with patch.object(Download, "running", return_value=True), \ + patch("ytdl.os.getpgid", return_value=4321), \ + patch("ytdl.os.killpg") as mock_killpg: + dl._kill_if_alive() + + mock_killpg.assert_called_once_with(4321, signal.SIGKILL) + + def test_kill_if_alive_noop_when_process_already_exited(self): + dl = _make_test_download() + dl.proc = types.SimpleNamespace(pid=4321) + + with patch.object(Download, "running", return_value=False), \ + patch("ytdl.os.killpg") as mock_killpg, \ + patch("ytdl.os.kill") as mock_kill: + dl._kill_if_alive() + + mock_killpg.assert_not_called() + mock_kill.assert_not_called() + class ConvertSrtToTxtTests(unittest.TestCase): def test_basic_conversion(self): diff --git a/app/ytdl.py b/app/ytdl.py index 20a0f4b..a44e02b 100644 --- a/app/ytdl.py +++ b/app/ytdl.py @@ -47,6 +47,10 @@ _MP_CTX = ( else multiprocessing.get_context() ) +# Grace period between SIGINT (lets yt-dlp/ffmpeg finalize the partial file, +# e.g. when cancelling a livestream) and SIGKILL escalation. +_CANCEL_GRACE_SECONDS = 15 + _LIVE_CHECK_INTERVAL = 60 _LIVE_MAX_CHECK_INTERVAL = 3600 # Consecutive probe failures (network blips, rate limits, transient extractor @@ -704,29 +708,50 @@ class Download: self.status_queue.put(None) await self.status_task + def _signal_group(self, sig): + """Send *sig* to the download's process group, falling back to the + process itself. Returns True if a signal was delivered. + + Only signal the whole group when the child actually became its own + group leader via os.setpgrp() in _download() — that sets its pgid + equal to its own pid. If it hasn't run setpgrp() yet, or setpgrp() + failed, its pgid is still the SERVER's group and killpg would signal + the entire MeTube process (PID 1 in Docker). Fall back to signalling + just the child process by pid in that case. + """ + try: + pgid = os.getpgid(self.proc.pid) + if pgid == self.proc.pid: + os.killpg(pgid, sig) + return True + except (OSError, AttributeError): + pass + try: + if sig == signal.SIGINT: + os.kill(self.proc.pid, sig) + else: + self.proc.kill() + return True + except Exception as e: + log.error(f"Error signalling process for {self.info.title}: {e}") + return False + + def _kill_if_alive(self): + if self.running(): + log.info(f"Escalating cancel to SIGKILL for: {self.info.title}") + self._signal_group(signal.SIGKILL) + def cancel(self): log.info(f"Cancelling download: {self.info.title}") if self.running(): - killed_group = False - try: - pgid = os.getpgid(self.proc.pid) - # Only kill the whole group (yt-dlp + any ffmpeg children) when - # the child actually became its own group leader via - # os.setpgrp() in _download() — that sets its pgid equal to its - # own pid. If it hasn't run setpgrp() yet, or setpgrp() failed, - # its pgid is still the SERVER's group and killpg would SIGKILL - # the entire MeTube process (PID 1 in Docker). Fall back to - # killing just the child process by pid in that case. - if pgid == self.proc.pid: - os.killpg(pgid, signal.SIGKILL) - killed_group = True - except (OSError, AttributeError): - pass - if not killed_group: - try: - self.proc.kill() - except Exception as e: - log.error(f"Error killing process for {self.info.title}: {e}") + # SIGINT first so yt-dlp/ffmpeg can finalize the partial file + # (livestream recordings stay playable); SIGKILL after a grace + # period if the process ignores it. + interrupted = self._signal_group(signal.SIGINT) + if interrupted and self.loop is not None: + self.loop.call_later(_CANCEL_GRACE_SECONDS, self._kill_if_alive) + else: + self._kill_if_alive() self.canceled = True if self.status_queue is not None: self.status_queue.put(None)