changeable key/value in-db blob size limit (#19669)

* changeable key/value in-db blob size limit

* tests

* test_get_keys_both_disk_and_db

* test_get_keys_values_both_disk_and_db

* fixup schema, like moving the index from the blob to the hash

Author: altendky

chia/_tests/core/data_layer/test_data_store.py

@@ -12,6 +12,7 @@
 import time
 from collections.abc import Awaitable
 from dataclasses import dataclass
+from hashlib import sha256
 from pathlib import Path
 from random import Random
 from typing import Any, BinaryIO, Callable, Optional
@@ -60,7 +61,7 @@
     "root": ["tree_id", "generation", "node_hash", "status"],
     "subscriptions": ["tree_id", "url", "ignore_till", "num_consecutive_failures", "from_wallet"],
     "schema": ["version_id", "applied_at"],
-    "ids": ["kv_id", "blob", "store_id"],
+    "ids": ["kv_id", "hash", "blob", "store_id"],
     "nodes": ["store_id", "hash", "root_hash", "generation", "idx"],
 }
 
@@ -2173,3 +2174,100 @@ async def test_get_existing_hashes(
     not_existing_hashes = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_keys)]
     result = await data_store.get_existing_hashes(existing_hashes + not_existing_hashes, store_id)
     assert result == set(existing_hashes)
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize(argnames="size_offset", argvalues=[-1, 0, 1])
+async def test_basic_key_value_db_vs_disk_cutoff(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+    size_offset: int,
+) -> None:
+    size = data_store.prefer_file_kv_blob_length + size_offset
+
+    blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+    blob_hash = bytes32(sha256(blob).digest())
+    async with data_store.db_wrapper.writer() as writer:
+        await data_store.add_kvid(blob=blob, store_id=store_id, writer=writer)
+
+    file_exists = data_store.get_key_value_path(store_id=store_id, blob_hash=blob_hash).exists()
+    async with data_store.db_wrapper.writer() as writer:
+        async with writer.execute(
+            "SELECT blob FROM ids WHERE hash = :blob_hash",
+            {"blob_hash": blob_hash},
+        ) as cursor:
+            row = await cursor.fetchone()
+            assert row is not None
+            db_blob: Optional[bytes] = row["blob"]
+
+    if size_offset <= 0:
+        assert not file_exists
+        assert db_blob == blob
+    else:
+        assert file_exists
+        assert db_blob is None
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize(argnames="size_offset", argvalues=[-1, 0, 1])
+@pytest.mark.parametrize(argnames="limit_change", argvalues=[-2, -1, 1, 2])
+async def test_changing_key_value_db_vs_disk_cutoff(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+    size_offset: int,
+    limit_change: int,
+) -> None:
+    size = data_store.prefer_file_kv_blob_length + size_offset
+
+    blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+    async with data_store.db_wrapper.writer() as writer:
+        kv_id = await data_store.add_kvid(blob=blob, store_id=store_id, writer=writer)
+
+    data_store.prefer_file_kv_blob_length += limit_change
+    retrieved_blob = await data_store.get_blob_from_kvid(kv_id=kv_id, store_id=store_id)
+
+    assert blob == retrieved_blob
+
+
+@pytest.mark.anyio
+async def test_get_keys_both_disk_and_db(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+) -> None:
+    inserted_keys: set[bytes] = set()
+
+    for size_offset in [-1, 0, 1]:
+        size = data_store.prefer_file_kv_blob_length + size_offset
+
+        blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        await data_store.insert(key=blob, value=b"", store_id=store_id, status=Status.COMMITTED)
+        inserted_keys.add(blob)
+
+    retrieved_keys = set(await data_store.get_keys(store_id=store_id))
+
+    assert retrieved_keys == inserted_keys
+
+
+@pytest.mark.anyio
+async def test_get_keys_values_both_disk_and_db(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+) -> None:
+    inserted_keys_values: dict[bytes, bytes] = {}
+
+    for size_offset in [-1, 0, 1]:
+        size = data_store.prefer_file_kv_blob_length + size_offset
+
+        key = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        value = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        await data_store.insert(key=key, value=value, store_id=store_id, status=Status.COMMITTED)
+        inserted_keys_values[key] = value
+
+    terminal_nodes = await data_store.get_keys_values(store_id=store_id)
+    retrieved_keys_values = {node.key: node.value for node in terminal_nodes}
+
+    assert retrieved_keys_values == inserted_keys_values

chia/data_layer/data_store.py

@@ -80,6 +80,7 @@ class DataStore:
     recent_merkle_blobs: LRUCache[bytes32, MerkleBlob]
     merkle_blobs_path: Path
     key_value_blobs_path: Path
+    prefer_file_kv_blob_length: int = 4096
 
     @classmethod
     @contextlib.asynccontextmanager
