fix some race conditions in PWM and caches

Author: CamDavidsonPilon

core/pioreactor/actions/led_intensity.py

@@ -2,6 +2,7 @@
 import os
 from contextlib import contextmanager
 from contextlib import nullcontext
+from time import time_ns
 from typing import Any
 from typing import Iterator
 
@@ -55,15 +56,19 @@ def change_leds_intensities_temporarily(
 
 @contextmanager
 def lock_leds_temporarily(channels: list[LedChannel]) -> Iterator[None]:
+    lock_id = f"{os.getpid()}:{time_ns()}"
+    acquired_channels: list[LedChannel] = []
     try:
         with local_intermittent_storage("led_locks") as cache:
             for c in channels:
-                cache[c] = os.getpid()
+                if cache.set_if_absent(c, lock_id):
+                    acquired_channels.append(c)
         yield
     finally:
         with local_intermittent_storage("led_locks") as cache:
-            for c in channels:
-                cache.pop(c)
+            for c in acquired_channels:
+                if cache.get(c) == lock_id:
+                    cache.pop(c)
 
 
 def is_led_channel_locked(channel: LedChannel) -> bool:

core/pioreactor/actions/pump.py

@@ -314,9 +314,15 @@ def _get_pump_action(pump_device: PumpCalibrationDevices) -> str:
                 )
                 return 0.0
 
-        with PWMPump(
-            unit, experiment, pin, calibration=calibration, mqtt_client=mqtt_client, logger=logger
-        ) as pump:
+        try:
+            pump_instance = PWMPump(
+                unit, experiment, pin, calibration=calibration, mqtt_client=mqtt_client, logger=logger
+            )
+        except exc.PWMError as e:
+            logger.error(str(e))
+            return 0.0
+
+        with pump_instance as pump:
             sub_duration = 0.5
             volume_moved_ml = 0.0
 

core/pioreactor/utils/__init__.py

@@ -380,6 +380,16 @@ def __setitem__(self, key, value):
     def set(self, key, value):
         return self.__setitem__(key, value)
 
+    def set_if_absent(self, key, value) -> bool:
+        self.cursor.execute(
+            f"""
+            INSERT OR IGNORE INTO {self.table_name} (key, value)
+            VALUES (?, ?)
+        """,
+            (key, value),
+        )
+        return self.cursor.rowcount == 1
+
     def get(self, key, default=None):
         self.cursor.execute(f"SELECT value FROM {self.table_name} WHERE key = ?", (key,))
         result = self.cursor.fetchone()

core/pioreactor/utils/pwm.py

@@ -123,8 +123,10 @@ def dc(self, dc: pt.FloatBetween0and100) -> None:
         if self._started:
             try:
                 lgpio.tx_pwm(self._handle, self.pin, self.frequency, self.dc)
-            except lgpio.error:
-                pass
+            except lgpio.error as e:
+                raise PWMError(
+                    f"Failed to set software PWM on GPIO-{self.pin} to {dc:.5g}% duty cycle at {self.frequency:.5g} Hz."
+                ) from e
         elif dc == 0:
             pass
         else:
