Merge pull request #99 from theseriff/perf/sqlite

perf(storage): replace threading lock with asyncio.Queue in SQLiteStorage

Author: theseriff

benchmarks/results.txt

@@ -1,32 +1,32 @@
-jobify: 0.9.0
-OS: Linux 6.18.13-arch1-1
+jobify: 0.9.2
+OS: Linux 6.19.8-arch1-1
 Architecture: x86_64
 CPU Physical cores: 4
 CPU Logical cores: 8
 CPU Max Frequency: 4700.00MHz
 CPU Min Frequency: 400.00MHz
-CPU Current Frequency: 1861.95MHz
+CPU Current Frequency: 933.53MHz
 
 ======================= Serializers ========================
 Config Name                         | Total Seconds         
 ----------------------------------- | ----------------------
-pickle                              |   0.564 sec
-extended_json                       |   3.225 sec
+pickle                              |   0.558 sec
+extended_json                       |   3.176 sec
 ============================================================
 
 ======================== Jobify APP ========================
 Config Name                         | Latency    | Throughput  
 ----------------------------------- | ---------- | ------------
-dummyDB+extended_json+none          |   0.024 ms |  61688.8 TPS
-dummyDB+extended_json+adaptix       |   0.024 ms |  61375.7 TPS
-dummyDB+json+adaptix                |   0.024 ms |  60827.8 TPS
-dummyDB+json+none                   |   0.029 ms |  60814.8 TPS
-dummyDB+pickle+adaptix              |   0.024 ms |  60488.4 TPS
-dummyDB+pickle+none                 |   0.026 ms |  59461.0 TPS
-sqlite+pickle+adaptix               |   0.135 ms |   7709.6 TPS
-sqlite+json+none                    |   0.133 ms |   7704.6 TPS
-sqlite+pickle+none                  |   0.140 ms |   7637.9 TPS
-sqlite+json+adaptix                 |   0.135 ms |   7447.3 TPS
-sqlite+extended_json+adaptix        |   0.140 ms |   7168.0 TPS
-sqlite+extended_json+none           |   0.147 ms |   7078.2 TPS
-============================================================
+dummyDB+extended_json+adaptix       |   0.024 ms |  61468.7 TPS
+dummyDB+pickle+adaptix              |   0.024 ms |  61399.0 TPS
+dummyDB+json+none                   |   0.024 ms |  61303.4 TPS
+dummyDB+extended_json+none          |   0.024 ms |  61219.5 TPS
+dummyDB+json+adaptix                |   0.024 ms |  61182.5 TPS
+dummyDB+pickle+none                 |   0.026 ms |  59433.8 TPS
+sqlite+pickle+adaptix               |   0.114 ms |   9646.6 TPS
+sqlite+json+none                    |   0.111 ms |   9508.7 TPS
+sqlite+json+adaptix                 |   0.114 ms |   9332.3 TPS
+sqlite+pickle+none                  |   0.121 ms |   9313.6 TPS
+sqlite+extended_json+adaptix        |   0.119 ms |   8864.6 TPS
+sqlite+extended_json+none           |   0.122 ms |   8804.9 TPS
+============================================================

src/jobify/_internal/storage/sqlite.py

@@ -1,15 +1,14 @@
 import asyncio
 import sqlite3
-import threading
 from collections.abc import Callable, Sequence
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
 from zoneinfo import ZoneInfo
 
 from typing_extensions import override
 
