157 lines
5.1 KiB
Python
157 lines
5.1 KiB
Python
"""Worker execution: success, failure, cooperative cancel, resume, fencing, and a
|
|
randomized concurrency stress with no duplicate completion or lost events."""
|
|
|
|
import threading
|
|
from collections import Counter
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import timedelta
|
|
|
|
import pytest
|
|
|
|
from photo_pipeline.db import create_db_engine, create_session_factory, run_migrations
|
|
from photo_pipeline.jobs.worker import Worker
|
|
from photo_pipeline.services.jobs import (
|
|
ItemState,
|
|
JobConflict,
|
|
JobService,
|
|
JobState,
|
|
_now,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def db_url(tmp_path):
|
|
url = f"sqlite:///{tmp_path / 'worker.db'}"
|
|
run_migrations(url)
|
|
return url
|
|
|
|
|
|
@pytest.fixture
|
|
def make_factory():
|
|
engines = []
|
|
|
|
def _make(url):
|
|
engine = create_db_engine(url)
|
|
engines.append(engine)
|
|
return create_session_factory(engine)
|
|
|
|
yield _make
|
|
for engine in engines:
|
|
engine.dispose()
|
|
|
|
|
|
@pytest.fixture
|
|
def sf(db_url, make_factory):
|
|
return make_factory(db_url)
|
|
|
|
|
|
@pytest.fixture
|
|
def jobs(sf):
|
|
return JobService(sf)
|
|
|
|
|
|
def test_worker_processes_all_items(sf, jobs):
|
|
handled = []
|
|
worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "w1")
|
|
job = jobs.enqueue("scan", items=["a", "b", "c"])
|
|
assert worker.run_once() == job["id"]
|
|
assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED
|
|
assert sorted(handled) == ["a", "b", "c"]
|
|
assert jobs.progress(job["id"])["done"] == 3
|
|
|
|
|
|
def test_handler_failure_fails_the_job(sf, jobs):
|
|
def handler(item, ctx):
|
|
if item == "bad":
|
|
raise ValueError("nope")
|
|
|
|
worker = Worker(sf, {"scan": handler}, "w1")
|
|
job = jobs.enqueue("scan", items=["ok", "bad"])
|
|
worker.run_once()
|
|
assert jobs.get(job["id"])["state"] == JobState.FAILED
|
|
assert jobs.progress(job["id"])["by_state"] == {
|
|
ItemState.SUCCEEDED: 1,
|
|
ItemState.FAILED: 1,
|
|
}
|
|
|
|
|
|
def test_cooperative_cancellation_leaves_items_resumable(sf, jobs):
|
|
# Cancel arrives while the first item runs; the worker stops before the next.
|
|
def handler(item, ctx):
|
|
if item == "a":
|
|
ctx.service.cancel(ctx.job_id)
|
|
|
|
worker = Worker(sf, {"scan": handler}, "w1")
|
|
job = jobs.enqueue("scan", items=["a", "b"])
|
|
worker.run_once()
|
|
assert jobs.get(job["id"])["state"] == JobState.CANCELLED
|
|
by_state = jobs.progress(job["id"])["by_state"]
|
|
assert by_state.get(ItemState.SUCCEEDED) == 1 # "a" finished
|
|
assert by_state.get(ItemState.QUEUED) == 1 # "b" left resumable
|
|
|
|
|
|
def test_fencing_rejects_superseded_worker(sf, jobs):
|
|
job = jobs.enqueue("scan", items=["a"])
|
|
stale = jobs.claim(["scan"], "old")
|
|
jobs.recover_stale(now=_now() + timedelta(hours=1))
|
|
fresh = jobs.claim(["scan"], "new")
|
|
assert fresh["fencing_token"] > stale["fencing_token"]
|
|
with pytest.raises(JobConflict):
|
|
jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=stale["fencing_token"])
|
|
with pytest.raises(JobConflict):
|
|
jobs.transition(
|
|
job["id"], JobState.SUCCEEDED, fencing_token=stale["fencing_token"]
|
|
)
|
|
|
|
|
|
def test_recovered_job_resumes_remaining_items(sf, jobs):
|
|
job = jobs.enqueue("scan", items=["a", "b"])
|
|
# Simulate a crashed attempt: "a" done, "b" left running, no finalize.
|
|
crashed = jobs.claim(["scan"], "dead")
|
|
token = crashed["fencing_token"]
|
|
jobs.set_item(job["id"], "a", ItemState.RUNNING, fencing_token=token)
|
|
jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=token)
|
|
jobs.set_item(job["id"], "b", ItemState.RUNNING, fencing_token=token)
|
|
jobs.recover_stale(now=_now() + timedelta(hours=1))
|
|
|
|
handled = []
|
|
fresh_worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "alive")
|
|
fresh_worker.run_once()
|
|
assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED
|
|
assert handled == ["b"] # "a" already succeeded, only "b" reprocessed
|
|
assert jobs.progress(job["id"])["by_state"][ItemState.SUCCEEDED] == 2
|
|
|
|
|
|
def test_concurrent_workers_no_duplicate_completion_or_lost_events(db_url, make_factory):
|
|
service = JobService(make_factory(db_url))
|
|
job_ids = [service.enqueue("scan", items=["x", "y"])["id"] for _ in range(12)]
|
|
|
|
calls = Counter()
|
|
lock = threading.Lock()
|
|
|
|
def handler(item, ctx):
|
|
with lock:
|
|
calls[(ctx.job_id, item)] += 1
|
|
|
|
def drain(worker_index):
|
|
worker = Worker(make_factory(db_url), {"scan": handler}, f"w{worker_index}")
|
|
empties = 0
|
|
for _ in range(200):
|
|
if worker.run_once() is None:
|
|
empties += 1
|
|
if empties >= 2:
|
|
break
|
|
else:
|
|
empties = 0
|
|
|
|
with ThreadPoolExecutor(max_workers=4) as pool:
|
|
list(pool.map(drain, range(4)))
|
|
|
|
# Every job completed exactly once; every item handled exactly once.
|
|
for job_id in job_ids:
|
|
assert service.get(job_id)["state"] == JobState.SUCCEEDED
|
|
types = [e["type"] for e in service.events(job_id)]
|
|
assert "claimed" in types and "state:succeeded" in types # no lost events
|
|
assert all(count == 1 for count in calls.values())
|
|
assert len(calls) == len(job_ids) * 2
|