Merge pull request #1326 from muzarski/enforce-request-coordinator

Enforce request coordinator

Author: wprzytula

docs/source/load-balancing/load-balancing.md

@@ -36,8 +36,12 @@ awareness enabled and latency-awareness disabled.
 
 ## Configuration
 
-Load balancing policies can be configured via execution profiles. In the code
-sample provided, a new execution profile is created using
+Load balancing policies can be configured in three different ways (sorted by descending precedence):
+1. directly on the `Statement`, `PreparedStatement` or `Batch` (`Statement::set_load_balancing_policy()`)
+2. execution profile set on the statement (`Statement::set_execution_profile_handle()`)
+3. default execution profile set on the session (`SessionBuilder::default_execution_profile_handle()`)
+
+ In the code sample provided, a new execution profile is created using
 `ExecutionProfile::builder()`, and the load balancing policy is set to the
 `DefaultPolicy` using `.load_balancing_policy(policy)`.
 

examples/Cargo.toml

@@ -98,6 +98,10 @@ path = "get_by_name.rs"
 name = "value_list"
 path = "value_list.rs"
 
+[[example]]
+name = "enforce_coordinator"
+path = "enforce_coordinator.rs"
+
 [[example]]
 name = "custom_load_balancing_policy"
 path = "custom_load_balancing_policy.rs"

examples/enforce_coordinator.rs

