diff --git a/sdk/cosmos/.cspell.json b/sdk/cosmos/.cspell.json index ab818543976..9c202c4ea4b 100644 --- a/sdk/cosmos/.cspell.json +++ b/sdk/cosmos/.cspell.json @@ -4,6 +4,7 @@ ], "ignoreWords": [ "accountname", + "acked", "activityid", "ALPN", "apacsoutheast", @@ -262,6 +263,7 @@ "RUPM", "rwcache", "sbyte", + "schemaversion", "serviceunavailable", "serviceversion", "sess", diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 38a3496ffb5..2ed418699b8 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -10,6 +10,8 @@ ### Other Changes +- Existing `ResponseHeaders` accessors now return Gateway 2.0 backend duration, quota, item-count, and local-LSN response metadata. ([#4797](https://github.com/Azure/azure-sdk-for-rust/pull/4797)) + ## 0.37.0 (2026-07-20) ### Features Added diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs index e4d19c51dce..7379c26ea51 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs @@ -546,25 +546,27 @@ async fn sdk_query_metadata_databases_and_containers() { let db_query = Query::from("SELECT * FROM c WHERE c.id = @id") .with_parameter("@id", db_name.as_str()) .unwrap(); - let emu_databases: Vec = backend - .emulator_client - .query_databases(db_query.clone(), None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); + let emu_databases: Vec = Box::pin( + backend + .emulator_client + .query_databases(db_query.clone(), None), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(emu_databases.len(), 1); assert_eq!(emu_databases[0].id.as_deref(), Some(db_name.as_str())); if let Some(ref real_client) = backend.real_client { - let real_databases: Vec = real_client - .query_databases(db_query, None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); + let real_databases: Vec = + Box::pin(real_client.query_databases(db_query, None)) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(real_databases.len(), emu_databases.len()); assert_eq!(real_databases[0].id, emu_databases[0].id); } @@ -572,27 +574,31 @@ async fn sdk_query_metadata_databases_and_containers() { let container_query = Query::from("SELECT * FROM c WHERE c.id = @id") .with_parameter("@id", "testcoll") .unwrap(); - let emu_containers: Vec = backend - .emulator_client - .database_client(&db_name) - .query_containers(container_query.clone(), None) + let emu_containers: Vec = Box::pin( + backend + .emulator_client + .database_client(&db_name) + .query_containers(container_query.clone(), None), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(emu_containers.len(), 1); + assert_eq!(emu_containers[0].id, "testcoll"); + + if let Some(ref real_client) = backend.real_client { + let real_containers: Vec = Box::pin( + real_client + .database_client(&db_name) + .query_containers(container_query, None), + ) .await .unwrap() .try_collect() .await .unwrap(); - assert_eq!(emu_containers.len(), 1); - assert_eq!(emu_containers[0].id, "testcoll"); - - if let Some(ref real_client) = backend.real_client { - let real_containers: Vec = real_client - .database_client(&db_name) - .query_containers(container_query, None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); assert_eq!(real_containers.len(), emu_containers.len()); assert_eq!(real_containers[0].id, emu_containers[0].id); } @@ -1085,25 +1091,25 @@ async fn sdk_query_items_with_filter_and_projection() { .unwrap() } - let emu_items: Vec = emu_container - .query_items(query(), FeedScope::partition("pk1"), None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); - assert_eq!(emu_items.len(), 2); - assert_eq!(emu_items[0].id, "query-1"); - assert_eq!(emu_items[1].id, "query-2"); - - if let Some(ref real) = real_container { - let real_items: Vec = real - .query_items(query(), FeedScope::partition("pk1"), None) + let emu_items: Vec = + Box::pin(emu_container.query_items(query(), FeedScope::partition("pk1"), None)) .await .unwrap() .try_collect() .await .unwrap(); + assert_eq!(emu_items.len(), 2); + assert_eq!(emu_items[0].id, "query-1"); + assert_eq!(emu_items[1].id, "query-2"); + + if let Some(ref real) = real_container { + let real_items: Vec = + Box::pin(real.query_items(query(), FeedScope::partition("pk1"), None)) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(real_items.len(), emu_items.len()); assert_eq!( real_items.iter().map(|i| &i.id).collect::>(), diff --git a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs index 3d2f5a5b290..cdc2a00b8c0 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs @@ -9,7 +9,10 @@ use azure_data_cosmos::fault_injection::{ FaultInjectionRuleBuilder, FaultOperationType, }; use azure_data_cosmos::models::{ContainerProperties, ThroughputProperties}; -use azure_data_cosmos::options::{ExcludedRegions, ItemReadOptions, OperationOptions}; +use azure_data_cosmos::options::{ + ExcludedRegions, ItemReadOptions, OperationOptions, OperationOptionsBuilder, + ThrottlingRetryOptionsBuilder, +}; use framework::{ assert_local_retry_attempted_on_region, assert_region_contacted_with_retry, assert_region_not_contacted, TestClient, TestOptions, HUB_REGION, SATELLITE_REGION, @@ -88,9 +91,20 @@ async fn verify_read_fails_with_injected_error( .expect("fault client should be available"); let fault_db_client = fault_client.database_client(db_client.id()); let fault_container_client = fault_db_client.container_client(&container_id).await?; + let read_options = (expected_status == StatusCode::TooManyRequests).then(|| { + ItemReadOptions::default().with_operation_options( + OperationOptionsBuilder::new() + .with_throttling_retry_options( + ThrottlingRetryOptionsBuilder::new() + .with_max_retry_count(0) + .build(), + ) + .build(), + ) + }); let result = run_context - .read_item(&fault_container_client, &pk, &item_id, None) + .read_item(&fault_container_client, &pk, &item_id, read_options) .await; let err = result.expect_err(&format!( diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index fa1d74748a4..aaa0f3ab3bd 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -10,6 +10,8 @@ ### Other Changes +- Gateway 2.0 responses now preserve backend duration, quota, item-count, quorum, replica, query, and physical-partition RNTBD metadata when converting to standard Cosmos response headers. ([#4797](https://github.com/Azure/azure-sdk-for-rust/pull/4797)) + ## 0.6.0 (2026-07-20) ### Features Added diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs index 3abefa55cc6..4991915de8f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs @@ -410,6 +410,84 @@ pub(crate) fn unwrap_response_for_gateway_v2( if let Some(charge) = response.request_charge { headers.insert(response_header_names::REQUEST_CHARGE, charge.to_string()); } + if let Some(duration_ms) = response.backend_request_duration_ms { + headers.insert( + response_header_names::SERVER_DURATION_MS, + duration_ms.to_string(), + ); + } + if let Some(value) = response.last_state_change_date_time { + headers.insert(response_header_names::LAST_STATE_CHANGE_UTC, value); + } + if let Some(value) = response.storage_max_resource_quota { + headers.insert(response_header_names::RESOURCE_QUOTA, value); + } + if let Some(value) = response.storage_resource_quota_usage { + headers.insert(response_header_names::RESOURCE_USAGE, value); + } + if let Some(value) = response.schema_version { + headers.insert(response_header_names::SCHEMA_VERSION, value); + } + if let Some(value) = response.item_count { + headers.insert(response_header_names::ITEM_COUNT, value.to_string()); + } + if let Some(value) = response.owner_id { + headers.insert(response_header_names::OWNER_ID, value); + } + if let Some(value) = response.quorum_acked_lsn { + headers.insert(response_header_names::QUORUM_ACKED_LSN, value.to_string()); + } + if let Some(value) = response.current_write_quorum { + headers.insert( + response_header_names::CURRENT_WRITE_QUORUM, + value.to_string(), + ); + } + if let Some(value) = response.current_replica_set_size { + headers.insert( + response_header_names::CURRENT_REPLICA_SET_SIZE, + value.to_string(), + ); + } + if let Some(value) = response.xp_role { + headers.insert(response_header_names::XP_ROLE, value.to_string()); + } + if let Some(value) = response.number_of_read_regions { + headers.insert( + response_header_names::NUMBER_OF_READ_REGIONS, + value.to_string(), + ); + } + if let Some(value) = response.local_lsn.filter(|value| *value >= 0) { + headers.insert(response_header_names::LOCAL_LSN, value.to_string()); + } + if let Some(value) = response.quorum_acked_local_lsn { + headers.insert( + response_header_names::QUORUM_ACKED_LOCAL_LSN, + value.to_string(), + ); + } + if let Some(value) = response.item_local_lsn.filter(|value| *value >= 0) { + headers.insert(response_header_names::ITEM_LOCAL_LSN, value.to_string()); + } + if let Some(value) = response.query_execution_info { + headers.insert(response_header_names::QUERY_EXECUTION_INFO, value); + } + if let Some(value) = response.pending_pk_delete { + headers.insert( + response_header_names::PENDING_PK_DELETE, + if value { "True" } else { "False" }, + ); + } + if let Some(value) = response.physical_partition_id { + headers.insert(response_header_names::PHYSICAL_PARTITION_ID, value); + } + if let Some(value) = response.conflict_resolved_timestamp { + headers.insert( + response_header_names::CONFLICT_RESOLVED_TIMESTAMP, + value.to_string(), + ); + } if let Some(token) = response.session_token { // GW2 surfaces only the vector portion of the session token, with the // partition key range id carried separately. Classic gateway emits @@ -1933,8 +2011,27 @@ mod tests { 404, activity_id, |tokens| { + write_small_string_token(tokens, 0x0002, "2026-07-21T00:00:00Z"); + write_string_token(tokens, 0x000E, "documentSize=10240;"); + write_string_token(tokens, 0x000F, "documentSize=1;"); + write_small_string_token(tokens, 0x0010, "1.0"); write_u32_token(tokens, 0x001C, 1002); + write_u32_token(tokens, 0x0014, 1); write_double_token(tokens, 0x0015, 3.5); + write_string_token(tokens, 0x0018, "owner-rid"); + write_i64_token(tokens, 0x001A, 40); + write_u32_token(tokens, 0x001E, 3); + write_u32_token(tokens, 0x001F, 4); + write_u32_token(tokens, 0x0026, 1); + write_u32_token(tokens, 0x0030, 2); + write_i64_token(tokens, 0x003A, 41); + write_i64_token(tokens, 0x003B, 40); + write_i64_token(tokens, 0x003C, 39); + write_string_token(tokens, 0x0045, "{\"reverseRidEnabled\":false}"); + write_double_token(tokens, 0x0051, 12.75); + write_byte_token(tokens, 0x0055, 0); + write_string_token(tokens, 0x0063, "physical-0"); + write_u64_token(tokens, 0x0087, 1_234_567); write_string_token(tokens, 0x003E, "1:2#3"); write_string_token(tokens, 0x0004, "\"etag\""); write_string_token(tokens, 0x0003, "continuation"); @@ -1967,6 +2064,47 @@ mod tests { .get_optional_str(&HeaderName::from_static("x-ms-request-charge")), Some("3.5") ); + assert_eq!( + unwrapped.headers.get_optional_str(&HeaderName::from_static( + response_header_names::SERVER_DURATION_MS + )), + Some("12.75") + ); + let parsed_headers = crate::models::CosmosResponseHeaders::from_headers(&unwrapped.headers); + assert_eq!(parsed_headers.server_duration_ms, Some(12.75)); + assert_eq!( + parsed_headers.last_state_change_utc.as_deref(), + Some("2026-07-21T00:00:00Z") + ); + assert_eq!( + parsed_headers.resource_quota.as_deref(), + Some("documentSize=10240;") + ); + assert_eq!( + parsed_headers.resource_usage.as_deref(), + Some("documentSize=1;") + ); + assert_eq!(parsed_headers.schema_version.as_deref(), Some("1.0")); + assert_eq!(parsed_headers.item_count, Some(1)); + assert_eq!(parsed_headers.owner_id.as_deref(), Some("owner-rid")); + assert_eq!(parsed_headers.quorum_acked_lsn, Some(40)); + assert_eq!(parsed_headers.current_write_quorum, Some(3)); + assert_eq!(parsed_headers.current_replica_set_size, Some(4)); + assert_eq!(parsed_headers.xp_role, Some(1)); + assert_eq!(parsed_headers.number_of_read_regions, Some(2)); + assert_eq!(parsed_headers.local_lsn, Some(41)); + assert_eq!(parsed_headers.quorum_acked_local_lsn, Some(40)); + assert_eq!(parsed_headers.item_local_lsn, Some(39)); + assert_eq!( + parsed_headers.query_execution_info.as_deref(), + Some("{\"reverseRidEnabled\":false}") + ); + assert_eq!(parsed_headers.pending_pk_delete, Some(false)); + assert_eq!( + parsed_headers.physical_partition_id.as_deref(), + Some("physical-0") + ); + assert_eq!(parsed_headers.conflict_resolved_timestamp, Some(1_234_567)); assert_eq!( unwrapped .headers @@ -2142,6 +2280,19 @@ mod tests { bytes.extend_from_slice(value.as_bytes()); } + fn write_small_string_token(bytes: &mut Vec, id: u16, value: &str) { + bytes.extend_from_slice(&id.to_le_bytes()); + bytes.push(0x07); + bytes.push(u8::try_from(value.len()).unwrap()); + bytes.extend_from_slice(value.as_bytes()); + } + + fn write_byte_token(bytes: &mut Vec, id: u16, value: u8) { + bytes.extend_from_slice(&id.to_le_bytes()); + bytes.push(0x00); + bytes.push(value); + } + fn write_u32_token(bytes: &mut Vec, id: u16, value: u32) { bytes.extend_from_slice(&id.to_le_bytes()); bytes.push(0x02); @@ -2154,6 +2305,12 @@ mod tests { bytes.extend_from_slice(&value.to_le_bytes()); } + fn write_u64_token(bytes: &mut Vec, id: u16, value: u64) { + bytes.extend_from_slice(&id.to_le_bytes()); + bytes.push(0x04); + bytes.extend_from_slice(&value.to_le_bytes()); + } + fn write_double_token(bytes: &mut Vec, id: u16, value: f64) { bytes.extend_from_slice(&id.to_le_bytes()); bytes.push(0x0E); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/response.rs index c8c8085ae97..bc9f3633597 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/response.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/response.rs @@ -30,22 +30,60 @@ pub(crate) struct RntbdResponse { pub(crate) etag: Option, /// Retry-after delay in milliseconds. pub(crate) retry_after_ms: Option, + /// Timestamp of the last replica state change. + pub(crate) last_state_change_date_time: Option, + /// Maximum resource quota values. + pub(crate) storage_max_resource_quota: Option, + /// Current resource quota usage. + pub(crate) storage_resource_quota_usage: Option, + /// Resource schema version. + pub(crate) schema_version: Option, /// Logical sequence number. pub(crate) lsn: Option, + /// Number of items returned. + pub(crate) item_count: Option, /// Request charge in request units. pub(crate) request_charge: Option, + /// Backend request processing duration in milliseconds. + pub(crate) backend_request_duration_ms: Option, /// Owner full name metadata. pub(crate) owner_full_name: Option, + /// Owner resource identifier. + pub(crate) owner_id: Option, + /// Quorum-acknowledged logical sequence number. + pub(crate) quorum_acked_lsn: Option, + /// Current write quorum. + pub(crate) current_write_quorum: Option, + /// Current replica set size. + pub(crate) current_replica_set_size: Option, /// Partition key range identifier. pub(crate) partition_key_range_id: Option, + /// Cross-partition role. + pub(crate) xp_role: Option, + /// Number of read regions. + pub(crate) number_of_read_regions: Option, /// Item logical sequence number. pub(crate) item_lsn: Option, /// Global committed logical sequence number. pub(crate) global_committed_lsn: Option, + /// Local logical sequence number. + pub(crate) local_lsn: Option, + /// Quorum-acknowledged local logical sequence number. + pub(crate) quorum_acked_local_lsn: Option, + /// Item local logical sequence number. + pub(crate) item_local_lsn: Option, /// Transport request identifier. pub(crate) transport_request_id: Option, /// Session token for session consistency. pub(crate) session_token: Option, + /// Query execution metadata. + pub(crate) query_execution_info: Option, + /// Whether partition-key deletion is pending. + pub(crate) pending_pk_delete: Option, + /// Physical partition identifier. + pub(crate) physical_partition_id: Option, + /// Conflict resolution timestamp. + pub(crate) conflict_resolved_timestamp: Option, } impl RntbdResponse { @@ -85,15 +123,34 @@ impl RntbdResponse { let mut continuation_token = None; let mut etag = None; let mut retry_after_ms = None; + let mut last_state_change_date_time = None; + let mut storage_max_resource_quota = None; + let mut storage_resource_quota_usage = None; + let mut schema_version = None; let mut lsn = None; + let mut item_count = None; let mut request_charge = None; + let mut backend_request_duration_ms = None; let mut owner_full_name = None; + let mut owner_id = None; + let mut quorum_acked_lsn = None; let mut sub_status = None; + let mut current_write_quorum = None; + let mut current_replica_set_size = None; let mut partition_key_range_id = None; + let mut xp_role = None; + let mut number_of_read_regions = None; let mut item_lsn = None; let mut global_committed_lsn = None; + let mut local_lsn = None; + let mut quorum_acked_local_lsn = None; + let mut item_local_lsn = None; let mut transport_request_id = None; let mut session_token = None; + let mut query_execution_info = None; + let mut pending_pk_delete = None; + let mut physical_partition_id = None; + let mut conflict_resolved_timestamp = None; while !frame.is_empty() { let token = Token::read_from(&mut frame)?; @@ -101,6 +158,10 @@ impl RntbdResponse { Ok(RntbdResponseToken::PayloadPresent) => { payload_present = expect_byte(token, "PayloadPresent")? != 0; } + Ok(RntbdResponseToken::LastStateChangeDateTime) => { + last_state_change_date_time = + Some(expect_small_string(token, "LastStateChangeDateTime")?); + } Ok(RntbdResponseToken::ContinuationToken) => { continuation_token = Some(expect_string(token, "ContinuationToken")?); } @@ -110,33 +171,91 @@ impl RntbdResponse { Ok(RntbdResponseToken::RetryAfterMilliseconds) => { retry_after_ms = Some(expect_u32(token, "RetryAfterMilliseconds")?); } + Ok(RntbdResponseToken::StorageMaxResourceQuota) => { + storage_max_resource_quota = + Some(expect_string(token, "StorageMaxResourceQuota")?); + } + Ok(RntbdResponseToken::StorageResourceQuotaUsage) => { + storage_resource_quota_usage = + Some(expect_string(token, "StorageResourceQuotaUsage")?); + } + Ok(RntbdResponseToken::SchemaVersion) => { + schema_version = Some(expect_small_string(token, "SchemaVersion")?); + } Ok(RntbdResponseToken::Lsn) => { lsn = Some(expect_i64(token, "LSN")?); } + Ok(RntbdResponseToken::ItemCount) => { + item_count = Some(expect_u32(token, "ItemCount")?); + } Ok(RntbdResponseToken::RequestCharge) => { request_charge = Some(expect_f64(token, "RequestCharge")?); } + Ok(RntbdResponseToken::BackendRequestDurationMilliseconds) => { + backend_request_duration_ms = + Some(expect_f64(token, "BackendRequestDurationMilliseconds")?); + } Ok(RntbdResponseToken::OwnerFullName) => { owner_full_name = Some(expect_string(token, "OwnerFullName")?); } + Ok(RntbdResponseToken::OwnerId) => { + owner_id = Some(expect_string(token, "OwnerId")?); + } + Ok(RntbdResponseToken::QuorumAckedLsn) => { + quorum_acked_lsn = Some(expect_i64(token, "QuorumAckedLSN")?); + } Ok(RntbdResponseToken::SubStatus) => { sub_status = Some(expect_u32(token, "SubStatus")?); } + Ok(RntbdResponseToken::CurrentWriteQuorum) => { + current_write_quorum = Some(expect_u32(token, "CurrentWriteQuorum")?); + } + Ok(RntbdResponseToken::CurrentReplicaSetSize) => { + current_replica_set_size = Some(expect_u32(token, "CurrentReplicaSetSize")?); + } Ok(RntbdResponseToken::PartitionKeyRangeId) => { partition_key_range_id = Some(expect_string(token, "PartitionKeyRangeId")?); } + Ok(RntbdResponseToken::XpRole) => { + xp_role = Some(expect_u32(token, "XPRole")?); + } + Ok(RntbdResponseToken::NumberOfReadRegions) => { + number_of_read_regions = Some(expect_u32(token, "NumberOfReadRegions")?); + } Ok(RntbdResponseToken::ItemLsn) => { item_lsn = Some(expect_i64(token, "ItemLSN")?); } Ok(RntbdResponseToken::GlobalCommittedLsn) => { global_committed_lsn = Some(expect_i64(token, "GlobalCommittedLSN")?); } + Ok(RntbdResponseToken::LocalLsn) => { + local_lsn = Some(expect_i64(token, "LocalLSN")?); + } + Ok(RntbdResponseToken::QuorumAckedLocalLsn) => { + quorum_acked_local_lsn = Some(expect_i64(token, "QuorumAckedLocalLSN")?); + } + Ok(RntbdResponseToken::ItemLocalLsn) => { + item_local_lsn = Some(expect_i64(token, "ItemLocalLSN")?); + } Ok(RntbdResponseToken::TransportRequestId) => { transport_request_id = Some(expect_u32(token, "TransportRequestID")?); } Ok(RntbdResponseToken::SessionToken) => { session_token = Some(expect_string(token, "SessionToken")?); } + Ok(RntbdResponseToken::QueryExecutionInfo) => { + query_execution_info = Some(expect_string(token, "QueryExecutionInfo")?); + } + Ok(RntbdResponseToken::PendingPkDelete) => { + pending_pk_delete = Some(expect_byte(token, "PendingPKDelete")? != 0); + } + Ok(RntbdResponseToken::PhysicalPartitionId) => { + physical_partition_id = Some(expect_string(token, "PhysicalPartitionId")?); + } + Ok(RntbdResponseToken::ConflictResolvedTimestamp) => { + conflict_resolved_timestamp = + Some(expect_u64(token, "ConflictResolvedTimestamp")?); + } Err(()) => {} } } @@ -162,14 +281,33 @@ impl RntbdResponse { continuation_token, etag, retry_after_ms, + last_state_change_date_time, + storage_max_resource_quota, + storage_resource_quota_usage, + schema_version, lsn, + item_count, request_charge, + backend_request_duration_ms, owner_full_name, + owner_id, + quorum_acked_lsn, + current_write_quorum, + current_replica_set_size, partition_key_range_id, + xp_role, + number_of_read_regions, item_lsn, global_committed_lsn, + local_lsn, + quorum_acked_local_lsn, + item_local_lsn, transport_request_id, session_token, + query_execution_info, + pending_pk_delete, + physical_partition_id, + conflict_resolved_timestamp, }) } } @@ -180,6 +318,12 @@ fn expect_string(token: Token, name: &str) -> azure_core::Result { .ok_or_else(|| unexpected_token_type(name)) } +fn expect_small_string(token: Token, name: &str) -> azure_core::Result { + token + .into_small_string() + .ok_or_else(|| unexpected_token_type(name)) +} + fn expect_byte(token: Token, name: &str) -> azure_core::Result { token.as_byte().ok_or_else(|| unexpected_token_type(name)) } @@ -192,6 +336,10 @@ fn expect_i64(token: Token, name: &str) -> azure_core::Result { token.as_i64().ok_or_else(|| unexpected_token_type(name)) } +fn expect_u64(token: Token, name: &str) -> azure_core::Result { + token.as_u64().ok_or_else(|| unexpected_token_type(name)) +} + fn expect_f64(token: Token, name: &str) -> azure_core::Result { token.as_f64().ok_or_else(|| unexpected_token_type(name)) } @@ -229,6 +377,116 @@ mod tests { assert!(response.body.is_empty()); } + #[test] + fn backend_request_duration_is_decoded() { + let mut frame = response_header(StatusCode::Ok); + Token::new(0x0051, TokenValue::Double(12.75)) + .write_to(&mut frame) + .unwrap(); + patch_total_len(&mut frame); + + let response = RntbdResponse::read(&frame).unwrap(); + + assert_eq!(response.backend_request_duration_ms, Some(12.75)); + } + + #[test] + fn backend_request_duration_rejects_wrong_token_type() { + let mut frame = response_header(StatusCode::Ok); + Token::new(0x0051, TokenValue::ULong(12)) + .write_to(&mut frame) + .unwrap(); + patch_total_len(&mut frame); + + let err = RntbdResponse::read(&frame).unwrap_err(); + + assert_eq!(*err.kind(), azure_core::error::ErrorKind::DataConversion); + } + + #[test] + fn known_token_with_wrong_type_is_rejected() { + let mut frame = response_header(StatusCode::Ok); + Token::new(0x0010, TokenValue::String("1.0".into())) + .write_to(&mut frame) + .unwrap(); + patch_total_len(&mut frame); + + let err = RntbdResponse::read(&frame).unwrap_err(); + + assert_eq!(*err.kind(), azure_core::error::ErrorKind::DataConversion); + } + + #[test] + fn live_observed_gateway_v2_metadata_is_decoded() { + let mut frame = response_header(StatusCode::Ok); + let tokens = [ + Token::new( + 0x0002, + TokenValue::SmallString("2026-07-21T00:00:00Z".into()), + ), + Token::new(0x000E, TokenValue::String("documentSize=10240;".into())), + Token::new(0x000F, TokenValue::String("documentSize=1;".into())), + Token::new(0x0010, TokenValue::SmallString("1.0".into())), + Token::new(0x0014, TokenValue::ULong(1)), + Token::new(0x0018, TokenValue::String("owner-rid".into())), + Token::new(0x001A, TokenValue::LongLong(40)), + Token::new(0x001E, TokenValue::ULong(3)), + Token::new(0x001F, TokenValue::ULong(4)), + Token::new(0x0026, TokenValue::ULong(1)), + Token::new(0x0030, TokenValue::ULong(2)), + Token::new(0x003A, TokenValue::LongLong(41)), + Token::new(0x003B, TokenValue::LongLong(40)), + Token::new(0x003C, TokenValue::LongLong(39)), + Token::new( + 0x0045, + TokenValue::String("{\"reverseRidEnabled\":false}".into()), + ), + Token::new(0x0055, TokenValue::Byte(0)), + Token::new(0x0063, TokenValue::String("physical-0".into())), + Token::new(0x0087, TokenValue::ULongLong(1_234_567)), + ]; + for token in tokens { + token.write_to(&mut frame).unwrap(); + } + patch_total_len(&mut frame); + + let response = RntbdResponse::read(&frame).unwrap(); + + assert_eq!( + response.last_state_change_date_time.as_deref(), + Some("2026-07-21T00:00:00Z") + ); + assert_eq!( + response.storage_max_resource_quota.as_deref(), + Some("documentSize=10240;") + ); + assert_eq!( + response.storage_resource_quota_usage.as_deref(), + Some("documentSize=1;") + ); + assert_eq!(response.schema_version.as_deref(), Some("1.0")); + assert_eq!(response.item_count, Some(1)); + assert_eq!(response.owner_id.as_deref(), Some("owner-rid")); + assert_eq!(response.quorum_acked_lsn, Some(40)); + assert_eq!(response.current_write_quorum, Some(3)); + assert_eq!(response.current_replica_set_size, Some(4)); + assert_eq!(response.xp_role, Some(1)); + assert_eq!(response.number_of_read_regions, Some(2)); + assert_eq!(response.local_lsn, Some(41)); + assert_eq!(response.quorum_acked_local_lsn, Some(40)); + assert_eq!(response.item_local_lsn, Some(39)); + assert_eq!( + response.query_execution_info.as_deref(), + Some("{\"reverseRidEnabled\":false}") + ); + assert_eq!(response.pending_pk_delete, Some(false)); + assert_eq!( + response.physical_partition_id.as_deref(), + Some("physical-0") + ); + assert_eq!(response.conflict_resolved_timestamp, Some(1_234_567)); + } + #[test] fn total_length_past_buffer_is_rejected() { let mut frame = response_header(StatusCode::Ok); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/tokens.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/tokens.rs index e66cd98cbe6..e2c702e23fb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/tokens.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/tokens.rs @@ -492,6 +492,14 @@ impl Token { } } + /// Returns the `ULongLong` (`u64`) value, or `None` if the token holds a different type. + pub(crate) fn as_u64(&self) -> Option { + match self.value { + TokenValue::ULongLong(value) => Some(value), + _ => None, + } + } + /// Returns the `Double` (`f64`) value, or `None` if the token holds a different type. pub(crate) fn as_f64(&self) -> Option { match self.value { @@ -509,6 +517,15 @@ impl Token { } } + /// Returns the owned `SmallString` value, consuming the token, or `None` + /// if the token holds a different type. + pub(crate) fn into_small_string(self) -> Option { + match self.value { + TokenValue::SmallString(value) => Some(value), + _ => None, + } + } + /// Returns the number of bytes this token occupies on the wire. pub(super) fn encoded_len(&self) -> usize { 2 + 1 + self.value.encoded_len() @@ -647,30 +664,68 @@ pub(super) enum RntbdResponseToken { /// Payload-present flag. When true, a u32 LE body-length prefix and that /// many body bytes follow the metadata section. PayloadPresent, + /// Timestamp of the last replica state change. + LastStateChangeDateTime, /// Continuation token. ContinuationToken, /// Entity tag. ETag, /// Retry-after delay in milliseconds. RetryAfterMilliseconds, + /// Maximum resource quota values. + StorageMaxResourceQuota, + /// Current resource quota usage. + StorageResourceQuotaUsage, + /// Resource schema version. + SchemaVersion, /// Logical sequence number. Lsn, + /// Number of items returned. + ItemCount, /// Request charge in request units. RequestCharge, /// Owner full name. OwnerFullName, + /// Owner resource identifier. + OwnerId, + /// Quorum-acknowledged logical sequence number. + QuorumAckedLsn, /// Cosmos DB sub-status code. SubStatus, + /// Current write quorum. + CurrentWriteQuorum, + /// Current replica set size. + CurrentReplicaSetSize, /// Partition key range identifier. PartitionKeyRangeId, + /// Cross-partition role. + XpRole, + /// Number of read regions. + NumberOfReadRegions, /// Item logical sequence number. ItemLsn, /// Global committed logical sequence number. GlobalCommittedLsn, + /// Local logical sequence number. + LocalLsn, + /// Quorum-acknowledged local logical sequence number. + QuorumAckedLocalLsn, + /// Item local logical sequence number. + ItemLocalLsn, /// Transport request identifier. TransportRequestId, /// Session token. SessionToken, + /// Query execution metadata. + QueryExecutionInfo, + /// Backend request processing duration in milliseconds. + BackendRequestDurationMilliseconds, + /// Whether partition-key deletion is pending. + PendingPkDelete, + /// Physical partition identifier. + PhysicalPartitionId, + /// Conflict resolution timestamp. + ConflictResolvedTimestamp, } impl TryFrom for RntbdResponseToken { @@ -679,18 +734,37 @@ impl TryFrom for RntbdResponseToken { fn try_from(value: u16) -> Result { match value { 0x0000 => Ok(Self::PayloadPresent), + 0x0002 => Ok(Self::LastStateChangeDateTime), 0x0003 => Ok(Self::ContinuationToken), 0x0004 => Ok(Self::ETag), 0x000C => Ok(Self::RetryAfterMilliseconds), + 0x000E => Ok(Self::StorageMaxResourceQuota), + 0x000F => Ok(Self::StorageResourceQuotaUsage), + 0x0010 => Ok(Self::SchemaVersion), 0x0013 => Ok(Self::Lsn), + 0x0014 => Ok(Self::ItemCount), 0x0015 => Ok(Self::RequestCharge), 0x0017 => Ok(Self::OwnerFullName), + 0x0018 => Ok(Self::OwnerId), + 0x001A => Ok(Self::QuorumAckedLsn), 0x001C => Ok(Self::SubStatus), + 0x001E => Ok(Self::CurrentWriteQuorum), + 0x001F => Ok(Self::CurrentReplicaSetSize), 0x0021 => Ok(Self::PartitionKeyRangeId), + 0x0026 => Ok(Self::XpRole), 0x0032 => Ok(Self::ItemLsn), 0x0029 => Ok(Self::GlobalCommittedLsn), + 0x0030 => Ok(Self::NumberOfReadRegions), 0x0035 => Ok(Self::TransportRequestId), + 0x003A => Ok(Self::LocalLsn), + 0x003B => Ok(Self::QuorumAckedLocalLsn), + 0x003C => Ok(Self::ItemLocalLsn), 0x003E => Ok(Self::SessionToken), + 0x0045 => Ok(Self::QueryExecutionInfo), + 0x0051 => Ok(Self::BackendRequestDurationMilliseconds), + 0x0055 => Ok(Self::PendingPkDelete), + 0x0063 => Ok(Self::PhysicalPartitionId), + 0x0087 => Ok(Self::ConflictResolvedTimestamp), _ => Err(()), } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs index 34e5235e18a..afd5c92661d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs @@ -179,6 +179,16 @@ pub(crate) mod response_header_names { pub const SERVICE_VERSION: &str = "x-ms-serviceversion"; pub const RESOURCE_QUOTA: &str = "x-ms-resource-quota"; pub const RESOURCE_USAGE: &str = "x-ms-resource-usage"; + pub const SCHEMA_VERSION: &str = "x-ms-schemaversion"; + pub const CURRENT_WRITE_QUORUM: &str = "x-ms-current-write-quorum"; + pub const CURRENT_REPLICA_SET_SIZE: &str = "x-ms-current-replica-set-size"; + pub const XP_ROLE: &str = "x-ms-xp-role"; + pub const QUERY_EXECUTION_INFO: &str = "x-ms-cosmos-query-execution-info"; + pub const PENDING_PK_DELETE: &str = "x-ms-cosmos-is-partition-key-delete-pending"; + pub const PHYSICAL_PARTITION_ID: &str = "x-ms-cosmos-physical-partition-id"; + /// Synthetic header matching .NET Direct's `WFConstants.BackendHeaders` + /// name for the RNTBD-only conflict resolution timestamp. + pub const CONFLICT_RESOLVED_TIMESTAMP: &str = "x-ms-cosmos-conflict-resolved-timestamp"; pub const HAS_TENTATIVE_WRITES: &str = "x-ms-cosmos-allow-tentative-writes"; pub const PARTITION_KEY_RANGE_ID: &str = "x-ms-documentdb-partitionkeyrangeid"; pub const INTERNAL_PARTITION_ID: &str = "x-ms-cosmos-internal-partition-id"; @@ -559,6 +569,30 @@ pub struct CosmosResponseHeaders { /// Resource usage information (`x-ms-resource-usage`). pub resource_usage: Option, + /// Resource schema version (`x-ms-schemaversion`). + pub(crate) schema_version: Option, + + /// Current write quorum (`x-ms-current-write-quorum`). + pub(crate) current_write_quorum: Option, + + /// Current replica set size (`x-ms-current-replica-set-size`). + pub(crate) current_replica_set_size: Option, + + /// Cross-partition role (`x-ms-xp-role`). + pub(crate) xp_role: Option, + + /// Query execution metadata (`x-ms-cosmos-query-execution-info`). + pub(crate) query_execution_info: Option, + + /// Whether partition-key deletion is pending (`x-ms-cosmos-is-partition-key-delete-pending`). + pub(crate) pending_pk_delete: Option, + + /// Physical partition identifier (`x-ms-cosmos-physical-partition-id`). + pub(crate) physical_partition_id: Option, + + /// Conflict resolution timestamp (`x-ms-cosmos-conflict-resolved-timestamp`). + pub(crate) conflict_resolved_timestamp: Option, + /// Whether the region has tentative (not yet committed) writes (`x-ms-cosmos-allow-tentative-writes`). pub has_tentative_writes: Option, @@ -720,6 +754,30 @@ impl CosmosResponseHeaders { response_header_names::RESOURCE_USAGE => { result.resource_usage = Some(value.as_str().to_owned()); } + response_header_names::SCHEMA_VERSION => { + result.schema_version = Some(value.as_str().to_owned()); + } + response_header_names::CURRENT_WRITE_QUORUM => { + result.current_write_quorum = value.as_str().parse().ok(); + } + response_header_names::CURRENT_REPLICA_SET_SIZE => { + result.current_replica_set_size = value.as_str().parse().ok(); + } + response_header_names::XP_ROLE => { + result.xp_role = value.as_str().parse().ok(); + } + response_header_names::QUERY_EXECUTION_INFO => { + result.query_execution_info = Some(value.as_str().to_owned()); + } + response_header_names::PENDING_PK_DELETE => { + result.pending_pk_delete = parse_bool_ci(value.as_str()); + } + response_header_names::PHYSICAL_PARTITION_ID => { + result.physical_partition_id = Some(value.as_str().to_owned()); + } + response_header_names::CONFLICT_RESOLVED_TIMESTAMP => { + result.conflict_resolved_timestamp = value.as_str().parse().ok(); + } response_header_names::HAS_TENTATIVE_WRITES => { result.has_tentative_writes = parse_bool_ci(value.as_str()); } @@ -898,6 +956,38 @@ impl CosmosResponseHeaders { response_header_names::RESOURCE_USAGE, self.resource_usage.clone(), ); + put_str( + response_header_names::SCHEMA_VERSION, + self.schema_version.clone(), + ); + put_str( + response_header_names::CURRENT_WRITE_QUORUM, + self.current_write_quorum.map(|v| v.to_string()), + ); + put_str( + response_header_names::CURRENT_REPLICA_SET_SIZE, + self.current_replica_set_size.map(|v| v.to_string()), + ); + put_str( + response_header_names::XP_ROLE, + self.xp_role.map(|v| v.to_string()), + ); + put_str( + response_header_names::QUERY_EXECUTION_INFO, + self.query_execution_info.clone(), + ); + put_str( + response_header_names::PENDING_PK_DELETE, + self.pending_pk_delete.map(|b| bool_to_wire(b).to_owned()), + ); + put_str( + response_header_names::PHYSICAL_PARTITION_ID, + self.physical_partition_id.clone(), + ); + put_str( + response_header_names::CONFLICT_RESOLVED_TIMESTAMP, + self.conflict_resolved_timestamp.map(|v| v.to_string()), + ); put_str( response_header_names::HAS_TENTATIVE_WRITES, self.has_tentative_writes @@ -1091,6 +1181,16 @@ mod tests { } } + #[test] + fn parses_canonical_schema_version_header() { + let mut headers = Headers::new(); + headers.insert("x-ms-schemaversion", "1.0"); + + let cosmos_headers = CosmosResponseHeaders::from_headers(&headers); + + assert_eq!(cosmos_headers.schema_version.as_deref(), Some("1.0")); + } + #[test] fn invalid_base64_index_metrics_returns_none() { let mut headers = Headers::new(); @@ -1401,6 +1501,14 @@ mod tests { service_version: Some("version 2.18.0".into()), resource_quota: Some("documentSize=10240;".into()), resource_usage: Some("documentSize=0;".into()), + schema_version: Some("1.0".into()), + current_write_quorum: Some(3), + current_replica_set_size: Some(4), + xp_role: Some(1), + query_execution_info: Some("{\"reverseRidEnabled\":false}".into()), + pending_pk_delete: Some(false), + physical_partition_id: Some("physical-0".into()), + conflict_resolved_timestamp: Some(1_234_567), has_tentative_writes: Some(false), partition_key_range_id: Some("0".into()), internal_partition_id: Some("internal-xyz".into()), @@ -1426,6 +1534,12 @@ mod tests { )), Some("False") ); + assert_eq!( + raw.get_optional_str(&HeaderName::from_static( + response_header_names::PENDING_PK_DELETE + )), + Some("False") + ); // Sub-status is emitted as the bare numeric value. assert_eq!( raw.get_optional_str(&HeaderName::from_static(response_header_names::SUBSTATUS)), @@ -1510,6 +1624,29 @@ mod tests { assert_eq!(round_tripped.service_version, original.service_version); assert_eq!(round_tripped.resource_quota, original.resource_quota); assert_eq!(round_tripped.resource_usage, original.resource_usage); + assert_eq!(round_tripped.schema_version, original.schema_version); + assert_eq!( + round_tripped.current_write_quorum, + original.current_write_quorum + ); + assert_eq!( + round_tripped.current_replica_set_size, + original.current_replica_set_size + ); + assert_eq!(round_tripped.xp_role, original.xp_role); + assert_eq!( + round_tripped.query_execution_info, + original.query_execution_info + ); + assert_eq!(round_tripped.pending_pk_delete, original.pending_pk_delete); + assert_eq!( + round_tripped.physical_partition_id, + original.physical_partition_id + ); + assert_eq!( + round_tripped.conflict_resolved_timestamp, + original.conflict_resolved_timestamp + ); assert_eq!( round_tripped.has_tentative_writes, original.has_tentative_writes diff --git a/sdk/cosmos/azure_data_cosmos_perf/README.md b/sdk/cosmos/azure_data_cosmos_perf/README.md index 88d9d1bd021..17822ee92c6 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/README.md +++ b/sdk/cosmos/azure_data_cosmos_perf/README.md @@ -1,8 +1,8 @@ # Azure Cosmos DB Performance Testing Tool A CLI tool for performance and scale testing the Azure Cosmos DB Rust SDK. It runs -point reads, single-partition queries, upserts, and creates concurrently and reports -latency statistics at configurable intervals. +point reads, single-partition queries, bounded change feed reads, upserts, and creates +concurrently and reports latency statistics at configurable intervals. ## Prerequisites @@ -64,12 +64,15 @@ cargo run -p azure_data_cosmos_perf -- \ | `--application-region` | *required* | Azure region where the application is running (e.g., `"East US 2"`) | | `--excluded-regions` | — | Comma-separated excluded regions | | `--exclude-regions-for` | `both` | Scope for excluded regions: `reads`, `writes`, or `both` | -| `--concurrency` | `50` | Number of concurrent operations | +| `--concurrency` | `50` | Initial seeding concurrency and persistent worker count in closed-loop mode; only the benchmark worker count is ignored when `--target-rate` is set | +| `--target-rate` | — | Target arrival rate in operations per second; enables open-loop mode when set | +| `--max-in-flight` | `100000` | Maximum in-flight requests in open-loop mode; excess issuances are skipped | | `--duration` | indefinite | Run duration in seconds | | `--seed-count` | `1000` | Number of items to pre-seed | | `--throughput` | `100000` | Throughput (RU/s) when creating the container | | `--default-ttl` | `3600` | Default TTL in seconds for items (0 to disable) | | `--report-interval` | `300` | Stats reporting interval in seconds | +| `--diagnostics-threshold-ms` | — | Log detailed diagnostics for successful operations slower than this threshold; disabled when omitted | | `--results-container` | `perfresults` | Container for storing perf results and error documents | | `--results-endpoint` | — | Cosmos DB endpoint for results (omit to use same account as `--endpoint`) | | `--results-database` | `perfdb` | Database name on the results account | @@ -82,6 +85,11 @@ cargo run -p azure_data_cosmos_perf -- \ | `--no-queries` | `false` | Disable query operations | | `--no-upserts` | `false` | Disable upsert operations | | `--no-creates` | `false` | Disable create operations | +| `--no-change-feed` | `false` | Disable per-feed-range change feed operations | +| `--change-feed-max-pages` | `4` | Maximum pages consumed by each change feed operation | +| `--no-feed-range-queries` | `false` | Disable per-feed-range query operations | +| `--feed-range-query-max-pages` | `4` | Maximum pages consumed by each per-feed-range query | +| `--feed-range-refresh-secs` | `60` | Interval between feed-range cache refreshes (0 disables refresh) | ### Examples @@ -92,6 +100,7 @@ cargo run -p azure_data_cosmos_perf -- \ --endpoint https://myaccount.documents.azure.com:443/ \ --auth key --key "$AZURE_COSMOS_KEY" \ --no-queries --no-upserts --no-creates \ + --no-change-feed --no-feed-range-queries \ --concurrency 100 --duration 60 --report-interval 10 ``` @@ -167,6 +176,14 @@ unique IDs and partition keys. Successfully created items are added to the shared item pool so they become targets for subsequent read, query, and upsert operations. +### Change Feed Operation + +When enabled (the default), the `ChangeFeed` operation rotates round-robin +through the container's physical feed ranges and reads each range's +latest-version change feed from the beginning. Because change feed streams do +not terminate when caught up, each execution consumes at most +`--change-feed-max-pages` pages (default `4`). + ### Multi-Process Launcher The `run_perf.sh` script launches multiple OS processes of the perf tool in diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/config.rs b/sdk/cosmos/azure_data_cosmos_perf/src/config.rs index b797a23b129..6cabb3817b6 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/config.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/config.rs @@ -63,14 +63,32 @@ pub struct Config { #[arg(long, default_value_t = false)] pub no_creates: bool, + /// Disable the per-feed-range change feed operation. + #[arg(long, default_value_t = false)] + pub no_change_feed: bool, + + /// Maximum number of pages consumed by each change feed operation. + /// + /// Change feed streams are infinite, so every operation must stop after + /// a bounded number of page polls. + #[arg(long, default_value_t = 4)] + pub change_feed_max_pages: usize, + /// Disable the per-feed-range query operation. /// /// When enabled (the default), the harness fans out - /// `SELECT VALUE COUNT(1) FROM c` across the container's physical - /// partitions, round-robin one range per call. + /// `SELECT * FROM c` across the container's physical partitions, + /// round-robin one range per call. #[arg(long, default_value_t = false)] pub no_feed_range_queries: bool, + /// Maximum number of pages consumed by each per-feed-range query. + /// + /// Bounds the work performed by one query as the container grows while + /// still exercising continuation-token handling across multiple pages. + #[arg(long, default_value_t = 4)] + pub feed_range_query_max_pages: usize, + /// Interval in seconds between background refreshes of the cached feed /// range list. Set to 0 to disable the background refresher (the cache /// then stays at the seed snapshot taken once at startup). @@ -82,9 +100,36 @@ pub struct Config { pub feed_range_refresh_secs: u64, /// Number of concurrent operations (minimum: 1). + /// + /// Drives the default closed-loop mode: the harness spawns this many + /// persistent worker tasks, each serially issuing one operation at a + /// time (so this is also the maximum number of in-flight requests). + /// When `--target-rate` switches the benchmark to open-loop issuance, + /// this value still controls initial seeding concurrency but does not + /// control the benchmark worker count. #[arg(long, default_value_t = 50)] pub concurrency: usize, + /// Target arrival rate in operations per second (open-loop mode). + /// + /// When set, requests are issued at this fixed rate regardless of how + /// fast prior requests complete, reducing the coordinated-omission + /// bias of the closed-loop `--concurrency` model. In-flight requests + /// are bounded by `--max-in-flight`; once that bound is hit, excess + /// issuances are skipped and counted (not buffered) so a slow backend + /// surfaces as visible omission rather than unbounded memory growth. + /// Omit to use the default closed-loop `--concurrency` model. + #[arg(long)] + pub target_rate: Option, + + /// Maximum number of in-flight requests in open-loop mode (minimum: 1). + /// + /// Only used when `--target-rate` is set. Bounds memory/socket growth + /// when the backend can't keep up with the target rate: issuances that + /// would exceed this cap are skipped and counted instead of queued. + #[arg(long, default_value_t = 100_000)] + pub max_in_flight: usize, + /// Run duration in seconds. Omit for indefinite. #[arg(long)] pub duration: Option, @@ -97,6 +142,12 @@ pub struct Config { #[arg(long, default_value_t = 300)] pub report_interval: u64, + /// Log detailed diagnostics for successful operations slower than this threshold in milliseconds. + /// + /// Omit to disable diagnostics logging for successful operations. + #[arg(long)] + pub diagnostics_threshold_ms: Option, + /// Throughput (RU/s) to provision when creating the container. #[arg(long, default_value_t = 100000)] pub throughput: usize, diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/main.rs b/sdk/cosmos/azure_data_cosmos_perf/src/main.rs index 2173f66eb86..0dcfd327463 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/main.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/main.rs @@ -34,6 +34,19 @@ fn create_aad_credential( .map_err(|e| e.into()) } +fn validate_max_in_flight(max_in_flight: usize) -> Result<(), String> { + if max_in_flight == 0 { + return Err("--max-in-flight must be at least 1 when --target-rate is set.".to_string()); + } + if max_in_flight > tokio::sync::Semaphore::MAX_PERMITS { + return Err(format!( + "--max-in-flight cannot exceed {}.", + tokio::sync::Semaphore::MAX_PERMITS + )); + } + Ok(()) +} + #[tokio::main] async fn main() -> Result<(), Box> { // Initialize tokio-console subscriber when the feature is enabled. @@ -103,6 +116,7 @@ async fn main() -> Result<(), Box> { use azure_data_cosmos::{ AccountEndpoint, AccountReference, CosmosClientBuilder, RoutingStrategy, }; + use azure_data_cosmos_driver::options::ConnectionPoolOptionsBuilder; use clap::Parser; use crate::config::{AuthMethod, Config}; @@ -117,6 +131,7 @@ async fn main() -> Result<(), Box> { && config.no_queries && config.no_upserts && config.no_creates + && config.no_change_feed && config.no_feed_range_queries { eprintln!("Error: all operations are disabled. Enable at least one."); @@ -130,10 +145,28 @@ async fn main() -> Result<(), Box> { eprintln!("Error: --concurrency cannot exceed {}.", u32::MAX); std::process::exit(1); } + if let Some(rate) = config.target_rate { + if rate == 0 { + eprintln!("Error: --target-rate must be at least 1."); + std::process::exit(1); + } + if let Err(message) = validate_max_in_flight(config.max_in_flight) { + eprintln!("Error: {message}"); + std::process::exit(1); + } + } if config.seed_count == 0 { eprintln!("Error: --seed-count must be at least 1."); std::process::exit(1); } + if config.feed_range_query_max_pages == 0 { + eprintln!("Error: --feed-range-query-max-pages must be at least 1."); + std::process::exit(1); + } + if config.change_feed_max_pages == 0 { + eprintln!("Error: --change-feed-max-pages must be at least 1."); + std::process::exit(1); + } // Resolve the (optional) user-agent suffix. Composes // `{configured}-{first-8-chars-of-workload_id}` so a single perf @@ -316,7 +349,17 @@ async fn main() -> Result<(), Box> { // Build config snapshot for Grafana dashboard visibility let config_snapshot = ConfigSnapshot { - concurrency: config.concurrency as u64, + concurrency: config + .target_rate + .is_none() + .then_some(config.concurrency as u64), + target_rate: config.target_rate, + max_in_flight: config + .target_rate + .is_some() + .then_some(config.max_in_flight as u64), + skipped_issuances: None, + skipped_postprocessing: None, application_region: config.application_region.clone(), excluded_regions: config.excluded_regions.join(", "), tokio_threads: tokio::runtime::Handle::current().metrics().num_workers() as u64, @@ -324,6 +367,10 @@ async fn main() -> Result<(), Box> { .ok() .and_then(|v| v.parse::().ok()) .unwrap_or(true), + gateway_v2_disabled: ConnectionPoolOptionsBuilder::new() + .build()? + .gateway_v2_disabled(), + diagnostics_threshold_ms: config.diagnostics_threshold_ms, pyroscope_enabled: std::env::var("PYROSCOPE_SERVER_URL") .map(|v| !v.is_empty()) .unwrap_or(false), @@ -346,6 +393,9 @@ async fn main() -> Result<(), Box> { feed_range_refresher, stats, concurrency: config.concurrency, + target_rate: config.target_rate, + max_in_flight: config.max_in_flight, + diagnostics_threshold: config.diagnostics_threshold_ms.map(Duration::from_millis), duration, report_interval: Duration::from_secs(config.report_interval), results_container, @@ -358,3 +408,20 @@ async fn main() -> Result<(), Box> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::validate_max_in_flight; + use tokio::sync::Semaphore; + + #[test] + fn max_in_flight_accepts_semaphore_limit() { + assert!(validate_max_in_flight(Semaphore::MAX_PERMITS).is_ok()); + } + + #[test] + fn max_in_flight_rejects_zero_and_values_above_semaphore_limit() { + assert!(validate_max_in_flight(0).is_err()); + assert!(validate_max_in_flight(Semaphore::MAX_PERMITS + 1).is_err()); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/change_feed.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/change_feed.rs new file mode 100644 index 00000000000..9dce964c8e2 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/change_feed.rs @@ -0,0 +1,109 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Bounded per-feed-range change feed operation. + +use std::{ + sync::atomic::{AtomicUsize, Ordering}, + sync::Arc, + time::Duration, +}; + +use async_trait::async_trait; +use azure_data_cosmos::clients::ContainerClient; +use azure_data_cosmos::feed::FeedScope; +use azure_data_cosmos::options::ChangeFeedStartFrom; +use azure_data_cosmos::CosmosStatus; +use azure_data_cosmos_driver::error::CosmosError as DriverCosmosError; +use futures::StreamExt; + +use super::{extract_backend_duration, FeedRangeCache, Operation, OperationResult, PerfItem}; + +/// Reads a bounded number of pages from one physical feed range. +pub struct ChangeFeedOperation { + cache: FeedRangeCache, + cursor: AtomicUsize, + max_pages: usize, +} + +impl ChangeFeedOperation { + pub fn new(cache: FeedRangeCache, max_pages: usize) -> Self { + Self { + cache, + cursor: AtomicUsize::new(0), + max_pages, + } + } +} + +#[async_trait] +impl Operation for ChangeFeedOperation { + fn name(&self) -> &'static str { + "ChangeFeed" + } + + async fn execute( + &self, + container: &ContainerClient, + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { + let snapshot = { + let guard = self.cache.read().expect("feed-range cache lock poisoned"); + Arc::clone(&guard) + }; + if snapshot.is_empty() { + return Err(DriverCosmosError::builder() + .with_status(CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) + .with_message("feed-range cache is empty") + .build() + .into()); + } + + let idx = self.cursor.fetch_add(1, Ordering::Relaxed) % snapshot.len(); + let range = snapshot[idx].clone(); + let iterator = Box::pin(container.query_change_feed::( + FeedScope::range(range), + ChangeFeedStartFrom::Beginning, + None, + )) + .await?; + let mut stream = Box::pin(iterator.take(self.max_pages)); + + let mut backend_total: Option = None; + let mut diagnostics = Vec::new(); + while let Some(result) = stream.next().await { + let page = result?; + if let Some(duration) = extract_backend_duration(page.headers()) { + backend_total = Some(backend_total.unwrap_or_default() + duration); + } + if capture_diagnostics { + diagnostics.push(page.diagnostics()); + } + } + + Ok(OperationResult::paged(backend_total, diagnostics)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::RwLock; + + use azure_data_cosmos::feed::FeedRange; + + use super::*; + + fn cache_with_len(len: usize) -> FeedRangeCache { + let ranges: Vec = (0..len).map(|_| FeedRange::full()).collect(); + Arc::new(RwLock::new(Arc::new(ranges))) + } + + #[test] + fn cursor_round_robins_across_ranges() { + let operation = ChangeFeedOperation::new(cache_with_len(3), 4); + let picks: Vec = (0..6) + .map(|_| operation.cursor.fetch_add(1, Ordering::Relaxed) % 3) + .collect(); + assert_eq!(picks, vec![0, 1, 2, 0, 1, 2]); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/create_item.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/create_item.rs index 19d9635ab7c..08ab328ea4a 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/create_item.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/create_item.rs @@ -11,7 +11,7 @@ use azure_data_cosmos::options::ItemWriteOptions; use rand::RngExt; use uuid::Uuid; -use super::{extract_backend_duration, Operation, PerfItem}; +use super::{extract_backend_duration, Operation, OperationResult, PerfItem}; use crate::seed::{SeededItem, SharedItems}; /// Creates a new item with a unique ID and partition key. @@ -39,7 +39,8 @@ impl Operation for CreateItemOperation { async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result> { + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { let id = Uuid::new_v4().to_string(); let partition_key = Uuid::new_v4().to_string(); let value = rand::rng().random_range(0..u64::MAX); @@ -55,8 +56,9 @@ impl Operation for CreateItemOperation { .create_item(&item.partition_key, &id, &item, self.options.clone()) .await?; let backend = extract_backend_duration(response.headers()); + let diagnostics = capture_diagnostics.then(|| response.diagnostics()); self.items.push(SeededItem { id, partition_key }); - Ok(backend) + Ok(OperationResult::single(backend, diagnostics)) } } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_query.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_query.rs index d131a9c0d92..2502a012eb2 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_query.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_query.rs @@ -6,8 +6,8 @@ //! Each `execute()` runs `SELECT * FROM c` against ONE `FeedRange`, //! round-robin'd from a cache shared with //! [`FeedRangeRefresher`](super::feed_range_refresher::FeedRangeRefresher), -//! and drains every page so the harness exercises continuation-token -//! handling across multi-page responses per feed range. +//! and drains a bounded number of pages so the harness exercises +//! continuation-token handling without query work growing with the container. //! The harness's existing worker pool provides concurrency — N workers //! each issue one query at a time, naturally matching the concurrency //! of every other operation. @@ -22,7 +22,7 @@ use azure_data_cosmos::{feed::FeedRange, CosmosStatus, FeedScope, Query}; use azure_data_cosmos_driver::error::CosmosError as DriverCosmosError; use futures::StreamExt; -use super::{extract_backend_duration, Operation}; +use super::{extract_backend_duration, Operation, OperationResult}; /// Shared, swappable feed-range cache. /// @@ -34,13 +34,15 @@ pub type FeedRangeCache = Arc>>>; pub struct FeedRangeQueryOperation { cache: FeedRangeCache, cursor: AtomicUsize, + max_pages: usize, } impl FeedRangeQueryOperation { - pub fn new(cache: FeedRangeCache) -> Self { + pub fn new(cache: FeedRangeCache, max_pages: usize) -> Self { Self { cache, cursor: AtomicUsize::new(0), + max_pages, } } } @@ -54,7 +56,8 @@ impl Operation for FeedRangeQueryOperation { async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result> { + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { // Snapshot the current ranges; release the lock immediately. let snapshot = { let guard = self.cache.read().expect("feed-range cache lock poisoned"); @@ -80,20 +83,25 @@ impl Operation for FeedRangeQueryOperation { let mut stream = Box::pin(container.query_items::(query, FeedScope::range(fr), None)) .await? - .into_pages(); + .into_pages() + .take(self.max_pages); // Sum backend durations across pages so a multi-page response // reports the total server processing time, matching the // client-observed elapsed which wraps the entire stream drain. let mut backend_total: Option = None; + let mut diagnostics = Vec::new(); while let Some(result) = stream.next().await { let page = result?; if let Some(d) = extract_backend_duration(page.headers()) { backend_total = Some(backend_total.unwrap_or_default() + d); } + if capture_diagnostics { + diagnostics.push(page.diagnostics()); + } } - Ok(backend_total) + Ok(OperationResult::paged(backend_total, diagnostics)) } } @@ -108,7 +116,7 @@ mod tests { #[test] fn cursor_round_robins_and_wraps() { - let op = FeedRangeQueryOperation::new(cache_with_len(3)); + let op = FeedRangeQueryOperation::new(cache_with_len(3), 4); // 6 fetches over a 3-element cache should visit each index twice // (0,1,2,0,1,2 modulo 3). let picks: Vec = (0..6) @@ -122,7 +130,7 @@ mod tests { // Construct against a zero-length cache and confirm snapshot // is empty (execute() would surface this as an error instead // of dividing by zero). - let op = FeedRangeQueryOperation::new(cache_with_len(0)); + let op = FeedRangeQueryOperation::new(cache_with_len(0), 4); let guard = op.cache.read().unwrap(); assert!(guard.is_empty()); } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_refresher.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_refresher.rs index 25e7a826194..9a00b49a4eb 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_refresher.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_refresher.rs @@ -20,6 +20,7 @@ use std::time::{Duration, Instant}; use azure_data_cosmos::clients::ContainerClient; use azure_data_cosmos::feed::FeedRange; +use azure_data_cosmos::options::ReadFeedRangesOptions; use crate::operations::feed_range_query::FeedRangeCache; use crate::stats::Stats; @@ -72,7 +73,11 @@ impl FeedRangeRefresher { let container = Arc::clone(&container); refresh_once(&stats, &cache, || async move { - container.read_feed_ranges(None).await + container + .read_feed_ranges(Some( + ReadFeedRangesOptions::default().with_force_refresh(true), + )) + .await }) .await; } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/mod.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/mod.rs index 3e2fc03a509..de64260aae4 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/mod.rs @@ -8,6 +8,7 @@ //! 2. Register it in [`create_operations`]. //! 3. Add a CLI flag in `config.rs` to enable/disable it. +mod change_feed; mod create_item; mod feed_range_query; mod feed_range_refresher; @@ -17,6 +18,7 @@ mod upsert_item; use async_trait::async_trait; use azure_data_cosmos::clients::ContainerClient; +use azure_data_cosmos::diagnostics::DiagnosticsContext; use azure_data_cosmos::models::ResponseHeaders; use azure_data_cosmos::options::Region; use azure_data_cosmos::options::{ @@ -27,6 +29,7 @@ use std::sync::{Arc, RwLock}; use std::time::Duration; use crate::config::{Config, ExcludeRegionsScope}; +pub use crate::operations::change_feed::ChangeFeedOperation; pub use crate::operations::create_item::CreateItemOperation; pub use crate::operations::feed_range_query::{FeedRangeCache, FeedRangeQueryOperation}; pub use crate::operations::feed_range_refresher::{FeedRangeRefresher, READ_FEED_RANGES_STAT}; @@ -57,16 +60,43 @@ pub trait Operation: Send + Sync { /// Executes one instance of the operation. /// - /// Returns `Ok(Some(d))` when the server reported a processing duration - /// via the `x-ms-request-duration-ms` response header (this is the - /// backend latency surfaced separately from the client-observed - /// wall-clock latency). Returns `Ok(None)` when no backend duration - /// could be observed (multi-page query streams may aggregate, see - /// individual implementations). + /// Returns backend latency and, when `capture_diagnostics` is true, + /// finalized diagnostics for the response or query pages. async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result>; + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result; +} + +/// Successful operation data consumed by the runner. +pub struct OperationResult { + /// Aggregated server-reported processing duration, when available. + pub backend_duration: Option, + /// Finalized diagnostics captured from successful responses. + pub diagnostics: Vec>, +} + +impl OperationResult { + pub fn single( + backend_duration: Option, + diagnostics: Option>, + ) -> Self { + Self { + backend_duration, + diagnostics: diagnostics.into_iter().collect(), + } + } + + pub fn paged( + backend_duration: Option, + diagnostics: Vec>, + ) -> Self { + Self { + backend_duration, + diagnostics, + } + } } /// The item type used for seeding, reading, querying, and upserting. @@ -118,12 +148,11 @@ pub async fn create_operations( } if !config.no_creates { ops.push(Arc::new(CreateItemOperation::new( - seeded_items, + seeded_items.clone(), write_options, ))); } - - let feed_range_refresher = if config.no_feed_range_queries { + let feed_range_refresher = if config.no_feed_range_queries && config.no_change_feed { None } else { let initial = container.read_feed_ranges(None).await?; @@ -135,7 +164,18 @@ pub async fn create_operations( .into()); } let cache: FeedRangeCache = Arc::new(RwLock::new(Arc::new(initial))); - ops.push(Arc::new(FeedRangeQueryOperation::new(cache.clone()))); + if !config.no_feed_range_queries { + ops.push(Arc::new(FeedRangeQueryOperation::new( + cache.clone(), + config.feed_range_query_max_pages, + ))); + } + if !config.no_change_feed { + ops.push(Arc::new(ChangeFeedOperation::new( + cache.clone(), + config.change_feed_max_pages, + ))); + } if config.feed_range_refresh_secs > 0 { Some(FeedRangeRefresher::new( diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/query_items.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/query_items.rs index 8cb7efe700e..cb39000f662 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/query_items.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/query_items.rs @@ -11,7 +11,7 @@ use azure_data_cosmos::Query; use azure_data_cosmos::{clients::ContainerClient, feed::FeedScope}; use futures::StreamExt; -use super::{extract_backend_duration, Operation}; +use super::{extract_backend_duration, Operation, OperationResult}; use crate::seed::SharedItems; /// Runs a single-partition query against a random seeded partition key. @@ -35,7 +35,8 @@ impl Operation for QueryItemsOperation { async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result> { + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { let item = self.items.random(); let pk = &item.partition_key; @@ -54,13 +55,17 @@ impl Operation for QueryItemsOperation { // the total server processing time, mirroring how the client-observed // elapsed wraps the entire stream consumption. let mut backend_total: Option = None; + let mut diagnostics = Vec::new(); while let Some(result) = stream.next().await { let page = result?; if let Some(d) = extract_backend_duration(page.headers()) { backend_total = Some(backend_total.unwrap_or_default() + d); } + if capture_diagnostics { + diagnostics.push(page.diagnostics()); + } } - Ok(backend_total) + Ok(OperationResult::paged(backend_total, diagnostics)) } } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/read_item.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/read_item.rs index f5e5ccf8a1a..59c97f39204 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/read_item.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/read_item.rs @@ -9,7 +9,7 @@ use async_trait::async_trait; use azure_data_cosmos::clients::ContainerClient; use azure_data_cosmos::options::ItemReadOptions; -use super::{extract_backend_duration, Operation}; +use super::{extract_backend_duration, Operation, OperationResult}; use crate::seed::SharedItems; /// Reads a random seeded item by ID and partition key. @@ -34,12 +34,16 @@ impl Operation for ReadItemOperation { async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result> { + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { let item = self.items.random(); let response = container .read_item(&item.partition_key, &item.id, self.options.clone()) .await?; - Ok(extract_backend_duration(response.headers())) + Ok(OperationResult::single( + extract_backend_duration(response.headers()), + capture_diagnostics.then(|| response.diagnostics()), + )) } } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/operations/upsert_item.rs b/sdk/cosmos/azure_data_cosmos_perf/src/operations/upsert_item.rs index 1e71b78f5c0..04bbbf62a73 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/operations/upsert_item.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/operations/upsert_item.rs @@ -10,7 +10,7 @@ use azure_data_cosmos::clients::ContainerClient; use azure_data_cosmos::options::ItemWriteOptions; use rand::RngExt; -use super::{extract_backend_duration, Operation, PerfItem}; +use super::{extract_backend_duration, Operation, OperationResult, PerfItem}; use crate::seed::SharedItems; /// Upserts an item into a random seeded partition. @@ -35,7 +35,8 @@ impl Operation for UpsertItemOperation { async fn execute( &self, container: &ContainerClient, - ) -> azure_data_cosmos::Result> { + capture_diagnostics: bool, + ) -> azure_data_cosmos::Result { let seeded = self.items.random(); let value = rand::rng().random_range(0..u64::MAX); @@ -49,6 +50,9 @@ impl Operation for UpsertItemOperation { let response = container .upsert_item(&item.partition_key, &seeded.id, &item, self.options.clone()) .await?; - Ok(extract_backend_duration(response.headers())) + Ok(OperationResult::single( + extract_backend_duration(response.headers()), + capture_diagnostics.then(|| response.diagnostics()), + )) } } diff --git a/sdk/cosmos/azure_data_cosmos_perf/src/runner.rs b/sdk/cosmos/azure_data_cosmos_perf/src/runner.rs index d193d9af0e4..e1d9fa23dea 100644 --- a/sdk/cosmos/azure_data_cosmos_perf/src/runner.rs +++ b/sdk/cosmos/azure_data_cosmos_perf/src/runner.rs @@ -3,7 +3,7 @@ //! Concurrent operation execution engine. -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -13,7 +13,9 @@ use azure_data_cosmos_driver::DiagnosticsVerbosity; use rand::RngExt; use serde::Serialize; use sysinfo::System; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::task::JoinSet; +use tokio::time::MissedTickBehavior; use uuid::Uuid; use crate::operations::{FeedRangeRefresher, Operation}; @@ -88,6 +90,14 @@ struct PerfResult { #[serde(skip_serializing_if = "Option::is_none")] config_concurrency: Option, #[serde(skip_serializing_if = "Option::is_none")] + config_target_rate: Option, + #[serde(skip_serializing_if = "Option::is_none")] + config_max_in_flight: Option, + #[serde(skip_serializing_if = "Option::is_none")] + config_skipped_issuances: Option, + #[serde(skip_serializing_if = "Option::is_none")] + config_skipped_postprocessing: Option, + #[serde(skip_serializing_if = "Option::is_none")] config_application_region: Option, #[serde(skip_serializing_if = "Option::is_none")] config_excluded_regions: Option, @@ -96,6 +106,10 @@ struct PerfResult { #[serde(skip_serializing_if = "Option::is_none")] config_ppcb_enabled: Option, #[serde(skip_serializing_if = "Option::is_none")] + config_gateway_v2_disabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + config_diagnostics_threshold_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] config_pyroscope_enabled: Option, #[serde(skip_serializing_if = "Option::is_none")] config_tokio_console_enabled: Option, @@ -155,6 +169,16 @@ pub struct RunConfig { pub feed_range_refresher: Option, pub stats: Arc, pub concurrency: usize, + /// Target arrival rate (ops/sec) for open-loop mode. When `Some`, the + /// runner issues requests at this fixed rate instead of using the + /// closed-loop `concurrency` worker pool. + pub target_rate: Option, + /// Maximum in-flight requests in open-loop mode. Ignored when + /// `target_rate` is `None`. + pub max_in_flight: usize, + /// Successful operations exceeding this latency emit detailed diagnostics. + /// `None` disables successful-operation diagnostics logging. + pub diagnostics_threshold: Option, pub duration: Option, pub report_interval: Duration, pub results_container: ContainerClient, @@ -167,23 +191,113 @@ pub struct RunConfig { /// Snapshot of runtime configuration emitted with each PerfResult document. #[derive(Clone)] pub struct ConfigSnapshot { - pub concurrency: u64, + pub concurrency: Option, + pub target_rate: Option, + pub max_in_flight: Option, + pub skipped_issuances: Option>, + pub skipped_postprocessing: Option>, pub application_region: String, pub excluded_regions: String, pub tokio_threads: u64, pub ppcb_enabled: bool, + pub gateway_v2_disabled: bool, + pub diagnostics_threshold_ms: Option, pub pyroscope_enabled: bool, pub tokio_console_enabled: bool, pub tokio_metrics_enabled: bool, pub valgrind_tool: String, } -/// Runs operations concurrently until cancelled or duration expires. +/// Shared metadata and reporting dependencies for one operation execution. +struct OperationRunContext<'a> { + stats: &'a Stats, + results_container: &'a ContainerClient, + workload_id: &'a str, + commit_sha: &'a str, + hostname: &'a str, + diagnostics_threshold: Option, + postprocessing_slots: Option>, + postprocessing_skipped: Option>, +} + +/// Executes a single operation against `container`, recording its latency on +/// success or its error (and an error document) on failure. Shared by both the +/// closed-loop worker pool and the open-loop rate scheduler. +async fn run_one( + op: &Arc, + container: &ContainerClient, + context: OperationRunContext<'_>, + permit: Option, +) { + let op_start = Instant::now(); + let result = op + .execute(container, context.diagnostics_threshold.is_some()) + .await; + drop(permit); + + match result { + Ok(result) => { + let elapsed = op_start.elapsed(); + context + .stats + .record_latency(op.name(), elapsed, result.backend_duration); + if diagnostics_threshold_exceeded(elapsed, context.diagnostics_threshold) { + if let Some(_postprocessing_permit) = + try_acquire_postprocessing_slot(&context.postprocessing_slots) + { + log_slow_operation_diagnostics(op.name(), elapsed, &result.diagnostics); + } else { + record_skipped_postprocessing(&context.postprocessing_skipped); + } + } + } + Err(e) => { + context.stats.record_error(op.name()); + if let Some(_postprocessing_permit) = + try_acquire_postprocessing_slot(&context.postprocessing_slots) + { + upsert_error( + context.results_container, + op.name(), + &e, + context.workload_id, + context.commit_sha, + context.hostname, + ) + .await; + } else { + record_skipped_postprocessing(&context.postprocessing_skipped); + } + } + } +} + +fn record_skipped_postprocessing(skipped: &Option>) { + if let Some(skipped) = skipped { + skipped.fetch_add(1, Ordering::Relaxed); + } +} + +fn try_acquire_postprocessing_slot( + slots: &Option>, +) -> Option> { + match slots { + Some(slots) => Arc::clone(slots).try_acquire_owned().ok().map(Some), + None => Some(None), + } +} + +/// Runs operations until cancelled or duration expires. +/// +/// In the default closed-loop mode, spawns `concurrency` worker tasks, each +/// continuously picking a random operation and executing it serially. When +/// `target_rate` is set, switches to open-loop mode: operations are issued at +/// a fixed arrival rate (bounded by `max_in_flight` concurrent requests), +/// decoupling issuance from completion to reduce coordinated-omission bias. /// -/// Spawns `concurrency` tasks, each continuously picking a random operation -/// from `operations` and executing it against `container`. Latency and errors -/// are recorded in `stats`. A background reporter prints summaries at the -/// given `report_interval` and upserts results into `results_container`. +/// Latency and errors are recorded in `stats`. A background reporter prints +/// summaries at the given `report_interval` and upserts results into +/// `results_container`. pub async fn run(config: RunConfig) { let RunConfig { container, @@ -191,15 +305,22 @@ pub async fn run(config: RunConfig) { feed_range_refresher, stats, concurrency, + target_rate, + max_in_flight, + diagnostics_threshold, duration, report_interval, results_container, workload_id, commit_sha, hostname, - config_snapshot, + mut config_snapshot, } = config; let cancelled = Arc::new(AtomicBool::new(false)); + let skipped = target_rate.map(|_| Arc::new(AtomicU64::new(0))); + let postprocessing_skipped = target_rate.map(|_| Arc::new(AtomicU64::new(0))); + config_snapshot.skipped_issuances = skipped.clone(); + config_snapshot.skipped_postprocessing = postprocessing_skipped.clone(); // Spawn the optional feed-range refresher BEFORE workers so its // first refresh has a head start on warming the cache. @@ -305,49 +426,167 @@ pub async fn run(config: RunConfig) { } }); - // Spawn fixed worker pool — each worker loops until cancelled + // Issue operations until cancelled — open-loop (fixed rate) or + // closed-loop (fixed worker pool), depending on `target_rate`. let start = Instant::now(); - let mut workers = JoinSet::new(); - - for _ in 0..concurrency { - let ops = operations.clone(); - let container = container.clone(); - let stats = stats.clone(); - let cancelled = cancelled.clone(); - let err_container = results_container.clone(); - let err_workload_id = workload_id.clone(); - let err_commit_sha = commit_sha.clone(); - let err_hostname = hostname.clone(); - - workers.spawn(async move { - while !cancelled.load(Ordering::Relaxed) { - let op_idx = rand::rng().random_range(0..ops.len()); - let op = &ops[op_idx]; - - let op_start = Instant::now(); - match op.execute(&container).await { - Ok(backend_duration) => { - stats.record_latency(op.name(), op_start.elapsed(), backend_duration); + + if let Some(rate) = target_rate { + // Open-loop: issue at a fixed arrival rate, decoupled from request + // completion. Borrow the id Strings as `Arc` (clone, not move) + // so the final report below can still use the originals. + let workload_id: Arc = Arc::from(workload_id.as_str()); + let commit_sha: Arc = Arc::from(commit_sha.as_str()); + let hostname: Arc = Arc::from(hostname.as_str()); + let semaphore = Arc::new(Semaphore::new(max_in_flight)); + let postprocessing_slots = Arc::new(Semaphore::new(max_in_flight)); + let postprocessing_skipped = + postprocessing_skipped.expect("open-loop postprocessing counter exists"); + let skipped = skipped.expect("open-loop skipped-issuance counter exists"); + let mut issued_tasks = JoinSet::new(); + + println!("Open-loop mode: target_rate={rate} ops/s, max_in_flight={max_in_flight}"); + + // 1ms ticker; each tick issues up to `elapsed * rate` so issuance + // self-corrects against wall-clock drift. Burst on missed ticks. + let mut ticker = tokio::time::interval(Duration::from_millis(1)); + ticker.set_missed_tick_behavior(MissedTickBehavior::Burst); + let mut issued: u64 = 0; + + while !cancelled.load(Ordering::Relaxed) { + ticker.tick().await; + while let Some(result) = issued_tasks.try_join_next() { + if let Err(error) = result { + eprintln!("Warning: open-loop operation task failed: {error}"); + } + } + let elapsed_secs = start.elapsed().as_secs_f64(); + let target = (elapsed_secs * rate as f64) as u64; + while issued < target { + if cancelled.load(Ordering::Relaxed) { + break; + } + // Count every scheduled issuance against `issued` (even + // skips) so saturation doesn't trigger a catch-up burst once + // the backend recovers. + issued += 1; + match Arc::clone(&semaphore).try_acquire_owned() { + Ok(permit) => { + let op = operations[rand::rng().random_range(0..operations.len())].clone(); + let container = container.clone(); + let stats = stats.clone(); + let err_container = results_container.clone(); + let workload_id = workload_id.clone(); + let commit_sha = commit_sha.clone(); + let hostname = hostname.clone(); + let postprocessing_slots = postprocessing_slots.clone(); + let postprocessing_skipped = postprocessing_skipped.clone(); + issued_tasks.spawn(async move { + run_one( + &op, + &container, + OperationRunContext { + stats: &stats, + results_container: &err_container, + workload_id: &workload_id, + commit_sha: &commit_sha, + hostname: &hostname, + diagnostics_threshold, + postprocessing_slots: Some(postprocessing_slots), + postprocessing_skipped: Some(postprocessing_skipped), + }, + Some(permit), + ) + .await; + }); } - Err(e) => { - stats.record_error(op.name()); - upsert_error( - &err_container, - op.name(), - &e, - &err_workload_id, - &err_commit_sha, - &err_hostname, - ) - .await; + Err(_) => { + let skipped_now = saturated_skip_count(issued, target); + issued = target; + let previous = skipped.fetch_add(skipped_now, Ordering::Relaxed); + let total = previous + skipped_now; + if previous / 10_000 != total / 10_000 { + println!( + "WARN: max_in_flight={max_in_flight} saturated, skipped {total} \ + issuance(s) (requests completing slower than target_rate)" + ); + } + break; } } } - }); - } + } - // Wait for all workers to finish - workers.join_all().await; + // Best-effort drain: permits track only the tested requests, while task + // completion also includes stats, diagnostics, and error persistence. + let drain_deadline = Instant::now() + Duration::from_secs(30); + while !issued_tasks.is_empty() && Instant::now() < drain_deadline { + match tokio::time::timeout(Duration::from_millis(50), issued_tasks.join_next()).await { + Ok(Some(Err(error))) => { + eprintln!("Warning: open-loop operation task failed: {error}"); + } + Ok(Some(Ok(()))) | Err(_) => {} + Ok(None) => break, + } + } + if !issued_tasks.is_empty() { + issued_tasks.abort_all(); + issued_tasks.join_all().await; + } + + let total_skipped = skipped.load(Ordering::Relaxed); + if total_skipped > 0 { + println!( + "Open-loop summary: skipped {total_skipped} issuance(s) due to \ + max_in_flight saturation" + ); + } + let total_postprocessing_skipped = postprocessing_skipped.load(Ordering::Relaxed); + if total_postprocessing_skipped > 0 { + eprintln!( + "Open-loop summary: skipped {total_postprocessing_skipped} diagnostics/error \ + persistence operation(s) due to postprocessing saturation" + ); + } + } else { + // Closed-loop: fixed worker pool — each worker loops until cancelled. + let mut workers = JoinSet::new(); + + for _ in 0..concurrency { + let ops = operations.clone(); + let container = container.clone(); + let stats = stats.clone(); + let cancelled = cancelled.clone(); + let err_container = results_container.clone(); + let err_workload_id = workload_id.clone(); + let err_commit_sha = commit_sha.clone(); + let err_hostname = hostname.clone(); + + workers.spawn(async move { + while !cancelled.load(Ordering::Relaxed) { + let op_idx = rand::rng().random_range(0..ops.len()); + run_one( + &ops[op_idx], + &container, + OperationRunContext { + stats: &stats, + results_container: &err_container, + workload_id: &err_workload_id, + commit_sha: &err_commit_sha, + hostname: &err_hostname, + diagnostics_threshold, + postprocessing_slots: None, + postprocessing_skipped: None, + }, + None, + ) + .await; + } + }); + } + + // Wait for all workers to finish + workers.join_all().await; + } // Print final report let total_elapsed = start.elapsed(); @@ -438,7 +677,17 @@ async fn upsert_results( tokio_busy_pct: tokio_fields.map(|t| t.busy_pct), tokio_park_count: tokio_fields.map(|t| t.park_count), tokio_queue_depth: tokio_fields.map(|t| t.queue_depth), - config_concurrency: Some(config.concurrency), + config_concurrency: config.concurrency, + config_target_rate: config.target_rate, + config_max_in_flight: config.max_in_flight, + config_skipped_issuances: config + .skipped_issuances + .as_ref() + .map(|value| value.load(Ordering::Relaxed)), + config_skipped_postprocessing: config + .skipped_postprocessing + .as_ref() + .map(|value| value.load(Ordering::Relaxed)), config_application_region: Some(config.application_region.clone()), config_excluded_regions: if config.excluded_regions.is_empty() { None @@ -447,6 +696,8 @@ async fn upsert_results( }, config_tokio_threads: Some(config.tokio_threads), config_ppcb_enabled: Some(config.ppcb_enabled), + config_gateway_v2_disabled: Some(config.gateway_v2_disabled), + config_diagnostics_threshold_ms: config.diagnostics_threshold_ms, config_pyroscope_enabled: Some(config.pyroscope_enabled), config_tokio_console_enabled: Some(config.tokio_console_enabled), config_tokio_metrics_enabled: Some(config.tokio_metrics_enabled), @@ -519,3 +770,63 @@ fn diagnostics_to_json(diagnostics: &DiagnosticsContext) -> serde_json::Value { serde_json::from_str(json_str) .expect("DiagnosticsContext::to_json_string should always produce valid JSON") } + +fn log_slow_operation_diagnostics( + operation: &str, + elapsed: Duration, + diagnostics: &[Arc], +) { + let diagnostics: Vec<_> = diagnostics + .iter() + .map(|context| diagnostics_to_json(context)) + .collect(); + let entry = serde_json::json!({ + "event": "slow_operation_diagnostics", + "operation": operation, + "elapsed_ms": elapsed.as_secs_f64() * 1000.0, + "diagnostics": diagnostics, + }); + eprintln!("{entry}"); +} + +fn diagnostics_threshold_exceeded(elapsed: Duration, threshold: Option) -> bool { + threshold.is_some_and(|threshold| elapsed > threshold) +} + +fn saturated_skip_count(issued_after_increment: u64, target: u64) -> u64 { + debug_assert!(issued_after_increment <= target); + target - issued_after_increment + 1 +} + +#[cfg(test)] +mod tests { + use super::{diagnostics_threshold_exceeded, saturated_skip_count}; + use std::time::Duration; + + #[test] + fn diagnostics_logging_is_disabled_without_threshold() { + assert!(!diagnostics_threshold_exceeded( + Duration::from_secs(1), + None + )); + } + + #[test] + fn diagnostics_logging_requires_latency_over_threshold() { + let threshold = Some(Duration::from_millis(100)); + assert!(!diagnostics_threshold_exceeded( + Duration::from_millis(100), + threshold + )); + assert!(diagnostics_threshold_exceeded( + Duration::from_millis(101), + threshold + )); + } + + #[test] + fn saturation_batches_current_and_remaining_arrivals() { + assert_eq!(saturated_skip_count(10, 10), 1); + assert_eq!(saturated_skip_count(10, 1_000_000), 999_991); + } +}