use fixed number of threads, hbase stress tests

Author: akhileshh

kvdbclient/base.py

@@ -1,4 +1,3 @@
-import os
 import sys
 import time
 import typing
@@ -77,7 +76,7 @@ def _init_common(self, logger_name, table_meta, lock_expiry, max_row_key_count):
         self._lock_expiry = lock_expiry
         self._version = None
         self._max_row_key_count = max_row_key_count
-        self._executor = ThreadPoolExecutor(max_workers=2 * (os.cpu_count() or 1))
+        self._executor = ThreadPoolExecutor(max_workers=8)
 
     # ── Abstract: backend-specific primitives ────────────────────────────
 

kvdbclient/hbase/client.py

@@ -1,6 +1,7 @@
 # pylint: disable=invalid-name, missing-docstring, line-too-long, too-many-arguments
 
 import base64
+import logging
 import typing
 import struct
 from collections import defaultdict
@@ -200,22 +201,9 @@ def _put_rows_with_retry(self, body: dict, slow_retry: bool = True):
 
     def lock_root(self, root_id: np.uint64, operation_id: np.uint64) -> bool:
         lock_column = attributes.Concurrency.Lock
-        indefinite_lock_column = attributes.Concurrency.IndefiniteLock
         lock_expiry = self._lock_expiry
         row_key = serialize_uint64(root_id)
 
-        # Check indefinite lock
-        indef_row = self._read_byte_row(row_key, columns=indefinite_lock_column)
-        if indef_row:
-            return False
-
-        # Check new parents
-        new_parents_col = attributes.Hierarchy.NewParent
-        np_row = self._read_byte_row(row_key, columns=new_parents_col)
-        if np_row:
-            return False
-
-        # Try checkAndPut: set lock only if column is empty
         lock_value = serialize_uint64(operation_id)
         timestamp = get_valid_timestamp(None)
         success = self._check_and_put(
@@ -229,26 +217,30 @@ def lock_root(self, root_id: np.uint64, operation_id: np.uint64) -> bool:
             put_timestamp=timestamp,
         )
 
-        if not success:
-            cells = self._read_byte_row(row_key, columns=lock_column)
-            if cells:
-                lock_ts = cells[0].timestamp
-                if lock_ts and datetime.now(timezone.utc) - lock_ts > lock_expiry:
-                    self._delete_cell(row_key, lock_column.family_id, lock_column.key)
-                    return self.lock_root(root_id, operation_id)
+        if success:
+            # Verify no indefinite lock (mirrors BigTable's RowFilterUnion check).
+            # HBase REST checkAndPut only supports single-column checks, so we
+            # acquire then verify, rolling back if an indefinite lock exists.
+            indefinite_lock_column = attributes.Concurrency.IndefiniteLock
+            if self._read_byte_row(row_key, columns=indefinite_lock_column):
+                self._delete_cell(row_key, lock_column.family_id, lock_column.key)
+                return False
+            return True
+
+        cells = self._read_byte_row(row_key, columns=lock_column)
+        if cells:
+            lock_ts = cells[0].timestamp
+            if lock_ts and datetime.now(timezone.utc) - lock_ts > lock_expiry:
+                self._delete_cell(row_key, lock_column.family_id, lock_column.key)
+                return self.lock_root(root_id, operation_id)
+            if self.logger.isEnabledFor(logging.DEBUG):
                 self.logger.debug(f"Locked operation ids: {[c.value for c in cells]}")
-            return False
-        return True
+        return False
 
     def lock_root_indefinitely(self, root_id: np.uint64, operation_id: np.uint64) -> bool:
         lock_column = attributes.Concurrency.IndefiniteLock
         row_key = serialize_uint64(root_id)
 
-        new_parents_col = attributes.Hierarchy.NewParent
-        np_row = self._read_byte_row(row_key, columns=new_parents_col)
-        if np_row:
-            return False
-
         lock_value = serialize_uint64(operation_id)
         timestamp = get_valid_timestamp(None)
         success = self._check_and_put(
@@ -262,9 +254,10 @@ def lock_root_indefinitely(self, root_id: np.uint64, operation_id: np.uint64) ->
             put_timestamp=timestamp,
         )
         if not success:
-            cells = self._read_byte_row(row_key, columns=lock_column)
-            if cells:
-                self.logger.debug(f"Indefinitely locked operation ids: {[c.value for c in cells]}")
+            if self.logger.isEnabledFor(logging.DEBUG):
+                cells = self._read_byte_row(row_key, columns=lock_column)
+                if cells:
+                    self.logger.debug(f"Indefinitely locked operation ids: {[c.value for c in cells]}")
             return False
         return True
 
@@ -295,11 +288,6 @@ def renew_lock(self, root_id: np.uint64, operation_id: np.uint64) -> bool:
         row_key = serialize_uint64(root_id)
         lock_value = serialize_uint64(operation_id)
 
