Serializing crawler queue transactions

Fix #867

Author: Yomguithereal

minet/crawl/queue.py

@@ -182,6 +182,10 @@ class BrokenCrawlerQueue(Exception):
 # because its #.get method can block, but will raise Empty when drained.
 # This is fine with `quenouille` if we don't forget to call #.unblock
 # on panic.
+# NOTE: sqlite needs to serialize writes, which means it is not particularly
+# useful in our case to allow for concurrent access to the queue because most
+# of its operations need to update the database somehow. This also means we
+# can rely on a single lock to make multithreaded transactions safe.
 class CrawlerQueue:
     # Params
     persistent: bool
@@ -192,15 +196,13 @@ class CrawlerQueue:
 
     # State
     tasks: Dict[str, int]
-    put_connection: sqlite3.Connection
-    task_connection: sqlite3.Connection
+    connection: sqlite3.Connection
     counter: int
     cleanup_interval: int
     vacuum_interval: int
     waiter: Condition
     currently_waiting_count: int
-    put_lock: Lock
-    task_lock: Lock
+    transaction_lock: Lock
     broken: bool
 
     def __init__(
@@ -257,29 +259,22 @@ def __init__(
 
         self.waiter = Condition()
         self.currently_waiting_count = 0
-        self.put_lock = Lock()
-        self.task_lock = Lock()
+        self.transaction_lock = Lock()
         self.broken = False
 
-        # NOTE: we need two connection if we are to allow concurrent
-        # put and task acquisition.
-        self.put_connection = sqlite3.connect(full_path, check_same_thread=False)
-        self.task_connection = sqlite3.connect(full_path, check_same_thread=False)
+        self.connection = sqlite3.connect(full_path, check_same_thread=False)
 
         # NOTE: it's seems it is safer and common practice to
         # reexecute pragmas each time because they might not
         # be stored persistently in some instances.
-        self.put_connection.executescript(SQL_PRAGMAS)
-        self.put_connection.commit()
-
-        self.task_connection.executescript(SQL_PRAGMAS)
-        self.task_connection.commit()
+        self.connection.executescript(SQL_PRAGMAS)
+        self.connection.commit()
 
         if inspect:
             return
 
         # Setup
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             if not self.resuming:
                 cursor.executescript(SQL_CREATE)
             else:
@@ -302,33 +297,11 @@ def __init__(
                 cursor.execute("VACUUM;")
 
     @contextmanager
-    def put_transaction(self):
-        cursor = None
-        try:
-            with self.put_lock, self.put_connection:
-                cursor = self.put_connection.cursor()
-                yield cursor
-        finally:
-            if cursor is not None:
-                cursor.close()
-
-    @contextmanager
-    def task_transaction(self):
-        cursor = None
-        try:
-            with self.task_lock, self.task_connection:
-                cursor = self.task_connection.cursor()
-                yield cursor
-        finally:
-            if cursor is not None:
-                cursor.close()
-
-    @contextmanager
-    def global_transaction(self):
+    def transaction(self):
         cursor = None
         try:
-            with self.put_lock, self.task_lock, self.task_connection:
-                cursor = self.task_connection.cursor()
+            with self.transaction_lock, self.connection:
+                cursor = self.connection.cursor()
                 yield cursor
         finally:
             if cursor is not None:
@@ -340,7 +313,7 @@ def explain_query_plan(self, sql: str) -> str:
 
         sql = sql.replace("?", "1")
 
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute("EXPLAIN QUERY PLAN %s" % sql)
 
             return "\n".join(row[3] for row in iterate_over_sqlite_cursor(cursor))
@@ -350,14 +323,14 @@ def __count(self, cursor: sqlite3.Cursor) -> int:
         return cursor.fetchone()[0]
 
     def qsize(self) -> int:
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             return self.__count(cursor)
 
     def __len__(self) -> int:
         return self.qsize()
 
     def put_many(self, jobs: Iterable[CrawlJob]) -> int:
-        with self.put_transaction() as cursor:
+        with self.transaction() as cursor:
             rows = []
 
             for job in jobs:
@@ -422,7 +395,7 @@ def get(self, block=True) -> CrawlJob:
                 if self.broken:
                     raise BrokenCrawlerQueue
 
-            with self.task_transaction() as cursor:
+            with self.transaction() as cursor:
                 cursor.execute(
                     SQL_GET_JOB % ("ASC" if not self.is_lifo else "DESC"),
                     (now(),),
@@ -497,7 +470,7 @@ def unblock(self) -> None:
     def worked_groups(self) -> Dict[str, Tuple[int, int]]:
         g = {}
 
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute(
                 'SELECT "group", "count", "allowed" FROM "parallelism" WHERE "count" > 0;'
             )
@@ -508,7 +481,7 @@ def worked_groups(self) -> Dict[str, Tuple[int, int]]:
         return g
 
     def dump(self) -> Iterator[CrawlerQueueRecord]:
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute(SQL_DUMP)
 
             for row in iterate_over_sqlite_cursor(cursor):
@@ -554,11 +527,11 @@ def __vacuum_and_analyze(self, cursor: sqlite3.Cursor) -> None:
         cursor.connection.commit()
 
     def cleanup(self) -> None:
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             self.__cleanup(cursor)
 
     def clear(self) -> None:
-        with self.global_transaction() as cursor:
+        with self.transaction() as cursor:
             self.__clear(cursor)
 
     # NOTE: there is a subtle difference between #.release_group
@@ -570,7 +543,7 @@ def clear(self) -> None:
     # task. This is important to ensure atomicity as well as possible.
     # Without creating backpressure on the queue's constraints.
     def release_group(self, job: CrawlJob) -> None:
-        with self.task_transaction() as cursor:
+        with self.transaction() as cursor:
             if job.id not in self.tasks:
                 raise RuntimeError("job is not being worked")
 
@@ -601,7 +574,7 @@ def group_releaser(self, job: CrawlJob):
             self.release_group(job)
 
     def task_done(self, job: CrawlJob) -> None:
-        with self.task_transaction() as cursor:
+        with self.transaction() as cursor:
             index = self.tasks.pop(job.id, None)
 
             if index is None:
@@ -620,8 +593,7 @@ def task_done(self, job: CrawlJob) -> None:
                 self.__vacuum_and_analyze(cursor)
 
     def close(self) -> None:
-        self.put_connection.close()
-        self.task_connection.close()
+        self.connection.close()
 
     def __del__(self) -> None:
         self.close()

minet/crawl/url_cache.py

@@ -32,22 +32,22 @@ class AtomicSet(Generic[T]):
     """
 
     __items: Set[T]
-    __lock: Lock
+    lock: Lock
 
     def __init__(self):
         self.__items = set()
-        self.__lock = Lock()
+        self.lock = Lock()
 
     def add(self, item: T) -> bool:
-        with self.__lock:
+        with self.lock:
             len_before = len(self.__items)
 
             self.__items.add(item)
 
             return len(self.__items) > len_before
 
     def add_many(self, items: Iterable[T]) -> int:
-        with self.__lock:
+        with self.lock:
             len_before = len(self.__items)
 
             for item in items:
@@ -66,7 +66,7 @@ def add_many_and_keep_new(self, items: Iterable[T], key: None = ...) -> List[T]:
         ...
 
     def add_many_and_keep_new(self, items, key=None):
-        with self.__lock:
+        with self.lock:
             new = []
 
             for item in items:
@@ -81,15 +81,15 @@ def add_many_and_keep_new(self, items, key=None):
             return new
 
     def __len__(self) -> int:
-        with self.__lock:
+        with self.lock:
             return len(self.__items)
 
     def __contains__(self, item: T) -> bool:
-        with self.__lock:
+        with self.lock:
             return item in self.__items
 
     def __iter__(self) -> Iterator[T]:
-        with self.__lock:
+        with self.lock:
             yield from self.__items
 
     def close(self) -> None:
@@ -112,30 +112,30 @@ def __init__(self, path: str, db_name: str = "set.db"):
 
         self.path = path
 
-        self.__connection = sqlite3.connect(
+        self.connection = sqlite3.connect(
             join(self.path, db_name), check_same_thread=False
         )
-        self.__lock = Lock()
+        self.lock = Lock()
 
         # Setup
         # NOTE: this is reexecuted on resume and this is fine
         # NOTE: it seems it's safer to reexecute pragmas anyway
-        self.__connection.executescript(SQL_CREATE)
-        self.__connection.commit()
+        self.connection.executescript(SQL_CREATE)
+        self.connection.commit()
 
     @contextmanager
-    def __transaction(self):
+    def transaction(self):
         cursor = None
         try:
-            with self.__lock, self.__connection:
-                cursor = self.__connection.cursor()
+            with self.lock, self.connection:
+                cursor = self.connection.cursor()
                 yield cursor
         finally:
             if cursor is not None:
                 cursor.close()
 
     def add(self, item: str) -> bool:
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             try:
                 cursor.execute('INSERT INTO "set" ("key") VALUES (?);', (item,))
             except sqlite3.IntegrityError:
@@ -144,7 +144,7 @@ def add(self, item: str) -> bool:
             return True
 
     def add_many(self, items: Iterable[str]) -> int:
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             rows = [(item,) for item in items]
             cursor.executemany('INSERT OR IGNORE INTO "set" ("key") VALUES (?);', rows)
             return cursor.rowcount
@@ -160,7 +160,7 @@ def add_many_and_keep_new(self, items: Iterable[str], key: None = ...) -> List[s
         ...
 
     def add_many_and_keep_new(self, items, key=None):
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             new = []
 
             for item in items:
@@ -175,31 +175,31 @@ def add_many_and_keep_new(self, items, key=None):
 
     # NOTE: I keep this for the ugly trick.
     # def contains_many(self, items: Iterable[str]) -> List[str]:
-    #     with self.__transaction() as cursor:
+    #     with self.transaction() as cursor:
     #         cursor.execute(
     #             'SELECT "key" FROM "set" WHERE "key" in (SELECT "value" FROM JSON_EACH(?));',
     #             (json.dumps(items),),
     #         )
 
     def __contains__(self, item: str) -> bool:
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute('SELECT 1 FROM "set" WHERE "key" = ?;', (item,))
             return cursor.fetchone() is not None
 
     def __len__(self) -> int:
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute('SELECT COUNT(*) FROM "set";')
             return cursor.fetchone()[0]
 
     def __iter__(self) -> Iterator[str]:
-        with self.__transaction() as cursor:
+        with self.transaction() as cursor:
             cursor.execute('SELECT "key" FROM "set";')
 
             for row in iterate_over_sqlite_cursor(cursor):
                 yield row[0]
 
     def close(self) -> None:
-        self.__connection.close()
+        self.connection.close()
 
     def __del__(self) -> None:
         self.close()