Make BlockStore.replace_proof own its write transaction (#21166)

* Make BlockStore.replace_proof own its write transaction

Co-authored-by: Cursor <cursoragent@cursor.com>

* Add failure-path test for _replace_proof; fix log wording

The except block was re-indented by the previous commit and had no
coverage. The log message said "error while adding block, rolling
back", but this path replaces a compact proof and the rollback happens
inside the store's transaction now.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

chia/_tests/core/full_node/test_full_node.py

@@ -3081,6 +3081,35 @@ async def test_compact_protocol_invalid_messages(
         assert not block.challenge_chain_ip_proof.normalized_to_identity
 
 
+@pytest.mark.limit_consensus_modes(reason="save time")
+@pytest.mark.anyio
+async def test_replace_proof_failure_logs_and_reraises(
+    one_node_one_block: tuple[FullNodeSimulator, ChiaServer, BlockTools],
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    full_node_api, _server, _bt = one_node_one_block
+    full_node = full_node_api.full_node
+    peak_block = await full_node.blockchain.get_full_peak()
+    assert peak_block is not None
+
+    async def failing_replace_proof(header_hash: bytes32, block: FullBlock) -> None:
+        raise RuntimeError("injected failure")
+
+    monkeypatch.setattr(full_node.block_store, "replace_proof", failing_replace_proof)
+
+    # the block's own vdf_info matches, so we reach the replace_proof call
+    with caplog.at_level(logging.ERROR):
+        with pytest.raises(RuntimeError, match="injected failure"):
+            await full_node._replace_proof(
+                peak_block.reward_chain_block.challenge_chain_ip_vdf,
+                peak_block.challenge_chain_ip_proof,
+                peak_block.header_hash,
+                CompressibleVDFField.CC_IP_VDF,
+            )
+    assert "error replacing proof" in caplog.text
+
+
 @pytest.mark.anyio
 @pytest.mark.parametrize("trusted", [True, False])
 async def test_unsolicited_compact_vdf(

chia/full_node/block_store.py

@@ -125,7 +125,8 @@ class BlockStore:
 
         self.block_cache.put(header_hash, block)
 
-        async with self.db_wrapper.writer_maybe_transaction() as conn:
+        # this method owns its write transaction, callers don't need to open one
+        async with self.db_wrapper.writer() as conn:
             await conn.execute(
                 "UPDATE full_blocks SET block=?,is_fully_compactified=? WHERE header_hash=?",
                 (

chia/full_node/full_node.py

@@ -3276,16 +3276,15 @@ class FullNode:
                 new_block = block.replace(challenge_chain_ip_proof=vdf_proof)
         if new_block is None:
             return False
-        async with self.db_wrapper.writer():
-            try:
-                await self.block_store.replace_proof(header_hash, new_block)
-                return True
-            except BaseException as e:
-                self.log.error(
-                    f"_replace_proof error while adding block {block.header_hash} height {block.height},"
-                    f" rolling back: {e} {traceback.format_exc()}"
-                )
-                raise
+        try:
+            await self.block_store.replace_proof(header_hash, new_block)
+            return True
+        except BaseException as e:
+            self.log.error(
+                f"_replace_proof error replacing proof for block {block.header_hash} height {block.height}:"
+                f" {e} {traceback.format_exc()}"
+            )
+            raise
 
     async def add_compact_proof_of_time(self, request: timelord_protocol.RespondCompactProofOfTime) -> None:
         peak = self.blockchain.get_peak()