Merge pull request #516 from cvybhu/retry_retry_policy

Fix RetryPolicy not retrying on Response::Error

Author: psarna

scylla-cql/src/frame/response/mod.rs

@@ -5,7 +5,7 @@ pub mod event;
 pub mod result;
 pub mod supported;
 
-use crate::frame::frame_errors::ParseError;
+use crate::{errors::QueryError, frame::frame_errors::ParseError};
 use num_enum::TryFromPrimitive;
 
 pub use error::Error;
@@ -57,4 +57,29 @@ impl Response {
 
         Ok(response)
     }
+
+    pub fn into_non_error_response(self) -> Result<NonErrorResponse, QueryError> {
+        Ok(match self {
+            Response::Error(err) => return Err(QueryError::from(err)),
+            Response::Ready => NonErrorResponse::Ready,
+            Response::Result(res) => NonErrorResponse::Result(res),
+            Response::Authenticate(auth) => NonErrorResponse::Authenticate(auth),
+            Response::AuthSuccess(auth_succ) => NonErrorResponse::AuthSuccess(auth_succ),
+            Response::AuthChallenge(auth_chal) => NonErrorResponse::AuthChallenge(auth_chal),
+            Response::Supported(sup) => NonErrorResponse::Supported(sup),
+            Response::Event(eve) => NonErrorResponse::Event(eve),
+        })
+    }
+}
+
+// A Response which can not be Response::Error
+#[derive(Debug)]
+pub enum NonErrorResponse {
+    Ready,
+    Result(result::Result),
+    Authenticate(authenticate::Authenticate),
+    AuthSuccess(authenticate::AuthSuccess),
+    AuthChallenge(authenticate::AuthChallenge),
+    Supported(Supported),
+    Event(event::Event),
 }

scylla/src/transport/connection.rs

