Skip to main content

cuprate_dandelion_tower/pool/
mod.rs

1//! # Dandelion++ Pool
2//!
3//! This module contains [`DandelionPoolManager`] which is a wrapper around a backing transaction store,
4//! which fully implements the dandelion++ protocol.
5//!
6//! The [`DandelionPoolManager`] is a middle man between a [preprocessing stage](#preprocessing-stage) and a dandelion router.
7//! It handles promoting transactions in the stem state to the fluff state and setting embargo timers on stem state transactions.
8//!
9//! ### Preprocessing stage
10//!
11//! The preprocessing stage (not handled in this crate) before giving the transaction to the [`DandelionPoolManager`]
12//! should handle:
13//!
14//! - verifying the tx.
15//! - checking if we have the tx in the pool already and giving that information to the [`IncomingTxBuilder`].
16//! - storing the tx in the pool, if it isn't there already.
17//!
18//! ### Keep Stem Transactions Hidden
19//!
20//! When using your handle to the backing store it must be remembered to keep transactions in the stem pool hidden.
21//! So handle any requests to the tx-pool like the stem side of the pool does not exist.
22use std::{
23    collections::HashMap,
24    hash::Hash,
25    marker::PhantomData,
26    task::{Context, Poll},
27};
28
29use futures::{future::BoxFuture, FutureExt};
30use rand_distr::Exp;
31use tokio::{
32    sync::{mpsc, oneshot},
33    task::{JoinHandle, JoinSet},
34};
35use tokio_util::{sync::PollSender, time::DelayQueue};
36use tower::Service;
37
38use crate::{
39    traits::{TxStoreRequest, TxStoreResponse},
40    DandelionConfig, DandelionRouteReq, DandelionRouterError, State,
41};
42
43mod incoming_tx;
44mod manager;
45
46pub use incoming_tx::{IncomingTx, IncomingTxBuilder};
47pub use manager::{DandelionPoolManager, DandelionPoolShutDown};
48
49/// Start the [`DandelionPoolManager`].
50///
51/// This function spawns the [`DandelionPoolManager`] and returns a
52/// [`DandelionPoolService`] and the manager's task handle.
53///
54/// ### Args
55///
56/// - `buffer_size` is the size of the channel's buffer between the [`DandelionPoolService`] and [`DandelionPoolManager`].
57/// - `dandelion_router` is the router service, kept generic instead of [`DandelionRouter`](crate::DandelionRouter) to allow
58///   user to customise routing functionality.
59/// - `backing_pool` is the backing transaction storage service
60/// - `config` is [`DandelionConfig`].
61pub fn start_dandelion_pool_manager<P, R, Tx, TxId, PeerId>(
62    buffer_size: usize,
63    dandelion_router: R,
64    backing_pool: P,
65    config: DandelionConfig,
66) -> (DandelionPoolService<Tx, TxId, PeerId>, JoinHandle<()>)
67where
68    Tx: Clone + Send + 'static,
69    TxId: Hash + Eq + Clone + Send + 'static,
70    PeerId: Hash + Eq + Clone + Send + 'static,
71    P: Service<TxStoreRequest<TxId>, Response = TxStoreResponse<Tx>, Error = tower::BoxError>
72        + Send
73        + 'static,
74    P::Future: Send + 'static,
75    R: Service<DandelionRouteReq<Tx, PeerId>, Response = State, Error = DandelionRouterError>
76        + Send
77        + 'static,
78    R::Future: Send + 'static,
79{
80    let (tx, rx) = mpsc::channel(buffer_size);
81
82    let pool = DandelionPoolManager {
83        dandelion_router,
84        backing_pool,
85        routing_set: JoinSet::new(),
86        stem_origins: HashMap::new(),
87        embargo_timers: DelayQueue::new(),
88        embargo_timer_keys: HashMap::new(),
89        embargo_dist: Exp::new(1.0 / config.average_embargo_timeout().as_secs_f64()).unwrap(),
90        config,
91        _tx: PhantomData,
92    };
93
94    let task = tokio::spawn(pool.run(rx));
95
96    (
97        DandelionPoolService {
98            tx: PollSender::new(tx),
99        },
100        task,
101    )
102}
103
104/// The dandelion pool manager service.
105///
106/// Used to send [`IncomingTx`]s to the [`DandelionPoolManager`]
107#[derive(Clone)]
108pub struct DandelionPoolService<Tx, TxId, PeerId> {
109    #[expect(clippy::type_complexity)]
110    /// The channel to [`DandelionPoolManager`].
111    tx: PollSender<(
112        (IncomingTx<Tx, TxId, PeerId>, tracing::Span),
113        oneshot::Sender<()>,
114    )>,
115}
116
117impl<Tx, TxId, PeerId> Service<IncomingTx<Tx, TxId, PeerId>>
118    for DandelionPoolService<Tx, TxId, PeerId>
119where
120    Tx: Clone + Send,
121    TxId: Hash + Eq + Clone + Send + 'static,
122    PeerId: Hash + Eq + Clone + Send + 'static,
123{
124    type Response = ();
125    type Error = DandelionPoolShutDown;
126    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
127
128    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
129        self.tx.poll_reserve(cx).map_err(|_| DandelionPoolShutDown)
130    }
131
132    fn call(&mut self, req: IncomingTx<Tx, TxId, PeerId>) -> Self::Future {
133        // although the channel isn't sending anything we want to wait for the request to be handled before continuing.
134        let (tx, rx) = oneshot::channel();
135
136        let res = self
137            .tx
138            .send_item(((req, tracing::Span::current()), tx))
139            .map_err(|_| DandelionPoolShutDown);
140
141        async move {
142            res?;
143            rx.await.map_err(|_| DandelionPoolShutDown)?;
144
145            Ok(())
146        }
147        .boxed()
148    }
149}