diff --git a/app/api/sse_manager.py b/app/api/sse_manager.py index d31d4a8..cf120c6 100644 --- a/app/api/sse_manager.py +++ b/app/api/sse_manager.py @@ -8,20 +8,24 @@ from app.external.sse_starlette import EventSourceResponse, ServerSentEvent class SSEManager: - _setup_finished = False _num_connections = 0 _connections = {} _sse_queue = asyncio.Queue() + _broadcast_task = None - def setup(self) -> None: - if self._setup_finished: - raise RuntimeError("SSEManager setup must not be called twice") - - loop = asyncio.get_event_loop() - loop.create_task(self._broadcast_data_sse()) - self._setup_finished = True + def _ensure_broadcast_task(self) -> None: + # Start the broadcast consumer lazily the first time it's needed from + # within a running loop. Doing this at import time relied on the + # deprecated asyncio.get_event_loop() and broke when the module was + # imported without a running loop (e.g. in a Celery worker). + if self._broadcast_task is not None: + return + self._broadcast_task = asyncio.get_running_loop().create_task( + self._broadcast_data_sse() + ) def add_connection(self, request: Request) -> Tuple[EventSourceResponse, int]: + self._ensure_broadcast_task() q = asyncio.Queue() id = self._num_connections self._num_connections += 1 @@ -33,6 +37,7 @@ class SSEManager: await self._connections[id].put(data) async def broadcast_to_all(self, data: ServerSentEvent): + self._ensure_broadcast_task() await self._sse_queue.put(data) async def _subscribe(self, request: Request, id: int, q: asyncio.Queue): diff --git a/app/api/utils.py b/app/api/utils.py index 45bb568..ade777f 100644 --- a/app/api/utils.py +++ b/app/api/utils.py @@ -21,7 +21,6 @@ from app.external.result_type.src.result import Err, Ok, Result from app.external.sse_starlette import ServerSentEvent sse_mgr = SSEManager() -sse_mgr.setup() def build_sse_event(event: str, json_data: Optional[Dict]): diff --git a/app/apps/router.py b/app/apps/router.py index 18d5c54..9fa7e7d 100644 --- a/app/apps/router.py +++ b/app/apps/router.py @@ -18,8 +18,7 @@ router = APIRouter(prefix=f"/{_PREFIX}", tags=["Apps"]) async def register_app_status_update_handlers(): # This handler watches for messages from the update app cache celery task # it is also responsible for notifying clients of the change - loop = asyncio.get_event_loop() - loop.create_task(watch_app_status_changes()) + asyncio.create_task(watch_app_status_changes()) @router.get( diff --git a/app/apps/service.py b/app/apps/service.py index 24e23e5..5b1e327 100644 --- a/app/apps/service.py +++ b/app/apps/service.py @@ -91,8 +91,7 @@ async def install_app(app_id: AppId): manage_app_task.delay(app_id, InstallMode.ON) # type: ignore - loop = asyncio.get_event_loop() - loop.create_task(_watcher(app_id, "install")) + asyncio.create_task(_watcher(app_id, "install")) async def uninstall_app(data: AppUninstallInput): @@ -111,8 +110,7 @@ async def uninstall_app(data: AppUninstallInput): manage_app_task.delay(app_id, InstallMode.OFF, data.keep_data) # type: ignore - loop = asyncio.get_event_loop() - loop.create_task(_watcher(app_id, "uninstall")) + asyncio.create_task(_watcher(app_id, "uninstall")) async def _watcher(app_id: AppId, action: str): diff --git a/app/auth/auth_handler.py b/app/auth/auth_handler.py index d942a76..fa942a7 100644 --- a/app/auth/auth_handler.py +++ b/app/auth/auth_handler.py @@ -81,5 +81,4 @@ def register_cookie_updater(): await asyncio.sleep(refresh_interval) handle_local_cookie() - loop = asyncio.get_event_loop() - loop.create_task(_cookie_updater()) + asyncio.create_task(_cookie_updater()) diff --git a/app/bitcoind/service.py b/app/bitcoind/service.py index 8e26df0..3c0af5c 100644 --- a/app/bitcoind/service.py +++ b/app/bitcoind/service.py @@ -192,8 +192,7 @@ async def handle_block_sub_redis(verbosity: int = 1) -> str: @logger.catch(exclude=(HTTPException,)) async def register_bitcoin_zmq_sub(): - loop = asyncio.get_event_loop() - loop.create_task(handle_block_sub_redis()) + asyncio.create_task(handle_block_sub_redis()) @logger.catch(exclude=(HTTPException,)) @@ -221,5 +220,4 @@ async def _handle_gather_bitcoin_status(): @logger.catch(exclude=(HTTPException,)) async def register_bitcoin_status_gatherer(): - loop = asyncio.get_event_loop() - loop.create_task(_handle_gather_bitcoin_status()) + asyncio.create_task(_handle_gather_bitcoin_status()) diff --git a/app/lightning/impl/lnd_grpc.py b/app/lightning/impl/lnd_grpc.py index 8256188..f47e0b8 100644 --- a/app/lightning/impl/lnd_grpc.py +++ b/app/lightning/impl/lnd_grpc.py @@ -250,8 +250,7 @@ This will show more debug information. logger.info("Trying to connect to LND daemon ...") - loop = asyncio.get_event_loop() - task = loop.create_task(self._check_lnd_status(sleep_time=2)) + task = asyncio.create_task(self._check_lnd_status(sleep_time=2)) while not self._initialized: res = await self._init_queue.get() # type: InitLnRepoUpdate @@ -267,7 +266,7 @@ This will show more debug information. # created the channel for us. self._create_stubs() - task = loop.create_task(self._check_lnd_status(sleep_time=0.5)) + task = asyncio.create_task(self._check_lnd_status(sleep_time=0.5)) elif res.state == LnInitState.DONE: self._initialized = True if not task.cancelled(): diff --git a/app/lightning/service.py b/app/lightning/service.py index 263eb7f..6edde58 100644 --- a/app/lightning/service.py +++ b/app/lightning/service.py @@ -202,10 +202,9 @@ async def register_lightning_listener(): await ln.get_ln_info() - loop = asyncio.get_event_loop() - loop.create_task(_handle_info_listener()) - loop.create_task(_handle_invoice_listener()) - loop.create_task(_handle_forward_event_listener()) + asyncio.create_task(_handle_info_listener()) + asyncio.create_task(_handle_invoice_listener()) + asyncio.create_task(_handle_forward_event_listener()) except NotImplementedError as r: raise HTTPException(status.HTTP_501_NOT_IMPLEMENTED, detail=r.args[0]) @@ -258,8 +257,7 @@ async def _handle_forward_event_listener(): _fwd_successes.append(i.model_dump()) if not _fwd_update_scheduled: - loop = asyncio.get_event_loop() - loop.create_task(_schedule_fwd_update()) + asyncio.create_task(_schedule_fwd_update()) _wallet_balance_update_scheduled = False @@ -279,5 +277,4 @@ def _schedule_wallet_balance_update(): global _wallet_balance_update_scheduled if not _wallet_balance_update_scheduled: - loop = asyncio.get_event_loop() - loop.create_task(_perform_update()) + asyncio.create_task(_perform_update()) diff --git a/app/main.py b/app/main.py index 174ac48..71d694d 100644 --- a/app/main.py +++ b/app/main.py @@ -98,9 +98,8 @@ async def lifespan(app: FastAPI): register_cookie_updater() await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) - loop = asyncio.get_event_loop() - btc_task = loop.create_task(_initialize_bitcoin()) - ln_task = loop.create_task(_initialize_lightning()) + btc_task = asyncio.create_task(_initialize_bitcoin()) + ln_task = asyncio.create_task(_initialize_lightning()) await register_all_handlers() handle_local_cookie() @@ -191,8 +190,7 @@ async def _set_startup_status( if lightning_msg is not None: api_startup_status.lightning_msg = lightning_msg - loop = asyncio.get_event_loop() - loop.create_task(warmup_new_connections()) + asyncio.create_task(warmup_new_connections()) await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) @@ -316,8 +314,7 @@ async def stream(request: Request): id, SSE.SYSTEM_STARTUP_INFO, jsonable_encoder(api_startup_status.model_dump()) ) - loop = asyncio.get_event_loop() - loop.create_task(warmup_new_connections()) + asyncio.create_task(warmup_new_connections()) return event_source diff --git a/app/system/service.py b/app/system/service.py index d3315a9..1e2df24 100644 --- a/app/system/service.py +++ b/app/system/service.py @@ -143,8 +143,7 @@ async def _handle_gather_hardware_info(): async def register_hardware_info_gatherer(): - loop = asyncio.get_event_loop() - loop.create_task(_handle_gather_hardware_info()) + asyncio.create_task(_handle_gather_hardware_info()) async def login(i: LoginInput) -> Dict[str, str]: diff --git a/tests/test_sse_manager.py b/tests/test_sse_manager.py new file mode 100644 index 0000000..2eb723d --- /dev/null +++ b/tests/test_sse_manager.py @@ -0,0 +1,25 @@ +""" +Regression test: SSEManager must not touch the event loop at import time. + +setup() ran at module import via asyncio.get_event_loop(), which is +deprecated without a running loop and breaks when the module is imported +outside a running loop (e.g. in a Celery worker). The broadcast consumer +must instead start lazily, the first time it's needed from within a loop. +""" + + +async def test_broadcast_task_starts_lazily_within_a_loop(): + from app.api.sse_manager import SSEManager + + mgr = SSEManager() + assert mgr._broadcast_task is None, "must not start a task before it's needed" + + mgr._ensure_broadcast_task() + assert mgr._broadcast_task is not None, "task should start within a running loop" + + # a second call must not spawn another task + first = mgr._broadcast_task + mgr._ensure_broadcast_task() + assert mgr._broadcast_task is first + + mgr._broadcast_task.cancel()