From 1730a2c32a44fb6245fddca0f6c104ffdf2af094 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:32:36 +0530 Subject: [PATCH] fix(gateway): serialize offloaded persists so concurrent flushes land in order Offloading the fsync to a worker thread (previous commit) lets two inbound messages have two persists in flight. Without ordering, a slow first flush lands after the second and the disk file loses the newer token/id/seen-id (in-memory state stayed correct, so the loss only surfaced after restart). The loop-side dict could also be mutated while the worker iterated it. - weixin ContextTokenStore / discord non-conversational tracker: snapshot the payload on the loop, hand it to the worker, serialize with a per-store asyncio.Lock. - feishu dedup: same asyncio.Lock around the offloaded persist (its snapshot was already taken under _dedup_lock; the write was not). - Regression tests: two concurrent persists with a slow first write must leave the union on disk. All three fail on the previous commit. --- gateway/platforms/weixin.py | 18 ++++++++++--- plugins/platforms/discord/adapter.py | 19 +++++++++++--- plugins/platforms/feishu/adapter.py | 10 +++++-- tests/gateway/test_discord_free_response.py | 27 +++++++++++++++++++ tests/gateway/test_feishu.py | 29 +++++++++++++++++++++ tests/gateway/test_weixin.py | 27 +++++++++++++++++++ 6 files changed, 121 insertions(+), 9 deletions(-) diff --git a/gateway/platforms/weixin.py b/gateway/platforms/weixin.py index 58f9d401ba..be698e4f71 100644 --- a/gateway/platforms/weixin.py +++ b/gateway/platforms/weixin.py @@ -302,6 +302,10 @@ class ContextTokenStore: def __init__(self, hermes_home: str): self._root = _account_dir(hermes_home) self._cache: Dict[str, str] = {} + # Serializes the offloaded flushes so two concurrent set() calls + # cannot land their writes out of order (last-writer-wins would drop + # the newer token from disk). + self._persist_lock = asyncio.Lock() def _path(self, account_id: str) -> Path: return self._root / f"{account_id}.context-tokens.json" @@ -334,16 +338,22 @@ class ContextTokenStore: # atomic_json_write() calls os.fsync(), which blocks until the write # reaches stable storage. _process_message runs on the event loop for # every inbound message, so offload the flush the same way #83906 did - # for the other gateway persist paths. - await asyncio.to_thread(self._persist, account_id) + # for the other gateway persist paths. The payload is snapshotted here, + # on the loop, so the worker never iterates ``_cache`` while another + # message task mutates it; the lock keeps flushes in mutation order. + async with self._persist_lock: + payload = self._payload(account_id) + await asyncio.to_thread(self._persist, account_id, payload) - def _persist(self, account_id: str) -> None: + def _payload(self, account_id: str) -> Dict[str, str]: prefix = f"{account_id}:" - payload = { + return { key[len(prefix) :]: value for key, value in self._cache.items() if key.startswith(prefix) } + + def _persist(self, account_id: str, payload: Dict[str, str]) -> None: try: atomic_json_write(self._path(account_id), payload) except Exception as exc: diff --git a/plugins/platforms/discord/adapter.py b/plugins/platforms/discord/adapter.py index 802b685f21..ee54212903 100644 --- a/plugins/platforms/discord/adapter.py +++ b/plugins/platforms/discord/adapter.py @@ -382,6 +382,10 @@ class _DiscordNonConversationalMessageTracker: def __init__(self, max_tracked: int = _MAX_TRACKED): self._max_tracked = max_tracked self._ids: dict[str, None] = dict.fromkeys(self._load()) + # Serializes the offloaded flushes so two concurrent mark_many() calls + # cannot land their writes out of order (last-writer-wins would drop + # the newer ids from disk). + self._persist_lock = asyncio.Lock() def _state_path(self) -> _Path: from hermes_constants import get_hermes_home @@ -404,11 +408,15 @@ class _DiscordNonConversationalMessageTracker: logger.debug("[%s] Failed to load non-conversational Discord IDs", "Discord") return [] - def _save(self) -> None: + def _snapshot(self) -> list[str]: + """Trim in-memory state and return the ids to persist (loop-side).""" ids = list(self._ids) if len(ids) > self._max_tracked: ids = ids[-self._max_tracked:] self._ids = dict.fromkeys(ids) + return ids + + def _save(self, ids: list[str]) -> None: try: atomic_json_write(self._state_path(), ids, indent=None) except Exception: @@ -425,8 +433,13 @@ class _DiscordNonConversationalMessageTracker: # atomic_json_write() calls os.fsync(), which blocks until the # write reaches stable storage. Both callers of mark_many() run # on the event loop, so offload the flush the same way #83906 - # did for the other gateway persist paths. - await asyncio.to_thread(self._save) + # did for the other gateway persist paths. The snapshot (and the + # trim that reassigns ``_ids``) stays on the loop so the worker + # never touches the dict while another task mutates it; the lock + # keeps flushes in mutation order. + async with self._persist_lock: + ids = self._snapshot() + await asyncio.to_thread(self._save, ids) def __contains__(self, message_id: str) -> bool: return str(message_id or "") in self._ids diff --git a/plugins/platforms/feishu/adapter.py b/plugins/platforms/feishu/adapter.py index 87207fcbc0..14ab109162 100644 --- a/plugins/platforms/feishu/adapter.py +++ b/plugins/platforms/feishu/adapter.py @@ -1614,6 +1614,9 @@ class FeishuAdapter(BasePlatformAdapter): self._seen_message_order: List[str] = [] self._dedup_state_path = get_hermes_home() / "feishu_seen_message_ids.json" self._dedup_lock = threading.Lock() + # Serializes the offloaded dedup-state flushes so two concurrent + # inbound messages cannot land their writes out of order. + self._dedup_persist_lock = asyncio.Lock() self._sender_name_cache: Dict[str, tuple[str, float]] = {} # sender_id → (name, expire_at) self._webhook_rate_counts: Dict[str, tuple[int, float]] = {} # rate_key → (count, window_start) self._webhook_anomaly_counts: Dict[str, tuple[int, str, float]] = {} # ip → (count, last_status, first_seen) @@ -4762,8 +4765,11 @@ class FeishuAdapter(BasePlatformAdapter): # atomic_json_write() calls os.fsync(), which blocks until the write # reaches stable storage. _handle_message_event_data runs on the # event loop for every inbound message, so offload the flush the - # same way #83906 did for the other gateway persist paths. - await asyncio.to_thread(self._persist_seen_message_ids) + # same way #83906 did for the other gateway persist paths. The lock + # keeps flushes in mutation order (the snapshot inside the worker is + # taken under _dedup_lock, but the write itself is not). + async with self._dedup_persist_lock: + await asyncio.to_thread(self._persist_seen_message_ids) return False # ========================================================================= diff --git a/tests/gateway/test_discord_free_response.py b/tests/gateway/test_discord_free_response.py index 75ee52ea19..fc58b982a3 100644 --- a/tests/gateway/test_discord_free_response.py +++ b/tests/gateway/test_discord_free_response.py @@ -854,4 +854,31 @@ class TestNonConversationalTrackerOffload: assert write_threads assert all(tid != loop_thread for tid in write_threads) + @pytest.mark.asyncio + async def test_concurrent_mark_many_persists_land_in_order(self): + """Two in-flight mark_many() calls (send() racing a history fetch) must + not let an older snapshot overwrite a newer one on disk.""" + import asyncio as _asyncio + import time + + tracker = discord_platform._DiscordNonConversationalMessageTracker() + tracker._ids = {} + writes = [] + calls = [0] + + def slow_first_write(path, data, *args, **kwargs): + idx = calls[0] + calls[0] += 1 + if idx == 0: + time.sleep(0.05) + writes.append(list(data)) + + with patch.object(discord_platform, "atomic_json_write", side_effect=slow_first_write): + first = _asyncio.create_task(tracker.mark_many(["1"])) + await _asyncio.sleep(0.005) + second = _asyncio.create_task(tracker.mark_many(["2"])) + await _asyncio.gather(first, second) + + assert sorted(writes[-1]) == ["1", "2"] + diff --git a/tests/gateway/test_feishu.py b/tests/gateway/test_feishu.py index c41a4ab0d5..2d1783e051 100644 --- a/tests/gateway/test_feishu.py +++ b/tests/gateway/test_feishu.py @@ -1682,6 +1682,35 @@ class TestDedupTTL(unittest.TestCase): self.assertTrue(write_threads) self.assertTrue(all(tid != loop_thread for tid in write_threads)) + @patch.dict(os.environ, {}, clear=True) + def test_concurrent_dedup_persists_land_in_order(self): + """Two in-flight _is_duplicate() calls (two chats) must not let an + older seen-ids snapshot overwrite a newer one on disk.""" + from gateway.config import PlatformConfig + from plugins.platforms.feishu.adapter import FeishuAdapter + + adapter = FeishuAdapter(PlatformConfig()) + writes = [] + calls = [0] + + def slow_first_write(path, data, *args, **kwargs): + idx = calls[0] + calls[0] += 1 + if idx == 0: + time.sleep(0.05) + writes.append(sorted(data["message_ids"])) + + async def run(): + first = asyncio.create_task(adapter._is_duplicate("om_a")) + await asyncio.sleep(0.005) + second = asyncio.create_task(adapter._is_duplicate("om_b")) + await asyncio.gather(first, second) + + with patch("plugins.platforms.feishu.adapter.atomic_json_write", side_effect=slow_first_write): + asyncio.run(run()) + + self.assertEqual(writes[-1], ["om_a", "om_b"]) + class TestGroupMentionAtAll(unittest.TestCase): """Tests for @_all (Feishu @everyone) group mention routing.""" diff --git a/tests/gateway/test_weixin.py b/tests/gateway/test_weixin.py index 1c44b1c204..936fef0f3e 100644 --- a/tests/gateway/test_weixin.py +++ b/tests/gateway/test_weixin.py @@ -191,6 +191,33 @@ class TestWeixinStatePersistence: assert write_threads assert all(tid != loop_thread for tid in write_threads) + @pytest.mark.asyncio + async def test_concurrent_context_token_persists_land_in_order(self, tmp_path): + """Two in-flight set() calls (two concurrent inbound messages) must not + let an older snapshot overwrite a newer one on disk. Without + serialization the first (slow) flush lands last and drops user-2.""" + import asyncio as _asyncio + import time + + store = ContextTokenStore(str(tmp_path)) + writes = [] + calls = [0] + + def slow_first_write(path, data, *args, **kwargs): + idx = calls[0] + calls[0] += 1 + if idx == 0: + time.sleep(0.05) + writes.append(dict(data)) + + with patch("gateway.platforms.weixin.atomic_json_write", side_effect=slow_first_write): + first = _asyncio.create_task(store.set("acct-1", "user-1", "t1")) + await _asyncio.sleep(0.005) + second = _asyncio.create_task(store.set("acct-1", "user-2", "t2")) + await _asyncio.gather(first, second) + + assert writes[-1] == {"user-1": "t1", "user-2": "t2"} + class TestWeixinQrLogin: @pytest.mark.asyncio