prepare multi threaded node

Cargo.lock

@@ -7490,6 +7490,7 @@ name = "nym-lp-data"
 version = "1.21.0"
 dependencies = [
  "bytes",
+ "dashmap",
  "num_enum",
  "nym-common",
  "nym-lp",

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

common/nym-lp-data/src/fragmentation/reconstruction.rs

@@ -1,12 +1,17 @@
 // Copyright 2026 - Nym Technologies SA <contact@nymtech.net>
 // 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<FragmentHashKey, MessageBuffer<Ts>>,
+    in_flight_messages: Arc<DashMap<FragmentHashKey, MessageBuffer<Ts>>>,
 
     /// 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<u8> {
-        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<u8>, 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<Instant, Duration> {
+    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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(60);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(10);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(10);
+        let rec = MessageReconstructor::<u64, u64>::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::<u64, u64>::new(10);
+        let rec = MessageReconstructor::<u64, u64>::new(10);
 
         assert!(
             rec.insert_new_fragment(make_fragment(1, 3, 0, SPHINX, vec![0xa]), 0)

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<SharedLpDataState>,
+
     /// Channel to receive incoming data
     input_rx: mpsc::Receiver<EncryptedLpPacket>,
 
@@ -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,
         }

nym-node/src/node/lp/data/handler/pipeline.rs

@@ -1,10 +1,7 @@
 // Copyright 2026 - Nym Technologies SA <contact@nymtech.net>
 // 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<R>
 where
     R: Rng,
 {
-    config: MixnodeDataPipelineConfig,
     /// Shared data state
-    state: SharedLpDataState,
-    fragment_reconstructor: MessageReconstructor<Instant, Duration>,
+    state: Arc<SharedLpDataState>,
     rng: R,
 }
 
@@ -47,13 +37,8 @@ impl<R> MixnodeDataPipeline<R>
 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<SharedLpDataState>, 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()

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<Instant, Duration>,
+
     /// 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,
         }