Add more benchmark tests

benchmark/db_class.py

@@ -7,6 +7,8 @@
 
 from rocksdict import Options, Rdict
 
+from flaxkv.pack import decode, encode
+
 
 class RocksDict:
     rdict_path = "./test_rocksdict"
@@ -77,3 +79,48 @@ def items(self):
     def destroy(self):
         self.sd.close()
         shutil.rmtree(self.root_path)
+
+
+class RedisDict:
+    def __init__(self):
+        import redis
+
+        self.client = redis.Redis(host="localhost", port=6379, db=0)
+
+    def __getitem__(self, key):
+        value = self.client.get(key)
+        return decode(value)
+
+    def __setitem__(self, key, value):
+        self.client.set(key, encode(value))
+
+    def __contains__(self, item):
+        return self.client.exists(item)
+
+    def __len__(self):
+        return self.client.dbsize()
+
+    def keys(self):
+        return self.client.keys('*')
+
+    def items(self):
+        ...
+
+    def destroy(self):
+        self.client.flushdb()
+        self.client.close()
+
+
+def wait_for_server_to_start(url, timeout=10):
+    import httpx
+
+    start_time = time.time()
+    while True:
+        try:
+            response = httpx.get(url)
+            response.raise_for_status()
+            break
+        except Exception:
+            time.sleep(0.5)
+            if time.time() - start_time > timeout:
+                raise RuntimeError("Server didn't start in time")

benchmark/run.py

@@ -1,17 +1,21 @@
 from __future__ import annotations
 
+import subprocess
+
 import numpy as np
 import pandas as pd
 import pytest
-from db_class import RocksDict, ShelveDict
+from db_class import RedisDict, RocksDict, ShelveDict, wait_for_server_to_start
 from rich import print
-from sparrow import MeasureTime
+from sparrow import MeasureTime  # pip install sparrow-python
 
 from flaxkv import FlaxKV
+from flaxkv._sqlite import SQLiteDict
+from flaxkv.core import BaseDBDict
 
 benchmark_info = {}
 
-N = 10_000
+N = 100
 
 
 def prepare_data(n):
@@ -21,8 +25,16 @@ def prepare_data(n):
     return d
 
 
-@pytest.fixture(scope="module", autouse=True)
-def print_info(request):
+@pytest.fixture(scope="session", autouse=True)
+def start_up(request):
+    process = subprocess.Popen(["flaxkv", "run"])
+    try:
+        wait_for_server_to_start(url="http://localhost:8000/healthz")
+        yield
+
+    finally:
+        process.kill()
+
     def plot(df: pd.DataFrame):
         import matplotlib.pyplot as plt
 
@@ -69,10 +81,13 @@ def print_result():
 @pytest.fixture(
     params=[
         "dict",
-        "flaxkv-LMDB",
-        "flaxkv-LevelDB",
+        "Redis",
         "RocksDict",
         "Shelve",
+        "Sqlite3",
+        "flaxkv-LMDB",
+        "flaxkv-LevelDB",
+        "flaxkv-REMOTE",
     ]
 )
 def temp_db(request):
@@ -81,10 +96,16 @@ def temp_db(request):
         db = FlaxKV('benchmark', backend='lmdb')
     elif request.param == "flaxkv-LevelDB":
         db = FlaxKV('benchmark', backend='leveldb')
+    elif request.param == "flaxkv-REMOTE":
+        db = FlaxKV('benchmark', "http://localhost:8000")
     elif request.param == "RocksDict":
         db = RocksDict()
     elif request.param == "Shelve":
         db = ShelveDict()
+    elif request.param == "Redis":
+        db = RedisDict()
+    elif request.param == "Sqlite3":
+        db = SQLiteDict('benchmark.db')
     elif request.param == "dict":
         db = {}
     else:
@@ -103,6 +124,8 @@ def benchmark(db, db_name, n=200):
     for i, (key, value) in enumerate(data.items()):
         db[key] = value
 
+    if isinstance(db, BaseDBDict):
+        db.write_immediately()
     write_cost = mt.show_interval(f"{db_name} write")
 
     keys = list(db.keys())

flaxkv/_sqlite.py

@@ -12,16 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import shutil
 import sqlite3
 
 from .pack import decode, decode_key, encode
 
 
 class SQLiteDict:
+    from pathlib import Path
+
+    root_path = Path("./SQLiteDB")
+    if not root_path.exists():
+        root_path.mkdir(parents=True, exist_ok=True)
+
     def __init__(self, filename=None):
-        self._filename = filename
-        self._conn = self._connect(filename)
+
+        self._filename = str(self.root_path / filename)
+        self._conn = self._connect(self._filename)
         self._cursor = self._conn.cursor()
         self._cursor.execute(
             "CREATE TABLE IF NOT EXISTS kv (key BLOB PRIMARY KEY, value BLOB)"
@@ -141,3 +149,6 @@ def update(self, E):
 
     def __repr__(self):
         return str(dict(self.items()))
+
+    def destroy(self):
+        self.__del__()