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 2b1344e commit 5ef967d
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 5ef967d

Please sign in to comment.