from __future__ import annotations import asyncio import dataclasses import logging from dataclasses import dataclass, field from queue import PriorityQueue, SimpleQueue from typing import ClassVar, Generic, TypeVar from chia_rs import SpendBundle from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import uint64 from chia.server.ws_connection import WSChiaConnection from chia.types.mempool_inclusion_status import MempoolInclusionStatus from chia.util.errors import Err T = TypeVar("T") class TransactionQueueFull(Exception): pass class ValuedEventSentinel: pass @dataclasses.dataclass class ValuedEvent(Generic[T]): _value_sentinel: ClassVar[ValuedEventSentinel] = ValuedEventSentinel() _event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event) _value: ValuedEventSentinel | T = _value_sentinel def set(self, value: T) -> None: if not isinstance(self._value, ValuedEventSentinel): raise Exception("Value already set") self._value = value self._event.set() async def wait(self) -> T: await self._event.wait() if isinstance(self._value, ValuedEventSentinel): raise Exception("Value not set despite event being set") return self._value @dataclasses.dataclass(frozen=True) class PeerWithTx: peer_host: str advertised_fee: uint64 advertised_cost: uint64 @dataclass(frozen=True, order=True) class TransactionQueueEntry: """ A transaction received from peer. This is put into a queue, and not yet in the mempool. """ transaction: SpendBundle = field(compare=False) transaction_bytes: bytes | None = field(compare=False) spend_name: bytes32 peer: WSChiaConnection | None = field(compare=False) test: bool = field(compare=False) # IDs of peers that advertised this transaction via new_transaction, along # with their hostname, fee and cost. peers_with_tx: dict[bytes32, PeerWithTx] = field(default_factory=dict, compare=False) done: ValuedEvent[tuple[MempoolInclusionStatus, Err | None]] = field( default_factory=ValuedEvent, compare=False, ) @dataclass class PeerTransactionsQueue: # Peer's priority queue of the form (negative fee per cost, entry). # We sort like this because PriorityQueue returns lowest first. priority_queue: PriorityQueue[tuple[float, TransactionQueueEntry]] = field(default_factory=PriorityQueue) # Peer's deficit in the context of deficit round robin algorithm. The unit # here is in CLVM cost. deficit: int = field(default=0, init=False) @dataclass class TransactionQueue: """ This class replaces one queue by using a high priority queue for local transactions and separate queues for peers. Local transactions are processed first. Then the next transaction is taken from the next non-empty queue after the last processed queue. (round-robin) This decreases the effects of one peer spamming your node with transactions. """ _list_cursor: int # this is which index _queue_length: asyncio.Semaphore _index_to_peer_map: list[bytes32] _peers_transactions_queues: dict[bytes32, PeerTransactionsQueue] _high_priority_queue: SimpleQueue[TransactionQueueEntry] peer_size_limit: int log: logging.Logger # Fallback cost for transactions without cost information _max_tx_clvm_cost: uint64 # Each 100 pops we do a cleanup of empty peer queues _cleanup_counter: int def __init__(self, peer_size_limit: int, log: logging.Logger, *, max_tx_clvm_cost: uint64) -> None: self._list_cursor = 0 self._queue_length = asyncio.Semaphore(0) # default is 1 self._index_to_peer_map = [] self._peers_transactions_queues = {} self._high_priority_queue = SimpleQueue() # we don't limit the number of high priority transactions self.peer_size_limit = peer_size_limit self.log = log self._max_tx_clvm_cost = max_tx_clvm_cost self._cleanup_counter = 0 def put(self, tx: TransactionQueueEntry, peer_id: bytes32 | None, high_priority: bool = False) -> None: if peer_id is None or high_priority: # when it's local there is no peer_id. self._high_priority_queue.put(tx) self._queue_length.release() return peer_queue = self._peers_transactions_queues.get(peer_id) if peer_queue is None: peer_queue = PeerTransactionsQueue() self._peers_transactions_queues[peer_id] = peer_queue self._index_to_peer_map.append(peer_id) if self._peers_transactions_queues[peer_id].priority_queue.qsize() >= self.peer_size_limit: self.log.warning(f"Transaction queue full for peer {peer_id}") raise TransactionQueueFull(f"Transaction queue full for peer {peer_id}") tx_info = tx.peers_with_tx.get(peer_id) if tx_info is not None and tx_info.advertised_cost > 0: fpc = tx_info.advertised_fee / tx_info.advertised_cost # PriorityQueue returns lowest first so we invert priority = -fpc else: # This peer didn't advertise cost and fee information for # this transaction (it sent a `RespondTransaction` message # instead of a `NewTransaction` one). priority = float("inf") peer_queue.priority_queue.put((priority, tx)) self._queue_length.release() # increment semaphore to indicate that we have a new item in the queue def _cleanup_peer_queues(self) -> None: """ Removes empty peer queues and updates the cursor accordingly. """ new_peer_map = [] for idx, peer_id in enumerate(self._index_to_peer_map): if self._peers_transactions_queues[peer_id].priority_queue.empty(): self._peers_transactions_queues.pop(peer_id, None) if idx < self._list_cursor: self._list_cursor -= 1 else: new_peer_map.append(peer_id) self._index_to_peer_map = new_peer_map if self._list_cursor >= len(self._index_to_peer_map): self._list_cursor = 0 async def pop(self) -> TransactionQueueEntry: await self._queue_length.acquire() if not self._high_priority_queue.empty(): return self._high_priority_queue.get() while True: # Map of peer ID to its top transaction's advertised cost. We want # to service transactions fairly between peers, based on cost, so # we need to find the lowest cost transaction among the top ones. top_txs_advertised_costs: dict[bytes32, uint64] = {} # Let's see if a peer can afford to send its top transaction num_peers = len(self._index_to_peer_map) assert num_peers != 0 start = self._list_cursor for offset in range(num_peers): peer_index = (start + offset) % num_peers peer_id = self._index_to_peer_map[peer_index] peer_queue = self._peers_transactions_queues[peer_id] if peer_queue.priority_queue.empty(): continue # There is no peek method so we access the internal `queue` _, entry = peer_queue.priority_queue.queue[0] tx_info = entry.peers_with_tx.get(peer_id) # If we don't know the cost information for this transaction # we fallback to the highest cost. if tx_info is not None: # At this point we have no transactions with zero cost assert tx_info.advertised_cost > 0 top_tx_advertised_cost = tx_info.advertised_cost else: top_tx_advertised_cost = self._max_tx_clvm_cost top_txs_advertised_costs[peer_id] = top_tx_advertised_cost if peer_queue.deficit >= top_tx_advertised_cost: # This peer can afford its top transaction _, entry = peer_queue.priority_queue.get() peer_queue.deficit -= top_tx_advertised_cost if peer_queue.priority_queue.empty(): peer_queue.deficit = 0 # Let's advance the cursor to the next peer self._list_cursor = (peer_index + 1) % num_peers # See if we need to perform the periodic cleanup self._cleanup_counter = (self._cleanup_counter + 1) % 100 if self._cleanup_counter == 0: self._cleanup_peer_queues() return entry # None of the peers could afford to send their top transactions, so # let's add the lowest cost among transactions to all the deficit # counters for the next iteration. assert len(top_txs_advertised_costs) != 0 lowest_cost_among_txs = min(top_txs_advertised_costs.values()) for peer_id in top_txs_advertised_costs: self._peers_transactions_queues[peer_id].deficit += lowest_cost_among_txs