session: remove get_tracing_info_custom

Remove Session::get_tracing_info_custom - a variant of
Session::get_tracing_info that allows for additional customization of
the fetch process (number of attempts, interval between attempts,
consistency level of fetching tracing info). This customization is now
moved to SessionConfig (as new parameters) and new helper functions
were introduced to SessionBuilder.

Since now this customization is on SessionConfig,
create_new_session_builder() in test_utils.rs returns a SessionBuilder
with pre-configured large number of TracingInfo fetch attempts, longer
delays (to account for a slow CI). Previously each tracing test had to
configure it on its own by using get_tracing_info_custom.

This change is a part of our effort to stabilize the API and reduce
the number of pub functions.

Refs #660

Author: avelanarius

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

@@ -24,6 +24,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;
@@ -56,7 +57,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;
@@ -130,6 +131,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
@@ -235,6 +239,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.
@@ -294,6 +314,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,
         }
     }
 
@@ -544,6 +567,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 {
@@ -1337,8 +1363,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
@@ -1356,35 +1398,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

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`]