@@ -156,6 +157,7 @@ async def managed(
                     """
                     CREATE TABLE IF NOT EXISTS ids(
                         kv_id INTEGER PRIMARY KEY,
+                        hash BLOB NOT NULL CHECK(length(store_id) == 32),
                         blob BLOB,
                         store_id BLOB NOT NULL CHECK(length(store_id) == 32)
                     )
@@ -175,7 +177,7 @@ async def managed(
                 )
                 await writer.execute(
                     """
-                    CREATE UNIQUE INDEX IF NOT EXISTS ids_blob_index ON ids(blob, store_id)
+                    CREATE UNIQUE INDEX IF NOT EXISTS ids_hash_index ON ids(hash, store_id)
                     """
                 )
                 await writer.execute(
@@ -562,20 +564,17 @@ async def insert_root_from_merkle_blob(
 
         return await self._insert_root(store_id, root_hash, status)
 
-    def _kvid_blob_is_file(self, blob: bytes) -> bool:
-        return len(blob) >= len(bytes32.zeros)
+    def _use_file_for_new_kv_blob(self, blob: bytes) -> bool:
+        return len(blob) > self.prefer_file_kv_blob_length
 
     async def get_kvid(self, blob: bytes, store_id: bytes32) -> Optional[KeyOrValueId]:
-        if self._kvid_blob_is_file(blob):
-            table_blob = sha256(blob).digest()
-        else:
-            table_blob = blob
+        blob_hash = bytes32(sha256(blob).digest())
 
         async with self.db_wrapper.reader() as reader:
             cursor = await reader.execute(
-                "SELECT kv_id FROM ids WHERE blob = ? AND store_id = ?",
+                "SELECT kv_id FROM ids WHERE hash = ? AND store_id = ?",
                 (
-                    table_blob,
+                    blob_hash,
                     store_id,
                 ),
             )
@@ -586,19 +585,15 @@ async def get_kvid(self, blob: bytes, store_id: bytes32) -> Optional[KeyOrValueI
 
             return KeyOrValueId(row[0])
 
-    def get_blob_from_table_blob(self, table_blob: bytes, store_id: bytes32) -> bytes:
-        if not self._kvid_blob_is_file(table_blob):
-            return table_blob
-
-        blob_hash = bytes32(table_blob)
+    def get_blob_from_file(self, blob_hash: bytes32, store_id: bytes32) -> bytes:
         # TODO: seems that zstd needs hinting
         # TODO: consider file-system based locking of either the file or the store directory
         return zstd.decompress(self.get_key_value_path(store_id=store_id, blob_hash=blob_hash).read_bytes())  # type: ignore[no-any-return]
 
     async def get_blob_from_kvid(self, kv_id: KeyOrValueId, store_id: bytes32) -> Optional[bytes]:
         async with self.db_wrapper.reader() as reader:
             cursor = await reader.execute(
-                "SELECT blob FROM ids WHERE kv_id = ? AND store_id = ?",
+                "SELECT hash, blob FROM ids WHERE kv_id = ? AND store_id = ?",
                 (
                     kv_id,
                     store_id,
@@ -609,7 +604,12 @@ async def get_blob_from_kvid(self, kv_id: KeyOrValueId, store_id: bytes32) -> Op
             if row is None:
                 return None
 
-        return self.get_blob_from_table_blob(bytes(row[0]), store_id)
+        blob: bytes = row["blob"]
+        if blob is not None:
+            return blob
+
+        blob_hash = bytes32(row["hash"])
+        return self.get_blob_from_file(blob_hash, store_id)
 
     async def get_terminal_node(self, kid: KeyId, vid: ValueId, store_id: bytes32) -> TerminalNode:
         key = await self.get_blob_from_kvid(kid.raw, store_id)
@@ -620,15 +620,17 @@ async def get_terminal_node(self, kid: KeyId, vid: ValueId, store_id: bytes32) -
         return TerminalNode(hash=leaf_hash(key, value), key=key, value=value)
 
     async def add_kvid(self, blob: bytes, store_id: bytes32, writer: aiosqlite.Connection) -> KeyOrValueId:
-        is_file = self._kvid_blob_is_file(blob)
-        if is_file:
-            table_blob = sha256(blob).digest()
+        use_file = self._use_file_for_new_kv_blob(blob)
+        blob_hash = bytes32(sha256(blob).digest())
+        if use_file:
+            table_blob = None
         else:
             table_blob = blob
         try:
             row = await writer.execute_insert(
-                "INSERT INTO ids (blob, store_id) VALUES (?, ?)",
+                "INSERT INTO ids (hash, blob, store_id) VALUES (?, ?, ?)",
                 (
+                    blob_hash,
                     table_blob,
                     store_id,
                 ),
@@ -644,14 +646,12 @@ async def add_kvid(self, blob: bytes, store_id: bytes32, writer: aiosqlite.Conne
 
         if row is None:
             raise Exception("Internal error")
-        kv_id = KeyOrValueId(row[0])
-        if is_file:
-            blob_hash = bytes32(table_blob)
+        if use_file:
             path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash)
             path.parent.mkdir(parents=True, exist_ok=True)
             # TODO: consider file-system based locking of either the file or the store directory
             path.write_bytes(zstd.compress(blob))
-        return kv_id
+        return KeyOrValueId(row[0])
 
     async def add_key_value(
         self, key: bytes, value: bytes, store_id: bytes32, writer: aiosqlite.Connection
@@ -1050,10 +1050,22 @@ async def get_internal_nodes(self, store_id: bytes32, root_hash: Optional[bytes3
         return internal_nodes
 
     def get_terminal_node_from_table_blobs(
-        self, kid: KeyId, vid: ValueId, table_blobs: dict[KeyOrValueId, bytes], store_id: bytes32
+        self,
+        kid: KeyId,
+        vid: ValueId,
+        table_blobs: dict[KeyOrValueId, tuple[bytes32, Optional[bytes]]],
+        store_id: bytes32,
     ) -> TerminalNode:
-        key = self.get_blob_from_table_blob(table_blobs[KeyOrValueId(kid.raw)], store_id)
-        value = self.get_blob_from_table_blob(table_blobs[KeyOrValueId(vid.raw)], store_id)
+        key = table_blobs[KeyOrValueId(kid.raw)][1]
+        if key is None:
+            key_hash = table_blobs[KeyOrValueId(kid.raw)][0]
+            key = self.get_blob_from_file(key_hash, store_id)
+
+        value = table_blobs[KeyOrValueId(vid.raw)][1]
+        if value is None:
+            value_hash = table_blobs[KeyOrValueId(vid.raw)][0]
+            value = self.get_blob_from_file(value_hash, store_id)
+
         return TerminalNode(hash=leaf_hash(key, value), key=key, value=value)
 
     async def get_keys_values(
@@ -1277,8 +1289,10 @@ async def get_keys(
             table_blobs = await self.get_table_blobs(raw_key_ids, store_id)
             keys: list[bytes] = []
             for kid in kv_ids.keys():
-                key = self.get_blob_from_table_blob(table_blobs[KeyOrValueId(kid.raw)], store_id)
-                keys.append(key)
+                blob_hash, blob = table_blobs[KeyOrValueId(kid.raw)]
+                if blob is None:
+                    blob = self.get_blob_from_file(blob_hash, store_id)
+                keys.append(blob)
 
         return keys
 
@@ -1653,22 +1667,25 @@ async def get_nodes_for_file(
 
     async def get_table_blobs(
         self, kv_ids_iter: Iterable[KeyOrValueId], store_id: bytes32
-    ) -> dict[KeyOrValueId, bytes]:
-        result: dict[KeyOrValueId, bytes] = {}
+    ) -> dict[KeyOrValueId, tuple[bytes32, Optional[bytes]]]:
+        result: dict[KeyOrValueId, tuple[bytes32, Optional[bytes]]] = {}
         batch_size = min(500, SQLITE_MAX_VARIABLE_NUMBER - 10)
         kv_ids = list(dict.fromkeys(kv_ids_iter))
 
         async with self.db_wrapper.reader() as reader:
             for i in range(0, len(kv_ids), batch_size):
                 chunk = kv_ids[i : i + batch_size]
                 placeholders = ",".join(["?"] * len(chunk))
-                query = (
-                    f"SELECT blob, kv_id FROM ids WHERE store_id = ? AND kv_id IN ({placeholders}) LIMIT {len(chunk)}"
-                )
+                query = f"""
+                    SELECT hash, blob, kv_id
+                    FROM ids
+                    WHERE store_id = ? AND kv_id IN ({placeholders})
+                    LIMIT {len(chunk)}
+                """
 
                 async with reader.execute(query, (store_id, *chunk)) as cursor:
                     rows = await cursor.fetchall()
-                    result.update({row["kv_id"]: row["blob"] for row in rows})
+                    result.update({row["kv_id"]: (row["hash"], row["blob"]) for row in rows})
 
         if len(result) != len(kv_ids):
             raise Exception("Cannot retrieve all the requested kv_ids")
@@ -1713,8 +1730,10 @@ async def write_tree_to_file(
                 blobs = []
                 for raw_id in (node.value1, node.value2):
                     id = KeyOrValueId.from_bytes(raw_id)
-                    id_table_blob = table_blobs[id]
-                    blobs.append(self.get_blob_from_table_blob(id_table_blob, store_id))
+                    blob_hash, blob = table_blobs[id]
+                    if blob is None:
+                        blob = self.get_blob_from_file(blob_hash, store_id)
+                    blobs.append(blob)
                 to_write = bytes(SerializedNode(True, blobs[0], blobs[1]))
             else:
                 to_write = bytes(node)