Skip to content

Commit

Permalink
session,iterator: record raw metadata&rows size
Browse files Browse the repository at this point in the history
Even though we can no longer record rows serialized size without
accounting metadata, we can record their serialized size together.
  • Loading branch information
wprzytula committed Nov 4, 2024
1 parent d7a3d7d commit a2fa253
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 1 deletion.
2 changes: 2 additions & 0 deletions scylla/src/transport/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions scylla/src/transport/query_result.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
20 changes: 19 additions & 1 deletion scylla/src/transport/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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))
}
Expand Down Expand Up @@ -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))
}
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -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>
Expand Down

0 comments on commit a2fa253

Please sign in to comment.