cuprate_dandelion_tower/pool/
mod.rs1use 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
49pub 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#[derive(Clone)]
108pub struct DandelionPoolService<Tx, TxId, PeerId> {
109 #[expect(clippy::type_complexity)]
110 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 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}