Fix queue claim concurrency Refs #6
This commit is contained in:
+17
-25
@@ -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,7 +694,7 @@ 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()
|
||||
|
||||
|
||||
|
||||
+26
-30
@@ -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,
|
||||
|
||||
@@ -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() == []
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user