Bound retired stream tombstones

src/transport.rs

@@ -30,6 +30,7 @@ use tokio::net::UdpSocket;
 pub const DEFAULT_INITIAL_WINDOW: usize = 1024 * 1024;
 pub const DEFAULT_RETRANSMIT_AFTER: Duration = Duration::from_millis(200);
 pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2);
+pub const DEFAULT_RETIRED_STREAM_TOMBSTONES: usize = 16 * 1024;
 pub const SERVICE_TARGET_PREFIX: &str = "@dosh-";
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -37,6 +38,7 @@ pub struct TransportConfig {
     pub initial_window: usize,
     pub retransmit_after: Duration,
     pub keepalive_after: Duration,
+    pub retired_stream_tombstones: usize,
 }
 
 impl Default for TransportConfig {
@@ -45,6 +47,7 @@ impl Default for TransportConfig {
             initial_window: DEFAULT_INITIAL_WINDOW,
             retransmit_after: DEFAULT_RETRANSMIT_AFTER,
             keepalive_after: DEFAULT_KEEPALIVE_AFTER,
+            retired_stream_tombstones: DEFAULT_RETIRED_STREAM_TOMBSTONES,
         }
     }
 }
@@ -174,6 +177,7 @@ pub struct StreamMux {
     next_recv_offset: HashMap<u64, u64>,
     recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
     retired_streams: HashSet<u64>,
+    retired_stream_order: VecDeque<u64>,
 }
 
 #[derive(Debug, Clone, Default)]
@@ -230,6 +234,7 @@ impl StreamMux {
             next_recv_offset: HashMap::new(),
             recv_pending: HashMap::new(),
             retired_streams: HashSet::new(),
+            retired_stream_order: VecDeque::new(),
         }
     }
 
@@ -687,7 +692,17 @@ impl StreamMux {
 
     fn retire_stream(&mut self, stream_id: u64) {
         self.cleanup_stream(stream_id);
-        self.retired_streams.insert(stream_id);
+        if self.config.retired_stream_tombstones == 0 {
+            return;
+        }
+        if self.retired_streams.insert(stream_id) {
+            self.retired_stream_order.push_back(stream_id);
+        }
+        while self.retired_stream_order.len() > self.config.retired_stream_tombstones {
+            if let Some(expired) = self.retired_stream_order.pop_front() {
+                self.retired_streams.remove(&expired);
+            }
+        }
     }
 
     fn cleanup_stream(&mut self, stream_id: u64) {
@@ -1318,6 +1333,35 @@ mod tests {
         assert!(mux.retired_streams.contains(&1));
     }
 
+    #[test]
+    fn retired_stream_tombstones_are_bounded() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retired_stream_tombstones: 2,
+            ..TransportConfig::default()
+        });
+        for stream_id in 1..=3 {
+            mux.open_stream(stream_id, "@dosh-test", 0).unwrap();
+            mux.close_stream(stream_id).unwrap();
+        }
+
+        assert_eq!(mux.retired_streams.len(), 2);
+        assert!(!mux.retired_streams.contains(&1));
+        assert!(mux.retired_streams.contains(&2));
+        assert!(mux.retired_streams.contains(&3));
+
+        let stale_open = protocol::to_body(&StreamOpen {
+            stream_id: 3,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+        assert_eq!(
+            mux.handle_packet(PacketKind::StreamOpen, &stale_open)
+                .unwrap(),
+            TransportEvent::Ignored { stream_id: 3 }
+        );
+    }
+
     #[tokio::test]
     async fn session_transport_partitions_client_and_server_stream_ids() {
         let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();