Merge pull request #1473 from Lorak-mmk/schema-agreement-dont-fail

Dont fail schema agreement on first error

Author: Lorak-mmk

scylla/src/client/session.rs

@@ -2200,42 +2200,13 @@ impl Session {
         last_error.map(Result::Err)
     }
 
-    /// Awaits schema agreement among all reachable nodes.
-    ///
-    /// Issues an agreement check each `Session::schema_agreement_interval`.
-    /// Loops indefinitely until the agreement is reached.
-    ///
-    /// If `required_node` is Some, only returns Ok if this node successfully
-    /// returned its schema version during the agreement process.
-    async fn await_schema_agreement_indefinitely(
-        &self,
-        required_node: Option<Uuid>,
-    ) -> Result<Uuid, SchemaAgreementError> {
-        loop {
-            tokio::time::sleep(self.schema_agreement_interval).await;
-            if let Some(agreed_version) = self
-                .check_schema_agreement_with_required_node(required_node)
-                .await?
-            {
-                return Ok(agreed_version);
-            }
-        }
-    }
-
     /// Awaits schema agreement among all reachable nodes.
     ///
     /// Issues an agreement check each `Session::schema_agreement_interval`.
     /// If agreement is not reached in `Session::schema_agreement_timeout`,
     /// `SchemaAgreementError::Timeout` is returned.
     pub async fn await_schema_agreement(&self) -> Result<Uuid, SchemaAgreementError> {
-        timeout(
-            self.schema_agreement_timeout,
-            self.await_schema_agreement_indefinitely(None),
-        )
-        .await
-        .unwrap_or(Err(SchemaAgreementError::Timeout(
-            self.schema_agreement_timeout,
-        )))
+        self.await_schema_agreement_with_required_node(None).await
     }
 
     /// Awaits schema agreement among all reachable nodes.
@@ -2250,14 +2221,34 @@ impl Session {
         &self,
         required_node: Option<Uuid>,
     ) -> Result<Uuid, SchemaAgreementError> {
-        timeout(
-            self.schema_agreement_timeout,
-            self.await_schema_agreement_indefinitely(required_node),
-        )
+        // None: no finished attempt recorded
+        // Some(Ok(())): Last attempt successful, without agreement
+        // Some(Err(_)): Last attempt failed
+        let mut last_agreement_failure: Option<Result<(), SchemaAgreementError>> = None;
+        timeout(self.schema_agreement_timeout, async {
+            loop {
+                let result = self
+                    .check_schema_agreement_with_required_node(required_node)
+                    .await;
+                match result {
+                    Ok(Some(agreed_version)) => return agreed_version,
+                    Ok(None) => last_agreement_failure = Some(Ok(())),
+                    Err(err) => last_agreement_failure = Some(Err(err)),
+                }
+                tokio::time::sleep(self.schema_agreement_interval).await;
+            }
+        })
         .await
-        .unwrap_or(Err(SchemaAgreementError::Timeout(
-            self.schema_agreement_timeout,
-        )))
+        .map_err(|_| {
+            match last_agreement_failure {
+                // There were no finished attempts - the only error we can return is Timeout.
+                None => SchemaAgreementError::Timeout(self.schema_agreement_timeout),
+                // If the last finished attempt resulted in an error, this error will be more informative than Timeout.
+                Some(Err(err)) => err,
+                // This is the canonical case for timeout - last attempt finished successfully, but without agreement.
+                Some(Ok(())) => SchemaAgreementError::Timeout(self.schema_agreement_timeout),
+            }
+        })
     }
 
     /// Checks if all reachable nodes have the same schema version.

scylla/tests/integration/session/schema_agreement.rs

@@ -1,16 +1,18 @@
 use std::sync::Arc;
 
 use assert_matches::assert_matches;
+use scylla::client::PoolSize;
 use scylla::client::session::Session;
 use scylla::client::session_builder::SessionBuilder;
-use scylla::errors::{ExecutionError, RequestAttemptError, SchemaAgreementError};
+use scylla::errors::{DbError, ExecutionError, RequestAttemptError, SchemaAgreementError};
 use scylla::policies::load_balancing::{NodeIdentifier, SingleTargetLoadBalancingPolicy};
 use scylla::response::query_result::QueryResult;
 use scylla::statement::Statement;
 use scylla_proxy::{
     Condition, ProxyError, Reaction, RequestOpcode, RequestReaction, RequestRule, RunningProxy,
     ShardAwareness, WorkerError,
 };
+use tracing::info;
 
 use crate::utils::{
     calculate_proxy_host_ids, setup_tracing, test_with_3_node_cluster, unique_keyspace_name,
@@ -157,3 +159,82 @@ async fn test_schema_await_with_unreachable_node() {
         Err(err) => panic!("{}", err),
     }
 }
+
+// Verifies that schema agreement process works correctly even if the first check fails.
+#[tokio::test]
+#[cfg_attr(scylla_cloud_tests, ignore)]
+async fn test_schema_await_with_transient_failure() {
+    setup_tracing();
+
+    let res = test_with_3_node_cluster(
+        ShardAwareness::QueryNode,
+        |proxy_uris, translation_map, mut running_proxy| async move {
+            // DB preparation phase
+            let session: Session = SessionBuilder::new()
+                .known_node(proxy_uris[0].as_str())
+                .address_translator(Arc::new(translation_map.clone()))
+                // Important in order to have a predictable amount of connections after session creation.
+                // Shard connections are created asynchronously, so it's hard to predict how many will be opened
+                // already when we check schema agreement.
+                .pool_size(PoolSize::PerHost(1.try_into().unwrap()))
+                .build()
+                .await
+                .unwrap();
+
+            let node_rules = Some(vec![RequestRule(
+                Condition::not(Condition::ConnectionRegisteredAnyEvent)
+                    .and(Condition::RequestOpcode(RequestOpcode::Query))
+                    .and(Condition::BodyContainsCaseSensitive(Box::new(
+                        *b"system.local",
+                    )))
+                    .and(Condition::TrueForLimitedTimes(1)),
+                // Use error that would prevent DefaultRetryPolicy from retrying.
+                // I don't think it is used for those queries, but it's additional future-proofing
+                // for the test.
+                RequestReaction::forge_with_error(DbError::SyntaxError),
+            )]);
+
+            // First, a sanity check for proxy rules.
+            // If for each node first request fails (and subsequent requests succeed),
+            // then first schema agreement check should return error, and second should succeed.
+            // Note that this is only true because we configured the session with 1-connection-per-node.
+            info!("Starting phase 1 - sanity check");
+            {
+                running_proxy
+                    .running_nodes
+                    .iter_mut()
+                    .for_each(|node| node.change_request_rules(node_rules.clone()));
+
+                // The important check: first call should error out, second one should succeed.
+                session.check_schema_agreement().await.unwrap_err();
+                session.check_schema_agreement().await.unwrap();
+
+                running_proxy
+                    .running_nodes
+                    .iter_mut()
+                    .for_each(|node| node.change_request_rules(Some(vec![])));
+            }
+
+            // Now let's check that awaiting schema agreement doesn't bail on error.
+            // I'll use the same proxy rules as before, so first check will error out.
+            info!("Starting phase 2 - main test");
+            {
+                running_proxy
+                    .running_nodes
+                    .iter_mut()
+                    .for_each(|node| node.change_request_rules(node_rules.clone()));
+
+                session.await_schema_agreement().await.unwrap();
+            }
+
+            running_proxy
+        },
+    )
+    .await;
+
+    match res {
+        Ok(()) => (),
+        Err(ProxyError::Worker(WorkerError::DriverDisconnected(_))) => (),
+        Err(err) => panic!("{}", err),
+    }
+}