@@ -0,0 +1,94 @@
+//! This example show how to enforce the target the request is sent to.
+
+use std::net::{IpAddr, SocketAddr};
+use std::sync::Arc;
+
+use anyhow::Result;
+use scylla::client::session::Session;
+use scylla::client::session_builder::SessionBuilder;
+use scylla::cluster::Node;
+use scylla::policies::load_balancing::{NodeIdentifier, SingleTargetLoadBalancingPolicy};
+use scylla::statement::prepared::PreparedStatement;
+
+/// Executes "SELECT host_id, rpc_address FROM system.local" query with `node` as the enforced target.
+/// Checks whether the result matches the expected values (i.e. ones stored in peers metadata).
+async fn query_system_local_and_verify(
+    session: &Session,
+    node: &Arc<Node>,
+    query_local: &PreparedStatement,
+) {
+    let (actual_host_id, actual_node_ip) = session
+        .execute_unpaged(query_local, ())
+        .await
+        .unwrap()
+        .into_rows_result()
+        .unwrap()
+        .single_row::<(uuid::Uuid, IpAddr)>()
+        .unwrap();
+
+    println!(
+        "queried host_id: {}; queried node_ip: {}",
+        actual_host_id, actual_node_ip
+    );
+    assert_eq!(node.host_id, actual_host_id);
+    assert_eq!(node.address.ip(), actual_node_ip);
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+
+    let session: Session = SessionBuilder::new().known_node(uri).build().await?;
+
+    let state = session.get_cluster_state();
+    let node = state
+        .get_nodes_info()
+        .first()
+        .ok_or_else(|| anyhow::anyhow!("No nodes in metadata!"))?;
+
+    let expected_host_id = node.host_id;
+    let expected_node_ip = node.address.ip();
+
+    let mut query_local = session
+        .prepare("SELECT host_id, rpc_address FROM system.local where key='local'")
+        .await?;
+
+    // Enforce the node using `Arc<Node>`.
+    {
+        let node_identifier = NodeIdentifier::Node(Arc::clone(node));
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    // Enforce the node using host_id.
+    {
+        let node_identifier = NodeIdentifier::HostId(expected_host_id);
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    // Enforce the node using **untranslated** node address.
+    {
+        let node_identifier =
+            NodeIdentifier::NodeAddress(SocketAddr::new(expected_node_ip, node.address.port()));
+        println!("Enforcing target using {:?}...", node_identifier);
+        query_local.set_load_balancing_policy(Some(SingleTargetLoadBalancingPolicy::new(
+            node_identifier,
+            None,
+        )));
+
+        query_system_local_and_verify(&session, node, &query_local).await;
+    }
+
+    Ok(())
+}

scylla/src/client/pager.rs

@@ -34,7 +34,7 @@ use crate::observability::driver_tracing::RequestSpan;
 use crate::observability::history::{self, HistoryListener};
 #[cfg(feature = "metrics")]
 use crate::observability::metrics::Metrics;
-use crate::policies::load_balancing::{self, RoutingInfo};
+use crate::policies::load_balancing::{self, LoadBalancingPolicy, RoutingInfo};
 use crate::policies::retry::{RequestInfo, RetryDecision, RetrySession};
 use crate::response::query_result::ColumnSpecs;
 use crate::response::{NonErrorQueryResponse, QueryResponse};
@@ -136,11 +136,11 @@ struct PagerWorker<'a, QueryFunc, SpanCreatorFunc> {
     // AsyncFn(Arc<Connection>, Option<Arc<[u8]>>) -> Result<QueryResponse, RequestAttemptError>
     page_query: QueryFunc,
 
+    load_balancing_policy: Arc<dyn LoadBalancingPolicy>,
     statement_info: RoutingInfo<'a>,
     query_is_idempotent: bool,
     query_consistency: Consistency,
     retry_session: Box<dyn RetrySession>,
-    execution_profile: Arc<ExecutionProfileInner>,
     #[cfg(feature = "metrics")]
     metrics: Arc<Metrics>,
 
@@ -162,7 +162,7 @@ where
 {
     // Contract: this function MUST send at least one item through self.sender
     async fn work(mut self, cluster_state: Arc<ClusterState>) -> PageSendAttemptedProof {
-        let load_balancer = self.execution_profile.load_balancing_policy.clone();
+        let load_balancer = Arc::clone(&self.load_balancing_policy);
         let statement_info = self.statement_info.clone();
         let query_plan =
             load_balancing::Plan::new(load_balancer.as_ref(), &statement_info, &cluster_state);
@@ -335,8 +335,7 @@ where
                 let _ = self.metrics.log_query_latency(elapsed.as_millis() as u64);
                 self.log_attempt_success();
                 self.log_request_success();
-                self.execution_profile
-                    .load_balancing_policy
+                self.load_balancing_policy
                     .on_request_success(&self.statement_info, elapsed, node);
 
                 request_span.record_raw_rows_fields(&rows);
@@ -369,9 +368,12 @@ where
             Err(err) => {
                 #[cfg(feature = "metrics")]
                 self.metrics.inc_failed_paged_queries();
-                self.execution_profile
-                    .load_balancing_policy
-                    .on_request_failure(&self.statement_info, elapsed, node, &err);
+                self.load_balancing_policy.on_request_failure(
+                    &self.statement_info,
+                    elapsed,
+                    node,
+                    &err,
+                );
                 Err(err)
             }
             Ok(NonErrorQueryResponse {
@@ -391,9 +393,12 @@ where
                 self.metrics.inc_failed_paged_queries();
                 let err =
                     RequestAttemptError::UnexpectedResponse(response.response.to_response_kind());
-                self.execution_profile
-                    .load_balancing_policy
-                    .on_request_failure(&self.statement_info, elapsed, node, &err);
+                self.load_balancing_policy.on_request_failure(
+                    &self.statement_info,
+                    elapsed,
+                    node,
+                    &err,
+                );
                 Err(err)
             }
         }
@@ -693,6 +698,12 @@ impl QueryPager {
             ..Default::default()
         };
 
+        let load_balancing_policy = Arc::clone(
+            statement
+                .get_load_balancing_policy()
+                .unwrap_or(&execution_profile.load_balancing_policy),
+        );
+
         let retry_session = statement
             .get_retry_policy()
             .map(|rp| &**rp)
@@ -733,8 +744,8 @@ impl QueryPager {
                 statement_info: routing_info,
                 query_is_idempotent: statement.config.is_idempotent,
                 query_consistency: consistency,
+                load_balancing_policy,
                 retry_session,
-                execution_profile,
                 #[cfg(feature = "metrics")]
                 metrics,
                 paging_state: PagingState::start(),
@@ -769,6 +780,13 @@ impl QueryPager {
 
         let page_size = config.prepared.get_validated_page_size();
 
+        let load_balancing_policy = Arc::clone(
+            config
+                .prepared
+                .get_load_balancing_policy()
+                .unwrap_or(&config.execution_profile.load_balancing_policy),
+        );
+
         let retry_session = config
             .prepared
             .get_retry_policy()
@@ -854,8 +872,8 @@ impl QueryPager {
                 statement_info,
                 query_is_idempotent: config.prepared.config.is_idempotent,
                 query_consistency: consistency,
+                load_balancing_policy,
                 retry_session,
-                execution_profile: config.execution_profile,
                 #[cfg(feature = "metrics")]
                 metrics: config.metrics,
                 paging_state: PagingState::start(),

scylla/src/client/session.rs

@@ -1801,12 +1801,15 @@ impl Session {
                 .as_ref()
                 .map(|hl| (&**hl, hl.log_request_start()));
 
-        let load_balancer = &execution_profile.load_balancing_policy;
+        let load_balancer = statement_config
+            .load_balancing_policy
+            .as_deref()
+            .unwrap_or(execution_profile.load_balancing_policy.as_ref());
 
         let runner = async {
             let cluster_state = self.cluster.get_state();
             let request_plan =
-                load_balancing::Plan::new(load_balancer.as_ref(), &statement_info, &cluster_state);
+                load_balancing::Plan::new(load_balancer, &statement_info, &cluster_state);
 
             // If a speculative execution policy is used to run request, request_plan has to be shared
             // between different async functions. This struct helps to wrap request_plan in mutex so it
@@ -1874,6 +1877,7 @@ impl Session {
                                 consistency_set_on_statement: statement_config.consistency,
                                 retry_session: retry_policy.new_session(),
                                 history_data,
+                                load_balancing_policy: load_balancer,
                                 query_info: &statement_info,
                                 request_span,
                             },
@@ -1910,6 +1914,7 @@ impl Session {
                             consistency_set_on_statement: statement_config.consistency,
                             retry_session: retry_policy.new_session(),
                             history_data,
+                            load_balancing_policy: load_balancer,
                             query_info: &statement_info,
                             request_span,
                         },
@@ -2007,7 +2012,7 @@ impl Session {
                         #[cfg(feature = "metrics")]
                         let _ = self.metrics.log_query_latency(elapsed.as_millis() as u64);
                         context.log_attempt_success(&attempt_id);
-                        execution_profile.load_balancing_policy.on_request_success(
+                        context.load_balancing_policy.on_request_success(
                             context.query_info,
                             elapsed,
                             node,
@@ -2022,7 +2027,7 @@ impl Session {
                         );
                         #[cfg(feature = "metrics")]
                         self.metrics.inc_failed_nonpaged_queries();
-                        execution_profile.load_balancing_policy.on_request_failure(
+                        context.load_balancing_policy.on_request_failure(
                             context.query_info,
                             elapsed,
                             node,
@@ -2135,6 +2140,7 @@ struct ExecuteRequestContext<'a> {
     consistency_set_on_statement: Option<Consistency>,
     retry_session: Box<dyn RetrySession>,
     history_data: Option<HistoryData<'a>>,
+    load_balancing_policy: &'a dyn load_balancing::LoadBalancingPolicy,
     query_info: &'a load_balancing::RoutingInfo<'a>,
     request_span: &'a RequestSpan,
 }

scylla/src/policies/load_balancing/mod.rs

@@ -13,8 +13,10 @@ use std::time::Duration;
 
 mod default;
 mod plan;
+mod single_target;
 pub use default::{DefaultPolicy, DefaultPolicyBuilder, LatencyAwarenessBuilder};
 pub use plan::Plan;
+pub use single_target::{NodeIdentifier, SingleTargetLoadBalancingPolicy};
 
 /// Represents info about statement that can be used by load balancing policies.
 #[derive(Default, Clone, Debug)]