Move batching and simplify protocol.

Author: richardkiss

chia/consensus/block_store_protocol.py

@@ -45,6 +45,6 @@ async def count_uncompactified_blocks(self) -> int: ...
     async def get_block_records_close_to_peak(
         self, blocks_n: int
     ) -> tuple[dict[bytes32, BlockRecord], Optional[bytes32]]: ...
-    def get_host_parameter_limit(self) -> int: ...
+
     async def get_prev_hash(self, header_hash: bytes32) -> bytes32: ...
     def transaction(self) -> AbstractAsyncContextManager[Any]: ...

chia/consensus/blockchain.py

@@ -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.get_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

@@ -13,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
@@ -190,9 +191,6 @@ async def get_sub_epoch_challenge_segments(
             return challenge_segments
         return None
 
-    def get_host_parameter_limit(self) -> int:
-        return self.db_wrapper.host_parameter_limit
-
     def transaction(self) -> AbstractAsyncContextManager[Any]:
         return self.db_wrapper.writer()
 
@@ -332,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: