diff --git a/importarr/main.py b/importarr/main.py index b77fda4..6a8f684 100644 --- a/importarr/main.py +++ b/importarr/main.py @@ -267,8 +267,12 @@ def ensure_worker_running() -> None: _worker_thread.start() -def stop_worker() -> None: +def stop_worker(*, wait: bool = False) -> bool: _worker_stop.set() + thread = _worker_thread + if wait and thread and thread.is_alive() and thread is not threading.current_thread(): + thread.join() + return not thread or not thread.is_alive() def _worker_loop() -> None: @@ -449,10 +453,9 @@ def queue_item_action(item_id: int, payload: QueueItemActionRequest, _: None = D if item is None: raise HTTPException(status_code=404, detail="queue item not found") if payload.action == "retry": - if item.get("claimed_by"): - raise HTTPException(status_code=409, detail="queue item is currently being imported") retry_state = "ready" if item["source_type"] in {"manual", "sab"} else "detected" - state.mark_queue_item_result(item["source_type"], item["source_id"], retry_state, "retry requested") + if not state.transition_queue_item_if_unclaimed(item_id, {"failed", "skipped"}, retry_state, "retry requested"): + raise HTTPException(status_code=409, detail="queue item is currently being imported or changed") elif payload.action == "run-now": claimed = state.claim_queue_item(item_id, WORKER_ID, {"ready", "failed", "retrying"}) if claimed is None: @@ -461,9 +464,8 @@ def queue_item_action(item_id: int, payload: QueueItemActionRequest, _: None = D updated = state.get_queue_item(item_id) return {"status": "imported" if imported else "updated", "imported": imported, "item": serialize_queue_item(updated or item)} elif payload.action == "ignore": - if item.get("claimed_by"): - raise HTTPException(status_code=409, detail="queue item is currently being imported") - state.mark_queue_item_result(item["source_type"], item["source_id"], "skipped", "ignored by user") + if not state.transition_queue_item_if_unclaimed(item_id, {str(item["state"])}, "skipped", "ignored by user"): + raise HTTPException(status_code=409, detail="queue item is currently being imported or changed") elif payload.action == "remove": if not state.delete_queue_item_if_unclaimed(item_id): raise HTTPException(status_code=409, detail="queue item is currently being imported") @@ -623,22 +625,10 @@ async def _import_ready_sab_jobs(importer: Importer, force: bool = False) -> int for video in scan_videos(readiness.storage): if consume_cancel_request(): return imported - set_current_job(str(video.path)) - try: - started = time.time() - result = importer.import_file(video.path, should_cancel=consume_cancel_request, on_progress=lambda copied, total, path=str(video.path), started=started: set_current_job(path, bytes_copied=copied, total_bytes=total, started_at=started)) - state.add_history(result.source, result.target, "imported", result.bytes) - state.mark_queue_item("sab", str(video.path), "imported") - imported += 1 - except ImportCancelled: - state.add_history(video.path, video.path, "cancelled", 0, "cancelled") - state.mark_queue_item("sab", str(video.path), "skipped", "cancelled") - return imported - except Exception as exc: - state.add_history(video.path, video.path, "failed", 0, exc.__class__.__name__) - state.mark_queue_item("sab", str(video.path), "failed", exc.__class__.__name__) - finally: - set_current_job(None) + row = state.upsert_queue_item(source_type="sab", source_id=str(video.path), source_path=video.path, name=video.path.name, state="ready", reason=readiness.reason, relative_path=str(video.relative_path), size=video.size) + claimed = state.claim_queue_item(row["id"], WORKER_ID, {"ready", "failed", "retrying"}) + if claimed is not None: + imported += _import_queue_item(claimed, importer, force=force, from_worker=True) return imported @@ -689,7 +679,9 @@ def _import_manual_batches(importer: Importer) -> int: for item in items: if consume_cancel_request(): return imported - imported += _import_queue_item(item, importer, force=True) + claimed = state.claim_queue_item(item["id"], WORKER_ID, {"ready", "failed", "retrying"}) + if claimed is not None: + imported += _import_queue_item(claimed, importer, force=True, from_worker=True) if not scan_videos(path) and not state.batch_has_active_items(batch["id"]): state.complete_manual_batch(batch["id"]) return imported @@ -702,8 +694,8 @@ def startup_queue_worker() -> None: @app.on_event("shutdown") def shutdown_queue_worker() -> None: - stop_worker() - state.release_stale_claims() + if stop_worker(wait=True): + state.release_stale_claims() def run() -> None: diff --git a/importarr/state.py b/importarr/state.py index 6029f9f..d191f04 100644 --- a/importarr/state.py +++ b/importarr/state.py @@ -207,32 +207,22 @@ class State: def claim_next_queue_item(self, worker_id: str) -> dict[str, Any] | None: with self._lock: row = self.conn.execute( - """ - select * from import_queue_items - where claimed_by is null - and ( - state = 'ready' - or (state = 'retrying' and (next_retry_at is null or next_retry_at <= current_timestamp)) - ) - order by - case state when 'ready' then 0 else 1 end, - updated_at asc, - id asc - limit 1 - """ - ).fetchone() - if row is None: - return None - self.conn.execute( """ update import_queue_items set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp - where id=? and claimed_by is null + where id = ( + select id from import_queue_items + where claimed_by is null + and (state = 'ready' or (state = 'retrying' and (next_retry_at is null or next_retry_at <= current_timestamp))) + order by case state when 'ready' then 0 else 1 end, updated_at asc, id asc + limit 1 + ) and claimed_by is null + returning * """, - (worker_id, row["id"]), - ) - claimed = self.conn.execute("select * from import_queue_items where id = ?", (row["id"],)).fetchone() - return dict(claimed) if claimed and claimed["claimed_by"] == worker_id else None + (worker_id,), + ).fetchone() + self.conn.commit() + return dict(row) if row else None def release_stale_claims(self) -> int: with self._lock: @@ -245,16 +235,22 @@ class State: def claim_queue_item(self, item_id: int, worker_id: str, allowed_states: set[str]) -> dict[str, Any] | None: placeholders = ",".join("?" for _ in allowed_states) with self._lock: - row = self.conn.execute(f"select * from import_queue_items where id = ? and state in ({placeholders}) and claimed_by is null", (item_id, *allowed_states)).fetchone() - if row is None: - return None - self.conn.execute( - "update import_queue_items set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp where id=? and claimed_by is null", - (worker_id, item_id), + row = self.conn.execute( + f"update import_queue_items set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp where id=? and state in ({placeholders}) and claimed_by is null returning *", + (worker_id, item_id, *allowed_states), + ).fetchone() + self.conn.commit() + return dict(row) if row else None + + def transition_queue_item_if_unclaimed(self, item_id: int, allowed_states: set[str], new_state: str, reason: str) -> bool: + placeholders = ",".join("?" for _ in allowed_states) + with self._lock: + cursor = self.conn.execute( + f"update import_queue_items set state=?, reason=?, updated_at=current_timestamp, next_retry_at=null, completed_at=case when ? in ('failed','skipped') then current_timestamp else null end where id=? and state in ({placeholders}) and claimed_by is null", + (new_state, reason, new_state, item_id, *allowed_states), ) self.conn.commit() - claimed = self.conn.execute("select * from import_queue_items where id = ?", (item_id,)).fetchone() - return dict(claimed) if claimed and claimed["claimed_by"] == worker_id else None + return cursor.rowcount > 0 def mark_queue_item_result( self, diff --git a/tests/test_manual_batches.py b/tests/test_manual_batches.py index 4e37851..b42b8f2 100644 --- a/tests/test_manual_batches.py +++ b/tests/test_manual_batches.py @@ -48,7 +48,7 @@ def test_manual_queue_items_are_persisted_and_imported(tmp_path, monkeypatch): queued = main.state.list_queue_items() assert len(queued) == 1 assert queued[0]["source_type"] == "manual" - assert queued[0]["state"] == "manual_batch" + assert queued[0]["state"] == "ready" assert main._import_manual_batches(main.Importer(movies, tv)) == 1 assert main.state.list_queue_items() == [] diff --git a/tests/test_queue_controls.py b/tests/test_queue_controls.py index 04e099e..6ce7dce 100644 --- a/tests/test_queue_controls.py +++ b/tests/test_queue_controls.py @@ -1,3 +1,5 @@ +import threading + from importarr.config import Settings from importarr.state import State @@ -259,6 +261,97 @@ def test_run_now_conflicts_when_item_is_claimed(tmp_path, monkeypatch): raise AssertionError("expected HTTPException") +def test_bulk_run_now_skips_item_claimed_by_worker(tmp_path, monkeypatch): + main, download, movies, tv = configure_main(tmp_path, monkeypatch) + batch = download / "Release" + batch.mkdir(parents=True) + source = batch / "A.mkv" + source.write_bytes(b"a") + created_batch = main.state.add_manual_batch(batch) + row = main.state.upsert_queue_item(source_type="manual", source_id=str(source), source_path=source, name=source.name, state="ready", batch_id=created_batch["id"]) + main.state.claim_queue_item(row["id"], "worker", {"ready"}) + + assert main._import_manual_batches(main.Importer(movies, tv)) == 0 + assert source.exists() + assert main.state.get_queue_item(row["id"])["claimed_by"] == "worker" + + +def test_retry_loses_atomic_race_with_worker_claim(tmp_path, monkeypatch): + main, _download, _movies, _tv = configure_main(tmp_path, monkeypatch) + row = main.state.upsert_queue_item(source_type="manual", source_id="a", name="A.mkv", state="failed") + original = main.state.transition_queue_item_if_unclaimed + + def claim_then_transition(*args, **kwargs): + main.state.claim_queue_item(row["id"], "worker", {"failed"}) + return original(*args, **kwargs) + + monkeypatch.setattr(main.state, "transition_queue_item_if_unclaimed", claim_then_transition) + + try: + main.queue_item_action(row["id"], main.QueueItemActionRequest(action="retry")) + except main.HTTPException as exc: + assert exc.status_code == 409 + else: + raise AssertionError("expected HTTPException") + assert main.state.get_queue_item(row["id"])["claimed_by"] == "worker" + + +def test_ignore_loses_atomic_race_with_worker_claim(tmp_path, monkeypatch): + main, _download, _movies, _tv = configure_main(tmp_path, monkeypatch) + row = main.state.upsert_queue_item(source_type="manual", source_id="a", name="A.mkv", state="ready") + original = main.state.transition_queue_item_if_unclaimed + + def claim_then_transition(*args, **kwargs): + main.state.claim_queue_item(row["id"], "worker", {"ready"}) + return original(*args, **kwargs) + + monkeypatch.setattr(main.state, "transition_queue_item_if_unclaimed", claim_then_transition) + + try: + main.queue_item_action(row["id"], main.QueueItemActionRequest(action="ignore")) + except main.HTTPException as exc: + assert exc.status_code == 409 + else: + raise AssertionError("expected HTTPException") + assert main.state.get_queue_item(row["id"])["claimed_by"] == "worker" + + +def test_shutdown_waits_for_active_import_before_releasing_claims(tmp_path, monkeypatch): + main, download, _movies, _tv = configure_main(tmp_path, monkeypatch) + source = download / "A.mkv" + source.parent.mkdir(parents=True) + source.write_bytes(b"a") + row = main.state.upsert_queue_item(source_type="manual", source_id=str(source), source_path=source, name=source.name, state="ready") + started = threading.Event() + finish = threading.Event() + + class BlockingImporter: + def import_file(self, *args, **kwargs): + started.set() + finish.wait() + raise RuntimeError("stopped") + + def active_worker(): + claimed = main.state.claim_queue_item(row["id"], main.WORKER_ID, {"ready"}) + main._import_queue_item(claimed, BlockingImporter(), from_worker=True) + + worker = threading.Thread(target=active_worker) + monkeypatch.setattr(main, "_worker_thread", worker) + worker.start() + assert started.wait(1) + shutdown = threading.Thread(target=main.shutdown_queue_worker) + shutdown.start() + + shutdown.join(0.05) + assert shutdown.is_alive() + assert main.state.get_queue_item(row["id"])["claimed_by"] == main.WORKER_ID + + finish.set() + shutdown.join(1) + assert not shutdown.is_alive() + assert main.state.get_queue_item(row["id"])["claimed_by"] is None + + def test_remove_conflicts_when_item_is_claimed(tmp_path, monkeypatch): main, _download, _movies, _tv = configure_main(tmp_path, monkeypatch) row = main.state.upsert_queue_item(source_type="manual", source_id="a", name="A.mkv", state="ready")