Skip to content

Commit dfe7fb3

Browse files
thedevbirbthedevbirb
authored
feat(msg-sim): configurable runtime (#180)
2 parents 4ac3712 + f60efc4 commit dfe7fb3

2 files changed

Lines changed: 108 additions & 42 deletions

File tree

‎msg-sim/src/namespace.rs‎

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use tokio::sync::oneshot;
88

99
use crate::dynch::{DynCh, DynRequestSender};
1010
use crate::namespace::helpers::current_netns;
11+
use crate::network::RuntimeFactory;
1112

1213
/// Base directory for named network namespaces.
1314
///
@@ -230,6 +231,7 @@ impl NetworkNamespaceInner {
230231
/// `setns(2)`, which is thread-local.
231232
pub fn spawn<Ctx: 'static>(
232233
self,
234+
runtime_factory: RuntimeFactory,
233235
make_ctx: impl FnOnce() -> Ctx + Send + 'static,
234236
) -> (std::thread::JoinHandle<Result<()>>, DynRequestSender<Ctx>) {
235237
let (tx, mut rx) = DynCh::<Ctx>::channel(8);
@@ -245,7 +247,7 @@ impl NetworkNamespaceInner {
245247
// Create mount namespace and remount /proc for namespace-specific sysctl access
246248
helpers::setup_mount_namespace()?;
247249

248-
let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
250+
let rt = runtime_factory();
249251

250252
tracing::debug!("started runtime");
251253
drop(_span);
@@ -286,6 +288,7 @@ pub struct NetworkNamespace<Ctx = ()> {
286288
impl NetworkNamespace {
287289
pub async fn new<Ctx: 'static>(
288290
name: impl Into<String>,
291+
runtime_factory: RuntimeFactory,
289292
make_ctx: impl FnOnce() -> Ctx + Send + 'static,
290293
) -> Result<NetworkNamespace<Ctx>> {
291294
let name = name.into();
@@ -296,7 +299,7 @@ impl NetworkNamespace {
296299
let file = tokio::fs::File::open(path).await?.into_std().await;
297300

298301
let inner = NetworkNamespaceInner { name, file };
299-
let (_receiver_task, task_sender) = inner.try_clone()?.spawn(make_ctx);
302+
let (_receiver_task, task_sender) = inner.try_clone()?.spawn(runtime_factory, make_ctx);
300303

301304
Ok(NetworkNamespace::<Ctx> { inner, task_sender, _receiver_task })
302305
}
@@ -349,11 +352,19 @@ mod tests {
349352

350353
const TCP_SLOW_START_AFTER_IDLE: &str = "/proc/sys/net/ipv4/tcp_slow_start_after_idle";
351354

355+
fn default_runtime() -> tokio::runtime::Runtime {
356+
tokio::runtime::Builder::new_multi_thread().enable_all().build().expect("to create runtime")
357+
}
358+
352359
#[tokio::test(flavor = "multi_thread")]
353360
async fn mount_namespace_isolates_proc() {
354361
// Create two namespaces
355-
let ns1 = NetworkNamespace::new("test-ns-mount-1", || ()).await.unwrap();
356-
let ns2 = NetworkNamespace::new("test-ns-mount-2", || ()).await.unwrap();
362+
let ns1 = NetworkNamespace::new("test-ns-mount-1", Box::new(default_runtime), || ())
363+
.await
364+
.unwrap();
365+
let ns2 = NetworkNamespace::new("test-ns-mount-2", Box::new(default_runtime), || ())
366+
.await
367+
.unwrap();
357368

358369
// Verify /proc is mounted in ns1 by checking /proc/self/ns/net exists
359370
let proc_mounted_ns1: bool = ns1
@@ -385,8 +396,12 @@ mod tests {
385396
#[tokio::test(flavor = "multi_thread")]
386397
async fn sysctl_values_are_namespace_specific() {
387398
// Create two namespaces
388-
let ns1 = NetworkNamespace::new("test-ns-sysctl-1", || ()).await.unwrap();
389-
let ns2 = NetworkNamespace::new("test-ns-sysctl-2", || ()).await.unwrap();
399+
let ns1 = NetworkNamespace::new("test-ns-sysctl-1", Box::new(default_runtime), || ())
400+
.await
401+
.unwrap();
402+
let ns2 = NetworkNamespace::new("test-ns-sysctl-2", Box::new(default_runtime), || ())
403+
.await
404+
.unwrap();
390405

391406
// Set different values in each namespace
392407
let write_result_ns1: std::io::Result<()> = ns1
@@ -446,24 +461,21 @@ mod tests {
446461

447462
assert_eq!(value_ns1, "0", "ns1 should have tcp_slow_start_after_idle=0");
448463
assert_eq!(value_ns2, "1", "ns2 should have tcp_slow_start_after_idle=1");
449-
assert_ne!(
450-
value_ns1, value_ns2,
451-
"sysctls should be isolated between namespaces"
452-
);
464+
assert_ne!(value_ns1, value_ns2, "sysctls should be isolated between namespaces");
453465
}
454466

455467
#[tokio::test(flavor = "multi_thread")]
456468
async fn namespace_has_isolated_network_identity() {
457469
// Create a namespace
458-
let ns = NetworkNamespace::new("test-ns-identity", || ()).await.unwrap();
470+
let ns = NetworkNamespace::new("test-ns-identity", Box::new(default_runtime), || ())
471+
.await
472+
.unwrap();
459473

460474
// Get the network namespace inode from inside the namespace
461475
let ns_inode_inside: u64 = ns
462476
.task_sender
463477
.submit(|_: &mut ()| -> DynFuture<'_, u64> {
464-
Box::pin(async {
465-
helpers::current_netns().map(|id| id.inode).unwrap_or(0)
466-
})
478+
Box::pin(async { helpers::current_netns().map(|id| id.inode).unwrap_or(0) })
467479
})
468480
.await
469481
.unwrap()
@@ -476,9 +488,6 @@ mod tests {
476488

477489
assert_ne!(ns_inode_inside, 0, "should get valid inode inside namespace");
478490
assert_ne!(host_inode, 0, "should get valid host inode");
479-
assert_ne!(
480-
ns_inode_inside, host_inode,
481-
"namespace inode should differ from host"
482-
);
491+
assert_ne!(ns_inode_inside, host_inode, "namespace inode should differ from host");
483492
}
484493
}

‎msg-sim/src/network.rs‎

Lines changed: 81 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@
5252
//! classes. See the [`crate::tc`] module for details on the qdisc hierarchy.
5353
5454
use std::{
55-
any::Any,
5655
collections::{HashMap, HashSet},
5756
fmt::{Debug, Display},
5857
io,
@@ -137,8 +136,8 @@ pub struct Link(pub PeerId, pub PeerId);
137136
impl Link {
138137
/// Create a new directed link from source to destination.
139138
#[inline]
140-
pub fn new(source: PeerId, destination: PeerId) -> Self {
141-
Link(source, destination)
139+
pub fn new(source: impl Into<PeerId>, destination: impl Into<PeerId>) -> Self {
140+
Link(source.into(), destination.into())
142141
}
143142

144143
/// Get the source peer (traffic originates here).
@@ -194,7 +193,7 @@ impl PeerTcState {
194193
}
195194

196195
/// Map from peer ID to peer instance.
197-
pub type PeerMap = HashMap<PeerId, Peer<Context>>;
196+
pub type PeerMap = HashMap<PeerId, Peer<PeerContext>>;
198197

199198
/// Map from peer ID to traffic control state.
200199
type TcStateMap = HashMap<PeerId, PeerTcState>;
@@ -217,21 +216,66 @@ impl Peer {
217216
}
218217
}
219218

219+
pub(crate) type RuntimeFactory = Box<dyn FnOnce() -> tokio::runtime::Runtime + Send>;
220+
221+
pub fn default_runtime_factory() -> RuntimeFactory {
222+
Box::new(|| {
223+
tokio::runtime::Builder::new_multi_thread().enable_all().build().expect("to create runtime")
224+
})
225+
}
226+
227+
/// Common context provided to all namespaces.
228+
///
229+
/// This context gives access to rtnetlink for network configuration.
230+
#[derive(Debug)]
231+
pub struct CommonContext {
232+
/// Handle for sending rtnetlink messages within this namespace.
233+
handle: rtnetlink::Handle,
234+
/// Background task processing rtnetlink responses.
235+
_connection_task: tokio::task::JoinHandle<()>,
236+
}
237+
220238
/// Context provided to tasks running within a peer's namespace.
221239
///
222240
/// This context gives access to rtnetlink for network configuration
223241
/// and metadata about the peer's position in the network.
224242
#[derive(Debug)]
225-
pub struct Context {
243+
pub struct PeerContext {
226244
/// Handle for sending rtnetlink messages within this namespace.
227-
handle: rtnetlink::Handle,
245+
pub handle: rtnetlink::Handle,
228246
/// Background task processing rtnetlink responses.
229247
_connection_task: tokio::task::JoinHandle<()>,
230-
231248
/// The subnet this network uses.
232-
subnet: Subnet,
249+
pub subnet: Subnet,
233250
/// This peer's ID.
234-
peer_id: usize,
251+
pub peer_id: PeerId,
252+
}
253+
254+
/// Options for configuring a peer.
255+
pub struct PeerOptions {
256+
runtime_factory: RuntimeFactory,
257+
}
258+
259+
impl Default for PeerOptions {
260+
fn default() -> Self {
261+
Self {
262+
runtime_factory: Box::new(|| {
263+
tokio::runtime::Builder::new_multi_thread()
264+
.enable_all()
265+
.build()
266+
.expect("to create runtime")
267+
}),
268+
}
269+
}
270+
}
271+
272+
impl PeerOptions {
273+
/// Create new peer options with a custom runtime factory.
274+
pub fn with_runtime(
275+
runtime_factory: impl FnOnce() -> tokio::runtime::Runtime + Send + 'static,
276+
) -> Self {
277+
Self { runtime_factory: Box::new(runtime_factory) }
278+
}
235279
}
236280

237281
// -------------------------------------------------------------------------------------
@@ -289,7 +333,7 @@ pub type Result<T> = std::result::Result<T, Error>;
289333
/// # Example
290334
///
291335
/// ```no_run
292-
/// use msg_sim::network::{Network, Link};
336+
/// use msg_sim::network::{Network, Link, PeerOptions};
293337
/// use msg_sim::tc::impairment::LinkImpairment;
294338
/// use msg_sim::ip::Subnet;
295339
/// use std::net::Ipv4Addr;
@@ -344,7 +388,7 @@ pub struct Network {
344388
subnet: Subnet,
345389

346390
/// The hub namespace containing the bridge device.
347-
network_hub_namespace: NetworkNamespace<Context>,
391+
network_hub_namespace: NetworkNamespace<CommonContext>,
348392

349393
/// Rtnetlink handle bound to the host namespace.
350394
///
@@ -379,11 +423,13 @@ impl Network {
379423
.map(|(connection, handle, _)| (handle, tokio::task::spawn(connection)))
380424
.unwrap();
381425

382-
Context { handle, subnet, peer_id: 0, _connection_task }
426+
CommonContext { handle, _connection_task }
383427
};
384428

385429
// Create the hub namespace that will host the bridge.
386-
let namespace_hub = NetworkNamespace::new(Self::hub_namespace_name(), make_ctx).await?;
430+
let namespace_hub =
431+
NetworkNamespace::new(Self::hub_namespace_name(), default_runtime_factory(), make_ctx)
432+
.await?;
387433
let fd = namespace_hub.fd();
388434

389435
let network = Self {
@@ -418,7 +464,7 @@ impl Network {
418464
/// 1. A new network namespace for the peer
419465
/// 2. A veth pair connecting the peer to the hub bridge
420466
/// 3. IP address assignment based on the subnet and peer ID
421-
pub async fn add_peer(&mut self) -> Result<PeerId> {
467+
pub async fn add_peer_with_options(&mut self, options: PeerOptions) -> Result<PeerId> {
422468
let peer_id = PEER_ID_NEXT.load(Ordering::Relaxed);
423469
let namespace_name = peer_id.namespace_name();
424470
let veth_name = Arc::new(peer_id.veth_name());
@@ -435,10 +481,12 @@ impl Network {
435481
.map(|(connection, handle, _)| (handle, tokio::task::spawn(connection)))
436482
.expect("to create rtnetlink socket");
437483

438-
Context { handle, peer_id, subnet, _connection_task }
484+
PeerContext { handle, _connection_task, subnet, peer_id }
439485
};
440486

441-
let network_namespace = NetworkNamespace::new(namespace_name.clone(), make_ctx).await?;
487+
let network_namespace =
488+
NetworkNamespace::new(namespace_name.clone(), options.runtime_factory, make_ctx)
489+
.await?;
442490

443491
// Step 1: Create the veth pair in the host namespace.
444492
// One end (veth_name) will go to the peer, the other (veth_br_name) to the bridge.
@@ -476,7 +524,7 @@ impl Network {
476524

477525
network_namespace
478526
.task_sender
479-
.submit(|ctx| {
527+
.submit(|ctx: &mut PeerContext| {
480528
Box::pin(async move {
481529
let address = ctx.peer_id.veth_address(ctx.subnet);
482530
let mask = ctx.subnet.netmask;
@@ -541,6 +589,11 @@ impl Network {
541589
Ok(peer_id)
542590
}
543591

592+
/// See [`Self::add_peer_with_options`].
593+
pub async fn add_peer(&mut self) -> Result<PeerId> {
594+
self.add_peer_with_options(PeerOptions::default()).await
595+
}
596+
544597
/// Run a task in a peer's network namespace.
545598
///
546599
/// The provided closure receives a mutable reference to the namespace's context,
@@ -560,7 +613,7 @@ impl Network {
560613
///
561614
/// ```no_run
562615
/// use msg_sim::ip::Subnet;
563-
/// use msg_sim::network::Network;
616+
/// use msg_sim::network::{Network, PeerOptions};
564617
/// use std::net::Ipv4Addr;
565618
/// use tokio::net::TcpListener;
566619
///
@@ -589,8 +642,8 @@ impl Network {
589642
fut: F,
590643
) -> Result<impl Future<Output = std::result::Result<T, oneshot::error::RecvError>>>
591644
where
592-
T: Any + Send + 'static,
593-
F: for<'a> FnOnce(&'a mut Context) -> DynFuture<'a, T> + Send + 'static,
645+
T: Send + 'static,
646+
F: for<'a> FnOnce(&'a mut PeerContext) -> DynFuture<'a, T> + Send + 'static,
594647
{
595648
let Some(peer) = self.peers.get(&peer_id) else {
596649
return Err(Error::PeerNotFound(peer_id));
@@ -630,7 +683,7 @@ impl Network {
630683
/// ```no_run
631684
/// use msg_sim::{
632685
/// ip::Subnet,
633-
/// network::{Link, Network},
686+
/// network::{Link, Network, PeerOptions},
634687
/// tc::impairment::LinkImpairment
635688
/// };
636689
///
@@ -681,7 +734,7 @@ impl Network {
681734
src_peer
682735
.namespace
683736
.task_sender
684-
.submit(move |ctx| {
737+
.submit(move |ctx: &mut PeerContext| {
685738
let span = tracing::debug_span!(
686739
"apply_impairment",
687740
link = %link,
@@ -996,13 +1049,17 @@ mod msg_sim_network {
9961049
req_socket_2.connect_sync(SocketAddr::new(address_2, port));
9971050
req_socket_3.connect_sync(SocketAddr::new(address_3, port));
9981051

999-
// Measure RTT to peer 2 (should be ~200ms round trip for 100ms one-way)
1052+
// Wait for both TCP connections to be established before starting
1053+
// measurements. Peer 3 has 500ms one-way latency, so TCP handshake takes ~1s.
1054+
tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
1055+
1056+
// Measure RTT to peer 2 (should be ~100ms for one-way latency)
10001057
let start = Instant::now();
10011058
let resp = req_socket_2.request("ping".into()).await.unwrap();
10021059
let rtt_2 = start.elapsed();
10031060
assert_eq!(resp.as_ref(), b"peer2");
10041061

1005-
// Measure RTT to peer 3 (should be ~1000ms round trip for 500ms one-way)
1062+
// Measure RTT to peer 3 (should be ~500ms for one-way latency)
10061063
let start = Instant::now();
10071064
let resp = req_socket_3.request("ping".into()).await.unwrap();
10081065
let rtt_3 = start.elapsed();

0 commit comments

Comments
 (0)