Merge pull request #3818 from jamshale/sonar-cloud-reports

Fix: Some asyncio task management and modernization

Author: swcurran

acapy_agent/admin/routes.py

@@ -1,6 +1,5 @@
 """Admin server routes."""
 
-import asyncio
 import re
 
 from aiohttp import web
@@ -219,7 +218,5 @@ async def shutdown_handler(request: web.BaseRequest):
 
     """
     request.app._state["ready"] = False
-    loop = asyncio.get_event_loop()
-    asyncio.ensure_future(request.app["conductor_stop"](), loop=loop)
-
+    await request.app["conductor_stop"]()
     return web.json_response({})

acapy_agent/admin/server.py

@@ -216,12 +216,23 @@ async def upgrade_middleware(request: web.BaseRequest, handler: Coroutine):
             # We need to check for completion (or fail) in another process
             in_progress_upgrades.set_wallet(context.profile.name)
             is_subwallet = context.metadata and "wallet_id" in context.metadata
-            asyncio.create_task(
+
+            # Create background task and store reference to prevent garbage collection
+            task = asyncio.create_task(
                 check_upgrade_completion_loop(
                     context.profile,
                     is_subwallet,
                 )
             )
+
+            # Store task reference on the app to prevent garbage collection
+            if not hasattr(request.app, "_background_tasks"):
+                request.app._background_tasks = set()
+            request.app._background_tasks.add(task)
+
+            # Remove task from set when it completes to prevent memory leaks
+            task.add_done_callback(request.app._background_tasks.discard)
+
             raise web.HTTPServiceUnavailable(reason="Upgrade in progress")
 
     return await handler(request)

acapy_agent/commands/start.py

@@ -1,14 +1,15 @@
 """Entrypoint."""
 
 import asyncio
-import functools
 import logging
 import signal
 import sys
-from typing import Coroutine, Sequence
+from typing import Sequence
 
 from configargparse import ArgumentParser
 
+from ..config.error import ArgsParseError
+
 try:
     import uvloop
 except ImportError:
@@ -24,89 +25,82 @@
 
 
 async def start_app(conductor: Conductor):
-    """Start up."""
+    """Start up the application."""
     await conductor.setup()
     await conductor.start()
 
 
 async def shutdown_app(conductor: Conductor):
-    """Shut down."""
+    """Shut down the application."""
     LOGGER.info("Shutting down")
     await conductor.stop()
 
+    # Cancel remaining tasks
+    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+    for task in tasks:
+        task.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+
 
 def init_argument_parser(parser: ArgumentParser):
     """Initialize an argument parser with the module's arguments."""
     return arg.load_argument_groups(parser, *arg.group.get_registered(arg.CAT_START))
 
 
-def execute(argv: Sequence[str] = None):
-    """Entrypoint."""
+async def run_app(argv: Sequence[str] = None):
+    """Main async runner for the app."""
     parser = arg.create_argument_parser(prog=PROG)
     parser.prog += " start"
     get_settings = init_argument_parser(parser)
     args = parser.parse_args(argv)
     settings = get_settings(args)
     common_config(settings)
 
-    # set ledger to read only if explicitly specified
+    # Set ledger to read-only if explicitly specified
     settings["ledger.read_only"] = settings.get("read_only_ledger", False)
 
-    # Create the Conductor instance
-    context_builder = DefaultContextBuilder(settings)
-    conductor = Conductor(context_builder)
-
-    # Run the application
     if uvloop:
         uvloop.install()
         LOGGER.info("uvloop installed")
-    run_loop(start_app(conductor), shutdown_app(conductor))
-
-
-def run_loop(startup: Coroutine, shutdown: Coroutine):
-    """Execute the application, handling signals and ctrl-c."""
-
-    async def init(cleanup):
-        """Perform startup, terminating if an exception occurs."""
-        try:
-            await startup
-        except Exception:
-            LOGGER.exception("Exception during startup:")
-            cleanup()
-
-    async def done():
-        """Run shutdown and clean up any outstanding tasks."""
-        await shutdown
-
-        if sys.version_info.major == 3 and sys.version_info.minor > 6:
-            all_tasks = asyncio.all_tasks()
-            current_task = asyncio.current_task()
-        else:
-            all_tasks = asyncio.Task.all_tasks()
-            current_task = asyncio.Task.current_task()
-
-        tasks = [task for task in all_tasks if task is not current_task]
-        for task in tasks:
-            task.cancel()
-        if tasks:
-            await asyncio.gather(*tasks, return_exceptions=True)
-        asyncio.get_event_loop().stop()
-
-    loop = asyncio.get_event_loop()
-    cleanup = functools.partial(asyncio.ensure_future, done(), loop=loop)
-    loop.add_signal_handler(signal.SIGTERM, cleanup)
-    asyncio.ensure_future(init(cleanup), loop=loop)
 
+    context_builder = DefaultContextBuilder(settings)
+    conductor = Conductor(context_builder)
+
+    loop = asyncio.get_running_loop()
+    shutdown_event = asyncio.Event()
+
+    def handle_signal():
+        LOGGER.info("Received stop signal")
+        shutdown_event.set()
+
+    loop.add_signal_handler(signal.SIGTERM, handle_signal)
+    loop.add_signal_handler(signal.SIGINT, handle_signal)
+
+    try:
+        await start_app(conductor)
+        await shutdown_event.wait()
+    finally:
+        await shutdown_app(conductor)
+
+
+def execute(argv: Sequence[str] = None):
+    """Entrypoint."""
     try:
-        loop.run_forever()
+        asyncio.run(run_app(argv))
+    except ArgsParseError as e:
+        LOGGER.error("Argument parsing error: %s", e)
+        raise e
     except KeyboardInterrupt:
-        loop.run_until_complete(done())
+        LOGGER.info("Interrupted by user")
+    except Exception:
+        LOGGER.exception("Unexpected exception during execution")
+        sys.exit(1)
 
 
 def main():
     """Execute the main line."""
-    if __name__ == "__main__":
-        execute()
+    execute()
 
 
-main()
+if __name__ == "__main__":
+    main()