Fix queue claim concurrency Refs #6

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"):
+        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")
+            raise HTTPException(status_code=409, detail="queue item is currently being imported or changed")
-        state.mark_queue_item_result(item["source_type"], item["source_id"], "skipped", "ignored by user")
     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))
+            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)
-            try:
+            claimed = state.claim_queue_item(row["id"], WORKER_ID, {"ready", "failed", "retrying"})
-                started = time.time()
+            if claimed is not None:
-                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))
+                imported += _import_queue_item(claimed, importer, force=force, from_worker=True)
-                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)
     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()
+    if stop_worker(wait=True):
-    state.release_stale_claims()
+        state.release_stale_claims()
 
 
 def run() -> None:

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"]),
+                (worker_id,),
-            )
+            ).fetchone()
-            claimed = self.conn.execute("select * from import_queue_items where id = ?", (row["id"],)).fetchone()
+            self.conn.commit()
-            return dict(claimed) if claimed and claimed["claimed_by"] == worker_id else None
+            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()
+            row = self.conn.execute(
-            if row is None:
+                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 *",
-                return None
+                (worker_id, item_id, *allowed_states),
-            self.conn.execute(
+            ).fetchone()
-                "update import_queue_items set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp where id=? and claimed_by is null",
+            self.conn.commit()
-                (worker_id, item_id),
+            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 cursor.rowcount > 0
-            return dict(claimed) if claimed and claimed["claimed_by"] == worker_id else None
 
     def mark_queue_item_result(
         self,

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() == []

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")