first crack at async context manager for ConsensusStore

Author: richardkiss

chia/consensus/blockchain.py

@@ -538,56 +538,57 @@ async def _reconsider_peak(
         else:
             records_to_add = await self.consensus_store.get_block_records_by_hash(fork_info.block_hashes)
 
-        for fetched_block_record in records_to_add:
-            if not fetched_block_record.is_transaction_block:
-                # Coins are only created in TX blocks so there are no state updates for this block
-                continue
-
-            height = fetched_block_record.height
-            # We need to recompute the additions and removals, since they are
-            # not stored on DB. We have all the additions and removals in the
-            # fork_info object, we just need to pick the ones belonging to each
-            # individual block height
-
-            # Apply the coin store changes for each block that is now in the blockchain
-            included_reward_coins = [
-                fork_add.coin
-                for fork_add in fork_info.additions_since_fork.values()
-                if fork_add.confirmed_height == height and fork_add.is_coinbase
-            ]
-            tx_additions = [
-                (coin_id, fork_add.coin, fork_add.same_as_parent)
-                for coin_id, fork_add in fork_info.additions_since_fork.items()
-                if fork_add.confirmed_height == height and not fork_add.is_coinbase
-            ]
-            tx_removals = [
-                coin_id for coin_id, fork_rem in fork_info.removals_since_fork.items() if fork_rem.height == height
-            ]
-            assert fetched_block_record.timestamp is not None
-            await self.consensus_store.new_block(
-                height,
-                fetched_block_record.timestamp,
-                included_reward_coins,
-                tx_additions,
-                tx_removals,
-            )
-            if self._log_coins and (len(tx_removals) > 0 or len(tx_additions) > 0):
-                log.info(
-                    f"adding new block to coin_store "
-                    f"(hh: {fetched_block_record.header_hash} "
-                    f"height: {fetched_block_record.height}), {len(tx_removals)} spends"
+        async with self.consensus_store as writer:
+            for fetched_block_record in records_to_add:
+                if not fetched_block_record.is_transaction_block:
+                    # Coins are only created in TX blocks so there are no state updates for this block
+                    continue
+
+                height = fetched_block_record.height
+                # We need to recompute the additions and removals, since they are
+                # not stored on DB. We have all the additions and removals in the
+                # fork_info object, we just need to pick the ones belonging to each
+                # individual block height
+
+                # Apply the coin store changes for each block that is now in the blockchain
+                included_reward_coins = [
+                    fork_add.coin
+                    for fork_add in fork_info.additions_since_fork.values()
+                    if fork_add.confirmed_height == height and fork_add.is_coinbase
+                ]
+                tx_additions = [
+                    (coin_id, fork_add.coin, fork_add.same_as_parent)
+                    for coin_id, fork_add in fork_info.additions_since_fork.items()
+                    if fork_add.confirmed_height == height and not fork_add.is_coinbase
+                ]
+                tx_removals = [
+                    coin_id for coin_id, fork_rem in fork_info.removals_since_fork.items() if fork_rem.height == height
+                ]
+                assert fetched_block_record.timestamp is not None
+                await writer.new_block(
+                    height,
+                    fetched_block_record.timestamp,
+                    included_reward_coins,
+                    tx_additions,
+                    tx_removals,
                 )
-                log.info("rewards: %s", ",".join([add.name().hex()[0:6] for add in included_reward_coins]))
-                log.info("additions: %s", ",".join([add[0].hex()[0:6] for add in tx_additions]))
-                log.info("removals: %s", ",".join([f"{rem}"[0:6] for rem in tx_removals]))
+                if self._log_coins and (len(tx_removals) > 0 or len(tx_additions) > 0):
+                    log.info(
+                        f"adding new block to coin_store "
+                        f"(hh: {fetched_block_record.header_hash} "
+                        f"height: {fetched_block_record.height}), {len(tx_removals)} spends"
+                    )
+                    log.info("rewards: %s", ",".join([add.name().hex()[0:6] for add in included_reward_coins]))
+                    log.info("additions: %s", ",".join([add[0].hex()[0:6] for add in tx_additions]))
+                    log.info("removals: %s", ",".join([f"{rem}"[0:6] for rem in tx_removals]))
 
-        # we made it to the end successfully
-        # Rollback sub_epoch_summaries
-        await self.consensus_store.rollback(fork_info.fork_height)
-        await self.consensus_store.set_in_chain([(br.header_hash,) for br in records_to_add])
+            # we made it to the end successfully
+            # Rollback sub_epoch_summaries
+            await writer.rollback(fork_info.fork_height)
+            await writer.set_in_chain([(br.header_hash,) for br in records_to_add])
 
-        # Changes the peak to be the new peak
-        await self.consensus_store.set_peak(block_record.header_hash)
+            # Changes the peak to be the new peak
+            await writer.set_peak(block_record.header_hash)
 
         return records_to_add, StateChangeSummary(
             block_record,

chia/consensus/consensus_store_protocol.py

@@ -1,10 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Collection
-from contextlib import AbstractAsyncContextManager
 from typing import Optional, Protocol
 
-import aiosqlite
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
@@ -13,54 +11,82 @@
 from chia.types.coin_record import CoinRecord
 
 
+class ConsensusStoreWriteProtocol(Protocol):
+    """
+    Protocol for performing mutating operations on the consensus store.
+
+    Instances implementing this protocol should be acquired via the async
+    context manager on ConsensusStoreProtocol to ensure atomic write
+    operations (e.g., wrapping all writes in a single DB transaction).
+    """
+
+    # Block store writes
+    async def add_full_block(self, header_hash: bytes32, block: FullBlock, block_record: BlockRecord) -> None: ...
+    def rollback_cache_block(self, header_hash: bytes32) -> None: ...
+    async def rollback(self, height: int) -> None: ...
+    async def set_in_chain(self, header_hashes: list[tuple[bytes32]]) -> None: ...
+    async def set_peak(self, header_hash: bytes32) -> None: ...
+    async def persist_sub_epoch_challenge_segments(
+        self, ses_block_hash: bytes32, segments: list[SubEpochChallengeSegment]
+    ) -> None: ...
+
+    # Coin store writes
+    async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]: ...
+    async def new_block(
+        self,
+        height: uint32,
+        timestamp: uint64,
+        included_reward_coins: Collection[Coin],
+        tx_additions: Collection[tuple[bytes32, Coin, bool]],
+        tx_removals: list[bytes32],
+    ) -> None: ...
+
+
 class ConsensusStoreProtocol(Protocol):
     """
-    Protocol for the consensus store, which provides methods to interact with
-    the consensus-related data in the blockchain.
+    Read-only protocol for the consensus store.
+
+    This protocol also acts as an async context manager. Entering the context
+    yields a ConsensusStoreWriteProtocol instance, which must be used for
+    performing write (mutating) operations. This ensures atomic writes and
+    makes it harder to accidentally perform writes outside a transaction.
+
+    Example usage:
+        async with store as writer:
+            await writer.add_full_block(...)
+            await writer.set_peak(...)
+
+        # Outside the context, only read methods are available
+        br = await store.get_block_record(header_hash)
     """
 
-    # Block store methods
-    def transaction(self) -> AbstractAsyncContextManager[aiosqlite.Connection]: ...
+    # Async context manager methods
+    async def __aenter__(self) -> ConsensusStoreWriteProtocol: ...
+    async def __aexit__(self, exc_type, exc, tb) -> Optional[bool]: ...
 
+    # Block store reads
     async def get_block_records_close_to_peak(
         self, blocks_n: int
     ) -> tuple[dict[bytes32, BlockRecord], Optional[bytes32]]: ...
     async def get_full_block(self, header_hash: bytes32) -> Optional[FullBlock]: ...
-    async def add_full_block(self, header_hash: bytes32, block: FullBlock, block_record: BlockRecord) -> None: ...
-    def rollback_cache_block(self, header_hash: bytes32) -> None: ...
     async def get_block_records_by_hash(self, header_hashes: list[bytes32]) -> list[BlockRecord]: ...
-    async def rollback(self, height: int) -> 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_in_range(self, start: int, stop: int) -> dict[bytes32, BlockRecord]: ...
     def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]: ...
     async def get_blocks_by_hash(self, header_hashes: list[bytes32]) -> list[FullBlock]: ...
     async def get_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]: ...
     async def get_prev_hash(self, header_hash: bytes32) -> bytes32: ...
-    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 get_generator(self, header_hash: bytes32) -> Optional[bytes]: ...
     async def get_generators_at(self, heights: set[uint32]) -> dict[uint32, bytes]: ...
 
-    # Coin store methods
+    # Coin store reads
     async def get_coin_records(self, names: Collection[bytes32]) -> list[CoinRecord]: ...
-    async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]: ...
-    async def new_block(
-        self,
-        height: uint32,
-        timestamp: uint64,
-        included_reward_coins: Collection[Coin],
-        tx_additions: Collection[tuple[bytes32, Coin, bool]],
-        tx_removals: list[bytes32],
-    ) -> None: ...
     async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]: ...
     async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]: ...
 
-    # Height map methods
+    # Height map methods (kept here for now; non-async and maybe_flush remain on read protocol)
     def get_ses_heights(self) -> list[uint32]: ...
     def get_ses(self, height: uint32) -> SubEpochSummary: ...
     def contains_height(self, height: uint32) -> bool: ...