context manager and subsystem work

appdaemon/context_manager.py

@@ -0,0 +1,153 @@
+import asyncio
+import functools
+import logging
+import signal
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack, contextmanager
+from threading import Event, Lock
+from time import sleep
+from typing import Any, Callable
+
+logger = logging.getLogger()
+
+
+def handler(signum, frame):
+    print('Signal handler called with signal', signum)
+    raise OSError("Couldn't open device!")
+
+
+class AppDaemonRunContext:
+    _stack: ExitStack
+    loop: asyncio.AbstractEventLoop
+    executor: ThreadPoolExecutor
+    stop_event: Event
+    shutdown_lock: Lock
+
+    def __init__(self):
+        self._stack = ExitStack()
+        self.stop_event = Event()
+        self.shutdown_lock = Lock()
+
+    def __enter__(self):
+        self.loop = self._stack.enter_context(self.asyncio_context())
+        logger.debug("Entered asyncio context")
+
+        signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
+        for s in signals:
+            self.loop.add_signal_handler(s, lambda s=s: self.loop.create_task(self.shutdown(s)))
+
+        # self.loop.add_signal_handler(signal.SIGINT, self.handle_signal, signal.SIGINT)
+        # self.loop.add_signal_handler(signal.SIGTERM, self.handle_signal)
+        # self.executor = self._stack.enter_context(ThreadPoolExecutor(max_workers=5))
+        self.executor = self._stack.enter_context(self.thread_context())
+        logger.debug("Entered threadpool context")
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._stack.__exit__(exc_type, exc_value, traceback)
+        logger.debug("Exited context")
+
+    async def shutdown(self, signal):
+        """Cleanup tasks tied to the service's shutdown.
+        
+        https://www.roguelynn.com/words/asyncio-graceful-shutdowns/
+        """
+        logger.info(f"Received exit signal {signal.name}...")
+        if not self.stop_event.is_set():
+            logger.debug('Setting stop event')
+            self.stop_event.set()
+        
+            tasks = [
+                t for t in asyncio.all_tasks()
+                if t is not asyncio.current_task()
+            ]
+
+            for task in tasks:
+                # logger.debug(f"Waiting on task to finish: {task.get_coro().__qualname__}")
+                logger.debug(f"Cancelling {task.get_name()}: {task.get_coro().__qualname__}")
+                task.cancel()
+
+            logger.debug('Waiting for tasks to finish...')
+            try:
+                await asyncio.gather(*tasks, return_exceptions=True)
+            except asyncio.CancelledError as e:
+                logger.debug(f'Cancelled: {e}')
+            logger.debug("Stopping event loop in context shutdown")
+            self.loop.stop()
+        else:
+            logger.warning('Already started shutdown')
+
+    @contextmanager
+    def asyncio_context(self):
+        try:
+            loop = asyncio.get_event_loop()
+            yield loop
+        finally:
+            loop.close()
+            if loop.is_closed():
+                logger.debug("Closed the event loop.")
+
+    @contextmanager
+    def thread_context(self):
+        with ThreadPoolExecutor(max_workers=5) as executor:
+            yield executor
+        if executor._shutdown:
+            logger.debug('Shut down the ThreadPoolExecutor')
+
+    async def run_in_executor(
+        self,
+        func: Callable,
+        *args,
+        timeout: float | None = None,
+        **kwargs
+    ) -> Any:
+        """Run the sync function using the ThreadPoolExecutor and await the result"""
+        timeout = timeout or 10.0
+
+        coro = self.loop.run_in_executor(
+            self.executor,
+            functools.partial(func, **kwargs),
+            *args,
+        )
+
+        try:
+            return await asyncio.wait_for(coro, timeout)
+        except asyncio.TimeoutError:
+            print('Timed out')
+
+
+if __name__ == "__main__":
+    import logging
+    import random
+    from uuid import uuid4
+
+    logging.basicConfig(level="DEBUG", format="{levelname:<8} {message}", style="{")
+
+    def dummy_function(delay: float):
+        id_ = uuid4().hex[:4]
+        logger.info(f'{id_} sleeping for {delay:.1f}s')
+        sleep(delay)
+        logger.info(f'{id_} Done')
+
+    async def async_dummy_function(delay: float):
+        id_ = uuid4().hex[:4]
+        logger.info(f'{id_} async sleeping for {delay:.1f}s')
+        await asyncio.sleep(delay)
+        logger.info(f'{id_} Done async')
+
+    with AppDaemonRunContext() as cm:
+        for _ in range(3):
+            logger.info('Submitting random dummy_functions')
+            cm.executor.submit(dummy_function, random.random() * 10.0)
+            cm.loop.create_task(async_dummy_function(random.random() * 5.0))
+        
+        try:
+            logger.info('Running until complete')
+            cm.loop.run_until_complete(asyncio.gather(*asyncio.all_tasks(cm.loop)))
+        except asyncio.CancelledError:
+            logger.error('Cancelled')
+    
+    if cm.loop.is_closed():
+        logger.debug('Loop is closed')
+    if cm.executor._shutdown:
+        logger.debug('Executor is shut down')