diff --git a/src/host.rs b/src/host.rs index 55d3dd7..c64e030 100644 --- a/src/host.rs +++ b/src/host.rs @@ -361,6 +361,10 @@ impl Tcp { rx } + pub(crate) fn stream_count(&self) -> usize { + self.sockets.len() + } + pub(crate) fn accept(&mut self, addr: SocketAddr) -> Option<(Syn, SocketAddr)> { self.binds[&addr.port()].deque.pop_front() } diff --git a/src/lib.rs b/src/lib.rs index 2eaecd7..dc61c30 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -246,3 +246,13 @@ pub fn partition(a: impl ToIpAddrs, b: impl ToIpAddrs) { pub fn repair(a: impl ToIpAddrs, b: impl ToIpAddrs) { World::current(|world| world.repair_many(a, b)) } + +/// Return the number of established tcp streams on the current host. +pub fn established_tcp_stream_count() -> usize { + World::current(|world| world.est_tcp_streams()) +} + +/// Return the number of established tcp streams on the given host. +pub fn established_tcp_stream_count_on(addr: impl ToIpAddr) -> usize { + World::current(|world| world.est_tcp_streams_on(addr)) +} diff --git a/src/world.rs b/src/world.rs index 2fd2698..fcda2ad 100644 --- a/src/world.rs +++ b/src/world.rs @@ -170,6 +170,18 @@ impl World { }); } + pub(crate) fn est_tcp_streams(&mut self) -> usize { + self.current_host().tcp.stream_count() + } + + pub(crate) fn est_tcp_streams_on(&mut self, addr: impl ToIpAddr) -> usize { + self.hosts + .get(&self.dns.lookup(addr)) + .unwrap() + .tcp + .stream_count() + } + /// Register a new host with the simulation. pub(crate) fn register( &mut self, diff --git a/tests/tcp.rs b/tests/tcp.rs index 42f8c5b..54baa37 100644 --- a/tests/tcp.rs +++ b/tests/tcp.rs @@ -548,8 +548,15 @@ fn hangup() -> Result { sim.client("client", async move { let s = TcpStream::connect(("server", PORT)).await?; + assert_eq!(1, turmoil::established_tcp_stream_count()); + assert_eq!(1, turmoil::established_tcp_stream_count_on("server")); + drop(s); + assert_eq!(0, turmoil::established_tcp_stream_count()); + assert_eq!(1, turmoil::established_tcp_stream_count_on("server")); // server sleeps in its loop + wait.notified().await; + assert_eq!(0, turmoil::established_tcp_stream_count_on("server")); Ok(()) });