P2P: Fix freeze in D++ (#204)

* Fix d++ router freeze

* update docs

* fix imports

* fix clippy

* Update p2p/dandelion-tower/src/lib.rs

Co-authored-by: hinto-janai <hinto.janai@protonmail.com>

---------

Co-authored-by: hinto-janai <hinto.janai@protonmail.com>
This commit is contained in:
Boog900 2024-07-01 19:23:40 +00:00 committed by GitHub
parent 7e9891de5b
commit f91be58a7f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 94 additions and 74 deletions

View file

@ -10,7 +10,7 @@ default = ["txpool"]
txpool = ["dep:rand_distr", "dep:tokio-util", "dep:tokio"]
[dependencies]
tower = { workspace = true, features = ["discover", "util"] }
tower = { workspace = true, features = ["util"] }
tracing = { workspace = true, features = ["std"] }
futures = { workspace = true, features = ["std"] }

View file

@ -26,9 +26,9 @@
//! The diffuse service should have a request of [`DiffuseRequest`](traits::DiffuseRequest) and it's error
//! should be [`tower::BoxError`].
//!
//! ## Outbound Peer Discoverer
//! ## Outbound Peer TryStream
//!
//! The outbound peer [`Discover`](tower::discover::Discover) should provide a stream of randomly selected outbound
//! The outbound peer [`TryStream`](futures::TryStream) should provide a stream of randomly selected outbound
//! peers, these peers will then be used to route stem txs to.
//!
//! The peers will not be returned anywhere, so it is recommended to wrap them in some sort of drop guard that returns
@ -37,10 +37,10 @@
//! ## Peer Service
//!
//! This service represents a connection to an individual peer, this should be returned from the Outbound Peer
//! Discover. This should immediately send the transaction to the peer when requested, i.e. it should _not_ set
//! TryStream. This should immediately send the transaction to the peer when requested, it should _not_ set
//! a timer.
//!
//! The diffuse service should have a request of [`StemRequest`](traits::StemRequest) and it's error
//! The peer service should have a request of [`StemRequest`](traits::StemRequest) and its error
//! should be [`tower::BoxError`].
//!
//! ## Backing Pool

View file

@ -6,11 +6,10 @@
//! ### What The Router Does Not Do
//!
//! It does not handle anything to do with keeping transactions long term, i.e. embargo timers and handling
//! loops in the stem. It is up to implementers to do this if they decide not top use [`DandelionPool`](crate::pool::DandelionPool)
//! loops in the stem. It is up to implementers to do this if they decide not to use [`DandelionPool`](crate::pool::DandelionPool)
//!
use std::{
collections::HashMap,
future::Future,
hash::Hash,
marker::PhantomData,
pin::Pin,
@ -18,12 +17,9 @@ use std::{
time::Instant,
};
use futures::TryFutureExt;
use futures::{future::BoxFuture, FutureExt, TryFutureExt, TryStream};
use rand::{distributions::Bernoulli, prelude::*, thread_rng};
use tower::{
discover::{Change, Discover},
Service,
};
use tower::Service;
use crate::{
traits::{DiffuseRequest, StemRequest},
@ -39,14 +35,22 @@ pub enum DandelionRouterError {
/// The broadcast service returned an error.
#[error("Broadcast service returned an err: {0}.")]
BroadcastError(tower::BoxError),
/// The outbound peer discoverer returned an error, this is critical.
#[error("The outbound peer discoverer returned an err: {0}.")]
OutboundPeerDiscoverError(tower::BoxError),
/// The outbound peer stream returned an error, this is critical.
#[error("The outbound peer stream returned an err: {0}.")]
OutboundPeerStreamError(tower::BoxError),
/// The outbound peer discoverer returned [`None`].
#[error("The outbound peer discoverer exited.")]
OutboundPeerDiscoverExited,
}
/// A response from an attempt to retrieve an outbound peer.
pub enum OutboundPeer<ID, T> {
/// A peer.
Peer(ID, T),
/// The peer store is exhausted and has no more to return.
Exhausted,
}
/// The dandelion++ state.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum State {
@ -116,9 +120,11 @@ pub struct DandelionRouter<P, B, ID, S, Tx> {
impl<Tx, ID, P, B, S> DandelionRouter<P, B, ID, S, Tx>
where
ID: Hash + Eq + Clone,
P: Discover<Key = ID, Service = S, Error = tower::BoxError>,
P: TryStream<Ok = OutboundPeer<ID, S>, Error = tower::BoxError>,
B: Service<DiffuseRequest<Tx>, Error = tower::BoxError>,
B::Future: Send + 'static,
S: Service<StemRequest<Tx>, Error = tower::BoxError>,
S::Future: Send + 'static,
{
/// Creates a new [`DandelionRouter`], with the provided services and config.
///
@ -165,15 +171,16 @@ where
match ready!(self
.outbound_peer_discover
.as_mut()
.poll_discover(cx)
.map_err(DandelionRouterError::OutboundPeerDiscoverError))
.try_poll_next(cx)
.map_err(DandelionRouterError::OutboundPeerStreamError))
.ok_or(DandelionRouterError::OutboundPeerDiscoverExited)??
{
Change::Insert(key, svc) => {
OutboundPeer::Peer(key, svc) => {
self.stem_peers.insert(key, svc);
}
Change::Remove(key) => {
self.stem_peers.remove(&key);
OutboundPeer::Exhausted => {
tracing::warn!("Failed to retrieve enough outbound peers for optimal dandelion++, privacy may be degraded.");
return Poll::Ready(Ok(()));
}
}
}
@ -181,11 +188,24 @@ where
Poll::Ready(Ok(()))
}
fn fluff_tx(&mut self, tx: Tx) -> B::Future {
self.broadcast_svc.call(DiffuseRequest(tx))
fn fluff_tx(&mut self, tx: Tx) -> BoxFuture<'static, Result<State, DandelionRouterError>> {
self.broadcast_svc
.call(DiffuseRequest(tx))
.map_ok(|_| State::Fluff)
.map_err(DandelionRouterError::BroadcastError)
.boxed()
}
fn stem_tx(&mut self, tx: Tx, from: ID) -> S::Future {
fn stem_tx(
&mut self,
tx: Tx,
from: ID,
) -> BoxFuture<'static, Result<State, DandelionRouterError>> {
if self.stem_peers.is_empty() {
tracing::debug!("Stem peers are empty, fluffing stem transaction.");
return self.fluff_tx(tx);
}
loop {
let stem_route = self.stem_routes.entry(from.clone()).or_insert_with(|| {
self.stem_peers
@ -201,11 +221,20 @@ where
continue;
};
return peer.call(StemRequest(tx));
return peer
.call(StemRequest(tx))
.map_ok(|_| State::Stem)
.map_err(DandelionRouterError::PeerError)
.boxed();
}
}
fn stem_local_tx(&mut self, tx: Tx) -> S::Future {
fn stem_local_tx(&mut self, tx: Tx) -> BoxFuture<'static, Result<State, DandelionRouterError>> {
if self.stem_peers.is_empty() {
tracing::warn!("Stem peers are empty, no outbound connections to stem local tx to, fluffing instead, privacy will be degraded.");
return self.fluff_tx(tx);
}
loop {
let stem_route = self.local_route.get_or_insert_with(|| {
self.stem_peers
@ -221,7 +250,11 @@ where
continue;
};
return peer.call(StemRequest(tx));
return peer
.call(StemRequest(tx))
.map_ok(|_| State::Stem)
.map_err(DandelionRouterError::PeerError)
.boxed();
}
}
}
@ -238,7 +271,7 @@ S: The Peer service - handles routing messages to a single node.
impl<Tx, ID, P, B, S> Service<DandelionRouteReq<Tx, ID>> for DandelionRouter<P, B, ID, S, Tx>
where
ID: Hash + Eq + Clone,
P: Discover<Key = ID, Service = S, Error = tower::BoxError>,
P: TryStream<Ok = OutboundPeer<ID, S>, Error = tower::BoxError>,
B: Service<DiffuseRequest<Tx>, Error = tower::BoxError>,
B::Future: Send + 'static,
S: Service<StemRequest<Tx>, Error = tower::BoxError>,
@ -246,8 +279,7 @@ where
{
type Response = State;
type Error = DandelionRouterError;
type Future =
Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
type Future = BoxFuture<'static, Result<State, DandelionRouterError>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if self.epoch_start.elapsed() > self.config.epoch_duration {
@ -309,39 +341,23 @@ where
tracing::trace!(parent: &self.span, "Handling route request.");
match req.state {
TxState::Fluff => Box::pin(
self.fluff_tx(req.tx)
.map_ok(|_| State::Fluff)
.map_err(DandelionRouterError::BroadcastError),
),
TxState::Fluff => self.fluff_tx(req.tx),
TxState::Stem { from } => match self.current_state {
State::Fluff => {
tracing::debug!(parent: &self.span, "Fluffing stem tx.");
Box::pin(
self.fluff_tx(req.tx)
.map_ok(|_| State::Fluff)
.map_err(DandelionRouterError::BroadcastError),
)
self.fluff_tx(req.tx)
}
State::Stem => {
tracing::trace!(parent: &self.span, "Steming transaction");
Box::pin(
self.stem_tx(req.tx, from)
.map_ok(|_| State::Stem)
.map_err(DandelionRouterError::PeerError),
)
self.stem_tx(req.tx, from)
}
},
TxState::Local => {
tracing::debug!(parent: &self.span, "Steming local tx.");
Box::pin(
self.stem_local_tx(req.tx)
.map_ok(|_| State::Stem)
.map_err(DandelionRouterError::PeerError),
)
self.stem_local_tx(req.tx)
}
}
}

View file

@ -3,44 +3,48 @@ mod router;
use std::{collections::HashMap, future::Future, hash::Hash, sync::Arc};
use futures::TryStreamExt;
use futures::{Stream, StreamExt, TryStreamExt};
use tokio::sync::mpsc::{self, UnboundedReceiver};
use tower::{
discover::{Discover, ServiceList},
util::service_fn,
Service, ServiceExt,
};
use tower::{util::service_fn, Service, ServiceExt};
use crate::{
traits::{TxStoreRequest, TxStoreResponse},
State,
OutboundPeer, State,
};
pub fn mock_discover_svc<Req: Send + 'static>() -> (
impl Discover<
Key = usize,
Service = impl Service<
Req,
Future = impl Future<Output = Result<(), tower::BoxError>> + Send + 'static,
Error = tower::BoxError,
> + Send
+ 'static,
Error = tower::BoxError,
impl Stream<
Item = Result<
OutboundPeer<
usize,
impl Service<
Req,
Future = impl Future<Output = Result<(), tower::BoxError>> + Send + 'static,
Error = tower::BoxError,
> + Send
+ 'static,
>,
tower::BoxError,
>,
>,
UnboundedReceiver<(u64, Req)>,
UnboundedReceiver<(usize, Req)>,
) {
let (tx, rx) = mpsc::unbounded_channel();
let discover = ServiceList::new((0..).map(move |i| {
let tx_2 = tx.clone();
let discover = futures::stream::iter(0_usize..1_000_000)
.map(move |i| {
let tx_2 = tx.clone();
service_fn(move |req| {
tx_2.send((i, req)).unwrap();
Ok::<_, tower::BoxError>(OutboundPeer::Peer(
i,
service_fn(move |req| {
tx_2.send((i, req)).unwrap();
async move { Ok::<(), tower::BoxError>(()) }
async move { Ok::<(), tower::BoxError>(()) }
}),
))
})
}))
.map_err(Into::into);
.map_err(Into::into);
(discover, rx)
}