Merge pull request #511 from wprzytula/debuggable-session

Made Session debuggable

Author: cvybhu

examples/custom_load_balancing_policy.rs

@@ -7,6 +7,7 @@ use scylla::{
 use std::{env, sync::Arc};
 
 /// Example load balancing policy that prefers nodes from favorite datacenter
+#[derive(Debug)]
 struct CustomLoadBalancingPolicy {
     fav_datacenter_name: String,
 }

scylla/src/transport/caching_session.rs

@@ -8,6 +8,7 @@ use bytes::Bytes;
 use dashmap::DashMap;
 
 /// Provides auto caching while executing queries
+#[derive(Debug)]
 pub struct CachingSession {
     pub session: Session,
     /// The prepared statement cache size

scylla/src/transport/cluster.rs

@@ -30,7 +30,19 @@ pub struct Cluster {
     _worker_handle: RemoteHandle<()>,
 }
 
-#[derive(Clone)]
+/// Enables printing [Cluster] struct in a neat way, by skipping the rather useless
+/// print of channels state and printing [ClusterData] neatly.
+pub struct ClusterNeatDebug<'a>(pub &'a Cluster);
+impl<'a> std::fmt::Debug for ClusterNeatDebug<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let cluster = self.0;
+        f.debug_struct("Cluster")
+            .field("data", &ClusterDataNeatDebug(&cluster.data.load()))
+            .finish_non_exhaustive()
+    }
+}
+
+#[derive(Clone, Debug)]
 pub struct Datacenter {
     pub nodes: Vec<Arc<Node>>,
     pub rack_count: usize,
@@ -45,6 +57,31 @@ pub struct ClusterData {
     pub(crate) datacenters: HashMap<String, Datacenter>,
 }
 
+/// Enables printing [ClusterData] struct in a neat way, skipping the clutter involved by
+/// [ClusterData::ring] being large and [Self::keyspaces] debug print being very verbose by default.
+pub struct ClusterDataNeatDebug<'a>(pub &'a Arc<ClusterData>);
+impl<'a> std::fmt::Debug for ClusterDataNeatDebug<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let cluster_data = &self.0;
+
+        f.debug_struct("ClusterData")
+            .field("known_peers", &cluster_data.known_peers)
+            .field("ring", {
+                struct RingSizePrinter(usize);
+                impl std::fmt::Debug for RingSizePrinter {
+                    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                        write!(f, "<size={}>", self.0)
+                    }
+                }
+                &RingSizePrinter(cluster_data.ring.len())
+            })
+            .field("keyspaces", &cluster_data.keyspaces.keys())
+            .field("all_nodes", &cluster_data.all_nodes)
+            .field("datacenters", &cluster_data.datacenters)
+            .finish_non_exhaustive()
+    }
+}
+
 // Works in the background to keep the cluster updated
 struct ClusterWorker {
     // Cluster data to keep updated:

scylla/src/transport/connection_pool.rs

@@ -144,6 +144,14 @@ pub struct NodeConnectionPool {
     pool_updated_notify: Arc<Notify>,
 }
 
+impl std::fmt::Debug for NodeConnectionPool {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("NodeConnectionPool")
+            .field("conns", &self.conns)
+            .finish_non_exhaustive()
+    }
+}
+
 impl NodeConnectionPool {
     pub fn new(
         address: IpAddr,

scylla/src/transport/load_balancing/dc_aware_round_robin.rs

@@ -7,6 +7,7 @@ use std::sync::{
 use tracing::trace;
 
 /// A data-center aware Round-robin load balancing policy.
+#[derive(Debug)]
 pub struct DcAwareRoundRobinPolicy {
     index: AtomicUsize,
     local_dc: String,

scylla/src/transport/load_balancing/mod.rs

@@ -35,7 +35,7 @@ impl<'a> Statement<'a> {
 pub type Plan<'a> = Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a>;
 
 /// Policy that decides which nodes to contact for each query
-pub trait LoadBalancingPolicy: Send + Sync {
+pub trait LoadBalancingPolicy: Send + Sync + std::fmt::Debug {
     /// It is used for each query to find which nodes to query first
     fn plan<'a>(&self, statement: &Statement, cluster: &'a ClusterData) -> Plan<'a>;
 

scylla/src/transport/load_balancing/round_robin.rs

@@ -7,6 +7,7 @@ use std::sync::{
 use tracing::trace;
 
 /// A Round-robin load balancing policy.
+#[derive(Debug)]
 pub struct RoundRobinPolicy {
     index: AtomicUsize,
 }

scylla/src/transport/load_balancing/token_aware.rs

@@ -9,6 +9,7 @@ use std::{collections::HashMap, sync::Arc};
 use tracing::trace;
 
 /// A wrapper load balancing policy that adds token awareness to a child policy.
+#[derive(Debug)]
 pub struct TokenAwarePolicy {
     child_policy: Box<dyn ChildLoadBalancingPolicy>,
 }
@@ -433,6 +434,7 @@ mod tests {
 
     // Used as child policy for TokenAwarePolicy tests
     // Forwards plan passed to it in apply_child_policy() method
+    #[derive(Debug)]
     struct DumbPolicy {}
 
     impl LoadBalancingPolicy for DumbPolicy {

scylla/src/transport/node.rs

@@ -15,6 +15,7 @@ use std::{
 };
 
 /// Node represents a cluster node along with it's data and connections
+#[derive(Debug)]
 pub struct Node {
     pub address: SocketAddr,
     pub datacenter: Option<String>,

scylla/src/transport/session.rs

@@ -25,7 +25,7 @@ use crate::query::Query;
 use crate::routing::Token;
 use crate::statement::{Consistency, SerialConsistency};
 use crate::tracing::{GetTracingConfig, TracingEvent, TracingInfo};
-use crate::transport::cluster::{Cluster, ClusterData};
+use crate::transport::cluster::{Cluster, ClusterData, ClusterNeatDebug};
 use crate::transport::connection::{
     BatchResult, Connection, ConnectionConfig, VerifiedKeyspaceName,
 };
@@ -65,6 +65,29 @@ pub struct Session {
     auto_await_schema_agreement_timeout: Option<Duration>,
 }
 
+/// This implementation deliberately omits some details from Cluster in order
+/// to avoid cluttering the print with much information of little usability.
+impl std::fmt::Debug for Session {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Session")
+            .field("cluster", &ClusterNeatDebug(&self.cluster))
+            .field("load_balancer", &self.load_balancer)
+            .field("schema_agreement_interval", &self.schema_agreement_interval)
+            .field("retry_policy", &self.retry_policy)
+            .field(
+                "speculative_execution_policy",
+                &self.speculative_execution_policy,
+            )
+            .field("metrics", &self.metrics)
+            .field("default_consistency", &self.default_consistency)
+            .field(
+                "auto_await_schema_agreement_timeout",
+                &self.auto_await_schema_agreement_timeout,
+            )
+            .finish()
+    }
+}
+
 /// Configuration options for [`Session`].
 /// Can be created manually, but usually it's easier to use
 /// [SessionBuilder](super::session_builder::SessionBuilder)