Fixed timer plugin, now checks every hour (with aiocron) for timers in the next hour and sets them

bot/plugins/timer.py

@@ -1,15 +1,17 @@
 # -*- coding: utf-8 -*-
+import re
 import asyncio
-from datetime import datetime, timedelta
+from datetime import datetime
 
 import irc3
+from aiocron import crontab
 from docopt import Dict as DocOptDict
 from irc3.plugins.command import command
 from irc3.utils import IrcString
 from psycopg2 import Error
+from psycopg2.extras import DictRow
 
 from . import DatabasePlugin
-from ..utils import time_delta
 
 
 @irc3.plugin
@@ -19,23 +21,25 @@ class Timer(DatabasePlugin):
 
     def __init__(self, bot: irc3.IrcBot):
         super().__init__(bot)
+        self.timers = set()
+        self.set_timers()
+        crontab('0 * * * *', func=self.set_timers)
 
-        # Fetch timers from database
+    def set_timers(self):
+        """Function which queries all timers in the next hour and schedules them."""
+        self.log.debug('Fetching timers')
         self.cur.execute('''
-        select
-          id, mask, target, message, delay, ends_at
-        from
+        SELECT
+          *
+        FROM
           timers
+        WHERE
+          ends_at >= now()
+          AND ends_at < now() + INTERVAL '1h'
         ''')
 
-        # Recreate timers
-        for res in self.cur.fetchall():
-            self.start_timer(IrcString(res['mask']),
-                             res['target'],
-                             res['message'],
-                             res['delay'],
-                             res['ends_at'] - datetime.now(),
-                             res['id'])
+        for timer in self.cur.fetchall():
+            asyncio.ensure_future(self.exec_timer(timer))
 
     @command
     def timer(self, mask: IrcString, target: IrcString, args: DocOptDict):
@@ -44,63 +48,51 @@ class Timer(DatabasePlugin):
         %%timer <delay> <message>...
         """
         delay = args['<delay>']
-        delta = time_delta(delay)
+        message = ' '.join(args['<message>'])
 
-        if not delta:
-            self.bot.privmsg(target, 'Invalid timer delay')
-        else:
-            message = ' '.join(args['<message>'])
-            values = [mask, target, message, delay]
+        if not re.match(r'\d+[smhdwy]|mon', delay):
+            return 'Invalid timer delay: {}'.format(delay)
 
-            try:
-                # Insert into database (add now + delta)
-                self.cur.execute('''
-                insert into
+        values = [mask, target, message, delay]
+
+        try:
+            self.cur.execute('''
+                INSERT INTO
                   timers (mask, target, message, delay, ends_at)
-                values
-                  (%s, %s, %s, %s, %s)
-                returning id
-                ''', values + [datetime.now() + delta])
-                self.con.commit()
+                VALUES
+                  (%s, %s, %s, %s, now() + INTERVAL %s)
+                RETURNING *
+                ''', values + [delay])
+            self.con.commit()
 
-                # Add delta and id from inserted and start timer
-                values.extend([delta, self.cur.fetchone()['id']])
-                self.start_timer(*values)
+            asyncio.ensure_future(self.exec_timer(self.cur.fetchone()))
 
-                # Send notice to user that timer has been set
-                self.bot.notice(mask.nick, 'Timer in {delay} set: {message}'
-                                .format(delay=delay, message=message))
-            except Error:
-                # Rollback transaction on error
-                self.con.rollback()
+            self.bot.notice(mask.nick, 'Timer in {delay} set: {message}'.format(delay=delay, message=message))
+        except Error as ex:
+            self.log.error(ex)
+            self.con.rollback()
 
-    def start_timer(self, mask: IrcString, target: IrcString, message: str,
-                    delay: str, delta: timedelta, row_id: int):
-        """Async function, sleeps for `delay` seconds and sends notification"""
+    async def exec_timer(self, timer: DictRow):
+        """Sets the actual timer (sleeps until it fires), sends the reminder and deletes the timer from database."""
+        if timer['id'] in self.timers:
+            return
 
-        async def callback():
-            # Sleep if necessary until timed
-            seconds = delta.total_seconds()
-            if seconds > 0:
-                await asyncio.sleep(seconds)
+        self.timers.add(timer['id'])
+        seconds = (timer['ends_at'] - datetime.now()).total_seconds()
+        if seconds > 0.0:
+            await asyncio.sleep(seconds)
 
-            # Send reminder
-            self.bot.privmsg(target, '\x02[Timer]\x0F {nick}: {message} '
-                                     '({delay})'.format(message=message,
-                                                        nick=mask.nick,
-                                                        delay=delay))
+        self.bot.privmsg(timer['target'], '\x02[Timer]\x02 {nick}: {message} ({delay})'.format(
+            message=timer['message'],
+            nick=IrcString(timer['mask']).nick,
+            delay=timer['delay'],
+        ))
 
-            try:
-                # Delete timer from database
-                self.cur.execute('''
-                delete from
-                  timers
-                where
-                  id = %s
-                ''', [row_id])
-                self.con.commit()
-            except Error:
-                # Rollback transaction on error
-                self.con.rollback()
-
-        asyncio.ensure_future(callback())
+        self.timers.remove(timer['id'])
+        self.cur.execute('''
+        DELETE FROM
+          timers
+        WHERE
+          id = %s
+        ''', [timer['id']])
+        self.con.commit()