@@ -213,6 +215,7 @@ def __init__(
         self.pin: GpioPin = pin
         self.hz = hz
         self.duty_cycle = 0.0
+        self._lock_id = f"{getpid()}:{id(self)}"
 
         if self.is_locked():
             msg = f"GPIO-{self.pin} is currently locked but a task is overwriting it. Either too many jobs are trying to access this pin, or a job didn't clean up properly. If your confident you can release it, use `pio cache clear pwm_locks {self.pin} --as-int` on the command line for {self.unit}."
@@ -318,11 +321,21 @@ def is_locked(self) -> bool:
 
     def lock(self) -> None:
         with local_intermittent_storage("pwm_locks") as pwm_locks:
-            pwm_locks[self.pin] = getpid()
+            if pwm_locks.set_if_absent(self.pin, self._lock_id):
+                return
+
+            raise PWMError(
+                f"GPIO-{self.pin} is currently locked but a task is overwriting it. Either too many jobs are trying to access this pin, or a job didn't clean up properly. If your confident you can release it, use `pio cache clear pwm_locks {self.pin} --as-int` on the command line for {self.unit}."
+            )
 
     def unlock(self) -> None:
         with local_intermittent_storage("pwm_locks") as pwm_locks:
-            pwm_locks.pop(self.pin)
+            owner_lock_id = pwm_locks.get(self.pin)
+            if owner_lock_id is None:
+                return
+
+            if owner_lock_id == self._lock_id:
+                pwm_locks.pop(self.pin)
 
     @contextmanager
     def lock_temporarily(self) -> Iterator[None]:

core/pioreactor/web/utils.py

@@ -156,8 +156,13 @@ def is_rate_limited(job: str, expire_time_seconds=1.0) -> bool:
     Check if the user has made a request within the debounce duration.
     """
     with local_intermittent_storage("debounce") as cache:
-        if cache.get(job) and (time() - cache.get(job)) < expire_time_seconds:
-            return True
-        else:
-            cache.set(job, time())
+        now = time()
+        if cache.set_if_absent(job, now):
             return False
+
+        last_request_time = cache.get(job)
+        if (last_request_time is not None) and ((now - float(last_request_time)) < expire_time_seconds):
+            return True
+
+        cache.set(job, now)
+        return False

core/tests/test_led_intensity.py

@@ -149,6 +149,18 @@ def test_is_led_channel_locked_directly() -> None:
     assert not is_led_channel_locked("A")
 
 
+def test_nested_lock_leds_temporarily_does_not_release_outer_lock() -> None:
+    assert not is_led_channel_locked("A")
+
+    with lock_leds_temporarily(["A"]):
+        assert is_led_channel_locked("A")
+        with lock_leds_temporarily(["A"]):
+            assert is_led_channel_locked("A")
+        assert is_led_channel_locked("A")
+
+    assert not is_led_channel_locked("A")
+
+
 def test_change_leds_intensities_temporarily_invalid_raises_and_state_unchanged() -> None:
     unit = get_unit_name()
     exp = "test_change_leds_invalid"

core/tests/test_pwms.py

@@ -1,12 +1,16 @@
 # -*- coding: utf-8 -*-
 # test_pwms
 import json
+import sys
 import time
+import types
 
 import pytest
 from pioreactor import pubsub
+from pioreactor.exc import PWMError
 from pioreactor.utils import local_intermittent_storage
 from pioreactor.utils.pwm import PWM
+from pioreactor.utils.pwm import SoftwarePWMOutputDevice
 from pioreactor.whoami import get_unit_name
 
 
@@ -98,3 +102,76 @@ def collect(msg) -> None:
 
     assert json.loads(mqtt_items[-1]).get("17", 0.0) == 0.0
     assert json.loads(mqtt_items[-1]).get("12", 0.0) == 0.0
+
+
+def _install_fake_lgpio(monkeypatch: pytest.MonkeyPatch) -> dict[str, float | list[float] | None]:
+    fake_lgpio = types.ModuleType("lgpio")
+
+    class FakeLgpioError(Exception):
+        pass
+
+    state: dict[str, float | list[float] | None] = {"raise_when_dc": None, "tx_calls": []}
+
+    def gpiochip_open(_chip: int) -> int:
+        return 1
+
+    def gpio_claim_output(_handle: int, _pin: int) -> None:
+        return
+
+    def tx_pwm(_handle: int, _pin: int, _frequency: float, duty_cycle: float) -> None:
+        tx_calls = state["tx_calls"]
+        assert isinstance(tx_calls, list)
+        tx_calls.append(duty_cycle)
+
+        raise_when_dc = state["raise_when_dc"]
+        if (raise_when_dc is not None) and (duty_cycle == raise_when_dc):
+            raise FakeLgpioError("simulated lgpio tx_pwm failure")
+
+    def gpiochip_close(_handle: int) -> None:
+        return
+
+    fake_lgpio.error = FakeLgpioError
+    fake_lgpio.gpiochip_open = gpiochip_open
+    fake_lgpio.gpio_claim_output = gpio_claim_output
+    fake_lgpio.tx_pwm = tx_pwm
+    fake_lgpio.gpiochip_close = gpiochip_close
+
+    monkeypatch.setitem(sys.modules, "lgpio", fake_lgpio)
+    return state
+
+
+def test_software_pwm_dc_errors_raise_pwm_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = _install_fake_lgpio(monkeypatch)
+    pwm = SoftwarePWMOutputDevice(pin=17, frequency=100)
+    pwm.start(0)
+
+    state["raise_when_dc"] = 25.0
+
+    with pytest.raises(PWMError, match="Failed to set software PWM"):
+        pwm.dc = 25.0
+
+
+def test_software_pwm_stop_errors_raise_pwm_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = _install_fake_lgpio(monkeypatch)
+    pwm = SoftwarePWMOutputDevice(pin=17, frequency=100)
+    pwm.start(20)
+
+    state["raise_when_dc"] = 0.0
+
+    with pytest.raises(PWMError, match="Failed to set software PWM"):
+        pwm.off()
+
+
+def test_lock_is_exclusive_after_creation() -> None:
+    exp = "test_lock_is_exclusive_after_creation"
+    unit = get_unit_name()
+
+    pwm_first = PWM(12, 10, experiment=exp, unit=unit)
+    pwm_second = PWM(12, 10, experiment=exp, unit=unit)
+
+    pwm_first.lock()
+    with pytest.raises(PWMError, match="GPIO-12 is currently locked"):
+        pwm_second.lock()
+
+    pwm_first.clean_up()
+    pwm_second.clean_up()

core/tests/test_utils.py

@@ -96,6 +96,18 @@ def test_caches_pop() -> None:
         assert cache.pop("C", default=3) == 3
 
 
+def test_cache_set_if_absent() -> None:
+    with local_intermittent_storage("test") as cache:
+        cache.empty()
+
+    with local_intermittent_storage("test") as cache:
+        assert cache.set_if_absent("A", "1")
+        assert not cache.set_if_absent("A", "2")
+
+    with local_intermittent_storage("test") as cache:
+        assert cache["A"] == "1"
+
+
 def test_caches_can_have_tuple_or_singleton_keys() -> None:
     with local_persistent_storage("test_caches_can_have_tuple_keys") as c:
         c[(1, 2)] = 1

core/tests/web/test_utils.py

@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
 import pytest
+from pioreactor.utils import local_intermittent_storage
+from pioreactor.web import utils as web_utils
+from pioreactor.web.utils import is_rate_limited
 from pioreactor.web.utils import is_valid_unix_filename
 from pioreactor.web.utils import scrub_to_valid
 
@@ -53,3 +56,25 @@ def test_valid_unix_filenames(name) -> None:
 )
 def test_invalid_unix_filenames(name) -> None:
     assert not is_valid_unix_filename(name)
+
+
+def test_is_rate_limited_blocks_second_request_within_window() -> None:
+    job_name = "test_rate_limit_second_blocked"
+    with local_intermittent_storage("debounce") as cache:
+        cache.pop(job_name)
+
+    assert not is_rate_limited(job_name, expire_time_seconds=10.0)
+    assert is_rate_limited(job_name, expire_time_seconds=10.0)
+
+
+def test_is_rate_limited_allows_after_expiry(monkeypatch: pytest.MonkeyPatch) -> None:
+    job_name = "test_rate_limit_allows_after_expiry"
+    with local_intermittent_storage("debounce") as cache:
+        cache.pop(job_name)
+
+    timeline = iter([1000.0, 1002.0, 1002.2])
+    monkeypatch.setattr(web_utils, "time", lambda: next(timeline))
+
+    assert not is_rate_limited(job_name, expire_time_seconds=1.0)
+    assert not is_rate_limited(job_name, expire_time_seconds=1.0)
+    assert is_rate_limited(job_name, expire_time_seconds=1.0)