From fc44a97925438c0f715d81f6baa816d098f188c5 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 23 Feb 2021 10:54:48 +1000
Subject: [PATCH] Revert "remove unnecessary Option around request timeout"

This reverts commit c3724031df237f9c58f2e13f8d89b3f97a77a922.
---
 zebra-network/src/peer/connection.rs | 27 ++++++++++++++++++---------
 zebra-network/src/peer/handshake.rs  |  1 +
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/zebra-network/src/peer/connection.rs b/zebra-network/src/peer/connection.rs
index 9bccfbce6..91990ac1d 100644
--- a/zebra-network/src/peer/connection.rs
+++ b/zebra-network/src/peer/connection.rs
@@ -324,7 +324,6 @@ pub(super) enum State {
         /// internal Response format.
         tx: MustUseOneshotSender<Result<Response, SharedPeerError>>,
         span: tracing::Span,
-        request_timer: Sleep,
     },
 }
 
@@ -404,12 +403,15 @@ impl State {
                 span,
                 mut tx,
                 mut handler,
-                request_timer,
             } => {
                 // we have to get rid of the span reference so we can tamper with the state
                 let span = span.clone();
                 trace!(parent: &span, "awaiting response to client request");
-                let cancel = future::select(request_timer, tx.cancellation());
+                let timer_ref = conn
+                    .request_timer
+                    .as_mut()
+                    .expect("timeout must be set while awaiting response");
+                let cancel = future::select(timer_ref, tx.cancellation());
                 match future::select(peer_rx.next(), cancel)
                     .instrument(span.clone())
                     .await
@@ -498,12 +500,9 @@ impl TryFrom<Transition> for State {
     fn try_from(trans: Transition) -> Result<Self, Self::Error> {
         match trans {
             Transition::AwaitRequest => Ok(State::AwaitingRequest),
-            Transition::AwaitResponse { handler, tx, span } => Ok(State::AwaitingResponse {
-                handler,
-                tx,
-                span,
-                request_timer: sleep(constants::REQUEST_TIMEOUT),
-            }),
+            Transition::AwaitResponse { handler, tx, span } => {
+                Ok(State::AwaitingResponse { handler, tx, span })
+            }
             Transition::ClientClose => Err(None),
             Transition::Close(e) => Err(Some(e)),
             Transition::CloseResponse { tx, e } => {
@@ -517,6 +516,12 @@ impl TryFrom<Transition> for State {
 /// The state associated with a peer connection.
 pub struct Connection<S, Tx> {
     pub(super) state: Option<State>,
+    /// A timeout for a client request. This is stored separately from
+    /// State so that we can move the future out of it independently of
+    /// other state handling.
+    // I don't think this is necessary, and will try moving it into `State` in
+    // the next commit TODO(jane)
+    pub(super) request_timer: Option<Sleep>,
     pub(super) svc: S,
     /// A `mpsc::Receiver<ClientRequest>` that converts its results to
     /// `InProgressClientRequest`
@@ -544,6 +549,10 @@ where
                 .step(&mut self, &mut peer_rx)
                 .await;
 
+            if matches!(transition, Transition::AwaitResponse { .. }) {
+                self.request_timer = Some(sleep(constants::REQUEST_TIMEOUT));
+            }
+
             self.state = match transition.try_into() {
                 Ok(state) => Some(state),
                 Err(None) => {
diff --git a/zebra-network/src/peer/handshake.rs b/zebra-network/src/peer/handshake.rs
index 7e2b3d1ae..92fc366dd 100644
--- a/zebra-network/src/peer/handshake.rs
+++ b/zebra-network/src/peer/handshake.rs
@@ -436,6 +436,7 @@ where
                 svc: inbound_service,
                 client_rx: server_rx.into(),
                 peer_tx,
+                request_timer: None,
             };
 
             tokio::spawn(