BlockStoreProtocol

Author: richardkiss

chia/_tests/blockchain/blockchain_test_utils.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, cast
 
 from chia_rs import FullBlock, SpendBundleConditions
 from chia_rs.sized_ints import uint32, uint64
@@ -10,19 +10,20 @@
 from chia.consensus.blockchain import AddBlockResult, Blockchain
 from chia.consensus.difficulty_adjustment import get_next_sub_slot_iters_and_difficulty
 from chia.consensus.multiprocess_validation import PreValidationResult, pre_validate_block
+from chia.full_node.block_store import BlockStore
 from chia.types.validation_state import ValidationState
 from chia.util.errors import Err
 
 
 async def check_block_store_invariant(bc: Blockchain):
-    db_wrapper = bc.block_store.db_wrapper
+    block_store = cast(BlockStore, bc.block_store)
 
-    if db_wrapper.db_version == 1:
+    if block_store.db_wrapper.db_version == 1:
         return
 
     in_chain = set()
     max_height = -1
-    async with db_wrapper.writer_maybe_transaction() as conn:
+    async with block_store.transaction() as conn:
         async with conn.execute("SELECT height, in_main_chain FROM full_blocks") as cursor:
             rows = await cursor.fetchall()
             for row in rows:

chia/_tests/blockchain/test_blockchain.py

@@ -3647,7 +3647,7 @@ async def test_get_blocks_at(self, empty_blockchain: Blockchain, default_1000_bl
             heights.append(block.height)
             await _validate_and_add_block(b, block)
 
-        blocks = await b.get_block_records_at(heights, batch_size=2)
+        blocks = await b.get_block_records_at(heights)
         assert blocks
         assert len(blocks) == 200
         assert blocks[-1].height == 199

chia/consensus/block_store_protocol.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Optional, Protocol
+
+from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint32
+
+
+class BlockStoreProtocol(Protocol):
+    async def add_full_block(self, header_hash: bytes32, block: FullBlock, block_record: BlockRecord) -> None: ...
+    def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]: ...
+    async def get_full_block(self, header_hash: bytes32) -> Optional[FullBlock]: ...
+    async def get_generator(self, header_hash: bytes32) -> Optional[bytes]: ...
+    async def get_generators_at(self, heights: set[uint32]) -> dict[uint32, bytes]: ...
+    async def get_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]: ...
+    async def get_block_records_in_range(self, start: int, stop: int) -> dict[bytes32, BlockRecord]: ...
+    async def get_block_records_by_hash(self, header_hashes: list[bytes32]) -> list[BlockRecord]: ...
+    async def get_blocks_by_hash(self, header_hashes: list[bytes32]) -> list[FullBlock]: ...
+    async def persist_sub_epoch_challenge_segments(
+        self, ses_block_hash: bytes32, segments: list[SubEpochChallengeSegment]
+    ) -> None: ...
+    async def get_sub_epoch_challenge_segments(
+        self,
+        ses_block_hash: bytes32,
+    ) -> Optional[list[SubEpochChallengeSegment]]: ...
+    async def rollback(self, height: int) -> None: ...
+    def rollback_cache_block(self, header_hash: bytes32) -> None: ...
+    async def set_in_chain(self, header_hashes: list[tuple[bytes32]]) -> None: ...
+    async def set_peak(self, header_hash: bytes32) -> None: ...
+    async def get_block_records_close_to_peak(
+        self, blocks_n: int
+    ) -> tuple[dict[bytes32, BlockRecord], Optional[bytes32]]: ...
+    async def get_prev_hash(self, header_hash: bytes32) -> bytes32: ...
+    def transaction(self) -> AbstractAsyncContextManager[Any]: ...

chia/consensus/blockchain.py

@@ -27,6 +27,7 @@
 from chia.consensus.block_body_validation import ForkInfo, validate_block_body
 from chia.consensus.block_header_validation import validate_unfinished_header_block
 from chia.consensus.block_height_map import BlockHeightMap
+from chia.consensus.block_store_protocol import BlockStoreProtocol
 from chia.consensus.coin_store_protocol import CoinStoreProtocol
 from chia.consensus.cost_calculator import NPCResult
 from chia.consensus.difficulty_adjustment import get_next_sub_slot_iters_and_difficulty
