From 0ed92013f9dae4dd7d7ea7fe1c293f433a496c52 Mon Sep 17 00:00:00 2001 From: agessaman Date: Sat, 18 Jul 2026 19:37:03 -0700 Subject: [PATCH] refactor(frame_server): enhance message persistence handling in CompanionFrameServer Updated the _persist_companion_message method to accept a queue_entry parameter for more precise removal of messages from the bridge queue. Introduced a new _remove_queue_entry method to ensure messages are removed by identity, preventing potential message loss during concurrent operations. Adjusted related tests to reflect these changes and verify correct behavior. --- repeater/companion/frame_server.py | 31 ++- tests/test_companion_bridge_frame_utils.py | 7 +- tests/test_companion_bridge_prefs.py | 2 +- tests/test_companion_message_persist_race.py | 206 +++++++++++++++++++ tests/test_companion_settings.py | 5 +- 5 files changed, 241 insertions(+), 10 deletions(-) create mode 100644 tests/test_companion_message_persist_race.py diff --git a/repeater/companion/frame_server.py b/repeater/companion/frame_server.py index b784365..7076fec 100644 --- a/repeater/companion/frame_server.py +++ b/repeater/companion/frame_server.py @@ -66,12 +66,19 @@ class CompanionFrameServer(_BaseFrameServer): # Persistence hook overrides # ----------------------------------------------------------------- - async def _persist_companion_message(self, msg_dict: dict) -> None: - """Persist message to SQLite and pop from bridge queue. + async def _persist_companion_message(self, msg_dict: dict, queue_entry=None) -> None: + """Persist message to SQLite and remove it from the bridge queue. The bridge's ``offline_queue_size`` (``message_queue.max_size``) doubles as the SQLite retention limit: 0 disables offline storage entirely, so the message is dropped instead of persisted. + + ``queue_entry`` is the exact in-memory entry this message came from. The + persisted entry is removed by identity (``message_queue.remove``) rather + than by ``pop_last``: pushes happen synchronously in sibling receive + tasks, so during the awaited ``to_thread`` another task can append a newer + entry and ``pop_last`` would remove that one instead — duplicating this + message and losing the newer one. """ if not self.sqlite_handler: return @@ -82,7 +89,7 @@ class CompanionFrameServer(_BaseFrameServer): getattr(self.bridge.message_queue, "_max_size", None), ) if retention == 0: - self.bridge.message_queue.pop_last() + self._remove_queue_entry(queue_entry) return persisted = await asyncio.to_thread( self.sqlite_handler.companion_push_message, @@ -91,13 +98,29 @@ class CompanionFrameServer(_BaseFrameServer): retention, ) if persisted: - self.bridge.message_queue.pop_last() + self._remove_queue_entry(queue_entry) else: logger.debug( "Companion %s: retaining message in memory after SQLite queue rejection", self.companion_hash, ) + def _remove_queue_entry(self, queue_entry) -> None: + """Remove exactly the persisted entry from the bridge queue by identity. + + Falling back to ``pop_last`` when ``queue_entry`` is None would reopen the + interleaving race (it could evict a newer, unpersisted entry), so an event + from an older core that carries no entry is left in memory: a possible + duplicate is preferable to losing a message. + """ + if queue_entry is None: + logger.debug( + "Companion %s: no queue entry on persisted message; leaving in memory", + self.companion_hash, + ) + return + self.bridge.message_queue.remove(queue_entry) + def _sync_next_from_persistence(self) -> Optional[QueuedMessage]: """Retrieve next message from SQLite when bridge queue is empty.""" if not self.sqlite_handler: diff --git a/tests/test_companion_bridge_frame_utils.py b/tests/test_companion_bridge_frame_utils.py index af12181..22133cd 100644 --- a/tests/test_companion_bridge_frame_utils.py +++ b/tests/test_companion_bridge_frame_utils.py @@ -128,7 +128,7 @@ async def test_frame_server_persistence_paths_and_stop(): companion_upsert_contact=MagicMock(), ) bridge = SimpleNamespace( - message_queue=SimpleNamespace(pop_last=MagicMock()), + message_queue=SimpleNamespace(remove=MagicMock(), pop_last=MagicMock()), sync_next_message=lambda: None, get_contacts=lambda: [], channels=SimpleNamespace(max_channels=2), @@ -147,9 +147,10 @@ async def test_frame_server_persistence_paths_and_stop(): srv._write_frame = MagicMock() srv._build_message_frame = MagicMock(return_value=b"frame") - await srv._persist_companion_message({"text": "x"}) + entry = object() + await srv._persist_companion_message({"text": "x"}, entry) sqlite.companion_push_message.assert_called_once_with("h", {"text": "x"}, None) - bridge.message_queue.pop_last.assert_called_once() + bridge.message_queue.remove.assert_called_once_with(entry) msg = srv._sync_next_from_persistence() assert msg is not None diff --git a/tests/test_companion_bridge_prefs.py b/tests/test_companion_bridge_prefs.py index 8d76308..aad2d8a 100644 --- a/tests/test_companion_bridge_prefs.py +++ b/tests/test_companion_bridge_prefs.py @@ -81,6 +81,6 @@ def test_bridge_accepts_host_radio_callbacks(identity): "coding_rate": 5, "tx_power_dbm": 19, "rx_delay_base": 0, - "airtime_factor": 0, + "airtime_factor": 1.0, } assert bridge.get_max_tx_power_dbm() == 20 diff --git a/tests/test_companion_message_persist_race.py b/tests/test_companion_message_persist_race.py new file mode 100644 index 0000000..6eecaf7 --- /dev/null +++ b/tests/test_companion_message_persist_race.py @@ -0,0 +1,206 @@ +"""Persistence must remove exactly the entry it persisted, not the queue tail. + +Offline-queue pushes happen synchronously in sibling receive tasks while +``_persist_companion_message`` is suspended inside its ``asyncio.to_thread`` +SQLite write. Removing the persisted entry with ``pop_last`` therefore races: a +newer entry pushed during the await becomes the tail and gets evicted instead, +duplicating the persisted message and losing the newer one. These tests pin the +identity-based removal that closes that window. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import patch + +import pytest +from openhop_core.companion.message_queue import MessageQueue +from openhop_core.companion.models import QueuedMessage +from openhop_core.protocol.constants import TXT_TYPE_PLAIN + +from repeater.companion.frame_server import CompanionFrameServer +from repeater.data_acquisition.sqlite_handler import SQLiteHandler + +_HASH = "0x01" +_TO_THREAD = "repeater.companion.frame_server.asyncio.to_thread" + + +class _Bridge: + """Minimal bridge exposing a real in-memory offline queue.""" + + def __init__(self, max_size: int = 100): + self.message_queue = MessageQueue(max_size=max_size) + + +def _server(handler: SQLiteHandler, bridge: _Bridge) -> CompanionFrameServer: + fs = CompanionFrameServer.__new__(CompanionFrameServer) + fs.sqlite_handler = handler + fs.companion_hash = _HASH + fs.bridge = bridge + return fs + + +def _msg(text: str, key_byte: bytes, packet_hash: str): + """A queue entry and the matching persistence dict for one direct message.""" + entry = QueuedMessage( + sender_key=key_byte * 32, + txt_type=TXT_TYPE_PLAIN, + timestamp=1000, + text=text, + is_channel=False, + path_len=0xFF, + ) + msg_dict = { + "sender_key": entry.sender_key, + "text": text, + "timestamp": entry.timestamp, + "txt_type": entry.txt_type, + "is_channel": False, + "channel_idx": 0, + "path_len": entry.path_len, + "packet_hash": packet_hash, + "snr": 0.0, + "rssi": 0, + "sender_prefix": b"", + } + return entry, msg_dict + + +def _memory_texts(queue: MessageQueue) -> set[str]: + """Drain the queue and return the set of retained message texts.""" + texts = set() + while True: + msg = queue.pop() + if msg is None: + break + texts.add(msg.text) + return texts + + +def _persisted_texts(handler: SQLiteHandler) -> set[str]: + return {m["text"] for m in handler.companion_load_messages(_HASH)} + + +class _Gate: + """Deterministic stand-in for ``asyncio.to_thread``. + + Each call registers a release event and suspends until the test releases it, + then runs the wrapped function inline. This lets the test interleave two + persistence coroutines and control which completes first. + """ + + def __init__(self): + self.releases: list[asyncio.Event] = [] + + async def to_thread(self, func, /, *args, **kwargs): + rel = asyncio.Event() + self.releases.append(rel) + await rel.wait() + return func(*args, **kwargs) + + async def release(self, idx: int) -> None: + self.releases[idx].set() + # Let the resumed coroutine run its synchronous remove/pop tail. + for _ in range(3): + await asyncio.sleep(0) + + +async def _settle() -> None: + """Yield enough for a freshly created task to reach its ``to_thread`` gate.""" + for _ in range(4): + await asyncio.sleep(0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("order", ["a_first", "b_first"]) +async def test_interleaved_persist_no_dup_no_loss(tmp_path, order): + """Two receive tasks persist concurrently; regardless of completion order the + older message ends up only in SQLite and the newer only in memory.""" + handler = SQLiteHandler(tmp_path) + bridge = _Bridge() + fs = _server(handler, bridge) + # Same packet_hash: the newer push is rejected as a SQLite duplicate, the + # damage case where pop_last would duplicate the older and drop the newer. + entry_old, dict_old = _msg("older", b"\x11", "DUP-HASH") + entry_new, dict_new = _msg("newer", b"\x22", "DUP-HASH") + + gate = _Gate() + with patch(_TO_THREAD, gate.to_thread): + # Task A: base_events pushed the older entry, then fired persist. + bridge.message_queue.push(entry_old) + task_a = asyncio.create_task(fs._persist_companion_message(dict_old, entry_old)) + await _settle() + # Task B (sibling): a newer message pushed during A's SQLite write. + bridge.message_queue.push(entry_new) + task_b = asyncio.create_task(fs._persist_companion_message(dict_new, entry_new)) + await _settle() + + if order == "a_first": + await gate.release(0) + await gate.release(1) + else: + await gate.release(1) + await gate.release(0) + await asyncio.gather(task_a, task_b) + + memory = _memory_texts(bridge.message_queue) + persisted = _persisted_texts(handler) + # Exactly one message persisted, the other retained: no message in both + # stores (no duplication) and none dropped (no loss). Which of the two wins + # the SQLite insert depends on completion order; the invariant does not. + assert len(persisted) == 1 + assert memory & persisted == set() + assert memory | persisted == {"older", "newer"} + + +@pytest.mark.asyncio +async def test_asymmetric_rejection_removes_persisted_not_newer(tmp_path): + """msg1 persists, msg2 is rejected: msg1 is removed by identity and msg2 stays, + even though msg2 is the queue tail pop_last would have taken.""" + handler = SQLiteHandler(tmp_path) + bridge = _Bridge() + fs = _server(handler, bridge) + entry1, dict1 = _msg("msg1", b"\x11", "SAME-HASH") + entry2, dict2 = _msg("msg2", b"\x22", "SAME-HASH") # duplicate hash -> rejected + + bridge.message_queue.push(entry1) + bridge.message_queue.push(entry2) + + await fs._persist_companion_message(dict1, entry1) # accepted + await fs._persist_companion_message(dict2, entry2) # rejected (duplicate) + + memory = _memory_texts(bridge.message_queue) + assert memory == {"msg2"} # newer retained, msg1 removed exactly + assert _persisted_texts(handler) == {"msg1"} # msg1 persisted once, not duplicated + + +@pytest.mark.asyncio +async def test_cancelled_sibling_entry_survives_other_persist(tmp_path): + """A sibling task cancelled after its push must keep its queue entry when + another task's persist completes.""" + handler = SQLiteHandler(tmp_path) + bridge = _Bridge() + fs = _server(handler, bridge) + entry1, dict1 = _msg("msg1", b"\x11", "HASH-1") + entry2, dict2 = _msg("msg2", b"\x22", "HASH-2") + + gate = _Gate() + with patch(_TO_THREAD, gate.to_thread): + bridge.message_queue.push(entry1) + task_a = asyncio.create_task(fs._persist_companion_message(dict1, entry1)) + await _settle() + # Sibling pushed its entry, then its persist task is cancelled mid-write. + bridge.message_queue.push(entry2) + task_b = asyncio.create_task(fs._persist_companion_message(dict2, entry2)) + await _settle() + task_b.cancel() + with pytest.raises(asyncio.CancelledError): + await task_b + # Task A now finishes its persist. + await gate.release(0) + await task_a + + memory = _memory_texts(bridge.message_queue) + assert "msg2" in memory # cancelled sibling's message retained + assert "msg1" not in memory # A removed only its own entry + assert _persisted_texts(handler) == {"msg1"} diff --git a/tests/test_companion_settings.py b/tests/test_companion_settings.py index b2fc2d3..1f49c85 100644 --- a/tests/test_companion_settings.py +++ b/tests/test_companion_settings.py @@ -531,10 +531,11 @@ class TestPersistSkipWhenOff: def test_skips_persistence_when_retention_zero(self): import asyncio + entry = object() fs = self._frame_server(0) - asyncio.run(fs._persist_companion_message({"text": "x"})) + asyncio.run(fs._persist_companion_message({"text": "x"}, entry)) fs.sqlite_handler.companion_push_message.assert_not_called() - fs.bridge.message_queue.pop_last.assert_called_once() + fs.bridge.message_queue.remove.assert_called_once_with(entry) def test_persists_with_retention(self): import asyncio