From ce50c9adaf2947ee4eaf10c33416543c4976b40e Mon Sep 17 00:00:00 2001
From: Tim <tim@dexorder.trade>
Date: Sat, 13 Apr 2024 01:46:47 -0400
Subject: [PATCH] shutdown fixes

---
 src/dexorder/bin/executable.py | 52 +++++++++++++++++++++-------------
 1 file changed, 33 insertions(+), 19 deletions(-)

diff --git a/src/dexorder/bin/executable.py b/src/dexorder/bin/executable.py
index 14d4763..5d3bd4f 100644
--- a/src/dexorder/bin/executable.py
+++ b/src/dexorder/bin/executable.py
@@ -1,3 +1,4 @@
+import inspect
 import logging
 import logging.config
 import tomllib
@@ -15,25 +16,16 @@ if __name__ == '__main__':
     raise Exception('this file is meant to be imported not executed')
 
 
-ignorable_exceptions = [CancelledError]
-
 log = logging.getLogger(__name__)
 
-async def _shutdown_coro(_sig, _loop, extra_shutdown):
-    log.info('shutting down')
-    if extra_shutdown is not None:
-        extra_shutdown()
-    tasks = [t for t in asyncio.all_tasks() if t is not
-             asyncio.current_task()]
-    for task in tasks:
-        task.cancel()
-    exceptions = await asyncio.gather(*tasks, return_exceptions=True)
-    for x in exceptions:
-        if x is not None and x.__class__ not in ignorable_exceptions:
-            print_exception(x)
-
+async def _shutdown_coro(_sig, _loop):
+    this_task = asyncio.current_task()
+    for task in asyncio.all_tasks():
+        if task is not this_task:
+            task.cancel()
 
 def execute(main:Coroutine, shutdown=None, *, parse_logging=True, parse_args=True):
+    # config
     configured = False
     if parse_logging:
         try:
@@ -51,10 +43,14 @@ def execute(main:Coroutine, shutdown=None, *, parse_logging=True, parse_args=Tru
         log.info('Logging configured to default')
     if parse_args:
         configuration.parse_args()
+
+    # loop setup
     loop = asyncio.get_event_loop()
     signals = Signals.SIGQUIT, Signals.SIGTERM, Signals.SIGINT
     for s in signals:
-        loop.add_signal_handler(s, lambda sig=s: asyncio.create_task(_shutdown_coro(sig, loop, shutdown), name=f'{s.name} handler'))
+        loop.add_signal_handler(s, lambda sig=s: asyncio.create_task(_shutdown_coro(sig, loop), name=f'{s.name} handler'))
+
+    # main
     task = loop.create_task(main, name='main')
     try:
         loop.run_until_complete(task)
@@ -62,13 +58,31 @@ def execute(main:Coroutine, shutdown=None, *, parse_logging=True, parse_args=Tru
         pass
     except Exception as x:
         print_exception(x)
+
+    # shutdown tasks
+    log.info('shutdown')
+    if shutdown is not None:
+        sd = shutdown()
+        if inspect.isawaitable(sd):
+            try:
+                loop.run_until_complete(sd)
+            except Exception as x:
+                log.error(f'Exception during shutdown: {x}')
+                print_exception(x)
+
+    # cancel anything remaining
     try:
         remaining_tasks = asyncio.all_tasks()
     except RuntimeError:
         pass
     else:
-        for t in remaining_tasks:
-            t.cancel()
-        loop.run_until_complete(asyncio.gather(*remaining_tasks))
+        for task in remaining_tasks:
+            task.cancel()
+        results = loop.run_until_complete(asyncio.gather(*remaining_tasks, return_exceptions=True))
+        for x in results:
+            if isinstance(x, Exception):
+                print_exception(x)
+
+    # end
     loop.stop()
     loop.close()