From 2deb9ee50bce75d2115c51829d19062cf5ceae73 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Przytu=C5=82a?= <wojciech.przytula@scylladb.com>
Date: Mon, 28 Oct 2024 08:52:27 +0100
Subject: [PATCH] session,iterator: record raw metadata&rows size

Even though we can no longer record rows serialized size without
accounting metadata, we can record their serialized size together.
---
 scylla-cql/src/frame/response/result.rs |  6 ++++++
 scylla/src/transport/iterator.rs        |  2 ++
 scylla/src/transport/query_result.rs    |  4 ++++
 scylla/src/transport/session.rs         | 20 +++++++++++++++++++-
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/scylla-cql/src/frame/response/result.rs b/scylla-cql/src/frame/response/result.rs
index f5cbbac9d..2fca59864 100644
--- a/scylla-cql/src/frame/response/result.rs
+++ b/scylla-cql/src/frame/response/result.rs
@@ -639,6 +639,12 @@ impl RawMetadataAndRawRows {
             cached_metadata: None,
         }
     }
+
+    /// Returns the serialized size of the raw metadata + raw rows.
+    #[inline]
+    pub fn metadata_and_rows_bytes_size(&self) -> usize {
+        self.raw_metadata_and_rows.len()
+    }
 }
 
 /// RESULT:Rows response, in partially serialized form.
diff --git a/scylla/src/transport/iterator.rs b/scylla/src/transport/iterator.rs
index 5fd037dc5..7d83ba4c1 100644
--- a/scylla/src/transport/iterator.rs
+++ b/scylla/src/transport/iterator.rs
@@ -328,6 +328,8 @@ where
                     .load_balancing_policy
                     .on_query_success(&self.statement_info, elapsed, node);
 
+                request_span.record_raw_rows_fields(&rows);
+
                 let received_page = ReceivedPage { rows, tracing_id };
 
                 // Send next page to RowIterator
diff --git a/scylla/src/transport/query_result.rs b/scylla/src/transport/query_result.rs
index ece0d2534..faf5496c2 100644
--- a/scylla/src/transport/query_result.rs
+++ b/scylla/src/transport/query_result.rs
@@ -172,6 +172,10 @@ impl QueryResult {
         }
     }
 
+    pub(crate) fn raw_metadata_and_rows(&self) -> Option<&RawMetadataAndRawRows> {
+        self.raw_rows.as_ref()
+    }
+
     /// Warnings emitted by the database.
     #[inline]
     pub fn warnings(&self) -> impl Iterator<Item = &str> {
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index bf875c8a0..8468183d2 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -18,6 +18,7 @@ use async_trait::async_trait;
 use futures::future::join_all;
 use futures::future::try_join_all;
 use itertools::{Either, Itertools};
+use scylla_cql::frame::response::result::RawMetadataAndRawRows;
 use scylla_cql::frame::response::result::{deser_cql_value, ColumnSpec};
 use scylla_cql::frame::response::NonErrorResponse;
 use scylla_cql::types::serialize::batch::BatchValues;
@@ -828,6 +829,7 @@ impl Session {
         self.handle_auto_await_schema_agreement(&response).await?;
 
         let (result, paging_state) = response.into_query_result_and_paging_state()?;
+        span.record_result_fields(&result);
         let result = result.into_legacy_result()?;
         Ok((result, paging_state))
     }
@@ -1265,6 +1267,7 @@ impl Session {
         self.handle_auto_await_schema_agreement(&response).await?;
 
         let (result, paging_state) = response.into_query_result_and_paging_state()?;
+        span.record_result_fields(&result);
         let result = result.into_legacy_result()?;
         Ok((result, paging_state))
     }
@@ -1460,8 +1463,12 @@ impl Session {
 
         let result = match run_query_result {
             RunQueryResult::IgnoredWriteError => LegacyQueryResult::mock_empty(),
-            RunQueryResult::Completed(response) => response.into_legacy_result()?,
+            RunQueryResult::Completed(result) => {
+                span.record_result_fields(&result);
+                result.into_legacy_result()?
+            }
         };
+
         Ok(result)
     }
 
@@ -2180,6 +2187,17 @@ impl RequestSpan {
         }
     }
 
+    pub(crate) fn record_raw_rows_fields(&self, raw_rows: &RawMetadataAndRawRows) {
+        self.span
+            .record("raw_result_size", raw_rows.metadata_and_rows_bytes_size());
+    }
+
+    pub(crate) fn record_result_fields(&self, query_result: &QueryResult) {
+        if let Some(raw_metadata_and_rows) = query_result.raw_metadata_and_rows() {
+            self.record_raw_rows_fields(raw_metadata_and_rows);
+        }
+    }
+
     pub(crate) fn record_replicas<'a>(&'a self, replicas: &'a [(impl Borrow<Arc<Node>>, Shard)]) {
         struct ReplicaIps<'a, N>(&'a [(N, Shard)]);
         impl<'a, N> Display for ReplicaIps<'a, N>