GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Also remove a bunch of panics in the native client

Browse Source
This commit is contained in:
Jon Häggblad
2025-02-17 10:46:19 +01:00
parent c6d38d3c4f
commit 1ccdd5d660
3 changed files with 77 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/native/src/client/mod.rs
+8 -3
View File
@@ -49,7 +49,7 @@ impl SocketClient {
client_output: ClientOutput,
client_state: ClientState,
self_address: &Recipient,
shutdown: nym_task::TaskClient,
task_client: nym_task::TaskClient,
packet_type: PacketType,
) {
info!("Starting websocket listener...");
@@ -77,10 +77,15 @@ impl SocketClient {
shared_lane_queue_lengths,
reply_controller_sender,
Some(packet_type),
task_client.fork("websocket_handler"),
);
websocket::Listener::new(config.socket.host, config.socket.listening_port)
.start(websocket_handler, shutdown);
websocket::Listener::new(
config.socket.host,
config.socket.listening_port,
task_client.with_suffix("websocket_listener"),
)
.start(websocket_handler);
}
/// blocking version of `start_socket` method. Will run forever (or until SIGINT is sent)
clients/native/src/websocket/handler.rs
+58 -39
View File
@@ -19,6 +19,7 @@ use nym_sphinx::receiver::ReconstructedMessage;
use nym_task::connections::{
ConnectionCommand, ConnectionCommandSender, ConnectionId, LaneQueueLengths, TransmissionLane,
};
use nym_task::TaskClient;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::time::Instant;
@@ -43,9 +44,11 @@ pub(crate) struct HandlerBuilder {
lane_queue_lengths: LaneQueueLengths,
reply_controller_sender: ReplyControllerSender,
packet_type: Option<PacketType>,
task_client: TaskClient,
}
impl HandlerBuilder {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
msg_input: InputMessageSender,
client_connection_tx: ConnectionCommandSender,
@@ -54,6 +57,7 @@ impl HandlerBuilder {
lane_queue_lengths: LaneQueueLengths,
reply_controller_sender: ReplyControllerSender,
packet_type: Option<PacketType>,
task_client: TaskClient,
) -> Self {
Self {
msg_input,
@@ -63,11 +67,14 @@ impl HandlerBuilder {
lane_queue_lengths,
reply_controller_sender,
packet_type,
task_client,
}
}
// TODO: make sure we only ever have one active handler
pub fn create_active_handler(&self) -> Handler {
let mut task_client = self.task_client.fork("active_handler");
task_client.disarm();
Handler {
msg_input: self.msg_input.clone(),
client_connection_tx: self.client_connection_tx.clone(),
@@ -78,6 +85,7 @@ impl HandlerBuilder {
lane_queue_lengths: self.lane_queue_lengths.clone(),
reply_controller_sender: self.reply_controller_sender.clone(),
packet_type: self.packet_type,
task_client,
}
}
}
@@ -92,16 +100,18 @@ pub(crate) struct Handler {
lane_queue_lengths: LaneQueueLengths,
reply_controller_sender: ReplyControllerSender,
packet_type: Option<PacketType>,
task_client: TaskClient,
}
impl Drop for Handler {
fn drop(&mut self) {
if self
if let Err(err) = self
.buffer_requester
.unbounded_send(ReceivedBufferMessage::ReceiverDisconnect)
.is_err()
{
error!("we failed to disconnect the receiver from the buffer! presumably the shutdown procedure has been initiated!")
if !self.task_client.is_shutdown_poll() {
error!("failed to disconnect the receiver from the buffer: {err}");
}
}
}
}
@@ -132,7 +142,13 @@ impl Handler {
{
Ok(length) => length,
Err(err) => {
error!("Failed to get reply queue length for connection {connection_id}: {err}");
if !self.task_client.is_shutdown_poll() {
error!(
"Failed to get reply queue length for connection {connection_id}: {err}"
);
}
// We're just going to assume that the queue is empty, and I think that's okay
// during shutdown.
0
}
};
@@ -175,10 +191,11 @@ impl Handler {
// the ack control is now responsible for chunking, etc.
let input_msg = InputMessage::new_regular(recipient, message, lane, self.packet_type);
self.msg_input
.send(input_msg)
.await
.expect("InputMessageReceiver has stopped receiving!");
if let Err(err) = self.msg_input.send(input_msg).await {
if !self.task_client.is_shutdown_poll() {
error!("Failed to send message to the input buffer: {err}");
}
}
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -207,10 +224,11 @@ impl Handler {
let input_msg =
InputMessage::new_anonymous(recipient, message, reply_surbs, lane, self.packet_type);
self.msg_input
.send(input_msg)
.await
.expect("InputMessageReceiver has stopped receiving!");
if let Err(err) = self.msg_input.send(input_msg).await {
if !self.task_client.is_shutdown_poll() {
error!("Failed to send anonymous message to the input buffer: {err}");
}
}
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -234,10 +252,11 @@ impl Handler {
});
let input_msg = InputMessage::new_reply(recipient_tag, message, lane, self.packet_type);
self.msg_input
.send(input_msg)
.await
.expect("InputMessageReceiver has stopped receiving!");
if let Err(err) = self.msg_input.send(input_msg).await {
if !self.task_client.is_shutdown_poll() {
error!("Failed to send reply message to the input buffer: {err}");
}
}
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -252,9 +271,14 @@ impl Handler {
}
fn handle_closed_connection(&self, connection_id: u64) -> Option<ServerResponse> {
self.client_connection_tx
if let Err(err) = self
.client_connection_tx
.unbounded_send(ConnectionCommand::Close(connection_id))
.unwrap();
{
if !self.task_client.is_shutdown_poll() {
error!("Failed to send close connection command: {err}");
}
}
None
}
@@ -369,11 +393,10 @@ impl Handler {
}
}
async fn listen_for_requests(
&mut self,
mut msg_receiver: ReconstructedMessagesReceiver,
mut task_client: nym_task::TaskClient,
) {
async fn listen_for_requests(&mut self, mut msg_receiver: ReconstructedMessagesReceiver) {
let mut task_client = self.task_client.fork("select");
task_client.disarm();
while !task_client.is_shutdown() {
tokio::select! {
// we can either get a client request from the websocket
@@ -422,15 +445,7 @@ impl Handler {
}
// consume self to make sure `drop` is called after this is done
pub(crate) async fn handle_connection(
mut self,
socket: TcpStream,
mut task_client: nym_task::TaskClient,
) {
// We don't want a crash in the connection handler to trigger a shutdown of the whole
// process.
task_client.disarm();
pub(crate) async fn handle_connection(mut self, socket: TcpStream) {
let ws_stream = match accept_async(socket).await {
Ok(ws_stream) => ws_stream,
Err(err) => {
@@ -443,14 +458,18 @@ impl Handler {
let (reconstructed_sender, reconstructed_receiver) = mpsc::unbounded();
// tell the buffer to start sending stuff to us
self.buffer_requester
.unbounded_send(ReceivedBufferMessage::ReceiverAnnounce(
reconstructed_sender,
))
.expect("the buffer request failed!");
if let Err(err) =
self.buffer_requester
.unbounded_send(ReceivedBufferMessage::ReceiverAnnounce(
reconstructed_sender,
))
{
if !self.task_client.is_shutdown_poll() {
error!("failed to announce the receiver to the buffer: {err}");
}
}
self.listen_for_requests(reconstructed_receiver, task_client)
.await;
self.listen_for_requests(reconstructed_receiver).await;
}
}
clients/native/src/websocket/listener.rs
+11 -17
View File
@@ -3,6 +3,7 @@
use super::handler::HandlerBuilder;
use log::*;
use nym_task::TaskClient;
use std::net::IpAddr;
use std::{net::SocketAddr, process, sync::Arc};
use tokio::io::AsyncWriteExt;
@@ -22,21 +23,19 @@ impl State {
pub(crate) struct Listener {
address: SocketAddr,
state: State,
task_client: TaskClient,
}
impl Listener {
pub(crate) fn new(host: IpAddr, port: u16) -> Self {
pub(crate) fn new(host: IpAddr, port: u16, task_client: TaskClient) -> Self {
Listener {
address: SocketAddr::new(host, port),
state: State::AwaitingConnection,
task_client,
}
}
pub(crate) async fn run(
&mut self,
handler: HandlerBuilder,
mut task_client: nym_task::TaskClient,
) {
pub(crate) async fn run(&mut self, handler: HandlerBuilder) {
let tcp_listener = match tokio::net::TcpListener::bind(self.address).await {
Ok(listener) => listener,
Err(err) => {
@@ -47,11 +46,11 @@ impl Listener {
let notify = Arc::new(Notify::new());
loop {
while !self.task_client.is_shutdown() {
tokio::select! {
// When the handler finishes we check if shutdown is signalled
_ = notify.notified() => {
if task_client.is_shutdown() {
if self.task_client.is_shutdown() {
log::trace!("Websocket listener: detected shutdown after connection closed");
break;
}
@@ -60,7 +59,7 @@ impl Listener {
}
// ... but when there is no connected client at the time of shutdown being
// signalled, we handle it here.
_ = task_client.recv() => {
_ = self.task_client.recv() => {
if !self.state.is_connected() {
log::trace!("Not connected: shutting down");
break;
@@ -88,9 +87,8 @@ impl Listener {
// hanging because the executor doesn't come back here
let notify_clone = Arc::clone(&notify);
let fresh_handler = handler.create_active_handler();
let task_client_handler = task_client.clone();
tokio::spawn(async move {
fresh_handler.handle_connection(socket, task_client_handler).await;
fresh_handler.handle_connection(socket).await;
notify_clone.notify_one();
});
self.state = State::Connected;
@@ -104,13 +102,9 @@ impl Listener {
log::debug!("Websocket listener: Exiting");
}
pub(crate) fn start(
mut self,
handler: HandlerBuilder,
shutdown: nym_task::TaskClient,
) -> JoinHandle<()> {
pub(crate) fn start(mut self, handler: HandlerBuilder) -> JoinHandle<()> {
info!("Running websocket on {:?}", self.address.to_string());
tokio::spawn(async move { self.run(handler, shutdown).await })
tokio::spawn(async move { self.run(handler).await })
}
}