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.

Author: wprzytula

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

scylla/src/transport/query_result.rs

@@ -167,6 +167,10 @@ impl QueryResult {
         }
     }
 
+    pub(crate) fn raw_metadata_and_rows(&self) -> Option<&RawMetadataAndRawRows> {
+        self.raw_metadata_and_rows.as_ref()
+    }
+
     /// Warnings emitted by the database.
     #[inline]
     pub fn warnings(&self) -> impl Iterator<Item = &str> {

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;
@@ -798,6 +799,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))
     }
@@ -1235,6 +1237,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))
     }
@@ -1430,8 +1433,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)
     }
 
@@ -2150,6 +2157,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>