CHIA-194: CHIP-0026 Mempool Updates (#17980)

* Initial draft of mempool updates [skip ci]

* More implementation work, but capability issue [skip ci]

* Temp (will probably revert)

* Get first test passing

* Setup spent coin tests

* Add unrelated test data to make tests better

* Refactor test

* Finish testing

* Update tests

* Include removal reason and block inclusion tests

* Remove unnecessary dict

* Add missing message to test

* Fix after rebase

* Broadcast support from node for mempool updates

* Remove redundant code

* Use new_block

* Filter and early return suggestions

* Reword comments on PeakPostProcessingResult

* Improve mempool fn name and doc comment

* Bump wallet protocol version

* Add ratelimits

* Add tests for querying mempool items

* Asserts and more tests

* Add peers for spend bundle tests

benchmarks/mempool.py

@@ -177,9 +177,9 @@ async def add_spend_bundles(spend_bundles: List[SpendBundle]) -> None:
                 spend_bundle_id = tx.name()
                 npc = await mempool.pre_validate_spendbundle(tx, None, spend_bundle_id)
                 assert npc is not None
-                _, status, error = await mempool.add_spend_bundle(tx, npc, spend_bundle_id, height)
-                assert status == MempoolInclusionStatus.SUCCESS
-                assert error is None
+                info = await mempool.add_spend_bundle(tx, npc, spend_bundle_id, height)
+                assert info.status == MempoolInclusionStatus.SUCCESS
+                assert info.error is None
 
         suffix = "st" if single_threaded else "mt"
 

chia/_tests/connection_utils.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 from pathlib import Path
-from typing import Set, Tuple
+from typing import List, Set, Tuple
 
 import aiohttp
 from cryptography import x509
@@ -39,15 +39,26 @@ async def disconnect_all_and_reconnect(server: ChiaServer, reconnect_to: ChiaSer
 
 
 async def add_dummy_connection(
-    server: ChiaServer, self_hostname: str, dummy_port: int, type: NodeType = NodeType.FULL_NODE
+    server: ChiaServer,
+    self_hostname: str,
+    dummy_port: int,
+    type: NodeType = NodeType.FULL_NODE,
+    *,
+    additional_capabilities: List[Tuple[uint16, str]] = [],
 ) -> Tuple[asyncio.Queue, bytes32]:
-    wsc, peer_id = await add_dummy_connection_wsc(server, self_hostname, dummy_port, type)
+    wsc, peer_id = await add_dummy_connection_wsc(
+        server, self_hostname, dummy_port, type, additional_capabilities=additional_capabilities
+    )
 
     return wsc.incoming_queue, peer_id
 
 
 async def add_dummy_connection_wsc(
-    server: ChiaServer, self_hostname: str, dummy_port: int, type: NodeType = NodeType.FULL_NODE
+    server: ChiaServer,
+    self_hostname: str,
+    dummy_port: int,
+    type: NodeType = NodeType.FULL_NODE,
+    additional_capabilities: List[Tuple[uint16, str]] = [],
 ) -> Tuple[WSChiaConnection, bytes32]:
     timeout = aiohttp.ClientTimeout(total=10)
     session = aiohttp.ClientSession(timeout=timeout)
@@ -86,7 +97,7 @@ async def add_dummy_connection_wsc(
         peer_id,
         100,
         30,
-        local_capabilities_for_handshake=default_capabilities[type],
+        local_capabilities_for_handshake=default_capabilities[type] + additional_capabilities,
     )
     await wsc.perform_handshake(server._network_id, dummy_port, type)
     if wsc.incoming_message_task is not None:

chia/_tests/core/full_node/test_subscriptions.py

@@ -1,10 +1,34 @@
 from __future__ import annotations
 
-from chia.full_node.subscriptions import PeerSubscriptions
+from chia_rs import AugSchemeMPL, Coin, CoinSpend, Program
+from chia_rs.sized_ints import uint32, uint64
+
+from chia.consensus.default_constants import DEFAULT_CONSTANTS
+from chia.full_node.bundle_tools import simple_solution_generator
+from chia.full_node.mempool_check_conditions import get_name_puzzle_conditions
+from chia.full_node.subscriptions import PeerSubscriptions, peers_for_spend_bundle
+from chia.types.blockchain_format.program import INFINITE_COST
 from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.spend_bundle import SpendBundle
+
+IDENTITY_PUZZLE = Program.to(1)
+IDENTITY_PUZZLE_HASH = IDENTITY_PUZZLE.get_tree_hash()
+
+OTHER_PUZZLE = Program.to(2)
+OTHER_PUZZLE_HASH = OTHER_PUZZLE.get_tree_hash()
+
+HINT_PUZZLE = Program.to(3)
+HINT_PUZZLE_HASH = HINT_PUZZLE.get_tree_hash()
+
+IDENTITY_COIN = Coin(bytes32(b"0" * 32), IDENTITY_PUZZLE_HASH, uint64(1000))
+OTHER_COIN = Coin(bytes32(b"3" * 32), OTHER_PUZZLE_HASH, uint64(1000))
+
+EMPTY_SIGNATURE = AugSchemeMPL.aggregate([])
 
 peer1 = bytes32(b"1" * 32)
 peer2 = bytes32(b"2" * 32)
+peer3 = bytes32(b"3" * 32)
+peer4 = bytes32(b"4" * 32)
 
 coin1 = bytes32(b"a" * 32)
 coin2 = bytes32(b"b" * 32)
@@ -420,3 +444,49 @@ def test_clear_subscriptions() -> None:
 
     subs.clear_puzzle_subscriptions(peer1)
     assert subs.peer_subscription_count(peer1) == 0
+
+
+def test_peers_for_spent_coin() -> None:
+    subs = PeerSubscriptions()
+
+    subs.add_puzzle_subscriptions(peer1, [IDENTITY_PUZZLE_HASH], 1)
+    subs.add_puzzle_subscriptions(peer2, [HINT_PUZZLE_HASH], 1)
+    subs.add_coin_subscriptions(peer3, [IDENTITY_COIN.name()], 1)
+    subs.add_coin_subscriptions(peer4, [OTHER_COIN.name()], 1)
+
+    coin_spends = [CoinSpend(IDENTITY_COIN, IDENTITY_PUZZLE, Program.to([]))]
+
+    spend_bundle = SpendBundle(coin_spends, AugSchemeMPL.aggregate([]))
+    generator = simple_solution_generator(spend_bundle)
+    npc_result = get_name_puzzle_conditions(
+        generator=generator, max_cost=INFINITE_COST, mempool_mode=True, height=uint32(0), constants=DEFAULT_CONSTANTS
+    )
+    assert npc_result.conds is not None
+
+    peers = peers_for_spend_bundle(subs, npc_result.conds, {HINT_PUZZLE_HASH})
+    assert peers == {peer1, peer2, peer3}
+
+
+def test_peers_for_created_coin() -> None:
+    subs = PeerSubscriptions()
+
+    new_coin = Coin(IDENTITY_COIN.name(), OTHER_PUZZLE_HASH, uint64(1000))
+
+    subs.add_puzzle_subscriptions(peer1, [OTHER_PUZZLE_HASH], 1)
+    subs.add_puzzle_subscriptions(peer2, [HINT_PUZZLE_HASH], 1)
+    subs.add_coin_subscriptions(peer3, [new_coin.name()], 1)
+    subs.add_coin_subscriptions(peer4, [OTHER_COIN.name()], 1)
+
+    coin_spends = [
+        CoinSpend(IDENTITY_COIN, IDENTITY_PUZZLE, Program.to([[51, OTHER_PUZZLE_HASH, 1000, [HINT_PUZZLE_HASH]]]))
+    ]
+
+    spend_bundle = SpendBundle(coin_spends, AugSchemeMPL.aggregate([]))
+    generator = simple_solution_generator(spend_bundle)
+    npc_result = get_name_puzzle_conditions(
+        generator=generator, max_cost=INFINITE_COST, mempool_mode=True, height=uint32(0), constants=DEFAULT_CONSTANTS
+    )
+    assert npc_result.conds is not None
+
+    peers = peers_for_spend_bundle(subs, npc_result.conds, set())
+    assert peers == {peer1, peer2, peer3}

chia/_tests/core/mempool/test_mempool.py

@@ -2863,7 +2863,7 @@ def test_limit_expiring_transactions(height: bool, items: List[int], expected: L
         invariant_check_mempool(mempool)
         if increase_fee:
             fee_rate += 0.1
-            assert ret is None
+            assert ret.error is None
         else:
             fee_rate -= 0.1
 

chia/_tests/core/mempool/test_mempool_item_queries.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+from typing import List
+
+from chia_rs import AugSchemeMPL, Coin, Program
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint32, uint64
+
+from chia._tests.core.mempool.test_mempool_manager import TEST_HEIGHT, make_bundle_spends_map_and_fee
+from chia.consensus.default_constants import DEFAULT_CONSTANTS
+from chia.full_node.bitcoin_fee_estimator import create_bitcoin_fee_estimator
+from chia.full_node.bundle_tools import simple_solution_generator
+from chia.full_node.fee_estimation import MempoolInfo
+from chia.full_node.mempool import Mempool
+from chia.full_node.mempool_check_conditions import get_name_puzzle_conditions
+from chia.types.blockchain_format.program import INFINITE_COST
+from chia.types.clvm_cost import CLVMCost
+from chia.types.coin_spend import CoinSpend
+from chia.types.fee_rate import FeeRate
+from chia.types.mempool_item import MempoolItem
+from chia.types.spend_bundle import SpendBundle
+
+MEMPOOL_INFO = MempoolInfo(
+    max_size_in_cost=CLVMCost(uint64(INFINITE_COST * 10)),
+    minimum_fee_per_cost_to_replace=FeeRate(uint64(5)),
+    max_block_clvm_cost=CLVMCost(uint64(INFINITE_COST)),
+)
+
+IDENTITY_PUZZLE = Program.to(1)
+IDENTITY_PUZZLE_HASH = IDENTITY_PUZZLE.get_tree_hash()
+
+OTHER_PUZZLE = Program.to(2)
+OTHER_PUZZLE_HASH = OTHER_PUZZLE.get_tree_hash()
+
+IDENTITY_COIN_1 = Coin(bytes32(b"0" * 32), IDENTITY_PUZZLE_HASH, uint64(1000))
+IDENTITY_COIN_2 = Coin(bytes32(b"1" * 32), IDENTITY_PUZZLE_HASH, uint64(1000))
+IDENTITY_COIN_3 = Coin(bytes32(b"2" * 32), IDENTITY_PUZZLE_HASH, uint64(1000))
+
+OTHER_COIN_1 = Coin(bytes32(b"3" * 32), OTHER_PUZZLE_HASH, uint64(1000))
+OTHER_COIN_2 = Coin(bytes32(b"4" * 32), OTHER_PUZZLE_HASH, uint64(1000))
+OTHER_COIN_3 = Coin(bytes32(b"5" * 32), OTHER_PUZZLE_HASH, uint64(1000))
+
+EMPTY_SIGNATURE = AugSchemeMPL.aggregate([])
+
+
+def make_item(coin_spends: List[CoinSpend]) -> MempoolItem:
+    spend_bundle = SpendBundle(coin_spends, EMPTY_SIGNATURE)
+    generator = simple_solution_generator(spend_bundle)
+    npc_result = get_name_puzzle_conditions(
+        generator=generator, max_cost=INFINITE_COST, mempool_mode=True, height=uint32(0), constants=DEFAULT_CONSTANTS
+    )
+    bundle_coin_spends, fee = make_bundle_spends_map_and_fee(spend_bundle, npc_result)
+    return MempoolItem(
+        spend_bundle=spend_bundle,
+        fee=fee,
+        npc_result=npc_result,
+        spend_bundle_name=spend_bundle.name(),
+        height_added_to_mempool=TEST_HEIGHT,
+        bundle_coin_spends=bundle_coin_spends,
+    )
+
+
+def test_empty_pool() -> None:
+    fee_estimator = create_bitcoin_fee_estimator(uint64(INFINITE_COST))
+    mempool = Mempool(MEMPOOL_INFO, fee_estimator)
+    assert mempool.items_with_coin_ids({IDENTITY_COIN_1.name()}) == []
+    assert mempool.items_with_puzzle_hashes({IDENTITY_PUZZLE_HASH}, False) == []
+
+
+def test_by_spent_coin_ids() -> None:
+    fee_estimator = create_bitcoin_fee_estimator(uint64(INFINITE_COST))
+    mempool = Mempool(MEMPOOL_INFO, fee_estimator)
+
+    # Add an item with both queried coins, to ensure there are no duplicates in the response.
+    item_1 = make_item(
+        [
+            CoinSpend(IDENTITY_COIN_1, IDENTITY_PUZZLE, Program.to([])),
+            CoinSpend(IDENTITY_COIN_2, IDENTITY_PUZZLE, Program.to([])),
+        ]
+    )
+    mempool.add_to_pool(item_1)
+
+    # Another coin with the same puzzle hash shouldn't match.
+    other = make_item(
+        [
+            CoinSpend(IDENTITY_COIN_3, IDENTITY_PUZZLE, Program.to([])),
+        ]
+    )
+    mempool.add_to_pool(other)
+
+    # And this coin is completely unrelated.
+    other = make_item([CoinSpend(OTHER_COIN_1, OTHER_PUZZLE, Program.to([[]]))])
+    mempool.add_to_pool(other)
+
+    # Only the first transaction includes these coins.
+    assert mempool.items_with_coin_ids({IDENTITY_COIN_1.name(), IDENTITY_COIN_2.name()}) == [item_1.spend_bundle_name]
+    assert mempool.items_with_coin_ids({IDENTITY_COIN_1.name()}) == [item_1.spend_bundle_name]
+    assert mempool.items_with_coin_ids({OTHER_COIN_2.name(), OTHER_COIN_3.name()}) == []
+
+
+def test_by_spend_puzzle_hashes() -> None:
+    fee_estimator = create_bitcoin_fee_estimator(uint64(INFINITE_COST))
+    mempool = Mempool(MEMPOOL_INFO, fee_estimator)
+
+    # Add a transaction with the queried puzzle hashes.
+    item_1 = make_item(
+        [
+            CoinSpend(IDENTITY_COIN_1, IDENTITY_PUZZLE, Program.to([])),
+            CoinSpend(IDENTITY_COIN_2, IDENTITY_PUZZLE, Program.to([])),
+        ]
+    )
+    mempool.add_to_pool(item_1)
+
+    # Another coin with the same puzzle hash should match.
+    item_2 = make_item(
+        [
+            CoinSpend(IDENTITY_COIN_3, IDENTITY_PUZZLE, Program.to([])),
+        ]
+    )
+    mempool.add_to_pool(item_2)
+
+    # But this coin has a different puzzle hash.
+    other = make_item([CoinSpend(OTHER_COIN_1, OTHER_PUZZLE, Program.to([[]]))])
+    mempool.add_to_pool(other)
+
+    # Only the first two transactions include the puzzle hash.
+    assert mempool.items_with_puzzle_hashes({IDENTITY_PUZZLE_HASH}, False) == [
+        item_1.spend_bundle_name,
+        item_2.spend_bundle_name,
+    ]
+
+    # Test the other puzzle hash as well.
+    assert mempool.items_with_puzzle_hashes({OTHER_PUZZLE_HASH}, False) == [
+        other.spend_bundle_name,
+    ]
+
+    # And an unrelated puzzle hash.
+    assert mempool.items_with_puzzle_hashes({bytes32(b"0" * 32)}, False) == []
+
+
+def test_by_created_coin_id() -> None:
+    fee_estimator = create_bitcoin_fee_estimator(uint64(INFINITE_COST))
+    mempool = Mempool(MEMPOOL_INFO, fee_estimator)
+
+    # Add a transaction that creates the queried coin id.
+    item = make_item(
+        [
+            CoinSpend(IDENTITY_COIN_1, IDENTITY_PUZZLE, Program.to([[51, IDENTITY_PUZZLE_HASH, 1000]])),
+        ]
+    )
+    mempool.add_to_pool(item)
+
+    # Test that the transaction is found.
+    assert mempool.items_with_coin_ids({Coin(IDENTITY_COIN_1.name(), IDENTITY_PUZZLE_HASH, uint64(1000)).name()}) == [
+        item.spend_bundle_name
+    ]
+
+
+def test_by_created_puzzle_hash() -> None:
+    fee_estimator = create_bitcoin_fee_estimator(uint64(INFINITE_COST))
+    mempool = Mempool(MEMPOOL_INFO, fee_estimator)
+
+    # Add a transaction that creates the queried puzzle hash.
+    item_1 = make_item(
+        [
+            CoinSpend(
+                IDENTITY_COIN_1,
+                IDENTITY_PUZZLE,
+                Program.to([[51, OTHER_PUZZLE_HASH, 400], [51, OTHER_PUZZLE_HASH, 600]]),
+            ),
+        ]
+    )
+    mempool.add_to_pool(item_1)
+
+    # This one is hinted.
+    item_2 = make_item(
+        [
+            CoinSpend(
+                IDENTITY_COIN_2,
+                IDENTITY_PUZZLE,
+                Program.to([[51, IDENTITY_PUZZLE_HASH, 1000, [OTHER_PUZZLE_HASH]]]),
+            ),
+        ]
+    )
+    mempool.add_to_pool(item_2)
+
+    # Test that the transactions are both found.
+    assert mempool.items_with_puzzle_hashes({OTHER_PUZZLE_HASH}, include_hints=True) == [
+        item_1.spend_bundle_name,
+        item_2.spend_bundle_name,
+    ]

chia/_tests/core/mempool/test_mempool_manager.py

@@ -388,7 +388,7 @@ async def add_spendbundle(
     npc_result = await mempool_manager.pre_validate_spendbundle(sb, None, sb_name)
     ret = await mempool_manager.add_spend_bundle(sb, npc_result, sb_name, TEST_HEIGHT)
     invariant_check_mempool(mempool_manager.mempool)
-    return ret
+    return ret.cost, ret.status, ret.error
 
 
 async def generate_and_add_spendbundle(

chia/_tests/util/build_network_protocol_files.py

@@ -107,6 +107,8 @@ def visit_wallet_protocol(visitor: Callable[[Any, str], None]) -> None:
     visitor(request_coin_state, "request_coin_state")
     visitor(respond_coin_state, "respond_coin_state")
     visitor(reject_coin_state, "reject_coin_state")
+    visitor(request_cost_info, "request_cost_info")
+    visitor(respond_cost_info, "respond_cost_info")
 
 
 def visit_harvester_protocol(visitor: Callable[[Any, str], None]) -> None:

chia/_tests/util/network_protocol_data.py

@@ -733,6 +733,27 @@
     uint8(wallet_protocol.RejectStateReason.EXCEEDED_SUBSCRIPTION_LIMIT)
 )
 
+removed_mempool_item = wallet_protocol.RemovedMempoolItem(
+    bytes32(bytes.fromhex("59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7")), uint8(1)
+)
+
+mempool_items_added = wallet_protocol.MempoolItemsAdded(
+    [bytes32(bytes.fromhex("59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7"))]
+)
+
+mempool_items_removed = wallet_protocol.MempoolItemsRemoved([removed_mempool_item])
+
+request_cost_info = wallet_protocol.RequestCostInfo()
+
+respond_cost_info = wallet_protocol.RespondCostInfo(
+    max_transaction_cost=uint64(100000),
+    max_block_cost=uint64(1000000),
+    max_mempool_cost=uint64(10000000),
+    mempool_cost=uint64(50000),
+    mempool_fee=uint64(500000),
+    bump_fee_per_cost=uint8(10),
+)
+
 
 ### HARVESTER PROTOCOL
 pool_difficulty = harvester_protocol.PoolDifficulty(