fix recovery of invalid sub epoch summaries cache (#16333)

chia/full_node/block_height_map.py

@@ -112,18 +112,18 @@ class BlockHeightMap:
 
         self.__first_dirty = height + 1
 
-        # if the peak hash is already in the height-to-hash map, we don't need
-        # to load anything more from the DB
         if self.get_hash(height) != peak:
             self.__set_hash(height, peak)
 
-            if row[3] is not None:
-                self.__sub_epoch_summaries[height] = row[3]
+        if row[3] is not None:
+            self.__sub_epoch_summaries[height] = row[3]
 
-            # prepopulate the height -> hash mapping
-            await self._load_blocks_from(height, prev_hash)
+        # prepopulate the height -> hash mapping
+        # run this unconditionally in to ensure both the height-to-hash and sub
+        # epoch summaries caches are in sync with the DB
+        await self._load_blocks_from(height, prev_hash)
 
-            await self.maybe_flush()
+        await self.maybe_flush()
 
         return self
 
@@ -163,14 +163,22 @@ class BlockHeightMap:
     # load height-to-hash map entries from the DB starting at height back in
     # time until we hit a match in the existing map, at which point we can
     # assume all previous blocks have already been populated
+    # the first iteration is mandatory on each startup, so we make it load fewer
+    # blocks to be fast. The common case is that the files are in sync with the
+    # DB so iteration can stop early.
     async def _load_blocks_from(self, height: uint32, prev_hash: bytes32) -> None:
+        # on mainnet, every 384th block has a sub-epoch summary. This should
+        # guarantee that we find at least one in the first iteration. If it
+        # matches, we're done reconciliating the cache with the DB.
+        window_size = 400
         while height > 0:
             # load 5000 blocks at a time
-            window_end = max(0, height - 5000)
+            window_end = max(0, height - window_size)
+            window_size = 5000
 
             query = (
                 "SELECT header_hash,prev_hash,height,sub_epoch_summary from full_blocks "
-                "INDEXED BY height WHERE height>=? AND height <?"
+                "INDEXED BY height WHERE in_main_chain=1 AND height>=? AND height <?"
             )
 
             async with self.db.reader_no_transaction() as conn:
@@ -195,6 +203,9 @@ class BlockHeightMap:
                         and height in self.__sub_epoch_summaries
                         and self.__sub_epoch_summaries[height] == entry[2]
                     ):
+                        # we only terminate the loop if we encounter a block
+                        # that has a sub epoch summary matching the cache and
+                        # the block hash matches the cache
                         return
                     self.__sub_epoch_summaries[height] = entry[2]
                 elif height in self.__sub_epoch_summaries:

tests/core/full_node/test_block_height_map.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import struct
 from pathlib import Path
 from typing import Optional
@@ -30,11 +31,12 @@ async def new_block(
 ) -> None:
     async with db.writer_maybe_transaction() as conn:
         cursor = await conn.execute(
-            "INSERT INTO full_blocks VALUES(?, ?, ?, ?)",
+            "INSERT INTO full_blocks VALUES(?, ?, ?, ?, ?)",
             (
                 block_hash,
                 parent,
                 height,
+                True,  # in_main_chain
                 # sub epoch summary
                 None if ses is None else bytes(ses),
             ),
@@ -52,6 +54,7 @@ async def setup_db(db: DBWrapper2) -> None:
             "header_hash blob PRIMARY KEY,"
             "prev_hash blob,"
             "height bigint,"
+            "in_main_chain tinyint,"
             "sub_epoch_summary blob)"
         )
         await conn.execute("CREATE TABLE IF NOT EXISTS current_peak(key int PRIMARY KEY, hash blob)")
