Merge pull request #768 from avelanarius/session_slimming_1

session: remove estimate_replicas_for_query, get_tracing_info_custom

Author: wprzytula

examples/compare-tokens.rs

@@ -1,8 +1,9 @@
 use anyhow::Result;
 use scylla::frame::value::ValueList;
+use scylla::routing::Token;
 use scylla::transport::partitioner::{Murmur3Partitioner, Partitioner};
 use scylla::transport::NodeAddr;
-use scylla::{load_balancing, Session, SessionBuilder};
+use scylla::{Session, SessionBuilder};
 use std::env;
 
 #[tokio::main]
@@ -32,16 +33,11 @@ async fn main() -> Result<()> {
         let serialized_pk = (pk,).serialized()?.into_owned();
         let t = Murmur3Partitioner::hash(&prepared.compute_partition_key(&serialized_pk)?).value;
 
-        let statement_info = load_balancing::RoutingInfo {
-            token: Some(scylla::routing::Token { value: t }),
-            keyspace: Some("ks"),
-            is_confirmed_lwt: false,
-            ..Default::default()
-        };
         println!(
-            "Estimated replicas for query: {:?}",
+            "Token endpoints for query: {:?}",
             session
-                .estimate_replicas_for_query(&statement_info)
+                .get_cluster_data()
+                .get_token_endpoints("ks", Token { value: t })
                 .iter()
                 .map(|n| n.address)
                 .collect::<Vec<NodeAddr>>()

examples/tracing.rs

@@ -7,7 +7,7 @@ use scylla::batch::Batch;
 use scylla::statement::{
     prepared_statement::PreparedStatement, query::Query, Consistency, SerialConsistency,
 };
-use scylla::tracing::{GetTracingConfig, TracingInfo};
+use scylla::tracing::TracingInfo;
 use scylla::transport::iterator::RowIterator;
 use scylla::QueryResult;
 use scylla::{Session, SessionBuilder};
@@ -21,7 +21,10 @@ async fn main() -> Result<()> {
     let uri = env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
 
     println!("Connecting to {} ...", uri);
-    let session: Session = SessionBuilder::new().known_node(uri).build().await?;
+    let session: Session = SessionBuilder::new()
+        .known_node(uri.as_str())
+        .build()
+        .await?;
 
     session.query("CREATE KEYSPACE IF NOT EXISTS ks WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}", &[]).await?;
 
@@ -106,19 +109,18 @@ async fn main() -> Result<()> {
     println!("Batch tracing id: {:?}\n", batch_result.tracing_id);
 
     // CUSTOM
-    // GetTracingConfig allows to specify a custom settings for querying tracing info
+    // Session configuration allows specifying custom settings for querying tracing info.
     // Tracing info might not immediately be available on queried node
     // so the driver performs a few attempts with sleeps in between.
-
-    let custom_config: GetTracingConfig = GetTracingConfig {
-        attempts: NonZeroU32::new(8).unwrap(),
-        interval: Duration::from_millis(100),
-        consistency: Consistency::One,
-    };
-
-    let _custom_info: TracingInfo = session
-        .get_tracing_info_custom(&query_tracing_id, &custom_config)
+    let session: Session = SessionBuilder::new()
+        .known_node(uri)
+        .tracing_info_fetch_attempts(NonZeroU32::new(8).unwrap())
+        .tracing_info_fetch_interval(Duration::from_millis(100))
+        .tracing_info_fetch_consistency(Consistency::One)
+        .build()
         .await?;
 
+    let _custom_info: TracingInfo = session.get_tracing_info(&query_tracing_id).await?;
+
     Ok(())
 }

scylla/src/tracing.rs

@@ -1,9 +1,6 @@
-use crate::statement::Consistency;
 use itertools::Itertools;
 use std::collections::HashMap;
 use std::net::IpAddr;
-use std::num::NonZeroU32;
-use std::time::Duration;
 use uuid::Uuid;
 
 use crate::cql_to_rust::{FromRow, FromRowError};
@@ -35,20 +32,6 @@ pub struct TracingEvent {
     pub thread: Option<String>,
 }
 
-/// Used to configure a custom retry strategy when querying tracing info
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct GetTracingConfig {
-    /// Number of attempts to be made before giving up.
-    /// Default value: 5
-    pub attempts: NonZeroU32,
-    /// Interval to wait between each attempt.
-    /// Default value: 3 milliseconds
-    pub interval: Duration,
-    /// Consistency to use in queries that read TracingInfo.
-    /// Default value: One
-    pub consistency: Consistency,
-}
-
 impl TracingInfo {
     /// Returns a list of unique nodes involved in the query
     pub fn nodes(&self) -> Vec<IpAddr> {
@@ -60,16 +43,6 @@ impl TracingInfo {
     }
 }
 
-impl Default for GetTracingConfig {
-    fn default() -> GetTracingConfig {
-        GetTracingConfig {
-            attempts: NonZeroU32::new(5).unwrap(),
-            interval: Duration::from_millis(3),
-            consistency: Consistency::One,
-        }
-    }
-}
-
 // A query used to query TracingInfo from system_traces.sessions
 pub(crate) const TRACES_SESSION_QUERY_STR: &str =
     "SELECT client, command, coordinator, duration, parameters, request, started_at \

scylla/src/transport/session.rs

@@ -27,6 +27,7 @@ use std::fmt::Display;
 use std::future::Future;
 use std::io;
 use std::net::SocketAddr;
+use std::num::NonZeroU32;
 use std::str::FromStr;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
@@ -59,7 +60,7 @@ use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
 use crate::query::Query;
 use crate::routing::Token;
 use crate::statement::{Consistency, SerialConsistency};
-use crate::tracing::{GetTracingConfig, TracingEvent, TracingInfo};
+use crate::tracing::{TracingEvent, TracingInfo};
 use crate::transport::cluster::{Cluster, ClusterData, ClusterNeatDebug};
 use crate::transport::connection::{Connection, ConnectionConfig, VerifiedKeyspaceName};
 use crate::transport::connection_pool::PoolConfig;
@@ -133,6 +134,9 @@ pub struct Session {
     auto_await_schema_agreement_timeout: Option<Duration>,
     refresh_metadata_on_auto_schema_agreement: bool,
     keyspace_name: ArcSwapOption<String>,
+    tracing_info_fetch_attempts: NonZeroU32,
+    tracing_info_fetch_interval: Duration,
+    tracing_info_fetch_consistency: Consistency,
 }
 
 /// This implementation deliberately omits some details from Cluster in order
@@ -238,6 +242,22 @@ pub struct SessionConfig {
     /// Please do performance measurements before committing to disabling
     /// this option.
     pub enable_write_coalescing: bool,
+
+    /// Number of attempts to fetch [`TracingInfo`]
+    /// in [`Session::get_tracing_info`]. Tracing info
+    /// might not be available immediately on queried node - that's why
+    /// the driver performs a few attempts with sleeps in between.
+    pub tracing_info_fetch_attempts: NonZeroU32,
+
+    /// Delay between attempts to fetch [`TracingInfo`]
+    /// in [`Session::get_tracing_info`]. Tracing info
+    /// might not be available immediately on queried node - that's why
+    /// the driver performs a few attempts with sleeps in between.
+    pub tracing_info_fetch_interval: Duration,
+
+    /// Consistency level of fetching [`TracingInfo`]
+    /// in [`Session::get_tracing_info`].
+    pub tracing_info_fetch_consistency: Consistency,
 }
 
 /// Describes database server known on Session startup.
@@ -297,6 +317,9 @@ impl SessionConfig {
             #[cfg(feature = "cloud")]
             cloud_config: None,
             enable_write_coalescing: true,
+            tracing_info_fetch_attempts: NonZeroU32::new(5).unwrap(),
+            tracing_info_fetch_interval: Duration::from_millis(3),
+            tracing_info_fetch_consistency: Consistency::One,
         }
     }
 
@@ -547,6 +570,9 @@ impl Session {
             refresh_metadata_on_auto_schema_agreement: config
                 .refresh_metadata_on_auto_schema_agreement,
             keyspace_name: ArcSwapOption::default(), // will be set by use_keyspace
+            tracing_info_fetch_attempts: config.tracing_info_fetch_attempts,
+            tracing_info_fetch_interval: config.tracing_info_fetch_interval,
+            tracing_info_fetch_consistency: config.tracing_info_fetch_consistency,
         };
 
         if let Some(keyspace_name) = config.used_keyspace {
@@ -1344,8 +1370,24 @@ impl Session {
     /// See [the book](https://rust-driver.docs.scylladb.com/stable/tracing/tracing.html)
     /// for more information about query tracing
     pub async fn get_tracing_info(&self, tracing_id: &Uuid) -> Result<TracingInfo, QueryError> {
-        self.get_tracing_info_custom(tracing_id, &GetTracingConfig::default())
-            .await
+        // tracing_info_fetch_attempts is NonZeroU32 so at least one attempt will be made
+        for _ in 0..self.tracing_info_fetch_attempts.get() {
+            let current_try: Option<TracingInfo> = self
+                .try_getting_tracing_info(tracing_id, Some(self.tracing_info_fetch_consistency))
+                .await?;
+
+            match current_try {
+                Some(tracing_info) => return Ok(tracing_info),
+                None => tokio::time::sleep(self.tracing_info_fetch_interval).await,
+            };
+        }
+
+        Err(QueryError::ProtocolError(
+            "All tracing queries returned an empty result, \
+            maybe the trace information didn't propagate yet. \
+            Consider configuring Session with \
+            a longer fetch interval (tracing_info_fetch_interval)",
+        ))
     }
 
     /// Gets the name of the keyspace that is currently set, or `None` if no
@@ -1363,35 +1405,6 @@ impl Session {
         self.keyspace_name.load_full()
     }
 
-    /// Queries tracing info with custom retry settings.\
-    /// Tracing info might not be available immediately on queried node -
-    /// that's why the driver performs a few attempts with sleeps in between.
-    /// [`GetTracingConfig`] allows to specify a custom querying strategy.
-    pub async fn get_tracing_info_custom(
-        &self,
-        tracing_id: &Uuid,
-        config: &GetTracingConfig,
-    ) -> Result<TracingInfo, QueryError> {
-        // config.attempts is NonZeroU32 so at least one attempt will be made
-        for _ in 0..config.attempts.get() {
-            let current_try: Option<TracingInfo> = self
-                .try_getting_tracing_info(tracing_id, Some(config.consistency))
-                .await?;
-
-            match current_try {
-                Some(tracing_info) => return Ok(tracing_info),
-                None => tokio::time::sleep(config.interval).await,
-            };
-        }
-
-        Err(QueryError::ProtocolError(
-            "All tracing queries returned an empty result, \
-            maybe information didn't reach this node yet. \
-            Consider using get_tracing_info_custom with \
-            bigger interval in GetTracingConfig",
-        ))
-    }
-
     // Tries getting the tracing info
     // If the queries return 0 rows then returns None - the information didn't reach this node yet
     // If there is some other error returns this error
@@ -1458,20 +1471,6 @@ impl Session {
         Ok(Some(tracing_info))
     }
 
-    // Returns which replicas are likely to take part in handling the query.
-    // If a list of replicas cannot be easily narrowed down, all available replicas
-    // will be returned.
-    pub fn estimate_replicas_for_query(&self, statement: &RoutingInfo) -> Vec<Arc<Node>> {
-        let cluster_data = self.cluster.get_data();
-        match statement.token {
-            Some(token) => {
-                let cluster_data = self.cluster.get_data();
-                cluster_data.get_token_endpoints(statement.keyspace.unwrap_or(""), token)
-            }
-            None => cluster_data.get_nodes_info().to_owned(),
-        }
-    }
-
     // This method allows to easily run a query using load balancing, retry policy etc.
     // Requires some information about the query and two closures
     // First closure is used to choose a connection

scylla/src/transport/session_builder.rs

@@ -10,11 +10,13 @@ use crate::cloud::{CloudConfig, CloudConfigError};
 #[cfg(feature = "cloud")]
 use crate::ExecutionProfile;
 
+use crate::statement::Consistency;
 use crate::transport::connection_pool::PoolSize;
 use crate::transport::host_filter::HostFilter;
 use std::borrow::Borrow;
 use std::marker::PhantomData;
 use std::net::SocketAddr;
+use std::num::NonZeroU32;
 #[cfg(feature = "cloud")]
 use std::path::Path;
 use std::sync::Arc;
@@ -828,6 +830,77 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self.config.refresh_metadata_on_auto_schema_agreement = refresh_metadata;
         self
     }
+
+    /// Set the number of attempts to fetch [TracingInfo](crate::tracing::TracingInfo)
+    /// in [`Session::get_tracing_info`].
+    /// The default is 5 attempts.
+    ///
+    /// Tracing info might not be available immediately on queried node - that's why
+    /// the driver performs a few attempts with sleeps in between.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use std::num::NonZeroU32;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .tracing_info_fetch_attempts(NonZeroU32::new(10).unwrap())
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn tracing_info_fetch_attempts(mut self, attempts: NonZeroU32) -> Self {
+        self.config.tracing_info_fetch_attempts = attempts;
+        self
+    }
+
+    /// Set the delay between attempts to fetch [TracingInfo](crate::tracing::TracingInfo)
+    /// in [`Session::get_tracing_info`].
+    /// The default is 3 milliseconds.
+    ///
+    /// Tracing info might not be available immediately on queried node - that's why
+    /// the driver performs a few attempts with sleeps in between.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use std::time::Duration;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .tracing_info_fetch_interval(Duration::from_millis(50))
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn tracing_info_fetch_interval(mut self, interval: Duration) -> Self {
+        self.config.tracing_info_fetch_interval = interval;
+        self
+    }
+
+    /// Set the consistency level of fetching [TracingInfo](crate::tracing::TracingInfo)
+    /// in [`Session::get_tracing_info`].
+    /// The default is [`Consistency::One`].
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder, statement::Consistency};
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .tracing_info_fetch_consistency(Consistency::One)
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn tracing_info_fetch_consistency(mut self, consistency: Consistency) -> Self {
+        self.config.tracing_info_fetch_consistency = consistency;
+        self
+    }
 }
 
 /// Creates a [`SessionBuilder`] with default configuration, same as [`SessionBuilder::new`]