From 0dd33d8339336cc548b317ed1d0e00bdfb197be0 Mon Sep 17 00:00:00 2001
From: november-pain <max.chikita@gmail.com>
Date: Tue, 19 Nov 2024 09:31:53 +0100
Subject: [PATCH] fix graceful shutdown

---
 src/main.py    | 30 ++++++++++++++++++++++++------
 src/monitor.py | 10 ++++++++--
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/main.py b/src/main.py
index ca137e4..c05671d 100644
--- a/src/main.py
+++ b/src/main.py
@@ -2,6 +2,7 @@
 import sys
 import asyncio
 import logging
+import signal
 from aiohttp import web
 from .config import Config
 from .monitor import TokenMonitor
@@ -18,7 +19,25 @@ async def main_async():
     config = None
     monitor = None
     runner = None
-    
+
+    async def shutdown():
+        logger.info("Received shutdown signal")
+        if monitor:
+            try:
+                await monitor.shutdown()
+            except Exception as e:
+                logger.error(f"Error during monitor shutdown: {e}")
+        if runner:
+            try:
+                logger.info("Cleaning up runner...")
+                await runner.cleanup()
+            except Exception as e:
+                logger.error(f"Error during runner cleanup: {e}")
+        logger.info("Finishing event loop...")
+
+    loop = asyncio.get_running_loop()
+    loop.add_signal_handler(signal.SIGTERM, lambda: asyncio.create_task(shutdown()))
+
     try:
         # Load configuration
         config_path = os.getenv('CONFIG_PATH', 'config.yaml')
@@ -42,10 +61,7 @@ async def main_async():
         
     except Exception as e:
         logger.error(f"Fatal error: {e}")
-        if monitor:
-            await monitor.shutdown()
-        if runner:
-            await runner.cleanup()
+        await shutdown()
         sys.exit(1)
 
 def main():
@@ -56,6 +72,8 @@ def main():
     except Exception as e:
         logger.error(f"Fatal error in main: {e}")
         sys.exit(1)
+    finally:
+        logger.info("Program terminated")
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/src/monitor.py b/src/monitor.py
index 770b004..6db0acf 100644
--- a/src/monitor.py
+++ b/src/monitor.py
@@ -198,6 +198,8 @@ async def collect_metrics(self):
                 logger.error(f"Error in collection loop: {e}")
 
             await asyncio.sleep(self.config.scrape_interval)
+        
+        logger.info("Metrics collection stopped")
     
     async def init_web3_connections(self):
         """Initialize Web3 connections and verify them"""
@@ -216,5 +218,9 @@ async def shutdown(self):
         self.running = False
         # Add a small delay to allow pending operations to complete
         await asyncio.sleep(1)
-        await self.app.shutdown()
-        await self.app.cleanup()
+        if self.app:
+            self.app.freeze()
+            await self.app.shutdown()
+            await self.app.cleanup()
+        else:
+            logger.warning("App is not initialized, skipping shutdown steps.")