@@ -80,7 +83,9 @@ async def setup_chain(
         peak_hash = gen_block_hash(height + chain_id * 65536)
 
     # we only set is_peak=1 for chain_id 0
-    await new_block(db, peak_hash, parent_hash, height, chain_id == 0, None)
+    if ses_every is not None and height % ses_every == 0:
+        ses = gen_ses(height)
+    await new_block(db, peak_hash, parent_hash, height, chain_id == 0, ses)
 
 
 class TestBlockHeightMap:
@@ -113,18 +118,19 @@ class TestBlockHeightMap:
             for height in reversed(range(10000)):
                 assert height_map.get_hash(uint32(height)) == gen_block_hash(height)
 
+    @pytest.mark.parametrize("ses_every", [20, 1])
     @pytest.mark.asyncio
-    async def test_save_restore(self, tmp_dir: Path, db_version: int) -> None:
+    async def test_save_restore(self, ses_every: int, tmp_dir: Path, db_version: int) -> None:
         async with DBConnection(db_version) as db_wrapper:
             await setup_db(db_wrapper)
-            await setup_chain(db_wrapper, 10000, ses_every=20)
+            await setup_chain(db_wrapper, 10000, ses_every=ses_every)
 
             height_map = await BlockHeightMap.create(tmp_dir, db_wrapper)
 
             for height in reversed(range(10000)):
                 assert height_map.contains_height(uint32(height))
                 assert height_map.get_hash(uint32(height)) == gen_block_hash(height)
-                if (height % 20) == 0:
+                if (height % ses_every) == 0:
                     assert height_map.get_ses(uint32(height)) == gen_ses(height)
                 else:
                     with pytest.raises(KeyError) as _:
@@ -142,13 +148,13 @@ class TestBlockHeightMap:
             async with db_wrapper.writer_maybe_transaction() as conn:
                 await conn.execute("DROP TABLE full_blocks")
             await setup_db(db_wrapper)
-            await setup_chain(db_wrapper, 10000, ses_every=20, start_height=9970)
+            await setup_chain(db_wrapper, 10000, ses_every=ses_every, start_height=9970)
             height_map = await BlockHeightMap.create(tmp_dir, db_wrapper)
 
             for height in reversed(range(10000)):
                 assert height_map.contains_height(uint32(height))
                 assert height_map.get_hash(uint32(height)) == gen_block_hash(height)
-                if (height % 20) == 0:
+                if (height % ses_every) == 0:
                     assert height_map.get_ses(uint32(height)) == gen_ses(height)
                 else:
                     with pytest.raises(KeyError) as _:
@@ -187,6 +193,38 @@ class TestBlockHeightMap:
                     with pytest.raises(KeyError) as _:
                         height_map.get_ses(uint32(height))
 
+    @pytest.mark.asyncio
+    async def test_restore_ses_only(self, tmp_dir: Path, db_version: int) -> None:
+        # this is a test where the height-to-hash is complete and correct but
+        # sub epoch summaries are missing. We need to be able to restore them in
+        # this case.
+        async with DBConnection(db_version) as db_wrapper:
+            await setup_db(db_wrapper)
+            await setup_chain(db_wrapper, 2000, ses_every=20)
+
+            height_map = await BlockHeightMap.create(tmp_dir, db_wrapper)
+            await height_map.maybe_flush()
+            del height_map
+
+            # corrupt the sub epoch cache
+            ses_cache = []
+            for i in range(0, 2000, 19):
+                ses_cache.append((uint32(i), bytes(gen_ses(i + 9999))))
+
+            await write_file_async(tmp_dir / "sub-epoch-summaries", bytes(SesCache(ses_cache)))
+
+            # the test starts here
+            height_map = await BlockHeightMap.create(tmp_dir, db_wrapper)
+
+            for height in reversed(range(2000)):
+                assert height_map.contains_height(uint32(height))
+                assert height_map.get_hash(uint32(height)) == gen_block_hash(height)
+                if (height % 20) == 0:
+                    assert height_map.get_ses(uint32(height)) == gen_ses(height)
+                else:
+                    with pytest.raises(KeyError) as _:
+                        height_map.get_ses(uint32(height))
+
     @pytest.mark.asyncio
     async def test_restore_extend(self, tmp_dir: Path, db_version: int) -> None:
         # test the case where the cache has fewer blocks than the DB, and that
@@ -450,6 +488,32 @@ class TestBlockHeightMap:
                 for idx in range(0, len(heights), 32):
                     assert new_heights[idx : idx + 32] == heights[idx : idx + 32]
 
+    @pytest.mark.asyncio
+    async def test_cache_file_truncate(self, tmp_dir: Path, db_version: int) -> None:
+        # Test the case where the cache has more blocks than the DB, the cache
+        # file will be truncated
+        async with DBConnection(db_version) as db_wrapper:
+            await setup_db(db_wrapper)
+            await setup_chain(db_wrapper, 2000, ses_every=20)
+            bh = await BlockHeightMap.create(tmp_dir, db_wrapper)
+            await bh.maybe_flush()
+
+            # extend the cache file
+            with open(tmp_dir / "height-to-hash", "r+b") as f:
+                f.truncate(32 * 4000)
+            assert os.path.getsize(tmp_dir / "height-to-hash") == 32 * 4000
+
+            bh = await BlockHeightMap.create(tmp_dir, db_wrapper)
+            await bh.maybe_flush()
+
+            with open(tmp_dir / "height-to-hash", "rb") as f:
+                new_heights = f.read()
+                assert len(new_heights) == 4000 * 32
+                # pytest doesn't behave very well comparing large buffers
+                # (when the test fails). Compare small portions at a time instead
+                for idx in range(0, 2000):
+                    assert new_heights[idx * 32 : idx * 32 + 32] == gen_block_hash(idx)
+
 
 @pytest.mark.asyncio
 async def test_unsupported_version(tmp_dir: Path) -> None: