From fb186cb034f0c53509de32c84b90a807bc384bd6 Mon Sep 17 00:00:00 2001
From: Marc Bowes <15209+marcbowes@users.noreply.github.com>
Date: Thu, 3 Oct 2024 17:03:04 -0700
Subject: [PATCH] Add established_tcp_stream_count{,_on}

This commit adds a utility method to count the number of established tcp
streams on a given host (or the current host). This can be used to
verify connections have been cleaned up, which is sometimes tricky to
reason about due to a combination of how tcp cleans up and how Drop
works.
---
 src/host.rs  |  4 ++++
 src/lib.rs   | 10 ++++++++++
 src/world.rs | 12 ++++++++++++
 tests/tcp.rs |  7 +++++++
 4 files changed, 33 insertions(+)

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