From 97b250433efc3e07873a891598a7689de697f8d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Tue, 22 Nov 2022 10:30:16 +0100 Subject: [PATCH] WIP --- clients/native/src/websocket/handler.rs | 15 +++++++-- .../network-requester/src/core.rs | 32 +++++++++++++++++-- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 07b196354d..e215179ac0 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -1,7 +1,7 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use client_connections::{ClosedConnectionSender, TransmissionLane}; +use client_connections::{ClosedConnectionSender, LaneQueueLengths, TransmissionLane}; use client_core::client::{ inbound_messages::{InputMessage, InputMessageSender}, received_buffer::{ @@ -40,6 +40,7 @@ pub(crate) struct Handler { self_full_address: Recipient, socket: Option>, received_response_type: ReceivedResponseType, + lane_queue_lengths: LaneQueueLengths, } // clone is used to use handler on a new connection, which initially is `None` @@ -52,6 +53,7 @@ impl Clone for Handler { self_full_address: self.self_full_address, socket: None, received_response_type: Default::default(), + lane_queue_lengths: self.lane_queue_lengths.clone(), } } } @@ -70,6 +72,7 @@ impl Handler { closed_connection_tx: ClosedConnectionSender, buffer_requester: ReceivedBufferRequestSender, self_full_address: &Recipient, + lane_queue_lengths: LaneQueueLengths, ) -> Self { Handler { msg_input, @@ -78,6 +81,7 @@ impl Handler { self_full_address: *self_full_address, socket: None, received_response_type: Default::default(), + lane_queue_lengths, } } @@ -95,7 +99,14 @@ impl Handler { panic!(); } - None + // WIP(JON) + let queue_lengh = self + .lane_queue_lengths + .lock() + .unwrap() + .get(&lane) + .unwrap_or(0); + Some(ServerResponse::LaneQueueLength(connection_id, queue_lengh)) } async fn handle_reply( diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs index 88d66d159d..d29d116264 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -7,7 +7,7 @@ use crate::error::NetworkRequesterError; use crate::statistics::ServiceStatisticsCollector; use crate::websocket; use crate::websocket::TSWebsocketStream; -use client_connections::ClosedConnectionReceiver; +use client_connections::{ClosedConnectionReceiver, LaneQueueLengths, TransmissionLane}; use futures::channel::mpsc; use futures::stream::{SplitSink, SplitStream}; use futures::{SinkExt, StreamExt}; @@ -119,6 +119,7 @@ impl ServiceProvider { async fn read_websocket_message( websocket_reader: &mut SplitStream, + lane_queue_lengths: LaneQueueLengths, ) -> Option { while let Some(msg) = websocket_reader.next().await { let data = msg @@ -139,6 +140,15 @@ impl ServiceProvider { let received = match deserialized_message { ServerResponse::Received(received) => received, + ServerResponse::LaneQueueLength(lane, queue_length) => { + log::info!( + "received LaneQueueLength lane: {lane}, queue_length: {queue_length}" + ); + let lane = TransmissionLane::ConnectionId(lane); + let mut guard = lane_queue_lengths.lock().unwrap(); + guard.map.insert(lane, queue_length); + continue; + } ServerResponse::Error(err) => { panic!("received error from native client! - {}", err) } @@ -155,6 +165,7 @@ impl ServiceProvider { return_address: Recipient, controller_sender: ControllerSender, mix_input_sender: MixProxySender<(Socks5Message, Recipient)>, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { let mut conn = match Connection::new(conn_id, remote_addr.clone(), return_address).await { @@ -196,7 +207,7 @@ impl ServiceProvider { ); // run the proxy on the connection - conn.run_proxy(mix_receiver, mix_input_sender, shutdown) + conn.run_proxy(mix_receiver, mix_input_sender, lane_queue_lengths, shutdown) .await; // proxy is done - remove the access channel from the controller @@ -212,6 +223,7 @@ impl ServiceProvider { ); } + #[allow(clippy::too_many_arguments)] async fn handle_proxy_connect( &mut self, controller_sender: &mut ControllerSender, @@ -219,6 +231,7 @@ impl ServiceProvider { conn_id: ConnectionId, remote_addr: String, return_address: Recipient, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { if !self.open_proxy && !self.outbound_request_filter.check(&remote_addr) { @@ -250,6 +263,7 @@ impl ServiceProvider { return_address, controller_sender_clone, mix_input_sender_clone, + lane_queue_lengths, shutdown, ) .await @@ -274,6 +288,7 @@ impl ServiceProvider { controller_sender: &mut ControllerSender, mix_input_sender: &MixProxySender<(Socks5Message, Recipient)>, stats_collector: Option, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { let deserialized_msg = match Socks5Message::try_from_bytes(raw_request) { @@ -299,6 +314,7 @@ impl ServiceProvider { req.conn_id, req.remote_addr, req.return_address, + lane_queue_lengths, shutdown, ) .await @@ -346,6 +362,10 @@ impl ServiceProvider { // `ClientRequest`. 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(); + // Controller for managing all active connections. // We provide it with a ShutdownListener since it requires it, even though for the network // requester shutdown signalling is not yet fully implemented. @@ -386,7 +406,12 @@ impl ServiceProvider { println!("\nAll systems go. Press CTRL-C to stop the server."); // for each incoming message from the websocket... (which in 99.99% cases is going to be a mix message) loop { - let received = match Self::read_websocket_message(&mut websocket_reader).await { + let received = match Self::read_websocket_message( + &mut websocket_reader, + shared_lane_queue_lengths.clone(), + ) + .await + { Some(msg) => msg, None => { error!("The websocket stream has finished!"); @@ -402,6 +427,7 @@ impl ServiceProvider { &mut controller_sender, &mix_input_sender, stats_collector.clone(), + shared_lane_queue_lengths.clone(), shutdown.subscribe(), ) .await;