From 4bba80b673eeb5a7121bdfdb4fbcda5a36dfe2a6 Mon Sep 17 00:00:00 2001
From: "Demi M. Obenour" <demi@parity.io>
Date: Sat, 7 Mar 2020 14:56:50 -0500
Subject: [PATCH] =?UTF-8?q?Implement=20=E2=80=98Connection::send=5Fstreams?=
 =?UTF-8?q?()=E2=80=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This allows users to know when all data has been acknowledged.
---
 quinn-proto/src/connection.rs |  5 +++++
 quinn-proto/src/streams.rs    | 12 ++++++++++++
 quinn-proto/src/tests/mod.rs  | 11 +++++++++++
 3 files changed, 28 insertions(+)

diff --git a/quinn-proto/src/connection.rs b/quinn-proto/src/connection.rs
index 4de7d2a1b..d96da39cc 100644
--- a/quinn-proto/src/connection.rs
+++ b/quinn-proto/src/connection.rs
@@ -2708,6 +2708,11 @@ where
         Some(id)
     }
 
+    /// The number of streams that may have unacknowledged data.
+    pub fn send_streams(&self) -> usize {
+        self.streams.send_streams()
+    }
+
     /// Finish a send stream, signalling that no more data will be sent
     pub fn finish(&mut self, id: StreamId) -> Result<(), FinishError> {
         let ss = self
diff --git a/quinn-proto/src/streams.rs b/quinn-proto/src/streams.rs
index 3c83e78e0..0eb0e34e2 100644
--- a/quinn-proto/src/streams.rs
+++ b/quinn-proto/src/streams.rs
@@ -21,6 +21,8 @@ pub(crate) struct Streams {
     pub next_remote: [u64; 2],
     // Next to report to the application, once opened
     next_reported_remote: [u64; 2],
+    // Outbound streams
+    send_streams: usize,
 }
 
 impl Streams {
@@ -33,6 +35,7 @@ impl Streams {
             max_remote: [max_remote_bi, max_remote_uni],
             next_remote: [0, 0],
             next_reported_remote: [0, 0],
+            send_streams: 0,
         };
 
         for dir in Dir::iter() {
@@ -52,9 +55,14 @@ impl Streams {
         self.next[dir as usize] += 1;
         let id = StreamId::new(side, dir, self.next[dir as usize] - 1);
         self.insert(false, id);
+        self.send_streams += 1;
         Some(id)
     }
 
+    pub fn send_streams(&self) -> usize {
+        self.send_streams
+    }
+
     pub fn alloc_remote_stream(&mut self, side: Side, dir: Dir) {
         self.max_remote[dir as usize] += 1;
         let id = StreamId::new(!side, dir, self.max_remote[dir as usize] - 1);
@@ -67,6 +75,9 @@ impl Streams {
         }
         let x = self.next_reported_remote[dir as usize];
         self.next_reported_remote[dir as usize] = x + 1;
+        if dir == Dir::Bi {
+            self.send_streams += 1;
+        }
         Some(StreamId::new(!side, dir, x))
     }
 
@@ -162,6 +173,7 @@ impl Streams {
             hash_map::Entry::Vacant(_) => {}
             hash_map::Entry::Occupied(e) => {
                 if e.get().is_closed() {
+                    self.send_streams -= 1;
                     e.remove_entry();
                 }
             }
diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
index 91d4faf4e..a0e1477ef 100644
--- a/quinn-proto/src/tests/mod.rs
+++ b/quinn-proto/src/tests/mod.rs
@@ -191,6 +191,7 @@ fn finish_stream_simple() {
 
     const MSG: &[u8] = b"hello";
     pair.client_conn_mut(client_ch).write(s, MSG).unwrap();
+    assert_eq!(pair.client_conn_mut(client_ch).send_streams(), 1);
     pair.client_conn_mut(client_ch).finish(s).unwrap();
     pair.drive();
 
@@ -199,10 +200,14 @@ fn finish_stream_simple() {
         Some(Event::StreamFinished { stream, stop_reason: None }) if stream == s
     );
     assert_matches!(pair.client_conn_mut(client_ch).poll(), None);
+    assert_eq!(pair.client_conn_mut(client_ch).send_streams(), 0);
+    assert_eq!(pair.server_conn_mut(client_ch).send_streams(), 0);
     assert_matches!(
         pair.server_conn_mut(server_ch).poll(),
         Some(Event::StreamOpened { dir: Dir::Uni })
     );
+    // Receive-only streams do not get `StreamFinished` events
+    assert_eq!(pair.server_conn_mut(client_ch).send_streams(), 0);
     assert_matches!(pair.server_conn_mut(server_ch).accept(Dir::Uni), Some(stream) if stream == s);
     assert_matches!(pair.server_conn_mut(server_ch).poll(), None);
     assert_matches!(
@@ -909,7 +914,9 @@ fn stop_opens_bidi() {
     let _guard = subscribe();
     let mut pair = Pair::default();
     let (client_conn, server_conn) = pair.connect();
+    assert_eq!(pair.client_conn_mut(client_conn).send_streams(), 0);
     let s = pair.client_conn_mut(client_conn).open(Dir::Bi).unwrap();
+    assert_eq!(pair.client_conn_mut(client_conn).send_streams(), 1);
     const ERROR: VarInt = VarInt(42);
     pair.client
         .connections
@@ -923,7 +930,9 @@ fn stop_opens_bidi() {
         pair.server_conn_mut(server_conn).poll(),
         Some(Event::StreamOpened { dir: Dir::Bi })
     );
+    assert_eq!(pair.server_conn_mut(client_conn).send_streams(), 0);
     assert_matches!(pair.server_conn_mut(server_conn).accept(Dir::Bi), Some(stream) if stream == s);
+    assert_eq!(pair.server_conn_mut(client_conn).send_streams(), 1);
     assert_matches!(
         pair.server_conn_mut(server_conn).read_unordered(s),
         Err(ReadError::Blocked)
@@ -932,6 +941,8 @@ fn stop_opens_bidi() {
         pair.server_conn_mut(server_conn).write(s, b"foo"),
         Err(WriteError::Stopped(ERROR))
     );
+    assert_eq!(pair.server_conn_mut(client_conn).send_streams(), 0);
+    assert_matches!(pair.server_conn_mut(server_conn).poll(), None);
 }
 
 #[test]