diff --git a/gateway/session.py b/gateway/session.py index 58b8368d7c6..fea3ba4a3c0 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -959,6 +959,13 @@ def build_session_key( return ":".join(key_parts) +class _SessionFlight: + def __init__(self) -> None: + self.event = threading.Event() + self.result: Optional["SessionEntry"] = None + self.error: Optional[BaseException] = None + + class AsyncSessionStore: """Async boundary for the synchronous, thread-safe SessionStore.""" @@ -995,6 +1002,10 @@ class SessionStore: # SQLite / fsync. Each writer snapshots the latest state only after # acquiring this lock, preventing stale delayed writes. self._save_lock = threading.Lock() + self._routing_generation = 0 + self._persisted_routing_generation = 0 + self._inflight_lock = threading.Lock() + self._inflight_sessions: Dict[str, _SessionFlight] = {} self._has_active_processes_fn = has_active_processes_fn # Whether to keep writing the legacy sessions.json mirror alongside # the primary gateway_routing table in state.db. Default True for @@ -1201,40 +1212,45 @@ class SessionStore: self._save() def _save(self) -> None: - """Persist the routing index (session key -> ID mapping). + """Persist the routing index while the caller holds ``_lock``.""" + data, generation = self._snapshot_routing_locked() + self._persist_routing_data(data, generation) - state.db's ``gateway_routing`` table is the primary store (#9006 - follow-up): the whole index is replaced atomically in one SQLite - transaction, mirroring the previous full-file JSON rewrite semantics. + def _snapshot_routing_locked(self) -> tuple[Dict[str, Any], int]: + """Capture immutable routing data and a monotonic generation.""" + self._routing_generation = getattr(self, "_routing_generation", 0) + 1 + return ( + {key: entry.to_dict() for key, entry in self._entries.items()}, + self._routing_generation, + ) - sessions.json is additionally written for backward compatibility - (external tooling, downgrade safety) unless the user disables it via - ``gateway.write_sessions_json: false`` in config.yaml. - """ - data = {key: entry.to_dict() for key, entry in self._entries.items()} - - # Primary: durable SQLite routing table. - db_saved = False - _db = getattr(self, "_db", None) - if _db: - replacer = getattr(_db, "replace_gateway_routing_entries", None) - if callable(replacer): - try: - replacer( - {k: json.dumps(v) for k, v in data.items()}, - scope=self._routing_scope(), - ) - db_saved = True - except Exception as exc: - logger.warning( - "gateway.session: state.db routing save failed: %s", exc - ) - - # Legacy mirror: sessions.json. Kept on by default for compat; when - # disabled we still fall back to it if the DB write failed, so the - # index is never lost entirely. - if getattr(self, "_write_sessions_json", True) or not db_saved: - self._save_sessions_json(data) + def _persist_routing_data(self, data: Dict[str, Any], generation: int) -> None: + """Serialize all whole-index writers through one durable write lock.""" + save_lock = getattr(self, "_save_lock", None) + if save_lock is None: + save_lock = threading.Lock() + self._save_lock = save_lock + with save_lock: + if generation <= getattr(self, "_persisted_routing_generation", 0): + return + db_saved = False + _db = getattr(self, "_db", None) + if _db: + replacer = getattr(_db, "replace_gateway_routing_entries", None) + if callable(replacer): + try: + replacer( + {k: json.dumps(v) for k, v in data.items()}, + scope=self._routing_scope(), + ) + db_saved = True + except Exception as exc: + logger.warning( + "gateway.session: state.db routing save failed: %s", exc + ) + if getattr(self, "_write_sessions_json", True) or not db_saved: + self._save_sessions_json(data) + self._persisted_routing_generation = generation def _save_sessions_json(self, data: Dict[str, Any]) -> None: """Write the legacy sessions.json mirror of the routing index.""" @@ -1277,35 +1293,10 @@ class SessionStore: raise def _save_entries(self) -> None: - """Persist the latest routing index without holding the state lock for I/O.""" - save_lock = getattr(self, "_save_lock", None) - if save_lock is None: - save_lock = threading.Lock() - self._save_lock = save_lock - with save_lock: - # Capture immutable dictionaries under the state lock, then release - # it before SQLite replacement and JSON fsync. - with self._lock: - data = {key: entry.to_dict() for key, entry in self._entries.items()} - - db_saved = False - _db = getattr(self, "_db", None) - if _db: - replacer = getattr(_db, "replace_gateway_routing_entries", None) - if callable(replacer): - try: - replacer( - {k: json.dumps(v) for k, v in data.items()}, - scope=self._routing_scope(), - ) - db_saved = True - except Exception as exc: - logger.warning( - "gateway.session: state.db routing save failed: %s", exc - ) - - if getattr(self, "_write_sessions_json", True) or not db_saved: - self._save_sessions_json(data) + """Snapshot latest state under ``_lock`` and persist after releasing it.""" + with self._lock: + data, generation = self._snapshot_routing_locked() + self._persist_routing_data(data, generation) def _resolve_profile_for_key(self, source: Optional[SessionSource] = None) -> Optional[str]: """Return the profile namespace for session keys, or None when off. @@ -1784,15 +1775,58 @@ class SessionStore: def get_or_create_session( self, source: SessionSource, - force_new: bool = False + force_new: bool = False, ) -> SessionEntry: - """Get an existing session or create a new one. + """Single-flight session lookup/create per routing key. - Evaluates reset policy to determine if the existing session is stale. - Creates a session record in SQLite when a new session starts. + Calls for different keys remain concurrent. Overlapping calls for the + same key share the owner's result, including concurrent ``force_new`` + deliveries, so only one routing transition and SQLite row is created. + """ + session_key = self._generate_session_key(source) + inflight_lock = getattr(self, "_inflight_lock", None) + if inflight_lock is None: + inflight_lock = threading.Lock() + self._inflight_lock = inflight_lock + self._inflight_sessions = {} + + with inflight_lock: + slot = self._inflight_sessions.get(session_key) + if slot is None: + slot = _SessionFlight() + self._inflight_sessions[session_key] = slot + owner = True + else: + owner = False + + if not owner: + slot.event.wait() + if slot.error is not None: + raise slot.error + assert slot.result is not None + return slot.result + + try: + result = self._get_or_create_session_impl(source, force_new=force_new) + slot.result = result + return result + except BaseException as exc: + slot.error = exc + raise + finally: + slot.event.set() + with inflight_lock: + self._inflight_sessions.pop(session_key, None) + + def _get_or_create_session_impl( + self, + source: SessionSource, + force_new: bool = False, + ) -> SessionEntry: + """Perform one session routing transition for the single-flight owner. All blocking I/O (SQLite SELECTs, routing-index rewrite + ``os.fsync``, - recovery DB queries) is performed *outside* ``self._lock``. The lock + recovery DB queries) is performed *outside* ``self._lock``. The lock protects only ``_entries`` / ``_loaded`` mutations. """ session_key = self._generate_session_key(source) @@ -1801,6 +1835,7 @@ class SessionStore: db_end_session_id = None db_create_kwargs = None existing_session_id = None + force_new_observed_entry = None # ---- Phase 0: lock read -- existing session_id for compression tip ---- if not force_new: @@ -1822,6 +1857,8 @@ class SessionStore: _entry_for_checks = None with self._lock: self._ensure_loaded_locked() + if force_new: + force_new_observed_entry = self._entries.get(session_key) if session_key in self._entries and not force_new: _entry_for_checks = self._entries[session_key] _stale_session_id = _entry_for_checks.session_id @@ -1903,7 +1940,7 @@ class SessionStore: _needs_recover = True # ---- Phase 3: no-lock I/O -- recovery + create + save + DB ops ---- - if _needs_recover: + if _needs_recover and db_end_session_id is None: recovered = self._query_recoverable_session( session_key=session_key, source=source, now=now, ) @@ -1934,10 +1971,16 @@ class SessionStore: reset_had_activity=reset_had_activity, ) with self._lock: - published = self._entries.get(session_key) if not force_new else None - if published is None: + current = self._entries.get(session_key) + may_publish = current is None or ( + force_new and current is force_new_observed_entry + ) + if may_publish: self._entries[session_key] = candidate published = candidate + else: + published = current + assert published is not None entry = published _needs_save = True if entry is candidate: diff --git a/tests/gateway/test_session_store_lock_io.py b/tests/gateway/test_session_store_lock_io.py index 64d6291c974..7b959f315c1 100644 --- a/tests/gateway/test_session_store_lock_io.py +++ b/tests/gateway/test_session_store_lock_io.py @@ -203,16 +203,22 @@ def test_concurrent_same_key_returns_one_published_session(tmp_path): source = _source() db = _db_with_rows({}) store = _make_store(tmp_path, db) - barrier = threading.Barrier(2) + owner_started = threading.Event() + release_owner = threading.Event() original_query = store._query_recoverable_session def synchronized_query(**kwargs): - barrier.wait(timeout=2) + owner_started.set() + assert release_owner.wait(timeout=2) return original_query(**kwargs) store._query_recoverable_session = synchronized_query # type: ignore[method-assign] with ThreadPoolExecutor(max_workers=2) as pool: - entries = list(pool.map(lambda _: store.get_or_create_session(source), range(2))) + owner = pool.submit(store.get_or_create_session, source) + assert owner_started.wait(timeout=2) + follower = pool.submit(store.get_or_create_session, source) + release_owner.set() + entries = [owner.result(timeout=2), follower.result(timeout=2)] key = store._generate_session_key(source) assert entries[0] is entries[1] @@ -221,6 +227,97 @@ def test_concurrent_same_key_returns_one_published_session(tmp_path): assert created_ids == {entries[0].session_id} +def test_concurrent_force_new_returns_one_published_session(tmp_path): + """Concurrent /new delivery must not create orphan SQLite sessions.""" + source = _source() + db = _db_with_rows({}) + store = _make_store(tmp_path, db) + owner_started = threading.Event() + release_owner = threading.Event() + original_impl = store._get_or_create_session_impl + + def synchronized_impl(*args, **kwargs): + owner_started.set() + assert release_owner.wait(timeout=2) + return original_impl(*args, **kwargs) + + store._get_or_create_session_impl = synchronized_impl # type: ignore[method-assign] + with ThreadPoolExecutor(max_workers=2) as pool: + owner = pool.submit(store.get_or_create_session, source, True) + assert owner_started.wait(timeout=2) + follower = pool.submit(store.get_or_create_session, source, True) + release_owner.set() + entries = [owner.result(timeout=2), follower.result(timeout=2)] + + assert entries[0] is entries[1] + created_ids = {call.kwargs["session_id"] for call in db.create_session.call_args_list} + assert created_ids == {entries[0].session_id} + + +def test_auto_reset_does_not_recover_session_being_ended(tmp_path): + source = _source() + db = _db_with_rows({}) + store = _make_store(tmp_path, db) + key = store._generate_session_key(source) + old = _seed_entry(store, key, "old-session") + old.suspended = True + db.find_latest_gateway_session_for_peer.return_value = { + "id": old.session_id, + "session_key": key, + "started_at": old.created_at.timestamp(), + } + + entry = store.get_or_create_session(source) + + assert entry.session_id != old.session_id + assert entry.was_auto_reset is True + db.reopen_session.assert_not_called() + db.end_session.assert_called_once_with(old.session_id, "session_reset") + + +def test_legacy_and_off_lock_saves_share_one_serialization_lock(tmp_path): + db = _db_with_rows({}) + persisted: dict[str, str] = {} + first_write_started = threading.Event() + release_first_write = threading.Event() + write_count = 0 + count_lock = threading.Lock() + + def replace(entries, *, scope): + nonlocal write_count, persisted + with count_lock: + write_count += 1 + call_number = write_count + if call_number == 1: + first_write_started.set() + assert release_first_write.wait(timeout=2) + persisted = dict(entries) + + db.replace_gateway_routing_entries.side_effect = replace + store = _make_store(tmp_path, db) + source_a = _source() + source_b = SessionSource( + platform=Platform.TELEGRAM, + chat_id="67890", + chat_type="dm", + user_id="67890", + ) + key_a = store._generate_session_key(source_a) + key_b = store._generate_session_key(source_b) + _seed_entry(store, key_a, "sid-a") + + with ThreadPoolExecutor(max_workers=2) as pool: + future_a = pool.submit(store._save_entries) + assert first_write_started.wait(timeout=2) + _seed_entry(store, key_b, "sid-b") + future_b = pool.submit(store._save) + release_first_write.set() + future_a.result(timeout=2) + future_b.result(timeout=2) + + assert set(persisted) == {key_a, key_b} + + def test_save_serialization_snapshots_latest_routing_index(tmp_path): """A delayed earlier writer must snapshot the state visible when it writes.""" db = _db_with_rows({})