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 Oct 28, 2024
1 parent 84abb72 commit 2deb9ee
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 1 deletion.
6 changes: 6 additions & 0 deletions scylla-cql/src/frame/response/result.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
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 @@ -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> {
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 2deb9ee

Please sign in to comment.