-from jobify._internal.common.constants import JobStatus
+from jobify._internal.common.constants import EMPTY, JobStatus
 from jobify._internal.storage.base import (
     ScheduledJob,
     Storage,
@@ -21,6 +20,7 @@
 
     from jobify._internal.common.types import LoopFactory
 
+
 CREATE_SCHEDULED_TABLE_QUERY = """
 CREATE TABLE IF NOT EXISTS {} (
     job_id TEXT PRIMARY KEY,
@@ -59,6 +59,12 @@
 
 ReturnT = TypeVar("ReturnT")
 
+_Callback = Callable[[sqlite3.Connection], ReturnT]
+_AsyncQueue: TypeAlias = asyncio.Queue[
+    tuple[_Callback[Any], asyncio.Future[Any]],
+]
+_STOP: Any = object()
+
 
 class SQLiteStorage(Storage):
     def __init__(
@@ -67,6 +73,7 @@ def __init__(
         *,
         table_name: str = "jobify_schedules",
         timeout: float = 20.0,
+        max_queue_size: int = 1024,
     ) -> None:
         validate_table_name(table_name)
         self.database: Path = (
@@ -77,8 +84,10 @@ def __init__(
         self.tz: ZoneInfo = ZoneInfo("UTC")
         self.getloop: LoopFactory = asyncio._get_running_loop
         self.threadpool: ThreadPoolExecutor | None = None
-        self._conn: sqlite3.Connection | None = None
-        self._lock: threading.Lock = threading.Lock()
+        self.max_queue_size: int = max_queue_size
+
+        self._queue: _AsyncQueue = EMPTY
+        self._worker_task: asyncio.Task[None] | None = None
 
         self.create_scheduled_table_query: str = (
             CREATE_SCHEDULED_TABLE_QUERY.format(table_name)
@@ -93,20 +102,36 @@ def __init__(
             table_name,
         )
 
-    @property
-    def conn(self) -> sqlite3.Connection:
-        if self._conn is None:
-            msg = "Database not initialized. Call startup() first."
-            raise RuntimeError(msg)
-        return self._conn
+    async def _worker(self, conn: sqlite3.Connection) -> None:
+        loop = self.getloop()
+        while True:
+            item = await self._queue.get()
+
+            if item is _STOP:
+                self._queue.task_done()
+                break
+
+            callback, future = item
+            try:
+                result = await loop.run_in_executor(
+                    self.threadpool,
+                    callback,
+                    conn,
+                )
+            except Exception as exc:  # noqa: BLE001
+                future.set_exception(exc)
+            else:
+                future.set_result(result)
+            finally:
+                self._queue.task_done()
 
-    async def _to_thread(self, func: Callable[[], ReturnT]) -> ReturnT:
-        def thread_safe() -> ReturnT:
-            with self._lock:
-                return func()
+        conn.close()
 
+    async def _execute(self, callback: _Callback[ReturnT]) -> ReturnT:
         loop = self.getloop()
-        return await loop.run_in_executor(self.threadpool, thread_safe)
+        future: asyncio.Future[ReturnT] = loop.create_future()
+        await self._queue.put((callback, future))
+        return await future
 
     @override
     async def startup(self) -> None:
@@ -119,19 +144,24 @@ async def startup(self) -> None:
         _ = conn.execute("PRAGMA synchronous=NORMAL;")
         _ = conn.execute(self.create_scheduled_table_query)
         conn.commit()
-        self._conn = conn
+        self._queue = asyncio.Queue(self.max_queue_size)
+        self._worker_task = asyncio.create_task(self._worker(conn))
 
     @override
     async def shutdown(self) -> None:
-        if self._conn is not None:
-            with self._lock:
-                self._conn.close()
-                self._conn = None
+        if self._queue is not EMPTY:
+            await self._queue.put(_STOP)
+            await self._queue.join()
+            self._queue = EMPTY
+
+        if self._worker_task is not None:
+            await self._worker_task
+            self._worker_task = None
 
     @override
     async def get_schedules(self) -> list[ScheduledJob]:
-        def get() -> list[ScheduledJob]:
-            cursor = self.conn.execute(self.select_schedules_query)
+        def get(conn: sqlite3.Connection) -> list[ScheduledJob]:
+            cursor = conn.execute(self.select_schedules_query)
             return [
                 ScheduledJob(
                     job_id=row[0],
@@ -145,35 +175,35 @@ def get() -> list[ScheduledJob]:
                 for row in cursor.fetchall()
             ]
 
-        return await self._to_thread(get)
+        return await self._execute(get)
 
     @override
     async def add_schedule(self, *scheduled: ScheduledJob) -> None:
-        def insert_many() -> None:
-            with self.conn as conn:
-                _ = conn.executemany(
-                    self.insert_schedule_query,
-                    [
-                        (
-                            sch.job_id,
-                            sch.name,
-                            sch.message,
-                            sch.status,
-                            sch.next_run_at.isoformat(),
-                        )
-                        for sch in scheduled
-                    ],
-                )
+        def insert_many(conn: sqlite3.Connection) -> None:
+            _ = conn.executemany(
+                self.insert_schedule_query,
+                [
+                    (
+                        sch.job_id,
+                        sch.name,
+                        sch.message,
+                        sch.status,
+                        sch.next_run_at.isoformat(),
+                    )
+                    for sch in scheduled
+                ],
+            )
+            conn.commit()
 
-        return await self._to_thread(insert_many)
+        return await self._execute(insert_many)
 
     @override
     async def delete_schedule(self, job_id: str) -> None:
-        def delete() -> None:
-            with self.conn as conn:
-                _ = conn.execute(self.delete_schedule_query, (job_id,))
+        def delete(conn: sqlite3.Connection) -> None:
+            _ = conn.execute(self.delete_schedule_query, (job_id,))
+            conn.commit()
 
-        return await self._to_thread(delete)
+        return await self._execute(delete)
 
     @override
     async def delete_schedule_many(self, job_ids: Sequence[str]) -> None:
@@ -182,17 +212,17 @@ async def delete_schedule_many(self, job_ids: Sequence[str]) -> None:
 
         job_ids = sorted(job_ids)
 
-        def delete_many() -> None:
+        def delete_many(conn: sqlite3.Connection) -> None:
             batch_size = 500
-            with self.conn as conn:
-                for i in range(0, len(job_ids), batch_size):
-                    batch = job_ids[i : i + batch_size]
-                    query = DELETE_SCHEDULE_MANY_QUERY.format_map(
-                        {
-                            "table_name": self.table_name,
-                            "placeholder": ",".join("?" * len(batch)),
-                        }
-                    )
-                    _ = conn.execute(query, batch)
+            for i in range(0, len(job_ids), batch_size):
+                batch = job_ids[i : i + batch_size]
+                query = DELETE_SCHEDULE_MANY_QUERY.format_map(
+                    {
+                        "table_name": self.table_name,
+                        "placeholder": ",".join("?" * len(batch)),
+                    }
+                )
+                _ = conn.execute(query, batch)
+            conn.commit()
 
-        return await self._to_thread(delete_many)
+        return await self._execute(delete_many)

tests/test_storage.py

@@ -38,8 +38,6 @@ async def test_sqlite(now: datetime) -> None:
     db = Path("test.db")
     storage = SQLiteStorage(database=db, table_name="test_table")
 
-    with pytest.raises(RuntimeError):
-        _ = storage.conn
     await storage.shutdown()
 
     scheduled = ScheduledJob(
@@ -529,3 +527,11 @@ def _(name: str) -> None:
         await app.wait_all()
 
     mock.assert_called_once_with("biba")
+
+
+async def test_sqlite_execute_with_raise() -> None:
+    storage = SQLiteStorage(":memory:")
+    await storage.startup()
+    with pytest.raises(TypeError):
+        await storage._execute(Mock(side_effect=TypeError))
+    await storage.shutdown()