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>

p2p/dandelion-tower/Cargo.toml

@@ -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"] }

p2p/dandelion-tower/src/lib.rs

@@ -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

p2p/dandelion-tower/src/router.rs

@@ -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::{
+use tower::Service;
-    discover::{Change, Discover},
-    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.
+    /// The outbound peer stream returned an error, this is critical.
-    #[error("The outbound peer discoverer returned an err: {0}.")]
+    #[error("The outbound peer stream returned an err: {0}.")]
-    OutboundPeerDiscoverError(tower::BoxError),
+    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)
+                .try_poll_next(cx)
-                .map_err(DandelionRouterError::OutboundPeerDiscoverError))
+                .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) => {
+                OutboundPeer::Exhausted => {
-                    self.stem_peers.remove(&key);
+                    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 {
+    fn fluff_tx(&mut self, tx: Tx) -> BoxFuture<'static, Result<State, DandelionRouterError>> {
-        self.broadcast_svc.call(DiffuseRequest(tx))
+        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 =
+    type Future = BoxFuture<'static, Result<State, DandelionRouterError>>;
-        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
 
     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(
+            TxState::Fluff => self.fluff_tx(req.tx),
-                self.fluff_tx(req.tx)
-                    .map_ok(|_| State::Fluff)
-                    .map_err(DandelionRouterError::BroadcastError),
-            ),
             TxState::Stem { from } => match self.current_state {
                 State::Fluff => {
                     tracing::debug!(parent: &self.span, "Fluffing stem tx.");
 
-                    Box::pin(
+                    self.fluff_tx(req.tx)
-                        self.fluff_tx(req.tx)
-                            .map_ok(|_| State::Fluff)
-                            .map_err(DandelionRouterError::BroadcastError),
-                    )
                 }
                 State::Stem => {
                     tracing::trace!(parent: &self.span, "Steming transaction");
 
-                    Box::pin(
+                    self.stem_tx(req.tx, from)
-                        self.stem_tx(req.tx, from)
-                            .map_ok(|_| State::Stem)
-                            .map_err(DandelionRouterError::PeerError),
-                    )
                 }
             },
             TxState::Local => {
                 tracing::debug!(parent: &self.span, "Steming local tx.");
 
-                Box::pin(
+                self.stem_local_tx(req.tx)
-                    self.stem_local_tx(req.tx)
-                        .map_ok(|_| State::Stem)
-                        .map_err(DandelionRouterError::PeerError),
-                )
             }
         }
     }

p2p/dandelion-tower/src/tests/mod.rs

@@ -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::{
+use tower::{util::service_fn, Service, ServiceExt};
-    discover::{Discover, ServiceList},
-    util::service_fn,
-    Service, ServiceExt,
-};
 
 use crate::{
     traits::{TxStoreRequest, TxStoreResponse},
-    State,
+    OutboundPeer, State,
 };
 
 pub fn mock_discover_svc<Req: Send + 'static>() -> (
-    impl Discover<
+    impl Stream<
-        Key = usize,
+        Item = Result<
-        Service = impl Service<
+            OutboundPeer<
-            Req,
+                usize,
-            Future = impl Future<Output = Result<(), tower::BoxError>> + Send + 'static,
+                impl Service<
-            Error = tower::BoxError,
+                        Req,
-        > + Send
+                        Future = impl Future<Output = Result<(), tower::BoxError>> + Send + 'static,
-                      + 'static,
+                        Error = tower::BoxError,
-        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 discover = futures::stream::iter(0_usize..1_000_000)
-        let tx_2 = tx.clone();
+        .map(move |i| {
+            let tx_2 = tx.clone();
 
-        service_fn(move |req| {
+            Ok::<_, tower::BoxError>(OutboundPeer::Peer(
-            tx_2.send((i, req)).unwrap();
+                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)
 }