diff --git a/custom_components/nestquest/completion.py b/custom_components/nestquest/completion.py index 87e6f93..d87bb31 100644 --- a/custom_components/nestquest/completion.py +++ b/custom_components/nestquest/completion.py @@ -25,9 +25,10 @@ is no permission check here (that is Feature 09). """ from __future__ import annotations +import asyncio import datetime -from .dao_children import _connection_lock +from .dao_children import _CONNECTION_LOCKS, _connection_lock from .dao_instances import ( EVENT_COMPLETED, EVENT_UNCOMPLETED, @@ -62,6 +63,63 @@ def _stamp_occurred_at(now: datetime.datetime | None) -> str: ) +class _TaskReentrantLock: + """Shared connection lock the owning task may re-enter.""" + + def __init__(self, inner: asyncio.Lock) -> None: + self._inner = inner + self._task: asyncio.Task | None = None + self._depth = 0 + + def locked(self) -> bool: + return self._inner.locked() + + async def acquire(self) -> bool: + task = asyncio.current_task() + if self._task is task: + self._depth += 1 + return True + await self._inner.acquire() + self._task = task + self._depth = 1 + return True + + def release(self) -> None: + if self._task is not asyncio.current_task(): + raise RuntimeError("lock released by a non-owner") + self._depth -= 1 + if self._depth == 0: + self._task = None + self._inner.release() + + async def __aenter__(self) -> _TaskReentrantLock: + await self.acquire() + return self + + async def __aexit__(self, *exc: object) -> None: + self.release() + + +def _reentrant_connection_lock(database) -> _TaskReentrantLock: + """Return the shared connection lock, made task-reentrant. + + ``CompletionEventsDao.append`` also acquires this lock, so + check+append in :func:`complete_instance` must re-enter on the + same task. The wrapper is stored in the shared dict so the DAO + sees the same object. + """ + loop = asyncio.get_running_loop() + key = (id(database), id(loop)) + lock = _CONNECTION_LOCKS.get(key) + if isinstance(lock, _TaskReentrantLock): + return lock + wrapped = _TaskReentrantLock( + lock if lock is not None else asyncio.Lock() + ) + _CONNECTION_LOCKS[key] = wrapped + return wrapped + + def _validate_int_id(value: object, field: str) -> int: """Reject non-int ids (bools included) before any lookup.""" if isinstance(value, bool) or not isinstance(value, int): @@ -214,7 +272,9 @@ async def complete_instance( Returns the derived state ``done``. An already-done instance is a no-op: nothing is appended. Unknown ``instance_id`` raises - ValueError naming the field. + ValueError naming the field. Existence, latest-event, and the + append run under one connection lock so a concurrent complete + cannot double-append. ``was_on_time`` is recorded as True — a placeholder until the on-time task computes it from ``now`` / ``today`` against the @@ -226,22 +286,31 @@ async def complete_instance( actor_source, actor_user_id, actor_child_id = _validate_actor( actor_source, actor_user_id, actor_child_id ) - instance = await QuestInstancesDao(database).get_by_id(instance_id) - if instance is None: - raise ValueError( - f"instance_id: quest instance {instance_id} does not exist" + async with _reentrant_connection_lock(database): + instance = await QuestInstancesDao(database).get_by_id(instance_id) + if instance is None: + raise ValueError( + f"instance_id: quest instance {instance_id} does not exist" + ) + latest = await CompletionEventsDao(database).get_latest_for_instance( + instance_id + ) + if latest is not None and latest.event_type == EVENT_COMPLETED: + return "done" + if latest is not None and latest.event_type != EVENT_UNCOMPLETED: + raise ValueError( + f"event_type must be '{EVENT_COMPLETED}' or " + f"'{EVENT_UNCOMPLETED}', got {latest.event_type!r}" + ) + await append_event( + database, + instance.id, + instance.child_id, + EVENT_COMPLETED, + actor_source=actor_source, + actor_user_id=actor_user_id, + actor_child_id=actor_child_id, + was_on_time=True, + now=now, ) - if await instance_state(database, instance_id) == "done": - return "done" - await append_event( - database, - instance.id, - instance.child_id, - EVENT_COMPLETED, - actor_source=actor_source, - actor_user_id=actor_user_id, - actor_child_id=actor_child_id, - was_on_time=True, - now=now, - ) return "done" diff --git a/tests/test_completion.py b/tests/test_completion.py index b7272d7..e9e2ff4 100644 --- a/tests/test_completion.py +++ b/tests/test_completion.py @@ -592,3 +592,95 @@ def test_complete_instance_now_pins_occurred_at(tmp_path) -> None: assert rows[0].occurred_at == "2026-09-18T08:00:00+00:00" _with_db(tmp_path, "complete-pinned.db")(_body) + + +def _gated_hass(armed: dict, gate_open: asyncio.Event, a_started: asyncio.Event): + """A hass mock whose executor gates the FIRST ``_execute`` job after + arming — the INSERT of whichever ``complete_instance`` starts first + — proving that caller holds the connection lock and the second + caller must queue behind it.""" + from unittest.mock import MagicMock + + hass = MagicMock() + + async def _gated_executor(fn, *args): + if armed["active"]: + if getattr(fn, "__name__", "") == "_execute" and not armed["gated"]: + armed["gated"] = True + a_started.set() + await gate_open.wait() + return fn(*args) + + hass.async_add_executor_job = _gated_executor + return hass + + +def test_complete_instance_concurrent_second_call_is_noop(tmp_path) -> None: + """Two racing completes cannot both append. + + Deterministic gate: task A is paused at its INSERT (it holds the + connection lock); task B is then started and must queue. A + finishes, B sees done and appends nothing. + """ + async def _main(): + armed = {"active": False, "gated": False} + gate_open = asyncio.Event() + a_started = asyncio.Event() + hass = _gated_hass(armed, gate_open, a_started) + database = NestQuestDatabase(hass) + await database.open(tmp_path / "complete-race.db") + try: + await apply_migrations(database) + children = ChildrenDao(database) + rules = ScheduleRulesDao(database) + definitions = QuestDefinitionsDao(database) + instances = QuestInstancesDao(database) + child = await children.create("Ada", _now_stamp()) + rule = await rules.create("daily", D1) + definition = await definitions.create( + "Brush teeth", + rule.id, + _now_stamp(), + assignee_child_ids=[child.id], + ) + instance = await instances.upsert( + definition.id, + child.id, + D1, + _now_stamp(), + window="morning", + ) + + armed["active"] = True + task_a = asyncio.ensure_future( + complete_instance( + database, + instance.id, + actor_source="user", + actor_user_id="user-1", + ) + ) + await a_started.wait() + armed["active"] = False + task_b = asyncio.ensure_future( + complete_instance( + database, + instance.id, + actor_source="user", + actor_user_id="user-2", + ) + ) + await asyncio.sleep(0) + gate_open.set() + state_a, state_b = await asyncio.gather(task_a, task_b) + assert state_a == "done" + assert state_b == "done" + rows = await CompletionEventsDao(database).list_by_instance( + instance.id + ) + assert len(rows) == 1 + assert rows[0].actor_user_id == "user-1" + finally: + await database.close() + + _run(_main())