@@ -32,7 +32,7 @@ use crate::batch::{Batch, BatchStatement};
 use crate::frame::{
     self,
     request::{self, batch, execute, query, register, Request},
-    response::{event::Event, result, Response, ResponseOpcode},
+    response::{event::Event, result, NonErrorResponse, Response, ResponseOpcode},
     server_event_type::EventType,
     value::{BatchValues, ValueList},
     FrameParams, SerializedRequest,
@@ -142,6 +142,13 @@ pub struct QueryResponse {
     pub warnings: Vec<String>,
 }
 
+// A QueryResponse in which response can not be Response::Error
+pub struct NonErrorQueryResponse {
+    pub response: NonErrorResponse,
+    pub tracing_id: Option<Uuid>,
+    pub warnings: Vec<String>,
+}
+
 /// Result of Session::batch(). Contains no rows, only some useful information.
 pub struct BatchResult {
     /// Warnings returned by the database
@@ -151,29 +158,42 @@ pub struct BatchResult {
 }
 
 impl QueryResponse {
+    pub fn into_non_error_query_response(self) -> Result<NonErrorQueryResponse, QueryError> {
+        Ok(NonErrorQueryResponse {
+            response: self.response.into_non_error_response()?,
+            tracing_id: self.tracing_id,
+            warnings: self.warnings,
+        })
+    }
+
+    pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
+        self.into_non_error_query_response()?.into_query_result()
+    }
+}
+
+impl NonErrorQueryResponse {
     pub fn as_set_keyspace(&self) -> Option<&result::SetKeyspace> {
         match &self.response {
-            Response::Result(result::Result::SetKeyspace(sk)) => Some(sk),
+            NonErrorResponse::Result(result::Result::SetKeyspace(sk)) => Some(sk),
             _ => None,
         }
     }
 
     pub fn as_schema_change(&self) -> Option<&result::SchemaChange> {
         match &self.response {
-            Response::Result(result::Result::SchemaChange(sc)) => Some(sc),
+            NonErrorResponse::Result(result::Result::SchemaChange(sc)) => Some(sc),
             _ => None,
         }
     }
 
     pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
         let (rows, paging_state, col_specs) = match self.response {
-            Response::Error(err) => return Err(err.into()),
-            Response::Result(result::Result::Rows(rs)) => (
+            NonErrorResponse::Result(result::Result::Rows(rs)) => (
                 Some(rs.rows),
                 rs.metadata.paging_state,
                 rs.metadata.col_specs,
             ),
-            Response::Result(_) => (None, None, vec![]),
+            NonErrorResponse::Result(_) => (None, None, vec![]),
             _ => {
                 return Err(QueryError::ProtocolError(
                     "Unexpected server response, expected Result or Error",

scylla/src/transport/session.rs

@@ -14,6 +14,7 @@ use tokio::time::timeout;
 use tracing::{debug, error, trace, trace_span, Instrument};
 use uuid::Uuid;
 
+use super::connection::NonErrorQueryResponse;
 use super::connection::QueryResponse;
 use super::errors::{BadQuery, NewSessionError, QueryError};
 use crate::cql_to_rust::FromRow;
@@ -470,6 +471,7 @@ impl Session {
                         connection
                             .query(query_ref, values_ref, paging_state_ref.clone())
                             .await
+                            .and_then(QueryResponse::into_non_error_query_response)
                     }
                 },
             )
@@ -484,7 +486,7 @@ impl Session {
 
     async fn handle_set_keyspace_response(
         &self,
-        response: &QueryResponse,
+        response: &NonErrorQueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(set_keyspace) = response.as_set_keyspace() {
             debug!(
@@ -501,7 +503,7 @@ impl Session {
     async fn handle_auto_await_schema_agreement(
         &self,
         contents: &str,
-        response: &QueryResponse,
+        response: &NonErrorQueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(timeout) = self.auto_await_schema_agreement_timeout {
             if response.as_schema_change().is_some()
@@ -748,7 +750,7 @@ impl Session {
             "Request",
             prepared_id = format!("{:X}", prepared.get_id()).as_str()
         );
-        let response = self
+        let response: NonErrorQueryResponse = self
             .run_query(
                 statement_info,
                 &prepared.config,
@@ -762,6 +764,7 @@ impl Session {
                     connection
                         .execute(prepared, values_ref, paging_state_ref.clone())
                         .await
+                        .and_then(QueryResponse::into_non_error_query_response)
                 },
             )
             .instrument(span)
@@ -1104,6 +1107,7 @@ impl Session {
     where
         ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
         QueryFut: Future<Output = Result<ResT, QueryError>>,
+        ResT: AllowedRunQueryResTType,
     {
         let cluster_data = self.cluster.get_data();
         let query_plan = self.load_balancer.plan(&statement_info, &cluster_data);
@@ -1196,6 +1200,7 @@ impl Session {
     where
         ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
         QueryFut: Future<Output = Result<ResT, QueryError>>,
+        ResT: AllowedRunQueryResTType,
     {
         let mut last_error: Option<QueryError> = None;
 
@@ -1303,6 +1308,7 @@ impl Session {
     ) -> Result<ResT, QueryError>
     where
         QueryFut: Future<Output = Result<ResT, QueryError>>,
+        ResT: AllowedRunQueryResTType,
     {
         let info = Statement::default();
         let config = StatementConfig {
@@ -1395,3 +1401,18 @@ async fn resolve_hostname(hostname: &str) -> Result<SocketAddr, NewSessionError>
 
     ret.ok_or(failed_err)
 }
+
+// run_query, execute_query, etc have a template type called ResT.
+// There was a bug where ResT was set to QueryResponse, which could
+// be an error response. This was not caught by retry policy which
+// assumed all errors would come from analyzing Result<ResT, QueryError>.
+// This trait is a guard to make sure that this mistake doesn't
+// happen again.
+// When using run_query make sure that the ResT type is NOT able
+// to contain any errors.
+// See https://github.com/scylladb/scylla-rust-driver/issues/501
+pub trait AllowedRunQueryResTType {}
+
+impl AllowedRunQueryResTType for Uuid {}
+impl AllowedRunQueryResTType for BatchResult {}
+impl AllowedRunQueryResTType for NonErrorQueryResponse {}