|
1 | 1 | import logging |
| 2 | +from contextlib import asynccontextmanager |
2 | 3 | from ipaddress import ip_address |
3 | 4 | from typing import Any |
4 | 5 |
|
@@ -100,6 +101,19 @@ def render(self, content: Any) -> bytes: |
100 | 101 | return orjson.dumps(content, option=orjson.OPT_SERIALIZE_NUMPY) |
101 | 102 |
|
102 | 103 |
|
| 104 | +@asynccontextmanager |
| 105 | +async def lifespan(app: FastAPI): |
| 106 | + # Startup logic |
| 107 | + if not ApiServer._message_stream: |
| 108 | + # Creates the MessageStream class on startup so it has access to the same event loop |
| 109 | + # as uvicorn |
| 110 | + ApiServer._message_stream = MessageStream() |
| 111 | + yield |
| 112 | + # Shutdown logic |
| 113 | + if ApiServer._message_stream: |
| 114 | + ApiServer._message_stream = None |
| 115 | + |
| 116 | + |
103 | 117 | class ApiServer(RPCHandler): |
104 | 118 | __instance = None |
105 | 119 | __initialized = False |
@@ -137,6 +151,7 @@ def __init__(self, config: Config, standalone: bool = False) -> None: |
137 | 151 | redoc_url=None, |
138 | 152 | default_response_class=FTJSONResponse, |
139 | 153 | openapi_tags=_OPENAPI_TAGS, |
| 154 | + lifespan=lifespan, |
140 | 155 | ) |
141 | 156 | self.configure_app(self.app, self._config) |
142 | 157 | self.start_api() |
@@ -261,24 +276,6 @@ def configure_app(self, app: FastAPI, config): |
261 | 276 | ) |
262 | 277 |
|
263 | 278 | app.add_exception_handler(RPCException, self.handle_rpc_exception) |
264 | | - app.add_event_handler(event_type="startup", func=self._api_startup_event) |
265 | | - app.add_event_handler(event_type="shutdown", func=self._api_shutdown_event) |
266 | | - |
267 | | - async def _api_startup_event(self): |
268 | | - """ |
269 | | - Creates the MessageStream class on startup |
270 | | - so it has access to the same event loop |
271 | | - as uvicorn |
272 | | - """ |
273 | | - if not ApiServer._message_stream: |
274 | | - ApiServer._message_stream = MessageStream() |
275 | | - |
276 | | - async def _api_shutdown_event(self): |
277 | | - """ |
278 | | - Removes the MessageStream class on shutdown |
279 | | - """ |
280 | | - if ApiServer._message_stream: |
281 | | - ApiServer._message_stream = None |
282 | 279 |
|
283 | 280 | def start_api(self): |
284 | 281 | """ |
|
0 commit comments