Add established_tcp_stream_count{,_on}

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()
     }

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))
+}

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,

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(())
     });