Add locks to avoid race conditions when adding the same to grab to the DB

Author: linuxdaemon

plugins/grab.py

@@ -1,7 +1,11 @@
+import logging
 import random
-
 from collections import defaultdict
+from threading import RLock
+
 from sqlalchemy import Table, Column, String
+from sqlalchemy.exc import SQLAlchemyError
+
 from cloudbot import hook
 from cloudbot.util import database
 
@@ -18,19 +22,25 @@
 )
 
 grab_cache = {}
+grab_locks = defaultdict(dict)
+grab_locks_lock = RLock()
+cache_lock = RLock()
+
+logger = logging.getLogger("cloudbot")
 
 
 @hook.on_start()
 def load_cache(db):
     """
     :type db: sqlalchemy.orm.Session
     """
-    grab_cache.clear()
-    for row in db.execute(table.select().order_by(table.c.time)):
-        name = row["name"].lower()
-        quote = row["quote"]
-        chan = row["chan"]
-        grab_cache.setdefault(chan, {}).setdefault(name, []).append(quote)
+    with cache_lock:
+        grab_cache.clear()
+        for row in db.execute(table.select().order_by(table.c.time)):
+            name = row["name"].lower()
+            quote = row["quote"]
+            chan = row["chan"]
+            grab_cache.setdefault(chan, {}).setdefault(name, []).append(quote)
 
 
 def two_lines(bigstring, chan):
@@ -90,27 +100,43 @@ def grab_add(nick, time, msg, chan, db, conn):
     load_cache(db)
 
 
+def get_latest_line(conn, chan, nick):
+    for item in reversed(conn.history[chan]):
+        name, timestamp, msg = item
+        if nick.casefold() == name.casefold():
+            return item
+
+    return None, None, None
+
+
 @hook.command()
 def grab(text, nick, chan, db, conn):
     """grab <nick> grabs the last message from the
     specified nick and adds it to the quote database"""
     if text.lower() == nick.lower():
         return "Didn't your mother teach you not to grab yourself?"
 
-    for item in conn.history[chan].__reversed__():
-        name, timestamp, msg = item
-        if text.lower() == name.lower():
-            # check to see if the quote has been added
-            if check_grabs(name.lower(), msg, chan):
-                return "I already have that quote from {} in the database".format(text)
-            else:
-                # the quote is new so add it to the db.
-                grab_add(name.lower(),timestamp, msg, chan, db, conn)
-                if check_grabs(name.lower(), msg, chan):
-                    return "the operation succeeded."
-                else:
-                    return "the operation failed"
-    return "I couldn't find anything from {} in recent history.".format(text)
+    with grab_locks_lock:
+        grab_lock = grab_locks[conn.name.casefold()].setdefault(chan.casefold(), RLock())
+
+    with grab_lock:
+        name, timestamp, msg = get_latest_line(conn, chan, text)
+        if not msg:
+            return "I couldn't find anything from {} in recent history.".format(text)
+
+        if check_grabs(text.casefold(), msg, chan):
+            return "I already have that quote from {} in the database".format(text)
+
+        try:
+            grab_add(name.casefold(), timestamp, msg, chan, db, conn)
+        except SQLAlchemyError:
+            logger.exception("Error occurred when grabbing %s in %s", name, chan)
+            return "Error occurred."
+
+        if check_grabs(name.casefold(), msg, chan):
+            return "the operation succeeded."
+        else:
+            return "the operation failed"
 
 
 def format_grab(name, quote):