Use FastAPI Lifespan Events

gns3server/core/tasks.py

@@ -17,8 +17,8 @@
 
 import asyncio
 
-from typing import Callable
 from fastapi import FastAPI
+from contextlib import asynccontextmanager
 
 from gns3server.controller import Controller
 from gns3server.config import Config
@@ -35,74 +35,75 @@ log = logging.getLogger(__name__)
 auto_discover_images_task_handle = None
 
 
-def create_startup_handler(app: FastAPI) -> Callable:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+
+    await startup(app)
+    yield
+    await shutdown(app)
+
+
+async def startup(app: FastAPI) -> None:
     """
     Tasks to be performed when the server is starting.
     """
 
-    async def start_app() -> None:
-        loop = asyncio.get_event_loop()
-        logger = logging.getLogger("asyncio")
-        logger.setLevel(logging.ERROR)
+    loop = asyncio.get_event_loop()
+    logger = logging.getLogger("asyncio")
+    logger.setLevel(logging.ERROR)
 
-        if log.getEffectiveLevel() == logging.DEBUG:
-            # On debug version we enable info that
-            # coroutine is not called in a way await/await
-            loop.set_debug(True)
+    if log.getEffectiveLevel() == logging.DEBUG:
+        # On debug version we enable info that
+        # coroutine is not called in a way await/await
+        loop.set_debug(True)
 
-        # connect to the database
-        await connect_to_db(app)
+    # connect to the database
+    await connect_to_db(app)
 
-        # retrieve the computes from the database
-        computes = await get_computes(app)
+    # retrieve the computes from the database
+    computes = await get_computes(app)
 
-        await Controller.instance().start(computes)
+    await Controller.instance().start(computes)
 
-        # Because with a large image collection
-        # without md5sum already computed we start the
-        # computing with server start
-        from gns3server.compute.qemu import Qemu
+    # Because with a large image collection
+    # without md5sum already computed we start the
+    # computing with server start
+    from gns3server.compute.qemu import Qemu
 
-        if Config.instance().settings.Server.auto_discover_images is True:
-            # Start the discovering new images on file system 5 seconds after the server has started
-            # to give it a chance to process API requests
-            global auto_discover_images_task_handle
-            auto_discover_images_task_handle = asyncio.get_event_loop().call_later(
-                5,
-                lambda: asyncio.create_task(discover_images_on_filesystem(app))
-            )
+    if Config.instance().settings.Server.auto_discover_images is True:
+        # Start the discovering new images on file system 5 seconds after the server has started
+        # to give it a chance to process API requests
+        global auto_discover_images_task_handle
+        auto_discover_images_task_handle = asyncio.get_event_loop().call_later(
+            5,
+            lambda: asyncio.create_task(discover_images_on_filesystem(app))
+        )
 
-        for module in MODULES:
-            log.debug(f"Loading module {module.__name__}")
-            m = module.instance()
-            m.port_manager = PortManager.instance()
-
-    return start_app
+    for module in MODULES:
+        log.debug(f"Loading module {module.__name__}")
+        m = module.instance()
+        m.port_manager = PortManager.instance()
 
 
-def create_shutdown_handler(app: FastAPI) -> Callable:
+async def shutdown(app: FastAPI) -> None:
     """
     Tasks to be performed when the server is exiting.
     """
 
-    async def shutdown_handler() -> None:
+    if auto_discover_images_task_handle is not None and not auto_discover_images_task_handle.cancelled():
+        auto_discover_images_task_handle.cancel()
+    await HTTPClient.close_session()
+    await Controller.instance().stop()
 
-        if auto_discover_images_task_handle is not None and not auto_discover_images_task_handle.cancelled():
-            auto_discover_images_task_handle.cancel()
-        await HTTPClient.close_session()
-        await Controller.instance().stop()
+    for module in MODULES:
+        log.debug(f"Unloading module {module.__name__}")
+        m = module.instance()
+        await m.unload()
 
-        for module in MODULES:
-            log.debug(f"Unloading module {module.__name__}")
-            m = module.instance()
-            await m.unload()
+    if PortManager.instance().tcp_ports:
+        log.warning(f"TCP ports are still used {PortManager.instance().tcp_ports}")
 
-        if PortManager.instance().tcp_ports:
-            log.warning(f"TCP ports are still used {PortManager.instance().tcp_ports}")
+    if PortManager.instance().udp_ports:
+        log.warning(f"UDP ports are still used {PortManager.instance().udp_ports}")
 
-        if PortManager.instance().udp_ports:
-            log.warning(f"UDP ports are still used {PortManager.instance().udp_ports}")
-
-        await disconnect_from_db(app)
-
-    return shutdown_handler
+    await disconnect_from_db(app)