diff --git a/shelfmark/core/queue.py b/shelfmark/core/queue.py index 950ed29f..293a58b1 100644 --- a/shelfmark/core/queue.py +++ b/shelfmark/core/queue.py @@ -145,28 +145,36 @@ class BookQueue: with self._lock: self._queue_hook = hook - def update_status(self, book_id: str, status: QueueStatus) -> None: - """Update status of a book in the queue.""" + def _apply_status_locked( + self, book_id: str, status: QueueStatus + ) -> tuple[Callable[[str, QueueStatus, DownloadTask], None] | None, DownloadTask | None]: + """Apply a status change; returns the terminal hook to run after releasing the lock.""" hook: Callable[[str, QueueStatus, DownloadTask], None] | None = None hook_task: DownloadTask | None = None + previous_status = self._status.get(book_id) + self._update_status(book_id, status) + + if ( + status in TERMINAL_QUEUE_STATUSES + and previous_status != status + and self._terminal_status_hook is not None + ): + current_task = self._task_data.get(book_id) + if current_task is not None: + hook = self._terminal_status_hook + hook_task = current_task + + # Clean up active download tracking when finished + if status in TERMINAL_QUEUE_STATUSES: + self._active_downloads.pop(book_id, None) + self._cancel_flags.pop(book_id, None) + + return hook, hook_task + + def update_status(self, book_id: str, status: QueueStatus) -> None: + """Update status of a book in the queue.""" with self._lock: - previous_status = self._status.get(book_id) - self._update_status(book_id, status) - - if ( - status in TERMINAL_QUEUE_STATUSES - and previous_status != status - and self._terminal_status_hook is not None - ): - current_task = self._task_data.get(book_id) - if current_task is not None: - hook = self._terminal_status_hook - hook_task = current_task - - # Clean up active download tracking when finished - if status in TERMINAL_QUEUE_STATUSES: - self._active_downloads.pop(book_id, None) - self._cancel_flags.pop(book_id, None) + hook, hook_task = self._apply_status_locked(book_id, status) if hook is not None and hook_task is not None: hook(book_id, status, hook_task) @@ -257,7 +265,11 @@ class BookQueue: # Not in a cancellable state return False - self.update_status(task_id, QueueStatus.CANCELLED) + # Write under the same lock so a download that finishes first is not overwritten + hook, hook_task = self._apply_status_locked(task_id, QueueStatus.CANCELLED) + + if hook is not None and hook_task is not None: + hook(task_id, QueueStatus.CANCELLED, hook_task) return True def set_priority(self, task_id: str, new_priority: int) -> bool: diff --git a/tests/core/test_queue.py b/tests/core/test_queue.py index ac663024..6320ce1d 100644 --- a/tests/core/test_queue.py +++ b/tests/core/test_queue.py @@ -1,8 +1,9 @@ -"""Tests for queue hook failure handling.""" +"""Tests for queue hook failure handling and cancellation.""" +import threading from unittest.mock import patch -from shelfmark.core.models import DownloadTask +from shelfmark.core.models import DownloadTask, QueueStatus from shelfmark.core.queue import BookQueue @@ -51,3 +52,43 @@ def test_enqueue_existing_logs_queue_hook_failures(): assert args[0] == "Queue hook failed while requeueing task %s: %s" assert args[1] == "task-2" assert str(args[2]) == "boom" + + +def test_cancel_does_not_overwrite_a_download_that_finished_first(): + """A download that completes while a cancel is in flight must stay complete.""" + queue = BookQueue() + assert queue.add(_make_task("race-task")) is True + assert queue.get_next() is not None + queue.update_status("race-task", QueueStatus.DOWNLOADING) + + terminal_events: list[QueueStatus] = [] + queue.set_terminal_status_hook(lambda _task_id, status, _task: terminal_events.append(status)) + + original_update_status = queue.update_status + cancel_yielded = threading.Event() + worker_finished = threading.Event() + + def update_status_yielding_to_the_worker(task_id: str, status: QueueStatus) -> None: + # Hand the worker the queue whenever the cancel path lets go of it before writing. + if status == QueueStatus.CANCELLED: + cancel_yielded.set() + worker_finished.wait(timeout=5) + original_update_status(task_id, status) + + queue.update_status = update_status_yielding_to_the_worker # type: ignore[method-assign] + + def finish_download() -> None: + cancel_yielded.wait(timeout=5) + original_update_status("race-task", QueueStatus.COMPLETE) + worker_finished.set() + + worker = threading.Thread(target=finish_download, daemon=True, name="TestDownloadWorker") + worker.start() + try: + queue.cancel_download("race-task") + finally: + cancel_yielded.set() + worker.join(timeout=5) + + assert queue.get_task_status("race-task") == QueueStatus.COMPLETE + assert terminal_events[-1] == QueueStatus.COMPLETE