diff --git a/Cargo.lock b/Cargo.lock index 2da304990a..6c382cce87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -581,6 +581,7 @@ name = "client-connections" version = "0.1.0" dependencies = [ "futures", + "log", ] [[package]] diff --git a/clients/client-core/src/client/real_messages_control/mod.rs b/clients/client-core/src/client/real_messages_control/mod.rs index a7c03dc1c0..a1e9290c64 100644 --- a/clients/client-core/src/client/real_messages_control/mod.rs +++ b/clients/client-core/src/client/real_messages_control/mod.rs @@ -8,13 +8,15 @@ use self::{ acknowledgement_control::AcknowledgementController, real_traffic_stream::OutQueueControl, }; -use crate::client::real_messages_control::acknowledgement_control::AcknowledgementControllerConnectors; -use crate::client::{ - inbound_messages::InputMessageReceiver, mix_traffic::BatchMixMessageSender, - topology_control::TopologyAccessor, +use crate::{ + client::{ + inbound_messages::InputMessageReceiver, mix_traffic::BatchMixMessageSender, + real_messages_control::acknowledgement_control::AcknowledgementControllerConnectors, + topology_control::TopologyAccessor, + }, + spawn_future, }; -use crate::spawn_future; -use client_connections::ClosedConnectionReceiver; +use client_connections::{ClosedConnectionReceiver, LaneQueueLengths}; use futures::channel::mpsc; use gateway_client::AcknowledgementReceiver; use log::*; @@ -104,6 +106,7 @@ where // obviously when we finally make shared rng that is on 'higher' level, this should become // generic `R` impl RealMessagesController { + #[allow(clippy::too_many_arguments)] pub fn new( config: Config, ack_receiver: AcknowledgementReceiver, @@ -111,6 +114,7 @@ impl RealMessagesController { mix_sender: BatchMixMessageSender, topology_access: TopologyAccessor, #[cfg(feature = "reply-surb")] reply_key_storage: ReplyKeyStorage, + lane_queue_lengths: LaneQueueLengths, closed_connection_rx: ClosedConnectionReceiver, ) -> Self { let rng = OsRng; @@ -161,6 +165,7 @@ impl RealMessagesController { rng, config.self_recipient, topology_access, + lane_queue_lengths, closed_connection_rx, ); diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs index a42c24892e..43be7524b3 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs @@ -4,7 +4,9 @@ use crate::client::mix_traffic::BatchMixMessageSender; use crate::client::real_messages_control::acknowledgement_control::SentPacketNotificationSender; use crate::client::topology_control::TopologyAccessor; -use client_connections::{ClosedConnectionReceiver, ConnectionId, TransmissionLane}; +use client_connections::{ + ClosedConnectionReceiver, ConnectionId, LaneQueueLengths, TransmissionLane, +}; use futures::channel::mpsc; use futures::task::{Context, Poll}; use futures::{Future, Stream, StreamExt}; @@ -134,6 +136,9 @@ where /// Incoming channel for being notified of closed connections, so that we can close lanes /// corresponding to connections. To avoid sending traffic unnecessary closed_connection_rx: ClosedConnectionReceiver, + + /// Report queue lengths so that upstream can backoff sending data, and keep connections open. + lane_queue_lengths: LaneQueueLengths, } pub(crate) struct RealMessage { @@ -177,6 +182,7 @@ where rng: R, our_full_destination: Recipient, topology_access: TopologyAccessor, + lane_queue_lengths: LaneQueueLengths, closed_connection_rx: ClosedConnectionReceiver, ) -> Self { OutQueueControl { @@ -192,6 +198,7 @@ where topology_access, transmission_buffer: Default::default(), closed_connection_rx, + lane_queue_lengths, } } @@ -309,6 +316,17 @@ where } } + fn pop_next_message(&mut self) -> Option { + // Pop the next message from the transmission buffer + let (lane, real_next) = self.transmission_buffer.pop_next_message_at_random()?; + + // Update the published queue length + let lane_length = self.transmission_buffer.lane_length(&lane); + self.lane_queue_lengths.set(&lane, lane_length); + + Some(real_next) + } + fn poll_poisson(&mut self, cx: &mut Context<'_>) -> Poll> { // The average delay could change depending on if backpressure in the downstream channel // (mix_tx) was detected. @@ -359,16 +377,13 @@ where log::trace!("handling real_messages: size: {}", real_messages.len()); self.transmission_buffer.store(&conn_id, real_messages); - let real_next = self - .transmission_buffer - .pop_next_message_at_random() - .expect("we just added one"); + let real_next = self.pop_next_message().expect("Just stored one"); Poll::Ready(Some(StreamMessage::Real(Box::new(real_next)))) } Poll::Pending => { - if let Some(real_next) = self.transmission_buffer.pop_next_message_at_random() { + if let Some(real_next) = self.pop_next_message() { Poll::Ready(Some(StreamMessage::Real(Box::new(real_next)))) } else { // otherwise construct a dummy one @@ -411,16 +426,13 @@ where // First store what we got for the given connection id self.transmission_buffer.store(&conn_id, real_messages); - let real_next = self - .transmission_buffer - .pop_next_message_at_random() - .expect("we just added one"); + let real_next = self.pop_next_message().expect("we just added one"); Poll::Ready(Some(StreamMessage::Real(Box::new(real_next)))) } Poll::Pending => { - if let Some(real_next) = self.transmission_buffer.pop_next_message_at_random() { + if let Some(real_next) = self.pop_next_message() { Poll::Ready(Some(StreamMessage::Real(Box::new(real_next)))) } else { Poll::Pending diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream/transmission_buffer.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream/transmission_buffer.rs index 2510a65e05..b47aaf470e 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream/transmission_buffer.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream/transmission_buffer.rs @@ -41,6 +41,10 @@ impl TransmissionBuffer { self.buffer.keys().count() } + pub(crate) fn lane_length(&self, lane: &TransmissionLane) -> Option { + self.buffer.get(lane).map(LaneBufferEntry::len) + } + #[allow(unused)] pub(crate) fn connections(&self) -> HashSet { self.buffer @@ -127,7 +131,7 @@ impl TransmissionBuffer { Some(real_next) } - pub(crate) fn pop_next_message_at_random(&mut self) -> Option { + pub(crate) fn pop_next_message_at_random(&mut self) -> Option<(TransmissionLane, RealMessage)> { if self.buffer.is_empty() { return None; } @@ -142,8 +146,9 @@ impl TransmissionBuffer { *self.pick_random_lane()? }; + let msg = self.pop_front_from_lane(&lane)?; log::trace!("picking to send from lane: {:?}", lane); - self.pop_front_from_lane(&lane) + Some((lane, msg)) } pub(crate) fn prune_stale_connections(&mut self) { @@ -196,7 +201,6 @@ impl LaneBufferEntry { > Duration::from_secs(MSG_CONSIDERED_STALE_AFTER_SECS) } - #[cfg(not(target_arch = "wasm32"))] fn len(&self) -> usize { self.real_messages.len() } diff --git a/clients/native/src/client/mod.rs b/clients/native/src/client/mod.rs index f2aee9829f..91fb13bcfa 100644 --- a/clients/native/src/client/mod.rs +++ b/clients/native/src/client/mod.rs @@ -1,7 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use client_connections::{ClosedConnectionReceiver, ClosedConnectionSender, TransmissionLane}; +use client_connections::{ + ClosedConnectionReceiver, ClosedConnectionSender, LaneQueueLengths, TransmissionLane, +}; use client_core::client::cover_traffic_stream::LoopCoverTrafficStream; use client_core::client::inbound_messages::{ InputMessage, InputMessageReceiver, InputMessageSender, @@ -119,6 +121,7 @@ impl NymClient { ack_receiver: AcknowledgementReceiver, input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, + lane_queue_lengths: LaneQueueLengths, closed_connection_rx: ClosedConnectionReceiver, shutdown: ShutdownListener, ) { @@ -149,6 +152,7 @@ impl NymClient { mix_sender, topology_accessor, reply_key_storage, + lane_queue_lengths, closed_connection_rx, ) .start_with_shutdown(shutdown); @@ -417,12 +421,17 @@ impl NymClient { // controller that connections are closed. let (closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + // Shared queue length data. Published by the `OutQueueController` in the client, and used + // primarily to throttle incoming connections (e.g socks5 for attached network-requesters) + let shared_lane_queue_length = LaneQueueLengths::new(); + self.start_real_traffic_controller( shared_topology_accessor.clone(), reply_key_storage, ack_receiver, input_receiver, sphinx_message_sender.clone(), + shared_lane_queue_length, closed_connection_rx, shutdown.subscribe(), ); diff --git a/clients/socks5/src/client/mod.rs b/clients/socks5/src/client/mod.rs index 5b8530e6d0..b35d7fd9fe 100644 --- a/clients/socks5/src/client/mod.rs +++ b/clients/socks5/src/client/mod.rs @@ -9,7 +9,7 @@ use crate::socks::{ authentication::{AuthenticationMethods, Authenticator, User}, server::SphinxSocksServer, }; -use client_connections::{ClosedConnectionReceiver, ClosedConnectionSender}; +use client_connections::{ClosedConnectionReceiver, ClosedConnectionSender, LaneQueueLengths}; use client_core::client::cover_traffic_stream::LoopCoverTrafficStream; use client_core::client::inbound_messages::{ InputMessage, InputMessageReceiver, InputMessageSender, @@ -120,6 +120,7 @@ impl NymClient { input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, closed_connection_rx: ClosedConnectionReceiver, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { let mut controller_config = client_core::client::real_messages_control::Config::new( @@ -149,6 +150,7 @@ impl NymClient { mix_sender, topology_accessor, reply_key_storage, + lane_queue_lengths, closed_connection_rx, ) .start_with_shutdown(shutdown); @@ -409,6 +411,10 @@ impl NymClient { // This will be forwarded to `OutQueueControl` let (closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + // Shared queue length data. Published by the `OutQueueController` in the client, and used + // primarily to throttle incoming connections + let shared_lane_queue_lengths = LaneQueueLengths::new(); + self.start_real_traffic_controller( shared_topology_accessor.clone(), reply_key_storage, @@ -416,6 +422,7 @@ impl NymClient { input_receiver, sphinx_message_sender.clone(), closed_connection_rx, + shared_lane_queue_lengths, shutdown.subscribe(), ); diff --git a/clients/webassembly/src/client/mod.rs b/clients/webassembly/src/client/mod.rs index 7aed6c9a3b..e49d6199db 100644 --- a/clients/webassembly/src/client/mod.rs +++ b/clients/webassembly/src/client/mod.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use self::config::Config; -use client_connections::{ClosedConnectionReceiver, TransmissionLane}; +use client_connections::{ClosedConnectionReceiver, LaneQueueLengths, TransmissionLane}; use client_core::client::{ cover_traffic_stream::LoopCoverTrafficStream, inbound_messages::{InputMessage, InputMessageReceiver, InputMessageSender}, @@ -129,6 +129,7 @@ impl NymClient { input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, closed_connection_rx: ClosedConnectionReceiver, + lane_queue_lengths: LaneQueueLengths, ) { let mut controller_config = real_messages_control::Config::new( self.key_manager.ack_key(), @@ -153,6 +154,7 @@ impl NymClient { input_receiver, mix_sender, topology_accessor, + lane_queue_lengths, closed_connection_rx, ) .start(); @@ -353,12 +355,17 @@ impl NymClient { // The MixTrafficController then sends the actual traffic let sphinx_message_sender = Self::start_mix_traffic_controller(gateway_client); + // Shared queue length data. Published by the `OutQueueController` in the client, and used + // primarily to throttle incoming connections + let shared_lane_queue_lengths = LaneQueueLengths::new(); + self.start_real_traffic_controller( shared_topology_accessor.clone(), ack_receiver, input_receiver, sphinx_message_sender.clone(), closed_connection_rx, + shared_lane_queue_lengths, ); if !self.config.debug.disable_loop_cover_traffic_stream { diff --git a/common/client-connections/Cargo.toml b/common/client-connections/Cargo.toml index 067621c28f..7438be79ba 100644 --- a/common/client-connections/Cargo.toml +++ b/common/client-connections/Cargo.toml @@ -7,3 +7,4 @@ edition = "2021" [dependencies] futures = "0.3" +log = "0.4.17" diff --git a/common/client-connections/src/lib.rs b/common/client-connections/src/lib.rs index 2aabce78b0..265e21f9fd 100644 --- a/common/client-connections/src/lib.rs +++ b/common/client-connections/src/lib.rs @@ -1,6 +1,8 @@ // Copyright 2022 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; + use futures::channel::mpsc; pub type ConnectionId = u64; @@ -19,3 +21,53 @@ pub enum TransmissionLane { /// they can forward this to the `OutQueueControl` (via `ClientRequest` for the network-requester) pub type ClosedConnectionSender = mpsc::UnboundedSender; pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver; + +// The `OutQueueControl` publishes the backlog per lane, primarily so that upstream can slow down +// if needed. +#[derive(Clone)] +pub struct LaneQueueLengths(std::sync::Arc>); + +impl LaneQueueLengths { + pub fn new() -> Self { + LaneQueueLengths(std::sync::Arc::new(std::sync::Mutex::new( + LaneQueueLengthsInner { + map: HashMap::new(), + }, + ))) + } + + pub fn set(&mut self, lane: &TransmissionLane, lane_length: Option) { + match self.0.lock() { + Ok(mut inner) => { + if let Some(length) = lane_length { + inner + .map + .entry(*lane) + .and_modify(|e| *e = length) + .or_insert(length); + } else { + inner.map.remove(lane); + } + } + Err(err) => log::warn!("Failed to set lane queue length: {err}"), + } + } +} + +impl Default for LaneQueueLengths { + fn default() -> Self { + Self::new() + } +} + +impl std::ops::Deref for LaneQueueLengths { + type Target = std::sync::Arc>; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +pub struct LaneQueueLengthsInner { + map: HashMap, +}