From 82e50f86d626b65160797bd2207b4e0abe04aa51 Mon Sep 17 00:00:00 2001
From: Alek Ratzloff <alekratz@gmail.com>
Date: Tue, 24 May 2022 19:16:15 -0700
Subject: [PATCH] Add graceful exits

This sets up a set of ropes and pulleys that signal the `Bot.keepalive`
function to clean things up after a quit signal has been sent. This
allows plugins to define an `on_unload` function to save any important
datas on intentional exit.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>
---
 omnibot/__main__.py | 25 +++++++++++++++----------
 omnibot/bot.py      | 13 ++++++++++---
 omnibot/plugin.py   |  3 +++
 3 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/omnibot/__main__.py b/omnibot/__main__.py
index 2bf69d5..554dbfe 100644
--- a/omnibot/__main__.py
+++ b/omnibot/__main__.py
@@ -1,5 +1,7 @@
 import asyncio
+from functools import partial
 import logging
+import signal
 
 from .config import ServerConfig
 from .bot import Bot
@@ -10,15 +12,18 @@ logging.basicConfig(
 )
 log = logging.getLogger(__name__)
 
+log.debug("Loading config")
+config = ServerConfig()
+config.load("config.toml")
+log.debug("Using configuration: %s", config)
 
-async def main():
-    log.debug("Loading config")
-    config = ServerConfig()
-    config.load("config.toml")
-    log.debug("Using configuration: %s", config)
+bot = Bot(config)
 
-    server = Bot(config)
-    await server.run()
-
-
-asyncio.run(main())
+try:
+    asyncio.run(bot.run())
+except KeyboardInterrupt:
+    log.info("Got ctrl-c")
+finally:
+    log.info("Quitting, press ctrl-c to quit immediately")
+    bot.quit()
+    asyncio.run(bot.keepalive())
diff --git a/omnibot/bot.py b/omnibot/bot.py
index f954da4..8f83d88 100644
--- a/omnibot/bot.py
+++ b/omnibot/bot.py
@@ -22,6 +22,7 @@ class Bot:
         ]
         # TODO - this may not be needed
         self.__channels: Set[str] = set()
+        self.__quitting = asyncio.Event()
 
     @property
     def server_config(self) -> ServerConfig:
@@ -31,6 +32,9 @@ class Bot:
     def plugins(self) -> Sequence[plugin.Plugin]:
         return self.__plugins
 
+    def quit(self):
+        self.__quitting.set()
+
     def channel_plugins(self, channel: str) -> Sequence[plugin.Plugin]:
         return [plugin for plugin in self.plugins if channel in plugin.channels]
 
@@ -127,6 +131,9 @@ class Bot:
 
     async def keepalive(self):
         # loop while we're connected, check every second
-        log.info("Starting keepalive loop")
-        while self.connection.connected:
-            await asyncio.sleep(1.0)
+        await self.__quitting.wait()
+        log.info("Shutting down gracefully")
+        # TODO: unload modules
+        await asyncio.gather(
+            *[plugin.on_unload(self.connection) for plugin in self.plugins]
+        )
diff --git a/omnibot/plugin.py b/omnibot/plugin.py
index 9fdd5ee..68784b5 100644
--- a/omnibot/plugin.py
+++ b/omnibot/plugin.py
@@ -51,6 +51,9 @@ class Plugin:
     async def on_message(self, conn: IrcProtocol, channel: str, who: Prefix, line: str):
         pass
 
+    async def on_unload(self, conn: IrcProtocol):
+        pass
+
 
 def load_plugin(server_config: ServerConfig, plugin_config: PluginConfig) -> Plugin:
     name = plugin_config["module"]