diff --git a/bot.py b/bot.py index 4249acb..abef0e9 100755 --- a/bot.py +++ b/bot.py @@ -9,7 +9,7 @@ import threading import traceback from datetime import datetime -from twisted.internet import protocol +from twisted.internet import protocol, reactor from twisted.words.protocols import irc import db @@ -183,6 +183,11 @@ class Fulvia(irc.IRCClient): self._user_joined.remove(func) + def add_shutdown(self, func): + """Adds the given function to the reactor's shutdown sequence.""" + reactor.addSystemEventTrigger('before','shutdown', func) + + def stillConnected(self): """Returns true if the bot is still connected to the server.""" if self._heartbeat: diff --git a/scheduler.py b/scheduler.py index 8e360e9..d31f284 100644 --- a/scheduler.py +++ b/scheduler.py @@ -2,11 +2,24 @@ """ Allows the bot to scheduler tasks to be performed later. """ -import time import pickle import threading from datetime import datetime +stop = threading.Event() + +class CronThread(threading.Thread): + def __init__(self, scheduler): + super().__init__() + self.scheduler = scheduler + + def run(self): + while not self.scheduler.bot.stillConnected(): + stop.wait(1) + while not stop.is_set(): + self.scheduler.run_tasks() + stop.wait(5) + class Scheduler: def __init__(self, bot): self.bot = bot @@ -14,25 +27,23 @@ class Scheduler: self.lock = threading.Lock() self.init_database() self.load_database() - self._t = threading.Thread(target=self.loop) - self._t.start() + self.bot.add_shutdown(stop.set) + self.loop = CronThread(self) + self.loop.start() + + def run_tasks(self): + self.lock.acquire() + tasks_due = [t for t in self.tasks if t[1] <= datetime.now()] + for task in tasks_due: + args = (self.bot,) + task[2] + t = threading.Thread(target=task[0], args=args) + t.start() + self.tasks.remove(task) + self.bot.db.execute( + "DELETE FROM scheduled_task WHERE dt = ?", + (pickle.dumps(task[1]),)) + self.lock.release() - def loop(self): - while not self.bot.stillConnected(): - time.sleep(1) - while True: - self.lock.acquire() - tasks_due = [t for t in self.tasks if t[1] <= datetime.now()] - for task in tasks_due: - args = (self.bot,) + task[2] - t = threading.Thread(target=task[0], args=args) - t.start() - self.tasks.remove(task) - self.bot.db.execute( - "DELETE FROM scheduled_task WHERE dt = ?", - (pickle.dumps(task[1]),)) - self.lock.release() - time.sleep(5) def add_task(self, func, dt, args): """