Still wait if there is none

mozilla · Apr 18, 2024 · 710c500 · 710c500
1 parent f7a8f6c
commit 710c500
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 29 deletions.
diff --git a/neqo-bin/src/client/http09.rs b/neqo-bin/src/client/http09.rs
@@ -95,10 +95,16 @@ impl<'a> super::Handler for Handler<'a> {
         Ok(false)
     }
 
-    fn take_token(&mut self, client: &mut Self::Client) -> Option<ResumptionToken> {
-        self.token
-            .take()
-            .or_else(|| client.take_resumption_token(Instant::now()))
+    fn take_token(&mut self) -> Option<ResumptionToken> {
+        self.token.take()
+    }
+
+    fn has_token(&mut self, client: &mut Self::Client) -> bool {
+        if self.token.is_some() {
+            return true;
+        }
+        self.token = client.take_resumption_token(Instant::now());
+        self.token.is_some()
     }
 }
 

diff --git a/neqo-bin/src/client/http3.rs b/neqo-bin/src/client/http3.rs
@@ -223,10 +223,16 @@ impl<'a> super::Handler for Handler<'a> {
         Ok(self.url_handler.done())
     }
 
-    fn take_token(&mut self, client: &mut Self::Client) -> Option<ResumptionToken> {
-        self.token
-            .take()
-            .or_else(|| client.take_resumption_token(Instant::now()))
+    fn take_token(&mut self) -> Option<ResumptionToken> {
+        self.token.take()
+    }
+
+    fn has_token(&mut self, client: &mut Self::Client) -> bool {
+        if self.token.is_some() {
+            return true;
+        }
+        self.token = client.take_resumption_token(Instant::now());
+        self.token.is_some()
     }
 }
 

diff --git a/neqo-bin/src/client/mod.rs b/neqo-bin/src/client/mod.rs
@@ -341,7 +341,8 @@ trait Handler {
     type Client: Client;
 
     fn handle(&mut self, client: &mut Self::Client) -> Res<bool>;
-    fn take_token(&mut self, client: &mut Self::Client) -> Option<ResumptionToken>;
+    fn take_token(&mut self) -> Option<ResumptionToken>;
+    fn has_token(&mut self, client: &mut Self::Client) -> bool;
 }
 
 /// Network client, e.g. [`neqo_transport::Connection`] or [`neqo_http3::Http3Client`].
@@ -373,20 +374,35 @@ struct Runner<'a, H: Handler> {
 
 impl<'a, H: Handler> Runner<'a, H> {
     async fn run(mut self) -> Res<Option<ResumptionToken>> {
-        let close_reason = loop {
+        loop {
             let handler_done = self.handler.handle(&mut self.client)?;
-            self.process_output().await?;
 
-            match (handler_done, self.client.is_closed()) {
-                // more work
-                (false, _) => {}
-                // no more work, closing connection
-                (true, None) => {
+            match (handler_done, self.args.resume, self.handler.has_token(&mut self.client)) {
+                // Handler isn't done. Continue.
+                (false, _, _) => {},
+                // Handler done. Resumption token needed but not present. Continue.
+                (true, true, false) => {
+                    qdebug!("Handler done. Waiting for resumption token.");
+                }
+                // Handler is done, no resumption token needed. Close.
+                (true, false, _) |
+                // Handler is done, resumption token needed and present. Close.
+                (true, true, true) => {
                     self.client.close(Instant::now(), 0, "kthxbye!");
-                    continue;
                 }
-                // no more work, connection closed, terminating
-                (true, Some(reason)) => break reason,
+            }
+
+            self.process_output().await?;
+
+            if let Some(reason) = self.client.is_closed() {
+                if self.args.stats {
+                    qinfo!("{:?}", self.client.stats());
+                }
+                return match reason {
+                    ConnectionError::Transport(TransportError::NoError)
+                    | ConnectionError::Application(0) => Ok(self.handler.take_token()),
+                    _ => Err(reason.into()),
+                };
             }
 
             if self.client.has_events() {
@@ -399,16 +415,6 @@ impl<'a, H: Handler> Runner<'a, H> {
                     self.timeout = None;
                 }
             }
-        };
-
-        if self.args.stats {
-            qinfo!("{:?}", self.client.stats());
-        }
-
-        match close_reason {
-            ConnectionError::Transport(TransportError::NoError)
-            | ConnectionError::Application(0) => Ok(self.handler.take_token(&mut self.client)),
-            _ => Err(close_reason.into()),
         }
     }