fix mempool TX cache cost (#8054)

* fix issue where the cost of the mempool TX cache (for spend bundles that can't be included in a block yet) would not be reset when the cache was emptied

* factor out the Pending TX cache from mempool, to allow unit testing

chia/full_node/mempool_manager.py

@@ -16,6 +16,7 @@ from chia.full_node.bundle_tools import simple_solution_generator
 from chia.full_node.coin_store import CoinStore
 from chia.full_node.mempool import Mempool
 from chia.full_node.mempool_check_conditions import mempool_check_conditions_dict, get_name_puzzle_conditions
+from chia.full_node.pending_tx_cache import PendingTxCache
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.program import SerializedProgram
 from chia.types.blockchain_format.sized_bytes import bytes32
@@ -52,8 +53,6 @@ class MempoolManager:
         self.constants: ConsensusConstants = consensus_constants
         self.constants_json = recurse_jsonify(dataclasses.asdict(self.constants))
 
-        # Transactions that were unable to enter mempool, used for retry. (they were invalid)
-        self.potential_txs: Dict[bytes32, MempoolItem] = {}
         # Keep track of seen spend_bundles
         self.seen_bundle_hashes: Dict[bytes32, bytes32] = {}
 
@@ -66,8 +65,9 @@ class MempoolManager:
 
         self.limit_factor = 0.5
         self.mempool_max_total_cost = int(self.constants.MAX_BLOCK_COST_CLVM * self.constants.MEMPOOL_BLOCK_BUFFER)
-        self.potential_cache_max_total_cost = int(self.constants.MAX_BLOCK_COST_CLVM * 5)
-        self.potential_cache_cost: int = 0
+
+        # Transactions that were unable to enter mempool, used for retry. (they were invalid)
+        self.potential_cache = PendingTxCache(self.constants.MAX_BLOCK_COST_CLVM * 5)
         self.seen_cache_size = 10000
         self.pool = ProcessPoolExecutor(max_workers=1)
 
@@ -366,7 +366,7 @@ class MempoolManager:
                 potential = MempoolItem(
                     new_spend, uint64(fees), npc_result, cost, spend_name, additions, removals, program
                 )
-                self.add_to_potential_tx_set(potential)
+                self.potential_cache.add(potential)
                 return (
                     uint64(cost),
                     MempoolInclusionStatus.PENDING,
@@ -411,7 +411,7 @@ class MempoolManager:
                     potential = MempoolItem(
                         new_spend, uint64(fees), npc_result, cost, spend_name, additions, removals, program
                     )
-                    self.add_to_potential_tx_set(potential)
+                    self.potential_cache.add(potential)
                     return uint64(cost), MempoolInclusionStatus.PENDING, error
                 break
 
@@ -467,22 +467,6 @@ class MempoolManager:
         # 5. If coins can be spent return list of unspents as we see them in local storage
         return None, []
 
-    def add_to_potential_tx_set(self, item: MempoolItem):
-        """
-        Adds SpendBundles that have failed to be added to the pool in potential tx set.
-        This is later used to retry to add them.
-        """
-        if item.spend_bundle_name in self.potential_txs:
-            return None
-
-        self.potential_txs[item.spend_bundle_name] = item
-        self.potential_cache_cost += item.cost
-
-        while self.potential_cache_cost > self.potential_cache_max_total_cost:
-            first_in = list(self.potential_txs.keys())[0]
-            self.potential_cache_max_total_cost -= self.potential_txs[first_in].cost
-            self.potential_txs.pop(first_in)
-
     def get_spendbundle(self, bundle_hash: bytes32) -> Optional[SpendBundle]:
         """Returns a full SpendBundle if it's inside one the mempools"""
         if bundle_hash in self.mempool.spends:
@@ -496,6 +480,7 @@ class MempoolManager:
         return None
 
     async def new_peak(self, new_peak: Optional[BlockRecord]) -> List[Tuple[SpendBundle, NPCResult, bytes32]]:
+        #        breakpoint()
         """
         Called when a new peak is available, we try to recreate a mempool for the new tip.
         """
@@ -523,10 +508,9 @@ class MempoolManager:
                 if result != MempoolInclusionStatus.SUCCESS:
                     self.remove_seen(item.spend_bundle_name)
 
-            potential_txs_copy = self.potential_txs.copy()
-            self.potential_txs = {}
+            potential_txs = self.potential_cache.drain()
             txs_added = []
-            for item in potential_txs_copy.values():
+            for item in potential_txs.values():
                 cost, status, error = await self.add_spendbundle(
                     item.spend_bundle, item.npc_result, item.spend_bundle_name, program=item.program
                 )

chia/full_node/pending_tx_cache.py

@@ -0,0 +1,40 @@
+from typing import Dict
+
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.mempool_item import MempoolItem
+
+
+class PendingTxCache:
+    _cache_max_total_cost: int
+    _cache_cost: int
+    _txs: Dict[bytes32, MempoolItem]
+
+    def __init__(self, cost_limit: int):
+        self._cache_max_total_cost = cost_limit
+        self._cache_cost = 0
+        self._txs = {}
+
+    def add(self, item: MempoolItem):
+        """
+        Adds SpendBundles that have failed to be added to the pool in potential tx set.
+        This is later used to retry to add them.
+        """
+        if item.spend_bundle_name in self._txs:
+            return None
+
+        self._txs[item.spend_bundle_name] = item
+        self._cache_cost += item.cost
+
+        while self._cache_cost > self._cache_max_total_cost:
+            first_in = list(self._txs.keys())[0]
+            self._cache_cost -= self._txs[first_in].cost
+            self._txs.pop(first_in)
+
+    def drain(self) -> Dict[bytes32, MempoolItem]:
+        ret = self._txs
+        self._txs = {}
+        self._cache_cost = 0
+        return ret
+
+    def cost(self) -> int:
+        return self._cache_cost

tests/core/full_node/test_mempool.py

@@ -19,6 +19,7 @@ from chia.types.coin_spend import CoinSpend
 from chia.types.condition_opcodes import ConditionOpcode
 from chia.types.condition_with_args import ConditionWithArgs
 from chia.types.spend_bundle import SpendBundle
+from chia.types.mempool_item import MempoolItem
 from chia.util.clvm import int_to_bytes
 from chia.util.condition_tools import conditions_for_solution
 from chia.util.errors import Err, ValidationError
@@ -27,6 +28,8 @@ from chia.util.hash import std_hash
 from chia.types.mempool_inclusion_status import MempoolInclusionStatus
 from chia.util.api_decorators import api_request, peer_required, bytes_required
 from chia.full_node.mempool_check_conditions import parse_condition_args, parse_condition, get_name_puzzle_conditions
+from chia.full_node.pending_tx_cache import PendingTxCache
+from blspy import G2Element
 
 from tests.connection_utils import connect_and_get_peer
 from tests.core.node_height import node_height_at_least
@@ -82,6 +85,79 @@ async def two_nodes():
         yield _
 
 
+def make_item(idx: int, cost: uint64 = uint64(80)) -> MempoolItem:
+    spend_bundle_name = bytes([idx] * 32)
+    return MempoolItem(
+        SpendBundle([], G2Element()),
+        uint64(0),
+        NPCResult(None, [], cost),
+        cost,
+        spend_bundle_name,
+        [],
+        [],
+        SerializedProgram(),
+    )
+
+
+class TestPendingTxCache:
+    def test_recall(self):
+        c = PendingTxCache(100)
+        item = make_item(1)
+        c.add(item)
+        tx = c.drain()
+        assert tx == {item.spend_bundle_name: item}
+
+    def test_fifo_limit(self):
+        c = PendingTxCache(200)
+        # each item has cost 80
+        items = [make_item(i) for i in range(1, 4)]
+        for i in items:
+            c.add(i)
+        # the max cost is 200, only two transactions will fit
+        # we evict items FIFO, so the to most recently added will be left
+        tx = c.drain()
+        assert tx == {items[-2].spend_bundle_name: items[-2], items[-1].spend_bundle_name: items[-1]}
+
+    def test_drain(self):
+        c = PendingTxCache(100)
+        item = make_item(1)
+        c.add(item)
+        tx = c.drain()
+        assert tx == {item.spend_bundle_name: item}
+
+        # drain will clear the cache, so a second call will be empty
+        tx = c.drain()
+        assert tx == {}
+
+    def test_cost(self):
+        c = PendingTxCache(200)
+        assert c.cost() == 0
+        item1 = make_item(1)
+        c.add(item1)
+        # each item has cost 80
+        assert c.cost() == 80
+
+        item2 = make_item(2)
+        c.add(item2)
+        assert c.cost() == 160
+
+        # the first item is evicted, so the cost stays the same
+        item3 = make_item(3)
+        c.add(item3)
+        assert c.cost() == 160
+
+        tx = c.drain()
+        assert tx == {item2.spend_bundle_name: item2, item3.spend_bundle_name: item3}
+
+        assert c.cost() == 0
+        item4 = make_item(4)
+        c.add(item4)
+        assert c.cost() == 80
+
+        tx = c.drain()
+        assert tx == {item4.spend_bundle_name: item4}
+
+
 class TestMempool:
     @pytest.mark.asyncio
     async def test_basic_mempool(self, two_nodes):
@@ -641,6 +717,21 @@ class TestMempoolManager:
         assert status == MempoolInclusionStatus.FAILED
         assert err == Err.ASSERT_SECONDS_ABSOLUTE_FAILED
 
+    @pytest.mark.asyncio
+    async def test_assert_height_pending(self, two_nodes):
+
+        full_node_1, full_node_2, server_1, server_2 = two_nodes
+        print(full_node_1.full_node.blockchain.get_peak())
+        current_height = full_node_1.full_node.blockchain.get_peak().height
+
+        cvp = ConditionWithArgs(ConditionOpcode.ASSERT_HEIGHT_ABSOLUTE, [int_to_bytes(current_height + 4)])
+        dic = {cvp.opcode: [cvp]}
+        blocks, spend_bundle1, peer, status, err = await self.condition_tester(two_nodes, dic)
+        sb1 = full_node_1.full_node.mempool_manager.get_spendbundle(spend_bundle1.name())
+        assert sb1 is None
+        assert status == MempoolInclusionStatus.PENDING
+        assert err == Err.ASSERT_HEIGHT_ABSOLUTE_FAILED
+
     @pytest.mark.asyncio
     async def test_assert_time_negative(self, two_nodes):