From df566933bab3e7c0cbdf36af581ed1bbcde5820e Mon Sep 17 00:00:00 2001 From: Simon Wicky Date: Tue, 19 May 2026 11:31:14 +0200 Subject: [PATCH] prepare multi threaded node --- Cargo.lock | 1 + common/nym-lp-data/Cargo.toml | 1 + .../src/fragmentation/reconstruction.rs | 94 +++++++++---------- nym-node/src/node/lp/data/handler/mod.rs | 18 ++-- nym-node/src/node/lp/data/handler/pipeline.rs | 28 ++---- nym-node/src/node/lp/data/shared.rs | 22 +---- 6 files changed, 69 insertions(+), 95 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8ab6e9a3b2..c19bd2426c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7490,6 +7490,7 @@ name = "nym-lp-data" version = "1.21.0" dependencies = [ "bytes", + "dashmap", "num_enum", "nym-common", "nym-lp", diff --git a/common/nym-lp-data/Cargo.toml b/common/nym-lp-data/Cargo.toml index 65971d9338..f197655938 100644 --- a/common/nym-lp-data/Cargo.toml +++ b/common/nym-lp-data/Cargo.toml @@ -14,6 +14,7 @@ documentation.workspace = true [dependencies] bytes.workspace = true +dashmap.workspace = true num_enum.workspace = true tracing.workspace = true thiserror.workspace = true diff --git a/common/nym-lp-data/src/fragmentation/reconstruction.rs b/common/nym-lp-data/src/fragmentation/reconstruction.rs index 45038beb90..52e89854f3 100644 --- a/common/nym-lp-data/src/fragmentation/reconstruction.rs +++ b/common/nym-lp-data/src/fragmentation/reconstruction.rs @@ -1,12 +1,17 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use tracing::{debug, trace, warn}; - use crate::fragmentation::fragment::{Fragment, FragmentHashKey, FragmentMetadata}; -use std::collections::HashMap; + +use dashmap::DashMap; +use dashmap::mapref::entry::Entry; use std::fmt::Debug; use std::ops::Add; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tracing::{debug, trace, warn}; + +pub const DEFAULT_FRAGMENT_TIMEOUT_DURATION: Duration = Duration::from_secs(30); /// Per-message buffer that collects every `Fragment` of a fragmented message /// and reassembles the original payload once they are all in. @@ -129,7 +134,7 @@ where /// In-flight messages keyed on `(id, frame_kind)`. The frame kind is /// part of the key so that a random-id collision between two unrelated /// kinds cannot accidentally route fragments into the same buffer. - in_flight_messages: HashMap>, + in_flight_messages: Arc>>, /// How long an incomplete message is allowed to sit before it is /// dropped on the next `cleanup_stale_buffers` pass. @@ -149,34 +154,13 @@ where } } - /// Whether the message at `key` is present and has all of its fragments. - fn is_message_fully_received(&self, key: FragmentHashKey) -> bool { - self.in_flight_messages - .get(&key) - .map(|buf| buf.is_complete) - .unwrap_or(false) - } - - /// Pop the buffer at `key` and reassemble its message bytes. The caller - /// must ensure `is_message_fully_received(key)` holds. - fn reconstruct_message(&mut self, key: FragmentHashKey) -> Vec { - debug_assert!(self.is_message_fully_received(key)); - - // SAFETY: the precondition above guarantees the entry exists, so - // `remove` returns `Some`. `is_message_fully_received` also implies - // the buffer's `is_complete` flag is set, which is the precondition - // `into_message` relies on. - #[allow(clippy::unwrap_used)] - self.in_flight_messages.remove(&key).unwrap().into_message() - } - /// Insert `fragment` into the buffer for its message and, if it was the /// last outstanding fragment, return the reassembled message bytes /// alongside the original metadata. /// /// Stale incomplete messages are evicted on every call. pub fn insert_new_fragment( - &mut self, + &self, fragment: Fragment, timestamp: Ts, ) -> Option<(Vec, FragmentMetadata)> { @@ -188,24 +172,34 @@ where // from any of the buffer's slots would be equivalent. let metadata = (fragment.frame_kind(), fragment.kind_metadata()).into(); - let buf = self - .in_flight_messages - .entry(key) - .or_insert_with(|| MessageBuffer::new(total_fragments, timestamp.clone())); - - buf.insert_fragment(fragment, timestamp.clone()); - let maybe_message = if self.is_message_fully_received(key) { - Some((self.reconstruct_message(key), metadata)) - } else { - None + let maybe_message = match self.in_flight_messages.entry(key) { + Entry::Occupied(mut entry) => { + entry.get_mut().insert_fragment(fragment, timestamp.clone()); + entry + .get() + .is_complete + .then(|| (entry.remove().into_message(), metadata)) + } + Entry::Vacant(entry) => { + let mut buf = MessageBuffer::new(total_fragments, timestamp.clone()); + buf.insert_fragment(fragment, timestamp.clone()); + if buf.is_complete { + Some((buf.into_message(), metadata)) + } else { + entry.insert(buf); + None + } + } }; + + // This might be a bit slow, keep an eye on it self.cleanup_stale_buffers(timestamp.clone()); maybe_message } /// Drop incomplete messages whose `last_fragment_timestamp` is older /// than `incomplete_message_timeout` ago. - pub fn cleanup_stale_buffers(&mut self, timestamp: Ts) { + pub fn cleanup_stale_buffers(&self, timestamp: Ts) { trace!("Cleaning up stale buffers"); self.in_flight_messages.retain(|_, buf| { let keep = buf.last_fragment_timestamp.clone() @@ -224,6 +218,12 @@ where } } +impl Default for MessageReconstructor { + fn default() -> Self { + MessageReconstructor::new(DEFAULT_FRAGMENT_TIMEOUT_DURATION) + } +} + #[cfg(test)] mod tests { #![allow(clippy::unwrap_used)] @@ -335,7 +335,7 @@ mod tests { let mut fragments = split(message, SPHINX, 64); assert_eq!(fragments.len(), 1); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); let out = rec.insert_new_fragment(fragments.pop().unwrap(), 0); let (bytes, metadata) = out.expect("single fragment must complete the message"); assert_eq!(bytes, message); @@ -348,7 +348,7 @@ mod tests { let fragments = split(&message, SPHINX, 16); assert!(fragments.len() > 1); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); let total = fragments.len(); let mut completed = None; for (i, f) in fragments.into_iter().enumerate() { @@ -371,7 +371,7 @@ mod tests { // Reverse arrival order. fragments.reverse(); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); let mut last = None; for (i, f) in fragments.into_iter().enumerate() { last = rec.insert_new_fragment(f, i as u64); @@ -384,7 +384,7 @@ mod tests { fn reconstructor_propagates_frame_kind() { let message = b"outfox bytes"; let mut fragments = split(message, OUTFOX, 64); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); let (_, metadata) = rec .insert_new_fragment(fragments.pop().unwrap(), 0) .expect("complete"); @@ -403,7 +403,7 @@ mod tests { make_fragment(2, 2, 1, SPHINX, vec![0xb2]), ]; - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); // Interleave. assert!(rec.insert_new_fragment(a.remove(0), 0).is_none()); assert!(rec.insert_new_fragment(b.remove(0), 1).is_none()); @@ -422,7 +422,7 @@ mod tests { let o1 = make_fragment(42, 2, 0, OUTFOX, vec![0x20]); let o2 = make_fragment(42, 2, 1, OUTFOX, vec![0x21]); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); assert!(rec.insert_new_fragment(s1, 0).is_none()); assert!(rec.insert_new_fragment(o1, 1).is_none()); let (s_msg, s_metadata) = rec.insert_new_fragment(s2, 2).unwrap(); @@ -437,7 +437,7 @@ mod tests { #[test] fn reconstructor_clears_buffer_after_emitting_message() { let f = make_fragment(99, 1, 0, SPHINX, vec![0xff]); - let mut rec = MessageReconstructor::::new(60); + let rec = MessageReconstructor::::new(60); let _ = rec.insert_new_fragment(f, 0).unwrap(); assert!( rec.in_flight_messages.is_empty(), @@ -450,7 +450,7 @@ mod tests { #[test] fn cleanup_evicts_buffers_older_than_timeout() { let f = make_fragment(1, 2, 0, SPHINX, vec![0]); - let mut rec = MessageReconstructor::::new(10); + let rec = MessageReconstructor::::new(10); // First (and only) fragment received at t=0; the message stays // incomplete. assert!(rec.insert_new_fragment(f, 0).is_none()); @@ -473,7 +473,7 @@ mod tests { let stale = make_fragment(1, 2, 0, SPHINX, vec![0]); let fresh = make_fragment(2, 1, 0, SPHINX, vec![0xff]); - let mut rec = MessageReconstructor::::new(10); + let rec = MessageReconstructor::::new(10); assert!(rec.insert_new_fragment(stale, 0).is_none()); assert_eq!(rec.in_flight_messages.len(), 1); @@ -489,7 +489,7 @@ mod tests { // A buffer that keeps receiving fragments must not be evicted // even if the absolute time exceeds the timeout, as long as the // gap between fragments stays under it. - let mut rec = MessageReconstructor::::new(10); + let rec = MessageReconstructor::::new(10); assert!( rec.insert_new_fragment(make_fragment(1, 3, 0, SPHINX, vec![0xa]), 0) diff --git a/nym-node/src/node/lp/data/handler/mod.rs b/nym-node/src/node/lp/data/handler/mod.rs index 356ab0b084..24ddcc068c 100644 --- a/nym-node/src/node/lp/data/handler/mod.rs +++ b/nym-node/src/node/lp/data/handler/mod.rs @@ -15,14 +15,14 @@ //! ``` //! -use crate::node::lp::data::handler::pipeline::{MixnodeDataPipeline, MixnodeDataPipelineConfig}; +use crate::node::lp::data::handler::pipeline::MixnodeDataPipeline; use crate::node::lp::data::shared::SharedLpDataState; use nym_lp_data::AddressedTimedData; use nym_lp_data::mixnodes::traits::MixnodeProcessingPipeline; use nym_lp_data::packet::EncryptedLpPacket; use nym_metrics::inc; use rand::rngs::OsRng; -use std::sync::mpsc; +use std::sync::{Arc, mpsc}; use std::time::Instant; use std::{net::SocketAddr, time::Duration}; use tokio::sync::mpsc::error::TrySendError; @@ -35,10 +35,12 @@ mod pipeline; mod processing; const PIPELINE_TICKING_DURATION: Duration = Duration::from_millis(1); -const FRAGMENT_TIMEOUT_DURATION: Duration = Duration::from_secs(30); /// LP Data Handler for UDP data plane, act as a pipeline driver and buffer for delaying packet pub struct LpDataHandler { + /// Shared state + _shared_state: Arc, + /// Channel to receive incoming data input_rx: mpsc::Receiver, @@ -60,16 +62,12 @@ impl LpDataHandler { output_tx: tokio::sync::mpsc::Sender<(EncryptedLpPacket, SocketAddr)>, shutdown: nym_task::ShutdownToken, ) -> Self { + let shared_state = Arc::new(state); Self { + _shared_state: shared_state.clone(), input_rx, output_tx, - pipeline: MixnodeDataPipeline::new( - state, - MixnodeDataPipelineConfig { - fragment_timeout: FRAGMENT_TIMEOUT_DURATION, - }, - OsRng, - ), + pipeline: MixnodeDataPipeline::new(shared_state, OsRng), outgoing_pkt_buffer: Vec::new(), shutdown, } diff --git a/nym-node/src/node/lp/data/handler/pipeline.rs b/nym-node/src/node/lp/data/handler/pipeline.rs index a7b094c2a8..521ce9cb55 100644 --- a/nym-node/src/node/lp/data/handler/pipeline.rs +++ b/nym-node/src/node/lp/data/handler/pipeline.rs @@ -1,10 +1,7 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use std::{ - net::SocketAddr, - time::{Duration, Instant}, -}; +use std::{net::SocketAddr, sync::Arc, time::Instant}; use nym_lp_data::{ AddressedTimedData, PipelinePayload, TimedData, TimedPayload, @@ -12,7 +9,7 @@ use nym_lp_data::{ Framing, FramingUnwrap, Transport, TransportUnwrap, WireUnwrappingPipeline, WireWrappingPipeline, }, - fragmentation::{fragment::fragment_payload, reconstruction::MessageReconstructor}, + fragmentation::fragment::fragment_payload, mixnodes::traits::MixnodeProcessingPipeline, packet::{ EncryptedLpPacket, LpFrame, LpHeader, LpPacket, MalformedLpPacketError, @@ -27,19 +24,12 @@ use crate::node::lp::data::{ shared::SharedLpDataState, }; -#[derive(Clone, Copy, Debug)] -pub struct MixnodeDataPipelineConfig { - pub fragment_timeout: Duration, -} - pub struct MixnodeDataPipeline where R: Rng, { - config: MixnodeDataPipelineConfig, /// Shared data state - state: SharedLpDataState, - fragment_reconstructor: MessageReconstructor, + state: Arc, rng: R, } @@ -47,13 +37,8 @@ impl MixnodeDataPipeline where R: Rng, { - pub fn new(state: SharedLpDataState, config: MixnodeDataPipelineConfig, rng: R) -> Self { - Self { - state, - config, - rng, - fragment_reconstructor: MessageReconstructor::new(config.fragment_timeout), - } + pub fn new(state: Arc, rng: R) -> Self { + Self { state, rng } } } @@ -190,7 +175,8 @@ where .inspect_err(|e| tracing::error!("Failed to recover a fragment : {e}")) .ok()?; let (payload, metadata) = self - .fragment_reconstructor + .state + .message_reconstructor .insert_new_fragment(fragment, frame.timestamp)?; let message_kind = metadata .try_into() diff --git a/nym-node/src/node/lp/data/shared.rs b/nym-node/src/node/lp/data/shared.rs index 2958076184..c1d00eed4e 100644 --- a/nym-node/src/node/lp/data/shared.rs +++ b/nym-node/src/node/lp/data/shared.rs @@ -6,6 +6,7 @@ use crate::config::LpConfig; use crate::node::key_rotation::active_keys::ActiveSphinxKeys; use crate::node::key_rotation::active_keys::SphinxKeyGuard; use crate::node::replay_protection::bloomfilter::ReplayProtectionBloomfilters; +use nym_lp_data::fragmentation::reconstruction::MessageReconstructor; use nym_node_metrics::NymNodeMetrics; use nym_node_metrics::mixnet::PacketKind; use nym_sphinx_framing::processing::{ @@ -15,35 +16,19 @@ use nym_sphinx_params::SphinxKeyRotation; use nym_task::ShutdownToken; use std::net::IpAddr; use std::time::Duration; +use std::time::Instant; use tracing::Span; use tracing::warn; #[derive(Clone, Copy)] pub(crate) struct ProcessingConfig { pub(crate) maximum_packet_delay: Duration, - // /// how long the task is willing to skip mutex acquisition before it will block the thread - // /// until it actually obtains it - // pub(crate) maximum_replay_detection_deferral: Duration, - - // /// how many packets the task is willing to queue before it will block the thread - // /// until it obtains the mutex - // pub(crate) maximum_replay_detection_pending_packets: usize, } impl ProcessingConfig { pub(crate) fn new(config: &Config) -> Self { ProcessingConfig { maximum_packet_delay: config.mixnet.debug.maximum_forward_packet_delay, - // maximum_replay_detection_deferral: config - // .mixnet - // .replay_protection - // .debug - // .maximum_replay_detection_deferral, - // maximum_replay_detection_pending_packets: config - // .mixnet - // .replay_protection - // .debug - // .maximum_replay_detection_pending_packets, } } } @@ -60,6 +45,8 @@ pub(crate) struct SharedLpDataState { pub replay_protection_filter: ReplayProtectionBloomfilters, + pub message_reconstructor: MessageReconstructor, + /// Metrics collection pub metrics: NymNodeMetrics, @@ -86,6 +73,7 @@ impl SharedLpDataState { lp_config: config.lp, sphinx_keys, replay_protection_filter, + message_reconstructor: Default::default(), metrics, shutdown_token, }