| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import os |
| 5 | import threading |
| 6 | import time |
| 7 | from dataclasses import dataclass, field |
| 8 | from typing import Any, TYPE_CHECKING |
| 9 | |
| 10 | from helpers import runtime |
| 11 | from helpers.print_style import PrintStyle |
| 12 | from helpers.state_snapshot import ( |
| 13 | StateRequestV1, |
| 14 | advance_state_request_after_snapshot, |
| 15 | build_snapshot_from_request, |
| 16 | ) |
| 17 | from helpers.ws import ConnectionIdentity, ConnectionNotFoundError, _ws_debug_enabled, ws_debug |
| 18 | from helpers.ws_manager import STATE_PUSH_EVENT |
| 19 | |
| 20 | if TYPE_CHECKING: # pragma: no cover - hints only |
| 21 | from helpers.ws_manager import WsManager |
| 22 | |
| 23 | |
| 24 | @dataclass |
| 25 | class ConnectionProjection: |
| 26 | namespace: str |
| 27 | sid: str |
| 28 | request: StateRequestV1 | None = None |
| 29 | seq: int = 0 |
| 30 | seq_base: int = 0 |
| 31 | # Incremented on every dirty signal. Used to coalesce bursts without delaying |
| 32 | # pushes indefinitely during continuous activity (throttled coalescing). |
| 33 | dirty_version: int = 0 |
| 34 | pushed_version: int = 0 |
| 35 | collections_dirty_version: int = 0 |
| 36 | # Development-only diagnostics - last known cause of the most recent dirty wave. |
| 37 | dirty_reason: str | None = None |
| 38 | dirty_wave_id: str | None = None |
| 39 | created_at: float = field(default_factory=time.time) |
| 40 | |
| 41 | |
| 42 | class StateMonitor: |
| 43 | """Per-sid dirty tracking with debounced snapshot push scheduling.""" |
| 44 | |
| 45 | def __init__(self, debounce_seconds: float = 0.1) -> None: |
| 46 | self.debounce_seconds = float(debounce_seconds) |
| 47 | self._lock = threading.RLock() |
| 48 | self._projections: dict[ConnectionIdentity, ConnectionProjection] = {} |
| 49 | self._debounce_handles: dict[ConnectionIdentity, asyncio.TimerHandle] = {} |
| 50 | self._push_tasks: dict[ConnectionIdentity, asyncio.Task[None]] = {} |
| 51 | self._manager: WsManager | None = None |
| 52 | self._emit_handler_id: str | None = None |
| 53 | self._dispatcher_loop: asyncio.AbstractEventLoop | None = None |
| 54 | self._dirty_wave_seq: int = 0 |
| 55 | |
| 56 | def bind_manager(self, manager: "WsManager", *, handler_id: str | None = None) -> None: |
| 57 | with self._lock: |
| 58 | self._manager = manager |
| 59 | if handler_id: |
| 60 | self._emit_handler_id = handler_id |
| 61 | # Use the manager's dispatcher loop for all scheduling so mark_dirty can be |
| 62 | # invoked safely from non-async contexts and other threads. |
| 63 | self._dispatcher_loop = getattr(manager, "_dispatcher_loop", None) |
| 64 | ws_debug( |
| 65 | f"[StateMonitor] bind_manager handler_id={handler_id or self._emit_handler_id}" |
| 66 | ) |
| 67 | |
| 68 | def register_sid(self, namespace: str, sid: str) -> None: |
| 69 | identity: ConnectionIdentity = (namespace, sid) |
| 70 | with self._lock: |
| 71 | self._projections.setdefault( |
| 72 | identity, ConnectionProjection(namespace=namespace, sid=sid) |
| 73 | ) |
| 74 | ws_debug(f"[StateMonitor] register_sid namespace={namespace} sid={sid}") |
| 75 | |
| 76 | def unregister_sid(self, namespace: str, sid: str) -> None: |
| 77 | identity: ConnectionIdentity = (namespace, sid) |
| 78 | with self._lock: |
| 79 | handle = self._debounce_handles.pop(identity, None) |
| 80 | if handle is not None: |
| 81 | handle.cancel() |
| 82 | task = self._push_tasks.pop(identity, None) |
| 83 | if task is not None: |
| 84 | task.cancel() |
| 85 | self._projections.pop(identity, None) |
| 86 | ws_debug(f"[StateMonitor] unregister_sid namespace={namespace} sid={sid}") |
| 87 | |
| 88 | def mark_dirty_all(self, *, reason: str | None = None) -> None: |
| 89 | wave_id = None |
| 90 | if _ws_debug_enabled(): |
| 91 | with self._lock: |
| 92 | self._dirty_wave_seq += 1 |
| 93 | wave_id = f"all_{self._dirty_wave_seq}" |
| 94 | with self._lock: |
| 95 | identities = list(self._projections.keys()) |
| 96 | for namespace, sid in identities: |
| 97 | self.mark_dirty(namespace, sid, reason=reason, wave_id=wave_id) |
| 98 | |
| 99 | def mark_dirty_for_context( |
| 100 | self, |
| 101 | context_id: str, |
| 102 | *, |
| 103 | reason: str | None = None, |
| 104 | include_collections: bool = True, |
| 105 | ) -> None: |
| 106 | if not isinstance(context_id, str) or not context_id.strip(): |
| 107 | return |
| 108 | target = context_id.strip() |
| 109 | wave_id = None |
| 110 | if _ws_debug_enabled(): |
| 111 | with self._lock: |
| 112 | self._dirty_wave_seq += 1 |
| 113 | wave_id = f"ctx_{self._dirty_wave_seq}" |
| 114 | with self._lock: |
| 115 | identities = [ |
| 116 | identity |
| 117 | for identity, projection in self._projections.items() |
| 118 | if projection.request is not None and projection.request.context == target |
| 119 | ] |
| 120 | for namespace, sid in identities: |
| 121 | self.mark_dirty( |
| 122 | namespace, |
| 123 | sid, |
| 124 | reason=reason, |
| 125 | wave_id=wave_id, |
| 126 | include_collections=include_collections, |
| 127 | ) |
| 128 | |
| 129 | def update_projection( |
| 130 | self, |
| 131 | namespace: str, |
| 132 | sid: str, |
| 133 | *, |
| 134 | request: StateRequestV1, |
| 135 | seq_base: int, |
| 136 | ) -> None: |
| 137 | identity: ConnectionIdentity = (namespace, sid) |
| 138 | with self._lock: |
| 139 | projection = self._projections.setdefault( |
| 140 | identity, ConnectionProjection(namespace=namespace, sid=sid) |
| 141 | ) |
| 142 | projection.request = request |
| 143 | projection.seq_base = seq_base |
| 144 | projection.seq = seq_base |
| 145 | ws_debug( |
| 146 | f"[StateMonitor] update_projection namespace={namespace} sid={sid} context={request.context!r} " |
| 147 | f"log_from={request.log_from} notifications_from={request.notifications_from} " |
| 148 | f"timezone={request.timezone!r} seq_base={seq_base}" |
| 149 | ) |
| 150 | |
| 151 | def mark_dirty( |
| 152 | self, |
| 153 | namespace: str, |
| 154 | sid: str, |
| 155 | *, |
| 156 | reason: str | None = None, |
| 157 | wave_id: str | None = None, |
| 158 | include_collections: bool = True, |
| 159 | ) -> None: |
| 160 | identity: ConnectionIdentity = (namespace, sid) |
| 161 | loop = self._dispatcher_loop |
| 162 | if loop is None or loop.is_closed(): |
| 163 | try: |
| 164 | loop = asyncio.get_running_loop() |
| 165 | except RuntimeError: |
| 166 | return |
| 167 | |
| 168 | try: |
| 169 | running_loop = asyncio.get_running_loop() |
| 170 | except RuntimeError: |
| 171 | running_loop = None |
| 172 | |
| 173 | if running_loop is loop: |
| 174 | self._mark_dirty_on_loop( |
| 175 | identity, |
| 176 | reason=reason, |
| 177 | wave_id=wave_id, |
| 178 | include_collections=include_collections, |
| 179 | ) |
| 180 | return |
| 181 | |
| 182 | loop.call_soon_threadsafe( |
| 183 | self._mark_dirty_on_loop, |
| 184 | identity, |
| 185 | reason, |
| 186 | wave_id, |
| 187 | include_collections, |
| 188 | ) |
| 189 | |
| 190 | def _mark_dirty_on_loop( |
| 191 | self, |
| 192 | identity: ConnectionIdentity, |
| 193 | reason: str | None = None, |
| 194 | wave_id: str | None = None, |
| 195 | include_collections: bool = True, |
| 196 | ) -> None: |
| 197 | with self._lock: |
| 198 | projection = self._projections.get(identity) |
| 199 | if projection is None: |
| 200 | return |
| 201 | projection.dirty_version += 1 |
| 202 | if include_collections: |
| 203 | projection.collections_dirty_version = projection.dirty_version |
| 204 | if runtime.is_development(): |
| 205 | projection.dirty_reason = ( |
| 206 | reason.strip() |
| 207 | if isinstance(reason, str) and reason.strip() |
| 208 | else "unknown" |
| 209 | ) |
| 210 | projection.dirty_wave_id = wave_id |
| 211 | self._schedule_debounce_on_loop(identity) |
| 212 | |
| 213 | def _schedule_debounce_on_loop(self, identity: ConnectionIdentity) -> None: |
| 214 | loop = asyncio.get_running_loop() |
| 215 | with self._lock: |
| 216 | projection = self._projections.get(identity) |
| 217 | if projection is None: |
| 218 | return |
| 219 | # INVARIANT.STATE.GATING: do not schedule pushes until a successful state_request |
| 220 | # established seq_base for this sid. |
| 221 | if projection.seq_base <= 0: |
| 222 | return |
| 223 | |
| 224 | # Throttled coalescing: schedule at most one push per debounce window. |
| 225 | # Do not postpone the scheduled push on subsequent dirties; this keeps |
| 226 | # streaming updates smooth while still capping to <= 1 push / 100ms / sid. |
| 227 | existing = self._debounce_handles.get(identity) |
| 228 | if existing is not None and not existing.cancelled(): |
| 229 | return |
| 230 | |
| 231 | running = self._push_tasks.get(identity) |
| 232 | if running is not None and not running.done(): |
| 233 | return |
| 234 | |
| 235 | handle = loop.call_later( |
| 236 | self.debounce_seconds, self._on_debounce_fire, identity |
| 237 | ) |
| 238 | self._debounce_handles[identity] = handle |
| 239 | ws_debug( |
| 240 | f"[StateMonitor] schedule_push namespace={projection.namespace} sid={projection.sid} " |
| 241 | f"delay_s={self.debounce_seconds} " |
| 242 | f"dirty={projection.dirty_version} pushed={projection.pushed_version} " |
| 243 | f"reason={projection.dirty_reason!r} wave={projection.dirty_wave_id!r}" |
| 244 | ) |
| 245 | |
| 246 | def _on_debounce_fire(self, identity: ConnectionIdentity) -> None: |
| 247 | with self._lock: |
| 248 | self._debounce_handles.pop(identity, None) |
| 249 | existing = self._push_tasks.get(identity) |
| 250 | if existing is not None and not existing.done(): |
| 251 | return |
| 252 | task = asyncio.create_task(self._flush_push(identity)) |
| 253 | self._push_tasks[identity] = task |
| 254 | |
| 255 | async def _flush_push(self, identity: ConnectionIdentity) -> None: |
| 256 | namespace, sid = identity |
| 257 | task = asyncio.current_task() |
| 258 | base_version = 0 |
| 259 | include_collections = True |
| 260 | dirty_reason: str | None = None |
| 261 | dirty_wave_id: str | None = None |
| 262 | try: |
| 263 | with self._lock: |
| 264 | projection = self._projections.get(identity) |
| 265 | manager = self._manager |
| 266 | handler_id = self._emit_handler_id |
| 267 | |
| 268 | if projection is None: |
| 269 | return |
| 270 | if manager is None: |
| 271 | # The handler binds the manager on connect; if not bound yet, |
| 272 | # we cannot emit. Keep dirty cleared to avoid infinite retry loops. |
| 273 | return |
| 274 | if projection.seq_base <= 0: |
| 275 | # INVARIANT.STATE.GATING: no push before a successful state_request. |
| 276 | return |
| 277 | |
| 278 | request = projection.request |
| 279 | if request is None: |
| 280 | return |
| 281 | base_version = projection.dirty_version |
| 282 | include_collections = ( |
| 283 | not request.collections_delta |
| 284 | or projection.collections_dirty_version > projection.pushed_version |
| 285 | ) |
| 286 | dirty_reason = projection.dirty_reason |
| 287 | dirty_wave_id = projection.dirty_wave_id |
| 288 | |
| 289 | snapshot = await build_snapshot_from_request( |
| 290 | request=request, |
| 291 | include_collections=include_collections, |
| 292 | ) |
| 293 | |
| 294 | with self._lock: |
| 295 | projection = self._projections.get(identity) |
| 296 | if projection is None: |
| 297 | return |
| 298 | if projection.request != request: |
| 299 | return |
| 300 | |
| 301 | # INVARIANT.STATE.SEQ_MONOTONIC + SEQ_RESET_ON_REQUEST |
| 302 | projection.seq += 1 |
| 303 | seq = projection.seq |
| 304 | |
| 305 | # Advance cursors after successful snapshot emission (incremental mode). |
| 306 | projection.request = advance_state_request_after_snapshot(request, snapshot) |
| 307 | |
| 308 | # Mark all dirties up to `base_version` as pushed. If new dirties |
| 309 | # arrived while building/emitting, a follow-up push will be scheduled. |
| 310 | projection.pushed_version = max(projection.pushed_version, base_version) |
| 311 | |
| 312 | payload = { |
| 313 | "runtime_epoch": runtime.get_runtime_id(), |
| 314 | "seq": seq, |
| 315 | "snapshot": snapshot, |
| 316 | } |
| 317 | |
| 318 | try: |
| 319 | logs_len = ( |
| 320 | len(snapshot.get("logs", [])) |
| 321 | if isinstance(snapshot.get("logs"), list) |
| 322 | else None |
| 323 | ) |
| 324 | ws_debug( |
| 325 | f"[StateMonitor] emit state_push namespace={namespace} sid={sid} seq={seq} " |
| 326 | f"context={request.context!r} logs_len={logs_len} " |
| 327 | f"include_collections={include_collections} " |
| 328 | f"reason={dirty_reason!r} wave={dirty_wave_id!r}" |
| 329 | ) |
| 330 | await manager.emit_to( |
| 331 | namespace, |
| 332 | sid, |
| 333 | STATE_PUSH_EVENT, |
| 334 | payload, |
| 335 | handler_id=handler_id, |
| 336 | ) |
| 337 | except ConnectionNotFoundError: |
| 338 | # Sid was removed before the emit; treat as benign. |
| 339 | ws_debug( |
| 340 | f"[StateMonitor] emit skipped: sid not found namespace={namespace} sid={sid}" |
| 341 | ) |
| 342 | return |
| 343 | except RuntimeError: |
| 344 | # Dispatcher loop may be closing (e.g., during shutdown or test teardown). |
| 345 | ws_debug( |
| 346 | f"[StateMonitor] emit skipped: dispatcher closing namespace={namespace} sid={sid}" |
| 347 | ) |
| 348 | return |
| 349 | finally: |
| 350 | follow_up = False |
| 351 | dirty_version = 0 |
| 352 | pushed_version = 0 |
| 353 | with self._lock: |
| 354 | if task is not None and self._push_tasks.get(identity) is task: |
| 355 | self._push_tasks.pop(identity, None) |
| 356 | projection = self._projections.get(identity) |
| 357 | if projection is not None: |
| 358 | dirty_version = projection.dirty_version |
| 359 | pushed_version = projection.pushed_version |
| 360 | follow_up = dirty_version > pushed_version |
| 361 | |
| 362 | # More dirties accumulated during push; schedule another coalesced push. |
| 363 | # IMPORTANT: this must not run from inside the `finally` block (a `return` in |
| 364 | # `finally` can swallow exceptions from the push task). |
| 365 | if not follow_up: |
| 366 | return |
| 367 | |
| 368 | ws_debug( |
| 369 | f"[StateMonitor] follow_up_push namespace={namespace} sid={sid} dirty={dirty_version} pushed={pushed_version}" |
| 370 | ) |
| 371 | try: |
| 372 | loop = self._dispatcher_loop or asyncio.get_running_loop() |
| 373 | except RuntimeError: |
| 374 | return |
| 375 | if loop.is_closed(): |
| 376 | return |
| 377 | loop.call_soon_threadsafe(self._schedule_debounce_on_loop, identity) |
| 378 | |
| 379 | # Testing hook: keep argument surface stable for future extensions |
| 380 | def _debug_state(self) -> dict[str, Any]: # pragma: no cover - helper |
| 381 | with self._lock: |
| 382 | return { |
| 383 | "identities": list(self._projections.keys()), |
| 384 | "handles": list(self._debounce_handles.keys()), |
| 385 | } |
| 386 | |
| 387 | |
| 388 | # Store singleton in a mutable container to avoid `global` assignment warnings while |
| 389 | # keeping a simple module-level accessor API. |
| 390 | _STATE_MONITOR_HOLDER: dict[str, StateMonitor | None] = {"monitor": None} |
| 391 | _STATE_MONITOR_LOCK = threading.RLock() |
| 392 | |
| 393 | |
| 394 | def get_state_monitor() -> StateMonitor: |
| 395 | with _STATE_MONITOR_LOCK: |
| 396 | monitor = _STATE_MONITOR_HOLDER.get("monitor") |
| 397 | if monitor is None: |
| 398 | monitor = StateMonitor() |
| 399 | _STATE_MONITOR_HOLDER["monitor"] = monitor |
| 400 | return monitor |
| 401 | |
| 402 | |
| 403 | def _reset_state_monitor_for_testing() -> None: # pragma: no cover - helper |
| 404 | with _STATE_MONITOR_LOCK: |
| 405 | _STATE_MONITOR_HOLDER["monitor"] = None |