fix: apply query_timeout to entire client/server exhange (#755)

Apply `query_timeout` to entire communication exchange between client
and server, not just individual server messages. Protects against
servers that stopped receiving mid transaction.

Author: levkk

pgdog/src/frontend/client/query_engine/multi_step/insert.rs

@@ -64,7 +64,7 @@ impl<'a> InsertMulti<'a> {
                 .await?;
 
             while self.engine.backend.has_more_messages() {
-                let message = self.engine.read_server_message(context).await?;
+                let message = self.engine.read_server_message().await?;
 
                 if self.state.forward(&message)? {
                     self.engine.process_server_message(context, message).await?;

pgdog/src/frontend/client/query_engine/multi_step/update.rs

@@ -172,7 +172,7 @@ impl<'a> UpdateMulti<'a> {
         let mut checker = ForwardCheck::new(context.client_request);
 
         while self.engine.backend.has_more_messages() {
-            let message = self.engine.read_server_message(context).await?;
+            let message = self.engine.read_server_message().await?;
             let code = message.code();
 
             if code == 'E' {
@@ -202,7 +202,7 @@ impl<'a> UpdateMulti<'a> {
             .await?;
 
         while self.engine.backend.has_more_messages() {
-            let message = self.engine.read_server_message(context).await?;
+            let message = self.engine.read_server_message().await?;
             self.engine.process_server_message(context, message).await?;
         }
 
@@ -236,7 +236,7 @@ impl<'a> UpdateMulti<'a> {
         let mut rows = 0;
 
         while self.engine.backend.has_more_messages() {
-            let message = self.engine.read_server_message(context).await?;
+            let message = self.engine.read_server_message().await?;
             match message.code() {
                 'D' => {
                     row.data_row = DataRow::try_from(message)?;

pgdog/src/frontend/client/query_engine/query.rs

@@ -58,6 +58,28 @@ impl QueryEngine {
             }
         }
 
+        match timeout(
+            context.timeouts.query_timeout(&State::Active),
+            self.client_server_exchange(context),
+        )
+        .await
+        {
+            Ok(response) => response?,
+            Err(err) => {
+                // Close the conn, it could be stuck executing a query
+                // or dead.
+                self.backend.force_close();
+                return Err(err.into());
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn client_server_exchange(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+    ) -> Result<(), Error> {
         match context.rewrite_result.take() {
             Some(RewriteResult::InsertSplit(requests)) => {
                 multi_step::InsertMulti::from_engine(self, requests)
@@ -75,7 +97,7 @@ impl QueryEngine {
                     && !self.streaming
                     && !self.test_mode.enabled
                 {
-                    let message = self.read_server_message(context).await?;
+                    let message = self.read_server_message().await?;
                     self.process_server_message(context, message).await?;
                 }
             }
@@ -90,26 +112,8 @@ impl QueryEngine {
         Ok(())
     }
 
-    pub async fn read_server_message(
-        &mut self,
-        context: &mut QueryEngineContext<'_>,
-    ) -> Result<Message, Error> {
-        let message = match timeout(
-            context.timeouts.query_timeout(&State::Active),
-            self.backend.read(),
-        )
-        .await
-        {
-            Ok(response) => response?,
-            Err(err) => {
-                // Close the conn, it could be stuck executing a query
-                // or dead.
-                self.backend.force_close();
-                return Err(err.into());
-            }
-        };
-
-        Ok(message)
+    pub async fn read_server_message(&mut self) -> Result<Message, Error> {
+        Ok(self.backend.read().await?)
     }
 
     pub async fn process_server_message(