Repository navigation
fix(backend): protect device websocket reconnect storms - #1384
Conversation
📝 WalkthroughWalkthroughThis PR enhances device connection resilience by adding rate limiting, registration-upsert debouncing with in-memory caching, stale-disconnect guards, and non-blocking background capability synchronization to the Socket.IO device namespace. Configuration constants, state tracking, and helper methods enable IP-keyed rate limiting, recent-registration caching to skip repeated database upserts, socket-id matching during disconnection, and scheduled background sync with configurable timeout. Four new tests validate each behavior in isolation. ChangesDevice Connection Resilience
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
backend/app/api/ws/device_namespace.py (1)
989-1000:⚠️ Potential issue | 🟠 Major | ⚡ Quick winThis “background” sync still runs blocking SQLAlchemy work on the event loop.
Moving the sync into
create_task()makeson_device_register()return sooner, but thewith _db_session()+build_desired_capabilities()block is still synchronous and now executes on the event-loop thread. Under a reconnect burst, that can still stall unrelated sockets while the payload is built. Please move the DB read behindrun_sync_in_executor()(or an async session) before awaitingsync_device_payload().Suggested shape
+def _build_desired_capabilities_sync(user_id: int) -> dict[str, Any]: + with _db_session() as db: + return device_capability_sync_service.build_desired_capabilities( + db, + user_id=user_id, + ) + async def _sync_global_capabilities_to_registered_device( self, *, user_id: int, device_id: str, ) -> None: """Best-effort desired-state sync when a device comes online.""" try: - with _db_session() as db: - payload = device_capability_sync_service.build_desired_capabilities( - db, - user_id=user_id, - ) + payload = await run_sync_in_executor( + _build_desired_capabilities_sync, + user_id, + ) result = await device_capability_sync_service.sync_device_payload( user_id=user_id, device_id=device_id, payload=payload, timeout_seconds=REGISTER_CAPABILITY_SYNC_TIMEOUT_SECONDS,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/app/api/ws/device_namespace.py` around lines 989 - 1000, The code currently runs synchronous DB work (with _db_session() and device_capability_sync_service.build_desired_capabilities) on the event loop; change it to run the DB read in a thread executor before awaiting sync_device_payload by moving the _db_session + build_desired_capabilities call into asyncio.get_running_loop().run_in_executor (or an equivalent helper) so the blocking call executes off the event loop and returns the payload; then pass that payload to device_capability_sync_service.sync_device_payload (with REGISTER_CAPABILITY_SYNC_TIMEOUT_SECONDS) as before. Ensure the executor-wrapped call acquires the session and calls build_desired_capabilities (not just build_desired_capabilities without session) so no blocking DB I/O runs on the event loop.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/app/api/ws/device_namespace.py`:
- Around line 968-974: The code currently schedules
_sync_global_capabilities_to_registered_device for every device:register (via
_schedule_background_task), causing repeated syncs on debounced reconnects;
update the registration handling to avoid scheduling when a recent sync is
in-flight or recently completed for the same (user_id, device_id): add a
short-lived in-memory lock/cache keyed by (user_id, device_id) (e.g.,
_capability_sync_inflight or _recent_capability_sync) that is set before
scheduling in _schedule_background_task and cleared after the task finishes (or
expires after N seconds), and only call _schedule_background_task for
_sync_global_capabilities_to_registered_device if the lock/cache is absent;
alternatively, only schedule when the registration actually changed (compare
previous registration state before enqueuing) inside the device:register
handling code that currently invokes _schedule_background_task.
In `@backend/tests/api/ws/test_device_reconnect_storm.py`:
- Around line 157-168: The monkeypatch targets the instance `namespace` but
`on_disconnect()` resolves `run_sync_in_executor` from the module-level
`device_namespace`, and `DeviceNamespace` has no instance attribute
`run_sync_in_executor`; update the test to patch
`device_namespace.run_sync_in_executor` (the module-level helper) instead of
`namespace`, i.e. replace the `monkeypatch.setattr(namespace,
"run_sync_in_executor", AsyncMock(...))` with
`monkeypatch.setattr(device_namespace, "run_sync_in_executor", AsyncMock(...))`;
keep the other patches and final assertions for `set_offline`,
`device_namespace.run_sync_in_executor`, and
`namespace._broadcast_device_offline` as-is so they assert the correct objects
called by `namespace.on_disconnect`.
---
Outside diff comments:
In `@backend/app/api/ws/device_namespace.py`:
- Around line 989-1000: The code currently runs synchronous DB work (with
_db_session() and device_capability_sync_service.build_desired_capabilities) on
the event loop; change it to run the DB read in a thread executor before
awaiting sync_device_payload by moving the _db_session +
build_desired_capabilities call into asyncio.get_running_loop().run_in_executor
(or an equivalent helper) so the blocking call executes off the event loop and
returns the payload; then pass that payload to
device_capability_sync_service.sync_device_payload (with
REGISTER_CAPABILITY_SYNC_TIMEOUT_SECONDS) as before. Ensure the executor-wrapped
call acquires the session and calls build_desired_capabilities (not just
build_desired_capabilities without session) so no blocking DB I/O runs on the
event loop.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9488e5a8-8787-42ab-a96a-a9d5637bac76
📒 Files selected for processing (3)
backend/app/api/ws/device_namespace.pybackend/app/services/device/capability_sync_service.pybackend/tests/api/ws/test_device_reconnect_storm.py
| persisted_display_name = self._get_recent_registration_display_name( | ||
| user_id, payload.device_id | ||
| ) | ||
| if not success: | ||
| return {"error": f"Registration failed: {error}"} | ||
| if persisted_display_name is None: | ||
| success, persisted_display_name, error = await run_sync_in_executor( | ||
| _register_device, | ||
| user_id, | ||
| payload.device_id, | ||
| payload.name, | ||
| payload.client_ip, | ||
| payload.device_type.value, | ||
| payload.bind_shell.value, | ||
| ) | ||
| if not success: | ||
| return {"error": f"Registration failed: {error}"} | ||
| self._remember_registration( | ||
| user_id, payload.device_id, persisted_display_name or payload.name | ||
| ) |
There was a problem hiding this comment.
Only debounce identical registration payloads.
This fast path keys the cache by (user_id, device_id) and only reuses a cached display name, so a second device:register within the debounce window skips _register_device() even when name, client_ip, device_type, or bind_shell changed. That leaves the CRD stale and can also overwrite the fresh name with the previous cached display name. Please include a fingerprint of the persisted fields in the debounce lookup and add a regression test for changed metadata.
| self._schedule_background_task( | ||
| self._sync_global_capabilities_to_registered_device( | ||
| user_id=user_id, | ||
| device_id=payload.device_id, | ||
| ), | ||
| "sync global capabilities after device registration", | ||
| ) |
There was a problem hiding this comment.
Debounced reconnects still fan out capability-sync storms.
Every successful device:register schedules _sync_global_capabilities_to_registered_device(), including the cached fast path above. A reconnect burst for one device now skips the DB upsert but still enqueues one capability-sync task per reconnect, which turns the storm into repeated DB reads and Socket.IO RPCs. Gate this with an in-flight/recent-sync key per (user_id, device_id) or only schedule after a real registration change.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/app/api/ws/device_namespace.py` around lines 968 - 974, The code
currently schedules _sync_global_capabilities_to_registered_device for every
device:register (via _schedule_background_task), causing repeated syncs on
debounced reconnects; update the registration handling to avoid scheduling when
a recent sync is in-flight or recently completed for the same (user_id,
device_id): add a short-lived in-memory lock/cache keyed by (user_id, device_id)
(e.g., _capability_sync_inflight or _recent_capability_sync) that is set before
scheduling in _schedule_background_task and cleared after the task finishes (or
expires after N seconds), and only call _schedule_background_task for
_sync_global_capabilities_to_registered_device if the lock/cache is absent;
alternatively, only schedule when the registration actually changed (compare
previous registration state before enqueuing) inside the device:register
handling code that currently invokes _schedule_background_task.
| monkeypatch.setattr( | ||
| device_namespace, | ||
| "run_sync_in_executor", | ||
| AsyncMock(return_value=[]), | ||
| ) | ||
| monkeypatch.setattr(namespace, "_broadcast_device_offline", AsyncMock()) | ||
|
|
||
| await namespace.on_disconnect("sid-old") | ||
|
|
||
| set_offline.assert_not_awaited() | ||
| device_namespace.run_sync_in_executor.assert_not_awaited() | ||
| namespace._broadcast_device_offline.assert_not_awaited() |
There was a problem hiding this comment.
Patch the module-level executor helper here, not the namespace instance.
DeviceNamespace does not expose a run_sync_in_executor attribute, so this monkeypatch.setattr(namespace, "run_sync_in_executor", ...) call fails before the test reaches on_disconnect(). Even if it did not, on_disconnect() resolves run_sync_in_executor from app.api.ws.device_namespace, so the final assertion is checking the wrong object.
Suggested fix
- monkeypatch.setattr(
- namespace,
- "run_sync_in_executor",
- AsyncMock(return_value=[]),
- )
+ run_sync_in_executor_mock = AsyncMock(return_value=[])
+ monkeypatch.setattr(
+ device_namespace,
+ "run_sync_in_executor",
+ run_sync_in_executor_mock,
+ )
monkeypatch.setattr(namespace, "_broadcast_device_offline", AsyncMock())
@@
- device_namespace.run_sync_in_executor.assert_not_awaited()
+ run_sync_in_executor_mock.assert_not_awaited()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/tests/api/ws/test_device_reconnect_storm.py` around lines 157 - 168,
The monkeypatch targets the instance `namespace` but `on_disconnect()` resolves
`run_sync_in_executor` from the module-level `device_namespace`, and
`DeviceNamespace` has no instance attribute `run_sync_in_executor`; update the
test to patch `device_namespace.run_sync_in_executor` (the module-level helper)
instead of `namespace`, i.e. replace the `monkeypatch.setattr(namespace,
"run_sync_in_executor", AsyncMock(...))` with
`monkeypatch.setattr(device_namespace, "run_sync_in_executor", AsyncMock(...))`;
keep the other patches and final assertions for `set_offline`,
`device_namespace.run_sync_in_executor`, and
`namespace._broadcast_device_offline` as-is so they assert the correct objects
called by `namespace.on_disconnect`.
2 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests