Require schema agreement response from DDL coordinator

This fixes the issue described in parent commit.

Internal schema agreement APIs now accept a `Option<Uuid>` that is a
host id of a node that must be a part of schema agreement. For
user-requested agreements it will be None, and for agreements after DDL
it will be the coordinator.

Author: Lorak-mmk

scylla/src/client/session.rs

@@ -1116,7 +1116,8 @@ impl Session {
         };
 
         self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
+        self.handle_auto_await_schema_agreement(&response, coordinator.node().host_id)
+            .await?;
 
         let (result, paging_state_response) =
             response.into_query_result_and_paging_state(coordinator)?;
@@ -1144,10 +1145,12 @@ impl Session {
     async fn handle_auto_await_schema_agreement(
         &self,
         response: &NonErrorQueryResponse,
+        coordinator_id: Uuid,
     ) -> Result<(), ExecutionError> {
         if self.schema_agreement_automatic_waiting {
             if response.as_schema_change().is_some() {
-                self.await_schema_agreement().await?;
+                self.await_schema_agreement_with_required_node(Some(coordinator_id))
+                    .await?;
             }
 
             if self.refresh_metadata_on_auto_schema_agreement
@@ -1489,7 +1492,8 @@ impl Session {
         };
 
         self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
+        self.handle_auto_await_schema_agreement(&response, coordinator.node().host_id)
+            .await?;
 
         let (result, paging_state_response) =
             response.into_query_result_and_paging_state(coordinator)?;
@@ -2160,10 +2164,19 @@ impl Session {
     ///
     /// Issues an agreement check each `Session::schema_agreement_interval`.
     /// Loops indefinitely until the agreement is reached.
-    async fn await_schema_agreement_indefinitely(&self) -> Result<Uuid, SchemaAgreementError> {
+    ///
+    /// 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().await? {
+            if let Some(agreed_version) = self
+                .check_schema_agreement_with_required_node(required_node)
+                .await?
+            {
                 return Ok(agreed_version);
             }
         }
@@ -2177,7 +2190,29 @@ impl Session {
     pub async fn await_schema_agreement(&self) -> Result<Uuid, SchemaAgreementError> {
         timeout(
             self.schema_agreement_timeout,
-            self.await_schema_agreement_indefinitely(),
+            self.await_schema_agreement_indefinitely(None),
+        )
+        .await
+        .unwrap_or(Err(SchemaAgreementError::Timeout(
+            self.schema_agreement_timeout,
+        )))
+    }
+
+    /// 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.
+    ///
+    /// 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_with_required_node(
+        &self,
+        required_node: Option<Uuid>,
+    ) -> Result<Uuid, SchemaAgreementError> {
+        timeout(
+            self.schema_agreement_timeout,
+            self.await_schema_agreement_indefinitely(required_node),
         )
         .await
         .unwrap_or(Err(SchemaAgreementError::Timeout(
@@ -2189,14 +2224,54 @@ impl Session {
     ///
     /// If so, returns that agreed upon version.
     pub async fn check_schema_agreement(&self) -> Result<Option<Uuid>, SchemaAgreementError> {
+        self.check_schema_agreement_with_required_node(None).await
+    }
+
+    /// Checks if all reachable nodes have the same schema version.
+    /// If so, returns that agreed upon version.
+    ///
+    /// If `required_node` is Some, only returns Ok if this node successfully
+    /// returned its schema version.
+    async fn check_schema_agreement_with_required_node(
+        &self,
+        required_node: Option<Uuid>,
+    ) -> Result<Option<Uuid>, SchemaAgreementError> {
         let cluster_state = self.get_cluster_state();
         // The iterator is guaranteed to be nonempty.
         let per_node_connections = cluster_state.iter_working_connections_per_node()?;
 
         // Therefore, this iterator is guaranteed to be nonempty, too.
-        let handles = per_node_connections.map(Session::read_node_schema_version);
+        let handles = per_node_connections.map(|(host_id, pool)| async move {
+            (host_id, Session::read_node_schema_version(pool).await)
+        });
         // Hence, this is nonempty, too.
-        let versions_results = try_join_all(handles).await?;
+        let versions_results = join_all(handles).await;
+
+        // Verify that required host is present, and returned success.
+        if let Some(required_node) = required_node {
+            match versions_results
+                .iter()
+                .find(|(host_id, _)| *host_id == required_node)
+            {
+                Some((_, Ok(SchemaNodeResult::Success(_version)))) => (),
+                // For other connections we can ignore Broken error, but for required
+                // host we need an actual schema version.
+                Some((_, Ok(SchemaNodeResult::BrokenConnection(e)))) => {
+                    return Err(SchemaAgreementError::RequestError(
+                        RequestAttemptError::BrokenConnectionError(e.clone()),
+                    ))
+                }
+                Some((_, Err(e))) => return Err(e.clone()),
+                None => return Err(SchemaAgreementError::RequiredHostAbsent(required_node)),
+            }
+        }
+
+        // Now we no longer need all the errors. We can return if there is
+        // irrecoverable one, and collect the Ok values otherwise.
+        let versions_results: Vec<_> = versions_results
+            .into_iter()
+            .map(|(_, result)| result)
+            .try_collect()?;
 
         // unwrap is safe because iterator is still not empty.
         let local_version = match versions_results

scylla/src/cluster/state.rs

@@ -324,17 +324,21 @@ impl ClusterState {
     /// Internal iterator iterates over working connections to all shards of given node.
     pub(crate) fn iter_working_connections_per_node(
         &self,
-    ) -> Result<impl Iterator<Item = impl Iterator<Item = Arc<Connection>>> + '_, ConnectionPoolError>
-    {
+    ) -> Result<
+        impl Iterator<Item = (Uuid, impl Iterator<Item = Arc<Connection>>)> + '_,
+        ConnectionPoolError,
+    > {
         // The returned iterator is nonempty by nonemptiness invariant of `self.known_peers`.
         assert!(!self.known_peers.is_empty());
         let nodes_iter = self.known_peers.values();
-        let mut connection_pool_per_node_iter =
-            nodes_iter.map(|node| node.get_working_connections());
+        let mut connection_pool_per_node_iter = nodes_iter.map(|node| {
+            node.get_working_connections()
+                .map(|pool| (node.host_id, pool))
+        });
 
         // First we try to find the first working pool of connections.
         // If none is found, return error.
-        let first_working_pool_or_error: Result<Vec<Arc<Connection>>, ConnectionPoolError> =
+        let first_working_pool_or_error: Result<(Uuid, Vec<Arc<Connection>>), ConnectionPoolError> =
             connection_pool_per_node_iter
                 .by_ref()
                 .find_or_first(Result::is_ok)
@@ -344,19 +348,19 @@ impl ClusterState {
         // 1. either consumed the whole iterator without success and got the first error,
         //    in which case we propagate it;
         // 2. or found the first working pool of connections.
-        let first_working_pool: Vec<Arc<Connection>> = first_working_pool_or_error?;
+        let first_working_pool: (Uuid, Vec<Arc<Connection>>) = first_working_pool_or_error?;
 
         // We retrieve connection pools for remaining nodes (those that are left in the iterator
         // once the first working pool has been found).
         let remaining_pools_iter = connection_pool_per_node_iter;
         // Errors (non-working pools) are filtered out.
         let remaining_working_pools_iter = remaining_pools_iter.filter_map(Result::ok);
-        // Pools are made iterators, so now we have `impl Iterator<Item = impl Iterator<Item = Arc<Connection>>>`.
-        let remaining_working_per_node_connections_iter =
-            remaining_working_pools_iter.map(IntoIterator::into_iter);
 
-        Ok(std::iter::once(first_working_pool.into_iter())
-            .chain(remaining_working_per_node_connections_iter))
+        // First pool is chained with the rest.
+        // Then, pools are made iterators, so now we have `impl Iterator<Item = (Uuid, impl Iterator<Item = Arc<Connection>>)>`.
+        Ok(std::iter::once(first_working_pool)
+            .chain(remaining_working_pools_iter)
+            .map(|(host_id, pool)| (host_id, IntoIterator::into_iter(pool))))
         // By an invariant `self.known_peers` is nonempty, so the returned iterator
         // is nonempty, too.
     }
@@ -366,7 +370,7 @@ impl ClusterState {
         &self,
     ) -> Result<impl Iterator<Item = Arc<Connection>> + '_, ConnectionPoolError> {
         self.iter_working_connections_per_node()
-            .map(|iter| iter.flatten())
+            .map(|outer_iter| outer_iter.flat_map(|(_, inner_iter)| inner_iter))
     }
 
     /// Returns nonempty iterator of working connections to all nodes.

scylla/src/errors.rs

@@ -6,6 +6,7 @@ use std::net::{AddrParseError, IpAddr, SocketAddr};
 use std::num::ParseIntError;
 use std::sync::Arc;
 use thiserror::Error;
+use uuid::Uuid;
 
 use crate::frame::response;
 
@@ -209,6 +210,12 @@ pub enum SchemaAgreementError {
     /// Schema agreement timed out.
     #[error("Schema agreement exceeded {}ms", std::time::Duration::as_millis(.0))]
     Timeout(std::time::Duration),
+
+    #[error(
+        "Host with id {} required for schema agreement is not present in connection pool",
+        0
+    )]
+    RequiredHostAbsent(Uuid),
 }
 
 /// An error that occurred during tracing info fetch.