@@ -35,7 +36,6 @@
 from chia.consensus.generator_tools import get_block_header
 from chia.consensus.get_block_generator import get_block_generator
 from chia.consensus.multiprocess_validation import PreValidationResult
-from chia.full_node.block_store import BlockStore
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.vdf import VDFInfo
 from chia.types.coin_record import CoinRecord
@@ -104,7 +104,7 @@ class Blockchain:
     # Unspent Store
     coin_store: CoinStoreProtocol
     # Store
-    block_store: BlockStore
+    block_store: BlockStoreProtocol
     # Used to verify blocks in parallel
     pool: Executor
     # Set holding seen compact proofs, in order to avoid duplicates.
@@ -122,7 +122,7 @@ class Blockchain:
     @staticmethod
     async def create(
         coin_store: CoinStoreProtocol,
-        block_store: BlockStore,
+        block_store: BlockStoreProtocol,
         height_map: BlockHeightMap,
         consensus_constants: ConsensusConstants,
         reserved_cores: int,
@@ -421,7 +421,7 @@ async def add_block(
 
         try:
             # Always add the block to the database
-            async with self.block_store.db_wrapper.writer():
+            async with self.block_store.transaction():
                 # Perform the DB operations to update the state, and rollback if something goes wrong
                 await self.block_store.add_full_block(header_hash, block, block_record)
                 records, state_change_summary = await self._reconsider_peak(block_record, genesis, fork_info)
@@ -883,7 +883,7 @@ async def get_header_blocks_in_range(
 
         blocks: list[FullBlock] = []
         for hash in hashes.copy():
-            block = self.block_store.block_cache.get(hash)
+            block = self.block_store.get_block_from_cache(hash)
             if block is not None:
                 blocks.append(block)
                 hashes.remove(hash)
@@ -926,27 +926,18 @@ async def get_header_block_by_height(
             return None
         return header_dict[header_hash]
 
-    async def get_block_records_at(self, heights: list[uint32], batch_size: int = 900) -> list[BlockRecord]:
+    async def get_block_records_at(self, heights: list[uint32]) -> list[BlockRecord]:
         """
         gets block records by height (only blocks that are part of the chain)
         """
-        records: list[BlockRecord] = []
         hashes: list[bytes32] = []
-        assert batch_size < self.block_store.db_wrapper.host_parameter_limit
         for height in heights:
             header_hash: Optional[bytes32] = self.height_to_hash(height)
             if header_hash is None:
                 raise ValueError(f"Do not have block at height {height}")
             hashes.append(header_hash)
-            if len(hashes) > batch_size:
-                res = await self.block_store.get_block_records_by_hash(hashes)
-                records.extend(res)
-                hashes = []
-
-        if len(hashes) > 0:
-            res = await self.block_store.get_block_records_by_hash(hashes)
-            records.extend(res)
-        return records
+
+        return await self.block_store.get_block_records_by_hash(hashes)
 
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
         if header_hash in self.__block_records:

chia/full_node/block_store.py

@@ -3,7 +3,8 @@
 import dataclasses
 import logging
 import sqlite3
-from typing import Optional
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Optional
 
 import typing_extensions
 import zstd
@@ -12,6 +13,7 @@
 from chia_rs.sized_ints import uint32
 
 from chia.full_node.full_block_utils import GeneratorBlockInfo, block_info_from_block, generator_from_block
+from chia.util.batches import to_batches
 from chia.util.db_wrapper import DBWrapper2, execute_fetchone
 from chia.util.errors import Err
 from chia.util.lru_cache import LRUCache
@@ -189,6 +191,12 @@ async def get_sub_epoch_challenge_segments(
             return challenge_segments
         return None
 
+    def transaction(self) -> AbstractAsyncContextManager[Any]:
+        return self.db_wrapper.writer()
+
+    def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]:
+        return self.block_cache.get(header_hash)
+
     def rollback_cache_block(self, header_hash: bytes32) -> None:
         try:
             self.block_cache.remove(header_hash)
@@ -322,20 +330,21 @@ async def get_block_records_by_hash(self, header_hashes: list[bytes32]) -> list[
         Returns a list of Block Records, ordered by the same order in which header_hashes are passed in.
         Throws an exception if the blocks are not present
         """
+
         if len(header_hashes) == 0:
             return []
 
         all_blocks: dict[bytes32, BlockRecord] = {}
-        async with self.db_wrapper.reader_no_transaction() as conn:
-            async with conn.execute(
-                "SELECT header_hash,block_record "
-                "FROM full_blocks "
-                f"WHERE header_hash in ({'?,' * (len(header_hashes) - 1)}?)",
-                header_hashes,
-            ) as cursor:
-                for row in await cursor.fetchall():
-                    block_rec = BlockRecord.from_bytes(row[1])
-                    all_blocks[block_rec.header_hash] = block_rec
+        for batch in to_batches(header_hashes, self.db_wrapper.host_parameter_limit):
+            async with self.db_wrapper.reader_no_transaction() as conn:
+                async with conn.execute(
+                    "SELECT header_hash,block_record FROM full_blocks "
+                    f"WHERE header_hash in ({'?,' * (len(batch.entries) - 1)}?)",
+                    batch.entries,
+                ) as cursor:
+                    for row in await cursor.fetchall():
+                        block_rec = BlockRecord.from_bytes(row[1])
+                        all_blocks[block_rec.header_hash] = block_rec
 
         ret: list[BlockRecord] = []
         for hh in header_hashes:
@@ -468,7 +477,8 @@ async def get_block_bytes_in_range(
         No orphan blocks.
         """
 
-        assert self.db_wrapper.db_version == 2
+        if self.db_wrapper.db_version != 2:
+            raise NotImplementedError("get_block_bytes_in_range requires DB version 2")
         async with self.db_wrapper.reader_no_transaction() as conn:
             async with conn.execute(
                 "SELECT block FROM full_blocks WHERE height >= ? AND height <= ? AND in_main_chain=1",

chia/full_node/full_node_api.py

@@ -1453,15 +1453,13 @@ async def request_block_headers(self, request: wallet_protocol.RequestBlockHeade
 
         if request.end_height < request.start_height or request.end_height - request.start_height > 128:
             return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
-        if self.full_node.block_store.db_wrapper.db_version == 2:
-            try:
-                blocks_bytes = await self.full_node.block_store.get_block_bytes_in_range(
-                    request.start_height, request.end_height
-                )
-            except ValueError:
-                return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
-
-        else:
+        try:
+            blocks_bytes = await self.full_node.block_store.get_block_bytes_in_range(
+                request.start_height, request.end_height
+            )
+        except NotImplementedError:
+            # The underlying block store may not support this optimized call
+            # (e.g. v1 DB). In this case, we fall back to the legacy approach
             height_to_hash = self.full_node.blockchain.height_to_hash
             header_hashes: list[bytes32] = []
             for i in range(request.start_height, request.end_height + 1):
@@ -1471,6 +1469,8 @@ async def request_block_headers(self, request: wallet_protocol.RequestBlockHeade
                 header_hashes.append(header_hash)
 
             blocks_bytes = await self.full_node.block_store.get_block_bytes_by_hash(header_hashes)
+        except ValueError:
+            return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
         if len(blocks_bytes) != (request.end_height - request.start_height + 1):  # +1 because interval is inclusive
             return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
         return_filter = request.return_filter

chia/simulator/full_node_simulator.py

@@ -167,7 +167,7 @@ async def revert_block_height(self, new_height: uint32) -> None:
                 raise ValueError("Cannot revert to a height less than 1.")
             block_record: BlockRecord = self.full_node.blockchain.height_to_block_record(new_height)
             # remove enough data to allow a bunch of blocks to be wiped.
-            async with self.full_node.block_store.db_wrapper.writer():
+            async with self.full_node.block_store.transaction():
                 # set coinstore
                 await self.full_node.coin_store.rollback_to_block(new_height)
                 # set blockstore to new height

tach.toml

@@ -22,7 +22,6 @@ path = "chia.consensus"
 depends_on = [
     "chia.types",
     "chia.util",
-    { path = "chia.full_node", deprecated = false },
 ]
 
 [[modules]]