diff --git a/pyproject.toml b/pyproject.toml index dcc0fd5e..3ad6b2d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ test = [ "pytest-qt", "pytest-asyncio", "pytest-timeout", + "fakeredis[lua]", "diff-cover", "ruff>=0.15" ] diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 68fd2ddc..2269fb8f 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import grp import logging import os diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 14db0716..060fd155 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -100,8 +100,8 @@ class BeamlineConfig: host = "localhost" else: host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch") - self.hw_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") self.redis = Redis(host=host, port=6379, db=0, decode_responses=True) + self.hw_lock = RedisLock(self.redis, name=f"{bl}:hardware_busy_lock") self.simulated_detector = bl is MXBeamline.SIMULATED self._initialize_optional_yaml_defaults() @@ -683,7 +683,7 @@ class BeamlineConfig: def reference_tools(self, data: SampleShortInfoList): self.redis.set(f"{self._bl}:reference-tools", data.model_dump_json()) - def listen_changes_reference_tools(self) -> redis.client.PubSub: + def listen_changes_reference_tools(self) -> PubSub: self.redis.config_set("notify-keyspace-events", "KEA") pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:reference-tools") @@ -711,7 +711,7 @@ class BeamlineConfig: if tmp is None: return BeamMarkCoeffModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return BeamMarkCoeffModel(**data_dict) @beam_mark_coeff.setter diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index b7c39029..bb96ea9b 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -95,8 +95,10 @@ async def _lock_hw(): if not hardware_busy_lock.acquire(blocking=False): raise BeamlineBusyException("Beamline hardware lock is held by another worker") logger.debug("Hardware lock acquired by process") - yield - hardware_busy_lock.release() + try: + yield + finally: + hardware_busy_lock.release() logger.debug("Hardware lock released by process") diff --git a/tests/unit/daq/test_server_hw_lock.py b/tests/unit/daq/test_server_hw_lock.py new file mode 100644 index 00000000..9113f162 --- /dev/null +++ b/tests/unit/daq/test_server_hw_lock.py @@ -0,0 +1,73 @@ +from unittest.mock import patch + +import fakeredis +import pytest +from redis.lock import Lock as RedisLock + +LOCK_NAME = "SIMULATED:hardware_busy_lock" +AUTH = {"Authorization": "Bearer fake-token"} + +THREAD_LOCAL = False + + +@pytest.fixture +def fake_redis(): + return fakeredis.FakeRedis(decode_responses=True) + + +@pytest.fixture +def hw_lock(server_module, client, fake_redis, monkeypatch): + """This worker's lock. + + Depends on ``client`` so that TestClient's lifespan (which overwrites + ``hardware_busy_lock``) has already run by the time we patch. + """ + lock = RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL) + monkeypatch.setattr(server_module, "hardware_busy_lock", lock) + return lock + + +@pytest.fixture +def other_worker(fake_redis): + """A second uvicorn worker's view of the same lock in Redis.""" + return RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL) + + +def test_endpoint_acquires_and_releases_the_lock(client, hw_lock, other_worker): + with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"): + response = client.put("/beamline/omega?val=10.5", headers=AUTH) + + assert response.status_code == 200 + # Released on the way out, so a second worker can take it. + assert not hw_lock.owned() + assert other_worker.acquire(blocking=False) + + +def test_endpoint_is_503_while_another_worker_holds_the_lock(client, hw_lock, other_worker): + assert other_worker.acquire(blocking=False) + + with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"): + response = client.put("/beamline/omega?val=10.5", headers=AUTH) + + assert response.status_code == 503 + assert response.json()["exception_class"] == "BeamlineBusyException" + # The rejected request must not have stolen or released the other worker's lock. + assert other_worker.owned() + assert not hw_lock.owned() + + +async def test_lock_hw_is_reentrant_for_the_owning_worker(server_module, hw_lock): + assert hw_lock.acquire(blocking=False) + + async with server_module._lock_hw(): + assert hw_lock.owned() + + assert hw_lock.owned(), "re-entrant exit released a lock it did not acquire" + + +async def test_lock_is_released_when_the_handler_raises(server_module, hw_lock): + with pytest.raises(RuntimeError): + async with server_module._lock_hw(): + raise RuntimeError("handler blew up") + + assert not hw_lock.owned(), "hardware lock leaked after an exception"