improve handling of UnfinishedBlocks

chia/full_node/full_node_store.py

@@ -38,6 +38,13 @@ class FullNodeStorePeakResult(Streamable):
     new_infusion_points: List[timelord_protocol.NewInfusionPointVDF]
 
 
+@dataclasses.dataclass
+class UnfinishedBlockEntry:
+    unfinished_block: UnfinishedBlock
+    result: PreValidationResult
+    height: uint32
+
+
 class FullNodeStore:
     constants: ConsensusConstants
 
@@ -52,7 +59,10 @@ class FullNodeStore:
     seen_unfinished_blocks: Dict[bytes32, None]
 
     # Unfinished blocks, keyed from reward hash
-    unfinished_blocks: Dict[bytes32, Tuple[uint32, UnfinishedBlock, PreValidationResult]]
+    # There may be multiple different unfinished blocks with the same partial
+    # hash (reward chain block hash). They are stored under their partial hash
+    # though. The inner dictionary uses the foliage hash as the key
+    unfinished_blocks: Dict[bytes32, Dict[Optional[bytes32], UnfinishedBlockEntry]]
 
     # Finished slots and sps from the peak's slot onwards
     # We store all 32 SPs for each slot, starting as 32 Nones and filling them as we go
@@ -164,32 +174,74 @@ def seen_unfinished_block(self, object_hash: bytes32) -> bool:
     def add_unfinished_block(
         self, height: uint32, unfinished_block: UnfinishedBlock, result: PreValidationResult
     ) -> None:
-        self.unfinished_blocks[unfinished_block.partial_hash] = (height, unfinished_block, result)
+        partial_hash = unfinished_block.partial_hash
+        entry = self.unfinished_blocks.setdefault(partial_hash, {})
+        entry[unfinished_block.foliage.foliage_transaction_block_hash] = UnfinishedBlockEntry(
+            unfinished_block, result, height
+        )
 
     def get_unfinished_block(self, unfinished_reward_hash: bytes32) -> Optional[UnfinishedBlock]:
         result = self.unfinished_blocks.get(unfinished_reward_hash, None)
         if result is None:
             return None
-        return result[1]
+        # The old API doesn't distinguish between duplicate UnfinishedBlocks,
+        # just return the first one
+        return next(iter(result.values())).unfinished_block
+
+    def get_unfinished_block2(
+        self, unfinished_reward_hash: bytes32, unfinished_foliage_hash: Optional[bytes32]
+    ) -> Tuple[Optional[UnfinishedBlock], int]:
+        result = self.unfinished_blocks.get(unfinished_reward_hash, None)
+        if result is None:
+            return None, 0
+        if unfinished_foliage_hash is None:
+            return next(iter(result.values())).unfinished_block, len(result)
+        else:
+            entry = result.get(unfinished_foliage_hash)
+            return (None if entry is None else entry.unfinished_block), len(result)
 
     def get_unfinished_block_result(self, unfinished_reward_hash: bytes32) -> Optional[PreValidationResult]:
         result = self.unfinished_blocks.get(unfinished_reward_hash, None)
         if result is None:
             return None
-        return result[2]
+        return next(iter(result.values())).result
+
+    def get_unfinished_block_result2(
+        self, unfinished_reward_hash: bytes32, unfinished_foliage_hash: Optional[bytes32]
+    ) -> Optional[PreValidationResult]:
+        result = self.unfinished_blocks.get(unfinished_reward_hash, None)
+        if result is None:
+            return None
+        if unfinished_foliage_hash is None:
+            return next(iter(result.values())).result
+        else:
+            entry = result.get(unfinished_foliage_hash)
+            return None if entry is None else entry.result
 
     # returns all unfinished blocks for the specified height
     def get_unfinished_blocks(self, height: uint32) -> List[UnfinishedBlock]:
-        return [block for ub_height, block, _ in self.unfinished_blocks.values() if ub_height == height]
+        ret: List[UnfinishedBlock] = []
+        for entry in self.unfinished_blocks.values():
+            for ube in entry.values():
+                if ube.height == height:
+                    ret.append(ube.unfinished_block)
+        return ret
 
     def clear_unfinished_blocks_below(self, height: uint32) -> None:
-        del_keys: List[bytes32] = []
-        for partial_reward_hash, (unf_height, unfinished_block, _) in self.unfinished_blocks.items():
-            if unf_height < height:
-                del_keys.append(partial_reward_hash)
-        for del_key in del_keys:
-            del self.unfinished_blocks[del_key]
-
+        del_partial: List[bytes32] = []
+        for partial_hash, entry in self.unfinished_blocks.items():
+            del_foliage: List[Optional[bytes32]] = []
+            for foliage_hash, ube in entry.items():
+                if ube.height < height:
+                    del_foliage.append(foliage_hash)
+            for fh in del_foliage:
+                del entry[fh]
+            if len(entry) == 0:
+                del_partial.append(partial_hash)
+        for ph in del_partial:
+            del self.unfinished_blocks[ph]
+
+    # TODO: this should be removed. It's only used by a test
     def remove_unfinished_block(self, partial_reward_hash: bytes32) -> None:
         if partial_reward_hash in self.unfinished_blocks:
             del self.unfinished_blocks[partial_reward_hash]