Handle nested context manager

Author: richardkiss

chia/_tests/blockchain/blockchain_test_utils.py

@@ -17,24 +17,18 @@
 async def check_block_store_invariant(bc: Blockchain):
     in_chain = set()
     max_height = -1
-    async with bc.consensus_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:
-                height = row[0]
-
-                # if this block is in-chain, ensure we haven't found another block
-                # at this height that's also in chain. That would be an invariant
-                # violation
-                if row[1]:
-                    # make sure we don't have any duplicate heights. Each block
-                    # height can only have a single block with in_main_chain set
-                    assert height not in in_chain
-                    in_chain.add(height)
-                    max_height = max(max_height, height)
-
-            # make sure every height is represented in the set
-            assert len(in_chain) == max_height + 1
+    async for height in bc.consensus_store.get_block_heights_in_main_chain():
+        # if this block is in-chain, ensure we haven't found another block
+        # at this height that's also in chain. That would be an invariant
+        # violation
+        # make sure we don't have any duplicate heights. Each block
+        # height can only have a single block with in_main_chain set
+        assert height not in in_chain
+        in_chain.add(height)
+        max_height = max(max_height, height)
+
+    # make sure every height is represented in the set
+    assert len(in_chain) == max_height + 1
 
 
 async def _validate_and_add_block(

chia/_tests/core/full_node/stores/test_block_store.py

@@ -410,7 +410,7 @@ def rand_vdf_proof() -> VDFProof:
 
             # make sure we get the same result when we hit the database
             # itself (and not just the block cache)
-            block_store.rollback_cache_block(block.header_hash)
+            consensus_store.rollback_cache_block(block.header_hash)
             b = await block_store.get_full_block(block.header_hash)
             assert b is not None
             assert b.challenge_chain_ip_proof == proof

chia/consensus/blockchain.py

@@ -413,9 +413,9 @@ async def add_block(
 
         try:
             # Always add the block to the database
-            async with self.consensus_store.transaction():
+            async with self.consensus_store as writer:
                 # Perform the DB operations to update the state, and rollback if something goes wrong
-                await self.consensus_store.add_full_block(header_hash, block, block_record)
+                await writer.add_full_block(header_hash, block, block_record)
                 records, state_change_summary = await self._reconsider_peak(block_record, genesis, fork_info)
 
                 # Then update the memory cache. It is important that this is not cancelled and does not throw

chia/consensus/consensus_store_protocol.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Collection
-from typing import Optional, Protocol
+from typing import AsyncIterator, Optional, Protocol
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
@@ -22,7 +22,6 @@ class ConsensusStoreWriteProtocol(Protocol):
 
     # 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: ...
@@ -80,6 +79,7 @@ async def get_sub_epoch_challenge_segments(
     ) -> 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]: ...
+    def get_block_heights_in_main_chain(self) -> AsyncIterator[int]: ...
 
     # Coin store reads
     async def get_coin_records(self, names: Collection[bytes32]) -> list[CoinRecord]: ...
@@ -94,3 +94,4 @@ def get_hash(self, height: uint32) -> bytes32: ...
     def rollback_height_map(self, height: uint32) -> None: ...
     def update_height_map(self, height: uint32, block_hash: bytes32, ses: Optional[SubEpochSummary]) -> None: ...
     async def maybe_flush_height_map(self) -> None: ...
+    def rollback_cache_block(self, header_hash: bytes32) -> None: ...

chia/full_node/consensus_store_sqlite3.py

@@ -3,7 +3,7 @@
 import dataclasses
 from collections.abc import Collection
 from contextlib import AbstractAsyncContextManager
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, AsyncIterator, Optional, TYPE_CHECKING
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
@@ -24,9 +24,6 @@ def __init__(self, block_store: BlockStore, coin_store: CoinStore):
     async def add_full_block(self, header_hash: bytes32, block: FullBlock, block_record: BlockRecord) -> None:
         await self._block_store.add_full_block(header_hash, block, block_record)
 
-    def rollback_cache_block(self, header_hash: bytes32) -> None:
-        self._block_store.rollback_cache_block(header_hash)
-
     async def rollback(self, height: int) -> None:
         await self._block_store.rollback(height)
 
@@ -54,6 +51,7 @@ async def new_block(
     ) -> None:
         await self._coin_store.new_block(height, timestamp, included_reward_coins, tx_additions, tx_removals)
 
+
 @dataclasses.dataclass
 class ConsensusStoreSQLite3:
     """
@@ -64,9 +62,10 @@ class ConsensusStoreSQLite3:
     coin_store: CoinStore
     height_map: BlockHeightMap
 
-    # Writer context and writer facade for transactional writes
+    # Writer context and writer facade for transactional writes (re-entrant via depth counter)
     _writer_ctx: Optional[AbstractAsyncContextManager[Any]] = None
     _writer: Optional[Any] = None
+    _txn_depth: int = 0
 
     @classmethod
     async def create(
@@ -88,22 +87,34 @@ async def create(
 
     # Async context manager yielding a writer for atomic writes
     async def __aenter__(self):
-        # Begin a transaction via the block_store. CoinStore shares the same DB.
-        self._writer_ctx = self.block_store.transaction()
-        await self._writer_ctx.__aenter__()
-        # Create and return a writer facade bound to this transaction
-        self._writer = ConsensusStoreSQLite3Writer(self.block_store, self.coin_store)
-        return self._writer
+        # Re-entrant async context manager:
+        # Begin a transaction only at the outermost level. CoinStore shares the same DB.
+        if self._txn_depth == 0:
+            self._writer_ctx = self.block_store.transaction()
+            await self._writer_ctx.__aenter__()
+            # Create writer facade bound to this transaction
+            self._writer = ConsensusStoreSQLite3Writer(self.block_store, self.coin_store)
+        self._txn_depth += 1
+        return self._writer  # Return the same writer for nested contexts
 
     async def __aexit__(self, exc_type, exc, tb):
-        # Commit on success, rollback on exception handled by transaction manager
         try:
-            if self._writer_ctx is not None:
-                return await self._writer_ctx.__aexit__(exc_type, exc, tb)
-            return None
+            # Check if we're at the outermost level before decrementing
+            if self._txn_depth == 1:
+                # This is the outermost context, handle transaction exit
+                if self._writer_ctx is not None:
+                    return await self._writer_ctx.__aexit__(exc_type, exc, tb)
+                return None
+            else:
+                # This is a nested context, just return None (don't suppress exceptions)
+                return None
         finally:
-            self._writer_ctx = None
-            self._writer = None
+            # Always decrement depth and clean up if we're at the outermost level
+            if self._txn_depth > 0:
+                self._txn_depth -= 1
+            if self._txn_depth == 0:
+                self._writer_ctx = None
+                self._writer = None
 
     # Block store methods
 
@@ -154,6 +165,16 @@ async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]:
     async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]:
         return await self.coin_store.get_coins_removed_at_height(height)
 
+    def get_block_heights_in_main_chain(self) -> AsyncIterator[int]:
+        async def gen():
+            async with self.block_store.transaction() as conn:
+                async with conn.execute("SELECT height, in_main_chain FROM full_blocks") as cursor:
+                    async for row in cursor:
+                        if row[1]:
+                            yield row[0]
+
+        return gen()
+
     # Height map methods
     def get_ses_heights(self) -> list[uint32]:
         return self.height_map.get_ses_heights()
@@ -179,9 +200,12 @@ async def maybe_flush_height_map(self) -> None:
         # BlockHeightMap.maybe_flush is asynchronous
         await self.height_map.maybe_flush()
 
-if TYPE_CHECKING:
+    def rollback_cache_block(self, header_hash: bytes32) -> None:
+        self.block_store.rollback_cache_block(header_hash)
 
-    from typing import  cast
+
+if TYPE_CHECKING:
+    from typing import cast
     from chia.consensus.consensus_store_protocol import ConsensusStoreProtocol
 
     def _protocol_check(o: ConsensusStoreProtocol) -> None: ...