Skip to content

Commit 0bb7113

Browse files
committed
Simplify BlockStore
1 parent b500593 commit 0bb7113

File tree

6 files changed

+45
-37
lines changed

6 files changed

+45
-37
lines changed

chia/_tests/blockchain/blockchain_test_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ async def check_block_store_invariant(bc: Blockchain):
2222

2323
in_chain = set()
2424
max_height = -1
25-
async with db_wrapper.writer_maybe_transaction() as conn:
25+
async with bc.block_store.transaction() as conn:
2626
async with conn.execute("SELECT height, in_main_chain FROM full_blocks") as cursor:
2727
rows = await cursor.fetchall()
2828
for row in rows:

chia/_tests/blockchain/test_blockchain.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3647,7 +3647,7 @@ async def test_get_blocks_at(self, empty_blockchain: Blockchain, default_1000_bl
36473647
heights.append(block.height)
36483648
await _validate_and_add_block(b, block)
36493649

3650-
blocks = await b.get_block_records_at(heights, batch_size=2)
3650+
blocks = await b.get_block_records_at(heights)
36513651
assert blocks
36523652
assert len(blocks) == 200
36533653
assert blocks[-1].height == 199

chia/consensus/blockchain.py

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -421,7 +421,7 @@ async def add_block(
421421

422422
try:
423423
# Always add the block to the database
424-
async with self.block_store.db_wrapper.writer():
424+
async with self.block_store.transaction():
425425
# Perform the DB operations to update the state, and rollback if something goes wrong
426426
await self.block_store.add_full_block(header_hash, block, block_record)
427427
records, state_change_summary = await self._reconsider_peak(block_record, genesis, fork_info)
@@ -883,7 +883,7 @@ async def get_header_blocks_in_range(
883883

884884
blocks: list[FullBlock] = []
885885
for hash in hashes.copy():
886-
block = self.block_store.block_cache.get(hash)
886+
block = self.block_store.get_block_from_cache(hash)
887887
if block is not None:
888888
blocks.append(block)
889889
hashes.remove(hash)
@@ -926,27 +926,18 @@ async def get_header_block_by_height(
926926
return None
927927
return header_dict[header_hash]
928928

929-
async def get_block_records_at(self, heights: list[uint32], batch_size: int = 900) -> list[BlockRecord]:
929+
async def get_block_records_at(self, heights: list[uint32]) -> list[BlockRecord]:
930930
"""
931931
gets block records by height (only blocks that are part of the chain)
932932
"""
933-
records: list[BlockRecord] = []
934933
hashes: list[bytes32] = []
935-
assert batch_size < self.block_store.db_wrapper.host_parameter_limit
936934
for height in heights:
937935
header_hash: Optional[bytes32] = self.height_to_hash(height)
938936
if header_hash is None:
939937
raise ValueError(f"Do not have block at height {height}")
940938
hashes.append(header_hash)
941-
if len(hashes) > batch_size:
942-
res = await self.block_store.get_block_records_by_hash(hashes)
943-
records.extend(res)
944-
hashes = []
945-
946-
if len(hashes) > 0:
947-
res = await self.block_store.get_block_records_by_hash(hashes)
948-
records.extend(res)
949-
return records
939+
940+
return await self.block_store.get_block_records_by_hash(hashes)
950941

951942
def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
952943
if header_hash in self.__block_records:

chia/full_node/block_store.py

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,27 @@
33
import dataclasses
44
import logging
55
import sqlite3
6+
from contextlib import AbstractAsyncContextManager
67
from typing import Optional
78

9+
import aiosqlite
810
import typing_extensions
911
import zstd
1012
from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSegments
1113
from chia_rs.sized_bytes import bytes32
1214
from chia_rs.sized_ints import uint32
1315

1416
from chia.full_node.full_block_utils import GeneratorBlockInfo, block_info_from_block, generator_from_block
17+
from chia.util.batches import to_batches
1518
from chia.util.db_wrapper import DBWrapper2, execute_fetchone
1619
from chia.util.errors import Err
1720
from chia.util.lru_cache import LRUCache
1821

22+
23+
class UnsupportedDatabaseVersionError(Exception):
24+
"""Raised when a method is called with an unsupported database version."""
25+
26+
1927
log = logging.getLogger(__name__)
2028

2129

@@ -189,6 +197,12 @@ async def get_sub_epoch_challenge_segments(
189197
return challenge_segments
190198
return None
191199

200+
def transaction(self) -> AbstractAsyncContextManager[aiosqlite.Connection]:
201+
return self.db_wrapper.writer()
202+
203+
def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]:
204+
return self.block_cache.get(header_hash)
205+
192206
def rollback_cache_block(self, header_hash: bytes32) -> None:
193207
try:
194208
self.block_cache.remove(header_hash)
@@ -322,20 +336,21 @@ async def get_block_records_by_hash(self, header_hashes: list[bytes32]) -> list[
322336
Returns a list of Block Records, ordered by the same order in which header_hashes are passed in.
323337
Throws an exception if the blocks are not present
324338
"""
339+
325340
if len(header_hashes) == 0:
326341
return []
327342

328343
all_blocks: dict[bytes32, BlockRecord] = {}
329-
async with self.db_wrapper.reader_no_transaction() as conn:
330-
async with conn.execute(
331-
"SELECT header_hash,block_record "
332-
"FROM full_blocks "
333-
f"WHERE header_hash in ({'?,' * (len(header_hashes) - 1)}?)",
334-
header_hashes,
335-
) as cursor:
336-
for row in await cursor.fetchall():
337-
block_rec = BlockRecord.from_bytes(row[1])
338-
all_blocks[block_rec.header_hash] = block_rec
344+
for batch in to_batches(header_hashes, self.db_wrapper.host_parameter_limit):
345+
async with self.db_wrapper.reader_no_transaction() as conn:
346+
async with conn.execute(
347+
"SELECT header_hash,block_record FROM full_blocks "
348+
f"WHERE header_hash in ({'?,' * (len(batch.entries) - 1)}?)",
349+
batch.entries,
350+
) as cursor:
351+
for row in await cursor.fetchall():
352+
block_rec = BlockRecord.from_bytes(row[1])
353+
all_blocks[block_rec.header_hash] = block_rec
339354

340355
ret: list[BlockRecord] = []
341356
for hh in header_hashes:
@@ -468,7 +483,8 @@ async def get_block_bytes_in_range(
468483
No orphan blocks.
469484
"""
470485

471-
assert self.db_wrapper.db_version == 2
486+
if self.db_wrapper.db_version != 2:
487+
raise UnsupportedDatabaseVersionError("get_block_bytes_in_range requires DB version 2")
472488
async with self.db_wrapper.reader_no_transaction() as conn:
473489
async with conn.execute(
474490
"SELECT block FROM full_blocks WHERE height >= ? AND height <= ? AND in_main_chain=1",

chia/full_node/full_node_api.py

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
from chia.consensus.get_block_generator import get_block_generator
4242
from chia.consensus.pot_iterations import calculate_ip_iters, calculate_iterations_quality, calculate_sp_iters
4343
from chia.consensus.signage_point import SignagePoint
44+
from chia.full_node.block_store import UnsupportedDatabaseVersionError
4445
from chia.full_node.coin_store import CoinStore
4546
from chia.full_node.fee_estimator_interface import FeeEstimatorInterface
4647
from chia.full_node.full_block_utils import get_height_and_tx_status_from_block, header_block_from_block
@@ -1453,15 +1454,15 @@ async def request_block_headers(self, request: wallet_protocol.RequestBlockHeade
14531454

14541455
if request.end_height < request.start_height or request.end_height - request.start_height > 128:
14551456
return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
1456-
if self.full_node.block_store.db_wrapper.db_version == 2:
1457-
try:
1458-
blocks_bytes = await self.full_node.block_store.get_block_bytes_in_range(
1459-
request.start_height, request.end_height
1460-
)
1461-
except ValueError:
1462-
return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
1463-
1464-
else:
1457+
try:
1458+
blocks_bytes = await self.full_node.block_store.get_block_bytes_in_range(
1459+
request.start_height, request.end_height
1460+
)
1461+
except ValueError:
1462+
return make_msg(ProtocolMessageTypes.reject_block_headers, reject)
1463+
except UnsupportedDatabaseVersionError:
1464+
# The underlying block store may not support this optimized call
1465+
# (e.g. v1 DB). In this case, we fall back to the legacy approach
14651466
height_to_hash = self.full_node.blockchain.height_to_hash
14661467
header_hashes: list[bytes32] = []
14671468
for i in range(request.start_height, request.end_height + 1):

chia/simulator/full_node_simulator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ async def revert_block_height(self, new_height: uint32) -> None:
167167
raise ValueError("Cannot revert to a height less than 1.")
168168
block_record: BlockRecord = self.full_node.blockchain.height_to_block_record(new_height)
169169
# remove enough data to allow a bunch of blocks to be wiped.
170-
async with self.full_node.block_store.db_wrapper.writer():
170+
async with self.full_node.block_store.transaction():
171171
# set coinstore
172172
await self.full_node.coin_store.rollback_to_block(new_height)
173173
# set blockstore to new height

0 commit comments

Comments
 (0)