client-core: expose lane queue length state to other components (#1777)

* client-core: publish lane queue lengths

* client-connections: rename to LaneQueueLenghts plural

* Fix clippy

* Fix wasm build

* rustfmt

* clippy
This commit is contained in:
Jon Häggblad
2022-11-21 11:21:28 +01:00
committed by GitHub
parent e30fd270a1
commit 1859ca0a30
9 changed files with 121 additions and 23 deletions
Generated
+1
View File
@@ -581,6 +581,7 @@ name = "client-connections"
version = "0.1.0"
dependencies = [
"futures",
"log",
]
[[package]]
@@ -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<OsRng> {
#[allow(clippy::too_many_arguments)]
pub fn new(
config: Config,
ack_receiver: AcknowledgementReceiver,
@@ -111,6 +114,7 @@ impl RealMessagesController<OsRng> {
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<OsRng> {
rng,
config.self_recipient,
topology_access,
lane_queue_lengths,
closed_connection_rx,
);
@@ -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<RealMessage> {
// 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<Option<StreamMessage>> {
// 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
@@ -41,6 +41,10 @@ impl TransmissionBuffer {
self.buffer.keys().count()
}
pub(crate) fn lane_length(&self, lane: &TransmissionLane) -> Option<usize> {
self.buffer.get(lane).map(LaneBufferEntry::len)
}
#[allow(unused)]
pub(crate) fn connections(&self) -> HashSet<u64> {
self.buffer
@@ -127,7 +131,7 @@ impl TransmissionBuffer {
Some(real_next)
}
pub(crate) fn pop_next_message_at_random(&mut self) -> Option<RealMessage> {
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()
}
+10 -1
View File
@@ -1,7 +1,9 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// 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(),
);
+8 -1
View File
@@ -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(),
);
+8 -1
View File
@@ -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 {
+1
View File
@@ -7,3 +7,4 @@ edition = "2021"
[dependencies]
futures = "0.3"
log = "0.4.17"
+52
View File
@@ -1,6 +1,8 @@
// Copyright 2022 - Nym Technologies SA <contact@nymtech.net>
// 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<ConnectionId>;
pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver<ConnectionId>;
// The `OutQueueControl` publishes the backlog per lane, primarily so that upstream can slow down
// if needed.
#[derive(Clone)]
pub struct LaneQueueLengths(std::sync::Arc<std::sync::Mutex<LaneQueueLengthsInner>>);
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<usize>) {
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<std::sync::Mutex<LaneQueueLengthsInner>>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
pub struct LaneQueueLengthsInner {
map: HashMap<TransmissionLane, usize>,
}