-        new_parents_col = attributes.Hierarchy.NewParent
-        np_row = self._read_byte_row(row_key, columns=new_parents_col)
-        if np_row:
-            return False
-
         return self._check_and_put(
             row_key=row_key,
             check_family=lock_column.family_id,
@@ -414,31 +402,21 @@ def _read_byte_row(self, row_key, columns=None, start_time=None, end_time=None,
         return row.get(row_key, [] if single_column else {})
 
     def _fetch_rows_by_keys(self, row_keys, col_specs, start_time, end_time, end_time_inclusive, single_column=False):
-        result = {}
-        for row_key in row_keys:
-            key_b64 = hbase_utils.encode_value(row_key)
-            col_path = ",".join(col_specs) if col_specs else ""
-            url = self._table_url(f"/{key_b64}")
-            if col_path:
-                url = self._table_url(f"/{key_b64}/{col_path}")
-
-            params = {"v": "1000"}
-            if start_time:
-                ts = hbase_utils.datetime_to_hbase_ts(
-                    hbase_utils.get_hbase_compatible_time_stamp(start_time)
-                )
-                params["ts.from"] = str(ts)
-            if end_time:
-                et = hbase_utils.get_hbase_compatible_time_stamp(end_time, round_up=end_time_inclusive)
-                params["ts.to"] = str(hbase_utils.datetime_to_hbase_ts(et))
-
-            resp = self._session.get(url, params=params)
-            if resp.status_code == 404:
-                continue
-            resp.raise_for_status()
-            parsed = hbase_utils.parse_cell_response(resp.json(), single_column=single_column)
-            result.update(parsed)
-        return result
+        if not row_keys:
+            return {}
+        sorted_keys = sorted(row_keys)
+        all_rows = self._fetch_rows_by_range(
+            start_key=sorted_keys[0],
+            end_key=sorted_keys[-1],
+            end_key_inclusive=True,
+            col_specs=col_specs,
+            start_time=start_time,
+            end_time=end_time,
+            end_time_inclusive=end_time_inclusive,
+            single_column=single_column,
+        )
+        key_set = set(row_keys)
+        return {k: v for k, v in all_rows.items() if k in key_set}
 
     def _fetch_rows_by_range(self, start_key, end_key, end_key_inclusive, col_specs, start_time, end_time, end_time_inclusive, single_column=False):
         scanner_spec = {

tests/conftest.py

@@ -4,11 +4,15 @@
 import subprocess
 import time
 import uuid
+from datetime import timedelta
 
 import pytest
 
 from kvdbclient.bigtable import BigTableConfig
 from kvdbclient.bigtable.client import Client
+from kvdbclient.hbase import HBaseConfig
+from kvdbclient.hbase.client import Client as HBaseClient
+from hbase_mock_server import start_hbase_mock_server
 
 
 EMULATOR_PROJECT = "test-project"
@@ -32,6 +36,9 @@ def _wait_for_port(host, port, timeout=30.0):
     raise RuntimeError(f"Emulator on {host}:{port} not ready within {timeout}s")
 
 
+# ── BigTable fixtures ────────────────────────────────────────────────────
+
+
 @pytest.fixture(scope="session")
 def bigtable_emulator():
     """Start the BigTable emulator or use one already running (CI)."""
@@ -104,3 +111,49 @@ def bt_client_small_batch(bigtable_emulator):
     client = Client(table_id=table_id, config=config)
     client.create_table(meta={"test": True}, version="0.0.1")
     yield client
+
+
+# ── HBase fixtures ───────────────────────────────────────────────────────
+
+
+@pytest.fixture(scope="session")
+def hbase_server():
+    _data, server, port = start_hbase_mock_server()
+    yield f"http://127.0.0.1:{port}"
+    server.shutdown()
+
+
+@pytest.fixture(scope="session")
+def hbase_config(hbase_server):
+    return HBaseConfig(BASE_URL=hbase_server)
+
+
+@pytest.fixture()
+def hbase_client(hbase_config):
+    table_id = f"test_{uuid.uuid4().hex[:12]}"
+    client = HBaseClient(table_id=table_id, config=hbase_config)
+    client.create_table(meta={"test": True}, version="0.0.1")
+    yield client
+
+
+@pytest.fixture()
+def hbase_client_no_table(hbase_config):
+    table_id = f"test_{uuid.uuid4().hex[:12]}"
+    client = HBaseClient(table_id=table_id, config=hbase_config)
+    yield client
+
+
+@pytest.fixture()
+def hbase_client_short_expiry(hbase_config):
+    table_id = f"test_{uuid.uuid4().hex[:12]}"
+    client = HBaseClient(table_id=table_id, config=hbase_config, lock_expiry=timedelta(seconds=1))
+    client.create_table(meta={"test": True}, version="0.0.1")
+    yield client
+
+
+@pytest.fixture()
+def bt_client_short_expiry(bt_config):
+    table_id = f"test_{uuid.uuid4().hex[:12]}"
+    client = Client(table_id=table_id, config=bt_config, lock_expiry=timedelta(seconds=1))
+    client.create_table(meta={"test": True}, version="0.0.1")
+    yield client

tests/test_bigtable_client.py

@@ -1,5 +1,6 @@
 # pylint: disable=missing-docstring, redefined-outer-name
 
+import time
 from datetime import datetime, timedelta, timezone
 
 import numpy as np
@@ -519,3 +520,21 @@ def test_no_lock_returns_none(self, bt_client):
     def test_wrong_op_returns_none(self, bt_client):
         bt_client.lock_root(np.uint64(100), np.uint64(1))
         assert bt_client.get_lock_timestamp(np.uint64(100), np.uint64(2)) is None
+
+
+# ── Lock Expiry ──────────────────────────────────────────────────────────
+
+
+class TestLockExpiry:
+    def test_expired_lock_allows_relock(self, bt_client_short_expiry):
+        client = bt_client_short_expiry
+        root, op1, op2 = np.uint64(100), np.uint64(1), np.uint64(2)
+        assert client.lock_root(root, op1) is True
+        time.sleep(1.5)
+        assert client.lock_root(root, op2) is True
+
+    def test_non_expired_lock_blocks(self, bt_client_short_expiry):
+        client = bt_client_short_expiry
+        root, op1, op2 = np.uint64(100), np.uint64(1), np.uint64(2)
+        assert client.lock_root(root, op1) is True
+        assert client.lock_root(root, op2) is False