Better USDT probes for ClickHouse protocol (#8622)

This adds more observability into the network protocol used to talk to
ClickHouse, in an attempt to help resolve #8595

Author: bnaecker

Cargo.lock

@@ -382,7 +382,7 @@ cargo_metadata = "0.20.0"
 chacha20poly1305 = "0.10.1"
 cfg-if = "1.0"
 chrono = { version = "0.4", features = [ "serde" ] }
-chrono-tz = "0.10.3"
+chrono-tz = { version = "0.10.3", features = [ "serde" ] }
 ciborium = "0.2.2"
 clap = { version = "4.5", features = ["cargo", "derive", "env", "wrap_help"] }
 clickana = { path = "dev-tools/clickana" }

Cargo.toml

@@ -37,7 +37,11 @@ async fn native_impl(
     conn: &Arc<Mutex<Connection>>,
     query: &str,
 ) {
-    conn.lock().await.query(query).await.expect("Expected to run query");
+    conn.lock()
+        .await
+        .query(uuid::Uuid::new_v4(), query)
+        .await
+        .expect("Expected to run query");
 }
 
 /// Setup the native query benchmark.

oximeter/db/benches/protocol.rs

@@ -4,7 +4,7 @@
 
 //! Rust client to ClickHouse database
 
-// Copyright 2024 Oxide Computer Company
+// Copyright 2025 Oxide Computer Company
 
 pub(crate) mod dbwrite;
 #[cfg(any(feature = "oxql", test))]
@@ -78,10 +78,16 @@ const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
 #[usdt::provider(provider = "clickhouse_client")]
 mod probes {
     /// Fires when a SQL query begins, with the query string.
-    fn sql__query__start(_: &usdt::UniqueId, sql: &str) {}
+    ///
+    /// The first argument is a UUID assigned to the query. This is sent to the
+    /// ClickHouse server along with the query string, where it's used as the
+    /// query ID in tables like the query log. This can be used to correlate
+    /// behavior on the client side with how the server itself processes the
+    /// query.
+    fn sql__query__start(query_id: &str, sql: &str) {}
 
     /// Fires when a SQL query ends, either in success or failure.
-    fn sql__query__done(_: &usdt::UniqueId) {}
+    fn sql__query__done(query_id: &str) {}
 }
 
 /// A `Client` to the ClickHouse metrics database.
@@ -1124,16 +1130,16 @@ impl Client {
             "n_rows" => block.n_rows(),
             "n_columns" => block.n_columns(),
         );
-        let id = usdt::UniqueId::new();
-        probes::sql__query__start!(|| (&id, sql));
+        let id = Uuid::new_v4();
+        probes::sql__query__start!(|| (id.to_string(), sql));
         let now = tokio::time::Instant::now();
         let result = tokio::time::timeout(
             self.request_timeout,
-            handle.insert(sql, block),
+            handle.insert(id, sql, block),
         )
         .await;
         let elapsed = now.elapsed();
-        probes::sql__query__done!(|| (&id));
+        probes::sql__query__done!(|| id.to_string());
         match result {
             Ok(result) => result.map_err(Error::from),
             Err(_) => Err(Error::DatabaseUnavailable(format!(
@@ -1166,13 +1172,14 @@ impl Client {
             "sql" => sql,
         );
 
-        let id = usdt::UniqueId::new();
-        probes::sql__query__start!(|| (&id, sql));
+        let id = Uuid::new_v4();
+        probes::sql__query__start!(|| (id.to_string(), sql));
         let now = tokio::time::Instant::now();
         let result =
-            tokio::time::timeout(self.request_timeout, handle.query(sql)).await;
+            tokio::time::timeout(self.request_timeout, handle.query(id, sql))
+                .await;
         let elapsed = now.elapsed();
-        probes::sql__query__done!(|| (&id));
+        probes::sql__query__done!(|| id.to_string());
         match result {
             Ok(result) => result.map_err(Error::from),
             Err(_) => Err(Error::DatabaseUnavailable(format!(

oximeter/db/src/client/mod.rs

@@ -2,7 +2,7 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 //
-// Copyright 2024 Oxide Computer Company
+// Copyright 2025 Oxide Computer Company
 
 //! Types for working with actual blocks and columns of data.
 
@@ -45,7 +45,7 @@ use uuid::Uuid;
 /// results of a query. It is not necessarily the entire set of data -- most
 /// queries result in sending more than one block. But every query (and some of
 /// the management or status commands) sends data as one or more block.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, serde::Serialize)]
 pub struct Block {
     /// A name for the block.
     ///
@@ -282,7 +282,7 @@ impl Block {
 ///
 /// This is only used for a few special kinds of queries. See the fields for
 /// details.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq, serde::Serialize)]
 pub struct BlockInfo {
     /// True if this is an "overflow" block, which is the case if:
     ///
@@ -325,7 +325,7 @@ impl Default for BlockInfo {
 ///
 /// This represents a single column of data fetched in a query or sent in an
 /// insert. It includes the type implicitly, in the value array it contains.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, serde::Serialize)]
 pub struct Column {
     /// The data values for this column.
     pub values: ValueArray,
@@ -370,7 +370,7 @@ impl Column {
 }
 
 /// An array of singly-typed data values from the server.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, serde::Serialize)]
 pub enum ValueArray {
     Bool(Vec<bool>),
     UInt8(Vec<u8>),
@@ -689,7 +689,8 @@ impl From<Vec<IpAddr>> for ValueArray {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq, serde::Serialize)]
+#[serde(transparent)]
 pub struct Precision(u8);
 
 impl TryFrom<u8> for Precision {
@@ -792,7 +793,7 @@ impl fmt::Display for Precision {
 }
 
 /// A type of a column of data.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, serde::Serialize)]
 pub enum DataType {
     Bool,
     UInt8,