diff --git a/Cargo.lock b/Cargo.lock index 8f818e46c21..458f9cd591d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -656,7 +656,7 @@ dependencies = [ [[package]] name = "azure_data_cosmos_driver_native" -version = "0.1.0" +version = "0.2.0" dependencies = [ "async-trait", "azure_core 1.1.0", diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index c0a0a8d8b67..c9ae09d759f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -8,6 +8,9 @@ ### Bugs Fixed +- Preserve admitted throttle retry settings during operation-triggered account metadata refresh, including regional and HTTP-version recovery paths. ([#5366](https://github.com/Azure/azure-sdk-for-rust/pull/5366)) +- Honor per-operation endpoint unavailability TTL and runtime TTL updates during routing without changing sibling operations' cooldown intervals. ([#5366](https://github.com/Azure/azure-sdk-for-rust/pull/5366)) + ### Other Changes ## 1.0.0-beta.1 (2026-09-28) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 0ff6455d076..9c860aaf4b7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -39,7 +39,7 @@ use crate::{ }, options::{ ConnectionPoolOptions, DriverOptions, OperationOptions, OperationOptionsView, - PartitionTopologyCacheMode, PlanOptions, ResolvedThroughputControl, + PartitionTopologyCacheMode, PlanOptions, ResolvedThroughputControl, ThrottlingRetryOptions, }, ActivityId, CosmosResponse, DiagnosticsContext, }; @@ -538,6 +538,7 @@ impl CosmosDriver { version: TransportHttpVersion, client_id: &azure_core::http::headers::HeaderValue, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result<(super::cache::AccountProperties, CosmosTransport)> { let endpoint = AccountEndpoint::from(account); let (transport, metadata_transport) = Self::build_metadata_transport_for_version( @@ -555,6 +556,7 @@ impl CosmosDriver { &user_agent, client_id, fault_injection_enabled, + throttling, ) .await?; Ok((props, transport)) @@ -576,6 +578,7 @@ impl CosmosDriver { >, client_id: &azure_core::http::headers::HeaderValue, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result { let endpoint = AccountEndpoint::from(account); let user_agent = Self::user_agent_header(runtime.user_agent()); @@ -604,6 +607,7 @@ impl CosmosDriver { &user_agent, client_id, fault_injection_enabled, + throttling, ) .await } @@ -699,6 +703,7 @@ impl CosmosDriver { TransportHttpVersion::Http11, client_id, fault_injection_enabled, + None, ) .await?; return Ok((TransportHttpVersion::Http11, props)); @@ -716,6 +721,7 @@ impl CosmosDriver { }, client_id, fault_injection_enabled, + None, ) .await { @@ -746,6 +752,7 @@ impl CosmosDriver { TransportHttpVersion::Http11, client_id, fault_injection_enabled, + None, ) .await?; Ok((TransportHttpVersion::Http11, props)) @@ -834,6 +841,10 @@ impl CosmosDriver { /// design (the driver / operation pipeline does not yet exist at bootstrap, nor for /// the 5-minute background refresh callback) but still produces a `DiagnosticsContext` /// matching the data-plane shape so error consumers see the same fields. + #[expect( + clippy::too_many_arguments, + reason = "metadata requests carry identity, diagnostics and an optional pinned retry policy" + )] async fn fetch_account_properties_with_transport( runtime: &CosmosDriverRuntime, transport: &super::transport::adaptive_transport::AdaptiveTransport, @@ -842,6 +853,7 @@ impl CosmosDriver { user_agent: &azure_core::http::headers::HeaderValue, client_id: &azure_core::http::headers::HeaderValue, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result { let endpoint = AccountEndpoint::from(account); let endpoint_url = endpoint.join_path("/"); @@ -874,30 +886,26 @@ impl CosmosDriver { // budget the metadata pipeline uses (previously bootstrap had no 429 // retry at all). let mut connectivity_retry_count = 0_u32; - // Resolve the caller-configured throttle limits the same way the - // metadata operation pipeline does, falling back to the metadata class - // defaults. This is a static helper that only receives the runtime, so - // only the runtime-level layers apply: the client-wide default set via - // `with_default_operation_options` (which takes precedence) and the - // environment (`AZURE_COSMOS_MAX_THROTTLE_RETRY_COUNT`). Without this a - // caller that disabled retries (`max_retry_count = 0`) would still see - // the nine-retry metadata default here, inconsistent with normal - // metadata operations. The per-retry delay cap has no caller override - // and stays at the metadata class default, matching the pipeline. - let bootstrap_options = OperationOptionsView::new( - Some(Arc::clone(runtime.env_operation_options())), - Some(runtime.default_operation_options()), - None, - None, - ); - let throttling_retry_options = bootstrap_options.throttling_retry_options(); + // Some(empty) pins absent fields too; only standalone bootstrap/background + // requests may resolve the current runtime generation. + let throttling_retry_options = throttling.cloned().unwrap_or_else(|| { + let view = OperationOptionsView::new( + Some(Arc::clone(runtime.env_operation_options())), + Some(runtime.default_operation_options()), + None, + None, + ); + let throttling = view.throttling_retry_options(); + ThrottlingRetryOptions { + max_retry_count: throttling.max_retry_count().copied(), + max_retry_wait_time: throttling.max_retry_wait_time().copied(), + } + }); let max_throttle_attempts = throttling_retry_options - .max_retry_count() - .copied() + .max_retry_count .unwrap_or(METADATA_MAX_THROTTLE_ATTEMPTS); let max_throttle_wait_time = throttling_retry_options - .max_retry_wait_time() - .copied() + .max_retry_wait_time .unwrap_or(METADATA_MAX_THROTTLE_WAIT); let mut throttle = ThrottleRetryState::with_limits( max_throttle_attempts, @@ -1183,6 +1191,7 @@ impl CosmosDriver { &self.client_id, None, fault_injection_enabled, + None, ) .await } @@ -1219,6 +1228,7 @@ impl CosmosDriver { client_id: &azure_core::http::headers::HeaderValue, previous_props: Option>, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result { let current_transport = transport_holder.load_full(); let current_version = current_transport.negotiated_version(); @@ -1234,6 +1244,7 @@ impl CosmosDriver { &user_agent_header, client_id, fault_injection_enabled, + throttling, ) .await { @@ -1247,6 +1258,7 @@ impl CosmosDriver { &endpoint, client_id, fault_injection_enabled, + throttling, ) .await; Ok(props) @@ -1262,6 +1274,7 @@ impl CosmosDriver { error, client_id, fault_injection_enabled, + throttling, ) .await { @@ -1278,6 +1291,7 @@ impl CosmosDriver { primary_error, previous_props, fault_injection_enabled, + throttling, ) .await } @@ -1302,6 +1316,7 @@ impl CosmosDriver { primary_error: crate::error::CosmosError, previous_props: Option>, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result { let Some(cached_props) = previous_props else { return Err(primary_error); @@ -1350,6 +1365,7 @@ impl CosmosDriver { &user_agent, client_id, fault_injection_enabled, + throttling, ) .await { @@ -1391,6 +1407,7 @@ impl CosmosDriver { endpoint: &AccountEndpoint, client_id: &azure_core::http::headers::HeaderValue, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) { if !matches!(current_version, TransportHttpVersion::Http11) || !runtime.connection_pool().is_http2_allowed() @@ -1409,6 +1426,7 @@ impl CosmosDriver { }, client_id, fault_injection_enabled, + throttling, ) .await { @@ -1458,6 +1476,7 @@ impl CosmosDriver { error: crate::error::CosmosError, client_id: &azure_core::http::headers::HeaderValue, fault_injection_enabled: bool, + throttling: Option<&ThrottlingRetryOptions>, ) -> crate::error::Result { if Self::should_downgrade_http2( current_version, @@ -1481,6 +1500,7 @@ impl CosmosDriver { fallback_version, client_id, fault_injection_enabled, + throttling, ) .await?; @@ -1740,7 +1760,8 @@ impl CosmosDriver { #[cfg(not(feature = "fault_injection"))] let fault_injection_for_callback = false; let refresh_callback = Arc::new( - move |previous_props: Option>| { + move |previous_props: Option>, + throttling: Option| { let runtime = Arc::clone(&runtime_for_callback); let account = account_for_callback.clone(); let transport_holder = Arc::clone(&transport_for_callback); @@ -1759,6 +1780,7 @@ impl CosmosDriver { &client_id, previous_props, fault_injection_enabled, + throttling.as_ref(), ) .await }); @@ -1774,13 +1796,8 @@ impl CosmosDriver { .or(runtime .default_operation_options() .endpoint_unavailability_ttl) - .unwrap_or_else(|| { - std::env::var("AZURE_COSMOS_ENDPOINT_UNAVAILABLE_TTL_MS") - .ok() - .and_then(|v| v.parse::().ok()) - .map(Duration::from_millis) - .unwrap_or(Duration::from_secs(60)) - }); + .or(runtime.env_operation_options().endpoint_unavailability_ttl) + .unwrap_or(Duration::from_secs(60)); // Wire the Gateway 2.0 connectivity probe. Before routing data-plane // traffic to a thin-client proxy endpoint, the store issues a @@ -2219,6 +2236,9 @@ impl CosmosDriver { &self, operation_options: &'a OperationOptions, ) -> OperationOptionsView<'a> { + if let Some(snapshot) = operation_options.resolution_snapshot_view() { + return snapshot; + } OperationOptionsView::new_with_override( Some(Arc::clone(self.runtime.env_override_operation_options())), Some(Arc::clone(self.runtime.env_operation_options())), @@ -5599,6 +5619,7 @@ mod tests { &TEST_CLIENT_ID, None, false, + None, ) .await .unwrap(); @@ -5643,6 +5664,7 @@ mod tests { &TEST_CLIENT_ID, None, false, + None, ) .await .unwrap(); @@ -5688,6 +5710,7 @@ mod tests { &TEST_CLIENT_ID, None, false, + None, ) .await .unwrap(); @@ -5925,6 +5948,7 @@ mod tests { &TEST_CLIENT_ID, Some(multi_region_previous_props()), false, + None, ) .await; @@ -5968,6 +5992,7 @@ mod tests { &TEST_CLIENT_ID, Some(multi_region_previous_props()), false, + None, ) .await; @@ -6006,6 +6031,7 @@ mod tests { &TEST_CLIENT_ID, None, false, + None, ) .await; @@ -6034,6 +6060,7 @@ mod tests { &user_agent, &client_id, false, + None, ) .await .expect_err( @@ -6090,6 +6117,136 @@ mod tests { ); } + #[tokio::test] + async fn operation_account_refresh_retains_admitted_throttle_policy() { + #[derive(Debug, Default)] + struct RecoveryClient { + operation_calls: std::sync::atomic::AtomicUsize, + refresh_calls: Mutex>, + } + + #[async_trait] + impl TransportClient for RecoveryClient { + async fn send(&self, request: &HttpRequest) -> Result { + let mut headers = Headers::new(); + if request.url.path() == "/dbs/db" { + if self.operation_calls.fetch_add(1, Ordering::SeqCst) == 0 { + headers.insert( + azure_core::http::headers::HeaderName::from_static("x-ms-substatus"), + azure_core::http::headers::HeaderValue::from_static("1008"), + ); + return Ok(HttpResponse { + status: 403, + headers, + body: Vec::new(), + }); + } + return Ok(HttpResponse { + status: 200, + headers, + body: b"{}".to_vec(), + }); + } + if request.url.path() == "/" && self.operation_calls.load(Ordering::SeqCst) > 0 { + *self + .refresh_calls + .lock() + .unwrap() + .entry(request.url.host_str().unwrap().to_owned()) + .or_default() += 1; + headers.insert( + azure_core::http::headers::HeaderName::from_static("x-ms-retry-after-ms"), + azure_core::http::headers::HeaderValue::from_static("1"), + ); + return Ok(HttpResponse { + status: 429, + headers, + body: Vec::new(), + }); + } + Ok(HttpResponse { + status: 200, + headers, + body: ACCOUNT_PROPERTIES_PAYLOAD.as_bytes().to_vec(), + }) + } + } + #[derive(Debug)] + struct RecoveryFactory(Arc); + impl HttpClientFactory for RecoveryFactory { + fn build( + &self, + _: &ConnectionPoolOptions, + _: HttpClientConfig, + ) -> crate::error::Result> { + Ok(self.0.clone()) + } + } + + let client = Arc::new(RecoveryClient::default()); + let runtime = CosmosDriverRuntimeBuilder::new() + .with_http_client_factory(Arc::new(RecoveryFactory(Arc::clone(&client)))) + .with_connection_pool( + ConnectionPoolOptions::builder() + .with_is_http2_allowed(false) + .build() + .unwrap(), + ) + .with_default_operation_options( + OperationOptionsBuilder::new() + .with_throttling_retry_options( + ThrottlingRetryOptionsBuilder::new() + .with_max_retry_count(0) + .build(), + ) + .build(), + ) + .build() + .await + .unwrap(); + let account = signed_test_account("https://test.documents.azure.com:443/"); + let driver = runtime + .create_driver(DriverOptions::builder(account.clone()).build()) + .await + .unwrap(); + let admitted = OperationOptions { + hedging_enabled: Some(false), + ..Default::default() + } + .with_resolution_snapshot( + Arc::clone(runtime.env_override_operation_options()), + Arc::clone(runtime.env_operation_options()), + runtime.default_operation_options(), + Arc::clone(driver.options.operation_options()), + ); + runtime.set_default_operation_options( + OperationOptionsBuilder::new() + .with_throttling_retry_options( + ThrottlingRetryOptionsBuilder::new() + .with_max_retry_count(9) + .build(), + ) + .build(), + ); + let response = driver + .execute_singleton_operation( + CosmosOperation::read_database(DatabaseReference::from_name(account, "db")), + admitted, + ) + .await + .unwrap(); + assert!(response.status().is_success()); + assert_eq!(client.operation_calls.load(Ordering::SeqCst), 2); + assert_eq!( + *client.refresh_calls.lock().unwrap(), + std::collections::HashMap::from([ + ("test.documents.azure.com".to_owned(), 1), + ("test-westus2.documents.azure.com".to_owned(), 1), + ]), + "both primary and regional refreshes must preserve the admitted zero-retry policy" + ); + } + #[tokio::test] async fn fetch_account_properties_retries_on_throttle_then_succeeds() { // Returns 429 (with x-ms-retry-after-ms) on the first attempt, then 200. @@ -6143,6 +6300,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect("bootstrap must retry the 429 and then succeed"); @@ -6212,6 +6370,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect_err("disabled throttle retries must surface the 429"); @@ -6307,6 +6466,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await } @@ -6333,6 +6493,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect("not-sent connectivity failures should be retried"); @@ -6362,6 +6523,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect_err("connectivity failures should surface after retry budget exhaustion"); @@ -6749,6 +6911,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect_err("transport-layer failure must surface as an error"); @@ -6821,6 +6984,7 @@ mod tests { &user_agent, &TEST_CLIENT_ID, false, + None, ) .await .expect_err("sign_request failure must surface as an error"); @@ -6961,6 +7125,7 @@ mod tests { &user_agent, &client_id, false, + None, ) .await; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs index d9269c2e7aa..24780e97929 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs @@ -499,6 +499,10 @@ pub(crate) async fn execute_operation_pipeline( let mut container_recreation_retry_attempted = false; let location_snapshot = location_state_store.snapshot(); let max_failover_retries = options.max_failover_retry_count().copied().unwrap_or(3); + let endpoint_unavailability_ttl = options + .endpoint_unavailability_ttl() + .copied() + .unwrap_or(Duration::from_secs(60)); // Throttle (HTTP 429) retry limits, resolved from the effective operation // options. These are the analogs of the .NET SDK's @@ -665,7 +669,7 @@ pub(crate) async fn execute_operation_pipeline( &location, pipeline_type.is_data_plane(), account_name.is_some(), - location_state_store.endpoint_unavailability_ttl(), + endpoint_unavailability_ttl, ), }; let attempt_read_consistency_strategy = @@ -1134,7 +1138,9 @@ pub(crate) async fn execute_operation_pipeline( effects, ); retry_state.pending_write_effects.extend(deferred_effects); - location_state_store.apply(&immediate_effects).await; + location_state_store + .apply_with_options(&immediate_effects, Some(options)) + .await; // ── STAGE 7: Act on the control-flow decision ────────────────── match action { @@ -1332,7 +1338,7 @@ pub(crate) async fn execute_operation_pipeline( &location, pipeline_type.is_data_plane(), account_name.is_some(), - location_state_store.endpoint_unavailability_ttl(), + endpoint_unavailability_ttl, ); // Re-evaluate hedge eligibility against the *post-advance* // primary. After `advance_to_next_attempt` rotates the @@ -3820,7 +3826,9 @@ async fn apply_hedge_leg_effects( transport_result, ); if !eval.effects.is_empty() { - ctx.location_state_store.apply(&eval.effects).await; + ctx.location_state_store + .apply_with_options(&eval.effects, Some(ctx.options)) + .await; } if eval.observed_session_unavailable { *race_observed_session_unavailable = true; @@ -8596,10 +8604,15 @@ mod tests { ), ); let loc = make_location(both); + assert_eq!( + super::resolve_endpoint(&read_op, &state, &loc, false, true, Duration::ZERO).endpoint, + r1, + "an operation with zero TTL can use the marked endpoint", + ); assert_eq!( resolve(&read_op, &loc), r2, - "a both-affecting mark on r1 must demote it for reads", + "another operation's TTL still demotes the same unmodified mark", ); assert_eq!( resolve(&write_op, &loc), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs index 3b4c23da6ad..81d7ba1e546 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs @@ -24,7 +24,7 @@ use crate::{ transport::connectivity_probe::{ConnectivityProbe, ProbeOutcome, ProbeRole}, }, models::AccountEndpoint, - options::{PartitionFailoverOptions, Region}, + options::{OperationOptionsView, PartitionFailoverOptions, Region, ThrottlingRetryOptions}, }; use super::{ @@ -68,6 +68,7 @@ impl LocationSnapshot { type AccountRefreshFn = Arc< dyn Fn( Option>, + Option, ) -> BoxFuture<'static, crate::error::Result> + Send + Sync, @@ -338,11 +339,6 @@ impl LocationStateStore { snapshot } - /// Returns the configured endpoint unavailability TTL. - pub fn endpoint_unavailability_ttl(&self) -> Duration { - self.endpoint_unavailability_ttl - } - /// Returns the latest account snapshot. #[allow(dead_code)] pub fn account_snapshot(&self) -> Arc { @@ -386,6 +382,15 @@ impl LocationStateStore { /// Applies location effects (endpoint unavailability and account refresh). pub async fn apply(&self, effects: &[LocationEffect]) { + self.apply_with_options(effects, None).await; + } + + /// Applies operation-triggered effects with the caller's pinned configuration. + pub(crate) async fn apply_with_options( + &self, + effects: &[LocationEffect], + options: Option<&OperationOptionsView<'_>>, + ) { for effect in effects { match effect { LocationEffect::MarkEndpointUnavailable { endpoint, reason } => { @@ -413,7 +418,14 @@ impl LocationStateStore { }); } LocationEffect::RefreshAccountProperties => { - self.refresh_account_properties_if_due().await; + let throttling = options.map(|options| { + let view = options.throttling_retry_options(); + ThrottlingRetryOptions { + max_retry_count: view.max_retry_count().copied(), + max_retry_wait_time: view.max_retry_wait_time().copied(), + } + }); + self.refresh_account_properties_if_due(throttling).await; } LocationEffect::CacheHubRegion { partition_key_range_id, @@ -507,7 +519,7 @@ impl LocationStateStore { } } - async fn refresh_account_properties_if_due(&self) { + async fn refresh_account_properties_if_due(&self, throttling: Option) { let now_ms = epoch_millis(); let refresh_after_ms = self.refresh_interval.as_millis() as u64; let last = self.last_refresh_epoch_ms.load(Ordering::Acquire); @@ -542,7 +554,7 @@ impl LocationStateStore { committed: false, }; - if self.refresh_account_properties_inner().await { + if self.refresh_account_properties_inner(throttling).await { claim.commit(); } } @@ -561,7 +573,7 @@ impl LocationStateStore { /// [`RefreshClaimGuard`] when its own fetch fails or is cancelled, so the /// same guarantee holds there. async fn force_refresh_account_properties(&self) { - let _ = self.refresh_account_properties_inner().await; + let _ = self.refresh_account_properties_inner(None).await; } /// Shared implementation of both `refresh_account_properties_if_due` @@ -581,7 +593,10 @@ impl LocationStateStore { /// /// Returns `true` only when a fresh snapshot was actually applied, so /// callers that pre-claimed the rate-limit clock can roll it back. - async fn refresh_account_properties_inner(&self) -> bool { + async fn refresh_account_properties_inner( + &self, + throttling: Option, + ) -> bool { // Capture the previous properties so the refresh callback can use // them for regional fallback if the primary endpoint fails. We // intentionally do NOT invalidate the cache here — concurrent @@ -594,7 +609,7 @@ impl LocationStateStore { .await; let refresh_fn = Arc::clone(&self.account_refresh_fn); - let fetched = (refresh_fn)(previous_props).await; + let fetched = (refresh_fn)(previous_props, throttling).await; let new_properties = match fetched { Ok(props) => props, @@ -1331,12 +1346,15 @@ mod tests { #[tokio::test] async fn apply_marks_endpoint_unavailable() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1364,12 +1382,15 @@ mod tests { #[test] fn global_database_account_name_prefers_metadata_id() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1418,12 +1439,15 @@ mod tests { #[tokio::test] async fn unavailable_endpoint_fails_back_only_after_successful_probe() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1476,12 +1500,15 @@ mod tests { #[tokio::test] async fn failed_probe_resets_cooldown_keeping_endpoint_out_of_rotation() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let make_store = || { LocationStateStore::new( @@ -1574,12 +1601,15 @@ mod tests { #[tokio::test] async fn endpoint_unavailability_lifecycle_mark_cooldown_probe_restore() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); // Non-zero cooldown so the "not yet due" gate is observable. let cooldown = Duration::from_secs(60); @@ -1675,12 +1705,15 @@ mod tests { #[test] fn account_sync_preserves_unavailable_marks_for_probe_loop() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1738,12 +1771,15 @@ mod tests { #[tokio::test] async fn probe_skips_endpoints_the_account_no_longer_advertises() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1829,12 +1865,15 @@ mod tests { #[tokio::test] async fn probe_does_not_fail_back_an_endpoint_removed_while_probing() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = Arc::new(LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1906,16 +1945,19 @@ mod tests { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); let refresh_calls = Arc::new(AtomicUsize::new(0)); let refresh_calls_clone = Arc::clone(&refresh_calls); - let refresh = Arc::new(move |_previous: Option>| { - let refresh_calls = Arc::clone(&refresh_calls_clone); - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { - refresh_calls.fetch_add(1, Ordering::SeqCst); - Ok(payload) - }); - fut - }); + let refresh = Arc::new( + move |_previous: Option>, + _throttling: Option| { + let refresh_calls = Arc::clone(&refresh_calls_clone); + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { + refresh_calls.fetch_add(1, Ordering::SeqCst); + Ok(payload) + }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -1952,27 +1994,30 @@ mod tests { let success_refreshes_clone = Arc::clone(&success_refreshes); let total_refreshes_clone = Arc::clone(&total_refreshes); // First call fails; subsequent calls succeed. - let refresh = Arc::new(move |_previous: Option>| { - let total = Arc::clone(&total_refreshes_clone); - let success = Arc::clone(&success_refreshes_clone); - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { - let n = total.fetch_add(1, Ordering::SeqCst); - if n == 0 { - Err(crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::new( - azure_core::http::StatusCode::BadRequest, - )) - .with_message("simulated network failure") - .build()) - } else { - success.fetch_add(1, Ordering::SeqCst); - Ok(payload) - } - }); - fut - }); + let refresh = Arc::new( + move |_previous: Option>, + _throttling: Option| { + let total = Arc::clone(&total_refreshes_clone); + let success = Arc::clone(&success_refreshes_clone); + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { + let n = total.fetch_add(1, Ordering::SeqCst); + if n == 0 { + Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::new( + azure_core::http::StatusCode::BadRequest, + )) + .with_message("simulated network failure") + .build()) + } else { + success.fetch_add(1, Ordering::SeqCst); + Ok(payload) + } + }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2019,27 +2064,30 @@ mod tests { let success_refreshes_clone = Arc::clone(&success_refreshes); let total_refreshes_clone = Arc::clone(&total_refreshes); // First call fails; subsequent calls succeed. - let refresh = Arc::new(move |_previous: Option>| { - let total = Arc::clone(&total_refreshes_clone); - let success = Arc::clone(&success_refreshes_clone); - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { - let n = total.fetch_add(1, Ordering::SeqCst); - if n == 0 { - Err(crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::new( - azure_core::http::StatusCode::BadRequest, - )) - .with_message("simulated network failure") - .build()) - } else { - success.fetch_add(1, Ordering::SeqCst); - Ok(payload) - } - }); - fut - }); + let refresh = Arc::new( + move |_previous: Option>, + _throttling: Option| { + let total = Arc::clone(&total_refreshes_clone); + let success = Arc::clone(&success_refreshes_clone); + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { + let n = total.fetch_add(1, Ordering::SeqCst); + if n == 0 { + Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::new( + azure_core::http::StatusCode::BadRequest, + )) + .with_message("simulated network failure") + .build()) + } else { + success.fetch_add(1, Ordering::SeqCst); + Ok(payload) + } + }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2111,16 +2159,19 @@ mod tests { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); let refreshes = Arc::new(AtomicUsize::new(0)); let refreshes_clone = Arc::clone(&refreshes); - let refresh = Arc::new(move |_previous: Option>| { - let count = Arc::clone(&refreshes_clone); - let payload = refresh_payload_with_g2(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { - count.fetch_add(1, Ordering::SeqCst); - Ok(payload) - }); - fut - }); + let refresh = Arc::new( + move |_previous: Option>, + _throttling: Option| { + let count = Arc::clone(&refreshes_clone); + let payload = refresh_payload_with_g2(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { + count.fetch_add(1, Ordering::SeqCst); + Ok(payload) + }); + fut + }, + ); let probe = Arc::new(BlockingProbe::default()); let store = Arc::new(LocationStateStore::new( @@ -2172,25 +2223,28 @@ mod tests { // Call 1 succeeds (bootstrap seeds the cache); calls 2+ surface a // typed 503 mirroring what `fetch_account_properties_with_transport` // now produces on a 5xx account-metadata response. - let refresh = Arc::new(move |_previous: Option>| { - let total = Arc::clone(&total_refreshes_clone); - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { - let n = total.fetch_add(1, Ordering::SeqCst); - if n == 0 { - Ok(payload) - } else { - Err(crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::new( - azure_core::http::StatusCode::ServiceUnavailable, - )) - .with_message("simulated 5xx on periodic account-metadata refresh") - .build()) - } - }); - fut - }); + let refresh = Arc::new( + move |_previous: Option>, + _throttling: Option| { + let total = Arc::clone(&total_refreshes_clone); + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { + let n = total.fetch_add(1, Ordering::SeqCst); + if n == 0 { + Ok(payload) + } else { + Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::new( + azure_core::http::StatusCode::ServiceUnavailable, + )) + .with_message("simulated 5xx on periodic account-metadata refresh") + .build()) + } + }); + fut + }, + ); let cache = Arc::new(AccountMetadataCache::new()); let store = LocationStateStore::new( @@ -2260,12 +2314,15 @@ mod tests { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); // The refresh fn is unused in this test — sync_account_properties // is called directly with explicit payloads. - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2400,12 +2457,15 @@ mod tests { #[test] fn sync_account_properties_adopts_gateway_v2_when_thin_client_locations_appear() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2528,12 +2588,15 @@ mod tests { // fix, the strong ref stays alive until the store itself is dropped. let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2704,12 +2767,15 @@ mod tests { async fn connectivity_probe_failure_suppresses_gateway_v2_then_recovers() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = refresh_payload_with_g2(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = refresh_payload_with_g2(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let probe = Arc::new(MockProbe::new(ProbeOutcome::Failed { failures: vec![( @@ -2778,12 +2844,15 @@ mod tests { async fn connectivity_probe_success_is_sticky_and_not_reprobed() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = refresh_payload_with_g2(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = refresh_payload_with_g2(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let probe = Arc::new(MockProbe::new(ProbeOutcome::AllHealthy)); @@ -2837,12 +2906,15 @@ mod tests { fn no_probe_wired_preserves_existing_gateway_v2_behavior() { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = refresh_payload_with_g2(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = refresh_payload_with_g2(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2869,12 +2941,15 @@ mod tests { fn build_store_with_two_regions() -> LocationStateStore { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_multi_region_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_multi_region_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); let store = LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -2951,12 +3026,15 @@ mod tests { fn build_store_for_ppaf_tests() -> LocationStateStore { let default_endpoint = CosmosEndpoint::global(test_endpoint().url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); LocationStateStore::new( Arc::new(AccountMetadataCache::new()), @@ -3054,12 +3132,15 @@ mod tests { .unwrap(), ); let default_endpoint = CosmosEndpoint::global(endpoint.url().clone()); - let refresh = Arc::new(|_previous: Option>| { - let payload = test_refresh_payload(); - let fut: BoxFuture<'static, crate::error::Result> = - Box::pin(async move { Ok(payload) }); - fut - }); + let refresh = Arc::new( + |_previous: Option>, + _throttling: Option| { + let payload = test_refresh_payload(); + let fut: BoxFuture<'static, crate::error::Result> = + Box::pin(async move { Ok(payload) }); + fut + }, + ); LocationStateStore::new( Arc::clone(&cache), endpoint, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs index 19521e73b33..ddb29ba92d9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs @@ -15,6 +15,8 @@ use serde::{Deserialize, Serialize}; use super::config::ContainerConfig; use super::dispatch::{OperationType, ParsedRequest}; use super::epk::{compute_epk, extract_pk_from_body, parse_partition_key_header, Epk}; +#[cfg(feature = "preview_dtx")] +use super::response::headers::ETAG; use super::response::headers::{ ACTIVITY_ID, CONTINUATION, GLOBAL_COMMITTED_LSN, INTERNAL_PARTITION_ID, ITEM_LOCAL_LSN, ITEM_LSN, LAST_STATE_CHANGE_UTC, LOCAL_LSN, NUMBER_OF_READ_REGIONS, PARTITION_KEY_RANGE_ID, @@ -22,7 +24,7 @@ use super::response::headers::{ TRANSPORT_REQUEST_ID, }; #[cfg(feature = "preview_dtx")] -use super::response::headers::{ETAG, REQUEST_CHARGE, SESSION_TOKEN, SUBSTATUS}; +use super::response::headers::{REQUEST_CHARGE, SESSION_TOKEN, SUBSTATUS}; use super::response::{ error_response, success_response, success_response_with_format, ResponseBuilder, ResponseFormat, }; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs index e9bf4c908d3..1f460406f2e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs @@ -4,6 +4,7 @@ //! Operation options that participate in runtime/account/operation resolution. use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use azure_core::http::headers::{HeaderName, HeaderValue}; @@ -38,6 +39,8 @@ use crate::options::{ #[options(layers(runtime, account, operation))] #[non_exhaustive] pub struct OperationOptions { + #[option(skip)] + pub(crate) resolution_snapshot: Option>, /// How PATCH operations are executed. /// /// `None` inherits from a lower layer (default: [`PatchStrategy::Auto`]). @@ -83,6 +86,7 @@ pub struct OperationOptions { pub max_failover_retry_count: Option, /// How long an endpoint is considered unavailable after a failure. + #[option(env = "AZURE_COSMOS_ENDPOINT_UNAVAILABLE_TTL_MS", parser = parse_milliseconds)] pub endpoint_unavailability_ttl: Option, /// Disables automatic session token management. @@ -172,6 +176,55 @@ pub struct OperationOptions { pub binary_encoding: Option, } +#[derive(Debug)] +pub(crate) struct OperationResolutionSnapshot { + env_override: Arc, + env: Arc, + runtime: Arc, + account: Arc, +} + +fn parse_milliseconds(value: &str) -> Option { + value.parse::().ok().map(Duration::from_millis) +} + +impl OperationOptions { + /// Pins the lower configuration layers for this operation and its clones. + /// + /// Intended for native bindings admitting work before driver initialization. + /// Nested fields still resolve independently; unset values cannot observe + /// subsequent runtime updates. + #[doc(hidden)] + pub fn with_resolution_snapshot( + mut self, + env_override: Arc, + env: Arc, + runtime: Arc, + account: Arc, + ) -> Self { + self.resolution_snapshot = Some(Arc::new(OperationResolutionSnapshot { + env_override, + env, + runtime, + account, + })); + self + } + + /// Returns the pinned layered view, when admission captured one. + #[doc(hidden)] + pub fn resolution_snapshot_view(&self) -> Option> { + let snapshot = self.resolution_snapshot.as_ref()?; + Some(OperationOptionsView::new_with_override( + Some(Arc::clone(&snapshot.env_override)), + Some(Arc::clone(&snapshot.env)), + Some(Arc::clone(&snapshot.runtime)), + Some(Arc::clone(&snapshot.account)), + Some(self), + )) + } +} + /// Retry behavior for requests throttled by the service (HTTP 429, /// rate-limited). /// @@ -259,6 +312,81 @@ pub struct ThroughputControlOptions { mod tests { use super::*; + #[test] + fn snapshot_preserves_all_layers_and_override_when_cloned_for_patch() { + let environment = Arc::new(OperationOptions::from_env_vars(|name| match name { + "AZURE_COSMOS_HEDGING_ENABLED" => Ok("true".into()), + "AZURE_COSMOS_MAX_SESSION_RETRY_COUNT" => Ok("3".into()), + "AZURE_COSMOS_ENDPOINT_UNAVAILABLE_TTL_MS" => Ok("12".into()), + _ => Err(std::env::VarError::NotPresent), + })); + let override_options = Arc::new(OperationOptions::from_env_override_vars(|name| { + (name == "AZURE_COSMOS_HEDGING_ENABLED_OVERRIDE") + .then(|| "false".into()) + .ok_or(std::env::VarError::NotPresent) + })); + let runtime = Arc::new(OperationOptions { + read_consistency_strategy: Some(ReadConsistencyStrategy::Eventual), + content_response_on_write: Some(ContentResponseOnWrite::Enabled), + max_failover_retry_count: Some(u32::MAX), + session_capturing_disabled: Some(true), + excluded_regions: Some(ExcludedRegions::new().with_region("East US")), + end_to_end_latency_policy: Some(Duration::from_secs(4).into()), + #[cfg(feature = "preview_patch")] + patch_strategy: Some(PatchStrategy::ServerSide), + custom_headers: Some(HashMap::from([( + HeaderName::from("x-test"), + HeaderValue::from("runtime"), + )])), + ..Default::default() + }); + let account = Arc::new(OperationOptions { + max_session_retry_count: Some(4), + #[cfg(feature = "preview_patch")] + patch_strategy: Some(PatchStrategy::Auto), + ..Default::default() + }); + let request = OperationOptions { + hedging_enabled: Some(true), + excluded_regions: Some(ExcludedRegions::new()), + custom_headers: Some(HashMap::new()), + availability_strategy: Some(AvailabilityStrategy::Disabled), + ..Default::default() + } + .with_resolution_snapshot(override_options, environment, runtime, account); + let mut patch_read = request.clone(); + patch_read.read_consistency_strategy = Some(ReadConsistencyStrategy::LatestCommitted); + let view = patch_read.resolution_snapshot_view().unwrap(); + assert_eq!(view.hedging_enabled(), Some(&false)); + assert_eq!(view.max_failover_retry_count(), Some(&u32::MAX)); + assert_eq!(view.max_session_retry_count(), Some(&4)); + assert_eq!( + view.endpoint_unavailability_ttl(), + Some(&Duration::from_millis(12)) + ); + #[cfg(feature = "preview_patch")] + assert_eq!(view.patch_strategy(), Some(&PatchStrategy::Auto)); + assert_eq!( + view.read_consistency_strategy(), + Some(&ReadConsistencyStrategy::LatestCommitted) + ); + assert_eq!( + view.content_response_on_write(), + Some(&ContentResponseOnWrite::Enabled) + ); + assert_eq!(view.session_capturing_disabled(), Some(&true)); + assert_eq!(view.excluded_regions(), Some(&ExcludedRegions::new())); + assert!(view.custom_headers().unwrap().is_empty()); + assert_eq!( + view.end_to_end_latency_policy().unwrap().timeout(), + Duration::from_secs(4) + ); + assert_eq!( + view.availability_strategy(), + Some(&AvailabilityStrategy::Disabled) + ); + } + #[test] fn default_operation_options() { let options = OperationOptions::default(); diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver_native/CHANGELOG.md new file mode 100644 index 00000000000..d7f3abeeae4 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver_native/CHANGELOG.md @@ -0,0 +1,17 @@ +# Release History + +## 0.2.0 (Unreleased) + +### Features Added + +- Added runtime, client, and request operation options to the native C FFI, including immutable admission snapshots for operation submission. ([#5366](https://github.com/Azure/azure-sdk-for-rust/pull/5366)) + +### Breaking Changes + +- Changed the pre-1.0 operation-options and request ABI layouts, including 64-bit retry-count fields; hosts using the 0.1.0 bootstrap ABI must rebuild for 0.2.0. ([#5366](https://github.com/Azure/azure-sdk-for-rust/pull/5366)) + +### Bugs Fixed + +### Other Changes + +- Began maintained native release history at 0.2.0; earlier 0.1.0 development is treated as the bootstrap baseline. diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/CMakeLists.txt b/sdk/cosmos/azure_data_cosmos_driver_native/CMakeLists.txt index 1ed763486de..7d0c6fb9e96 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/CMakeLists.txt +++ b/sdk/cosmos/azure_data_cosmos_driver_native/CMakeLists.txt @@ -61,7 +61,7 @@ configure_file( if(NOT WIN32) set(PKG_CONFIG_NAME "azurecosmosdriver") set(PKG_CONFIG_DESCRIPTION "Azure Cosmos DB driver C API") - set(PKG_CONFIG_VERSION "0.1.0") + set(PKG_CONFIG_VERSION "0.2.0") set(PKG_CONFIG_LIBS "-L\${libdir} -lazurecosmosdriver") set(PKG_CONFIG_CFLAGS "-I\${includedir}") diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml b/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml index acf016d7713..cfc2c4de1d1 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml +++ b/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml @@ -4,7 +4,7 @@ [package] name = "azure_data_cosmos_driver_native" description = "C ABI wrapper for the Azure Cosmos DB driver crate (azure_data_cosmos_driver). Exposes a schema-agnostic completion-queue-style FFI for cross-language SDK reuse (.NET, Java, Go, Python, native C/C++)." -version = "0.1.0" +version = "0.2.0" authors.workspace = true edition.workspace = true license.workspace = true diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/README.md b/sdk/cosmos/azure_data_cosmos_driver_native/README.md index 02674c6d0ea..7d99afb918d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/README.md +++ b/sdk/cosmos/azure_data_cosmos_driver_native/README.md @@ -40,7 +40,7 @@ for the full design. | Resource-token authentication | ⏳ follow-up | | Sync driver creation (`_blocking`) | ✅ | | Async driver creation (`_submit`) | ✅ | -| Cache-hit advisory (`5001 OPTIONS_IGNORED_ON_CACHE_HIT`) | ⏳ needs driver-side `was_cached` signal | +| Same-endpoint clients sharing cached container references | ✅ runtime-owned container cache | | Sync + async `resolve_container` | ✅ | | Single + hierarchical partition keys | ✅ | | Item-CRUD operations (read / create / upsert / replace / delete) | ✅ | @@ -60,6 +60,65 @@ for the full design. ## Building +Version **0.2.0** is a breaking C ABI revision. Check `cosmos_version()` against +`AZURECOSMOSDRIVER_H_VERSION` before any layout-sensitive call, including options +returned by value. Rebuild host bindings and vendor the generated header together +with the matching library; do not mix 0.1.x headers and 0.2.x libraries. + +### Operation defaults and admission snapshots + +The same `cosmos_operation_options_t` supplies request, driver and runtime defaults. +Initialize it with `cosmos_operation_options_default()`, not zero initialization: +negative numeric values inherit, while zero retry counts and durations are explicit. +Retry counts and throughput buckets accept the complete `uint32_t` range through +checked `int64_t` fields. Non-NULL, zero-length regions/headers clear inherited +collections. Throughput and throttle members inherit independently; binary flags +replace their whole group. Hedging environment overrides retain driver precedence. + +Set `cosmos_runtime_options_t.operation_options` during construction or atomically +replace the complete runtime group with `cosmos_runtime_set_operation_options`. +NULL resets the defaults; rejected updates do not change them. Inputs are copied. + +Before lazy initialization, call `cosmos_operation_options_snapshot_create` with +the runtime, client defaults and request overrides. It returns a resolved timeout +in milliseconds (-1 if absent) and an owned snapshot. Put the handle in +`cosmos_operation_request_t.options_snapshot`; the shared fields in `request.options` +are then ignored, but its query-plan mode still applies. Submit copies the snapshot, +so it may be freed immediately afterward. A driver from another runtime is rejected. +All native retries and client-side patch stages retain the same configuration +generation, including unset values and environment overrides. +Account metadata refreshes triggered during recovery also retain the admitted +throttle retry count and wait budget, including regional and HTTP-version fallback. +Independently scheduled background refreshes use current runtime defaults. + +Cursor requests can carry the same admission snapshot in their common operation +fields. The cursor retains it across opening and page requests: the original +deadline includes time the host spends between pages, rather than restarting +for each page. A snapshot from another runtime is rejected before admission. + +The snapshot starts its deadline at capture, deducting time spent in host +initialization before execution; it never restarts the full budget at submit. +Rust clamps configured latency below one second. Hosts may impose a stricter +context deadline independently. Bootstrap/metadata initialization remains separate +work, governed by the host's same deadline and the driver's metadata policy. + +Endpoint-unavailability TTL is resolved per operation during routing without +mutating shared cooldown settings. The existing TTL environment variable is +captured at runtime construction; background probes retain their construction policy. + +Drivers sharing a runtime reuse its container-reference cache, including name/RID +lookups and recreation refreshes. Cached references retain the credentials used +to create them, so same-account clients on one runtime must use compatible +credentials. Use separate runtimes when credential isolation is required. + +Host-requested cancellation is not supported: dropping an in-flight driver future +can discard its diagnostics. The CANCELLED outcome and handle-state values remain +reserved. Hosts must retain operation cookies and drain terminal completions even +if their own callers stop waiting. Admission deadlines remain enforced separately; +an admission timeout does not guarantee partial driver diagnostics. + +### Commands + ```bash # Rust side (produces the cdylib / staticlib and regenerates the header). cargo build --release -p azure_data_cosmos_driver_native @@ -485,6 +544,7 @@ internal static class Cosmos public StringView patch_tracking_id; // UUID, NULL/0 = generate public ushort patch_tracking_capacity; // 0 = driver default public uint patch_tracking_retention_seconds; // 0 = driver default + public IntPtr options_snapshot; // NULL = no admission snapshot } // A drained completion. All pointers are borrowed until free_completions. @@ -789,7 +849,8 @@ public final class CosmosSample { STRING_VIEW.withName("patch_tracking_id"), JAVA_SHORT.withName("patch_tracking_capacity"), MemoryLayout.paddingLayout(2), - JAVA_INT.withName("patch_tracking_retention_seconds")); + JAVA_INT.withName("patch_tracking_retention_seconds"), + ADDRESS.withName("options_snapshot")); // Layout of cosmos_completion_t. Pointers and intptr_t/uintptr_t are 8 bytes. static final GroupLayout COMPLETION = MemoryLayout.structLayout( @@ -1354,6 +1415,7 @@ class CosmosOperationRequest(ctypes.Structure): ("patch_tracking_id", CosmosStringView), ("patch_tracking_capacity", ctypes.c_uint16), ("patch_tracking_retention_seconds", ctypes.c_uint32), + ("options_snapshot", void_p), ] @@ -1595,11 +1657,10 @@ if __name__ == "__main__": retry storm the per-attempt list can be a compacted subset — `is_compacted` / `retained_request_count` report that, while `request_count` and `total_request_charge` stay exact. -7. **Single-runtime caching.** Drivers are cached by endpoint URL on the - `cosmos_runtime_t` that created them. Multiple `cosmos_runtime_t` - instances do **not** share their caches — see - [section 4.4.1 in the spec](https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/cosmos/docs/specs/0019-native-wrapper.md) - for the full contract. +7. **Shared runtime caches.** Every construction creates a fresh driver, but + container references, account metadata and runtime transport resources are + shared. Cached container references can reuse another same-account client's + credentials; use separate runtimes when that is not intended. ## Repository archaeology — files removed by PR #4103 diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/api.md b/sdk/cosmos/azure_data_cosmos_driver_native/api.md index a470c15d398..517e5bfdd2f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/api.md +++ b/sdk/cosmos/azure_data_cosmos_driver_native/api.md @@ -507,8 +507,8 @@ pub mod op_request { pub content_response_on_write: i32, pub patch_strategy: i32, pub session_capturing_disabled: i8, - pub max_failover_retry_count: i32, - pub max_session_retry_count: i32, + pub max_failover_retry_count: i64, + pub max_session_retry_count: i64, pub end_to_end_timeout_ms: i64, pub endpoint_unavailability_ttl_ms: i64, pub excluded_regions: *const crate::string::CosmosStringView, @@ -518,6 +518,13 @@ pub mod op_request { pub binary_encoding_enabled: i8, pub binary_encoding_request_text_response: i8, pub query_plan_mode: i32, + pub throughput_bucket: i64, + pub priority_level: i32, + pub max_throttle_retry_count: i64, + pub max_throttle_retry_wait_time_ms: i64, + pub hedging_enabled: i8, + pub availability_strategy: i32, + pub hedge_threshold_ms: i64, } #[repr(C)] pub struct CosmosOperationRequest { @@ -547,6 +554,7 @@ pub mod op_request { pub patch_tracking_id: crate::string::CosmosStringView, pub patch_tracking_capacity: u16, pub patch_tracking_retention_seconds: u32, + pub options_snapshot: *const crate::options_snapshot::OperationOptionsSnapshot, } pub struct OptOf; #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -618,6 +626,17 @@ pub mod op_request { CosmosReadConsistencyStrategyLatestCommitted = 5, } } +pub mod options_snapshot { + #[no_mangle] + pub unsafe extern "C" fn cosmos_operation_options_snapshot_create(runtime: *const crate::runtime::RuntimeContext, client_options: *const crate::op_request::CosmosOperationOptions, request_options: *const crate::op_request::CosmosOperationOptions, out_snapshot: *mut *mut OperationOptionsSnapshot, out_timeout_ms: *mut i64) -> crate::error::CosmosStatusCode; + #[no_mangle] + pub unsafe extern "C" fn cosmos_operation_options_snapshot_free(snapshot: *mut OperationOptionsSnapshot); + #[no_mangle] + pub unsafe extern "C" fn cosmos_runtime_set_operation_options(runtime: *const crate::runtime::RuntimeContext, options: *const crate::op_request::CosmosOperationOptions) -> crate::error::CosmosStatusCode; + #[derive(Clone)] + pub struct OperationOptionsSnapshot { + } +} pub mod partition_key { pub use azurecosmosdriver::string::CosmosStringView; #[no_mangle] @@ -736,6 +755,7 @@ pub mod runtime_builder { pub user_agent_suffix: crate::string::CosmosStringView, pub wrapping_sdk_identifier: crate::string::CosmosStringView, pub cpu_refresh_interval_ms: u64, + pub operation_options: *const crate::op_request::CosmosOperationOptions, } } #[attr = MacroUse {arguments:UseAll}] diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/build.rs b/sdk/cosmos/azure_data_cosmos_driver_native/build.rs index 2048be2b3cd..01ed1660c62 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/build.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/build.rs @@ -116,6 +116,10 @@ fn generate_c_header() { "cursor_completion_backing_t".into(), ), ("RuntimeContext".into(), "runtime_t".into()), + ( + "OperationOptionsSnapshot".into(), + "operation_options_snapshot_t".into(), + ), ("CosmosDriver".into(), "driver_t".into()), ("DriverHandle".into(), "driver_t".into()), ("AccountReference".into(), "account_ref_t".into()), diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/account_and_driver_options.c b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/account_and_driver_options.c index 458b52e372d..0a37e60598b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/account_and_driver_options.c +++ b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/account_and_driver_options.c @@ -15,9 +15,8 @@ // // `cosmos_driver_get_or_create_blocking` is not exercised here because // it touches the network. The emulator-backed scenario (stand up a driver, -// free it, recreate it, observe the cached instance, observe the -// OPTIONS_IGNORED_ON_CACHE_HIT advisory) is intentionally deferred to a -// CI-side integration test once the advisory itself lands. The Rust-side +// free it, recreate it, and verify independent credentials/defaults) belongs +// in the network integration suite. The Rust-side // integration test `blocking_against_invalid_endpoint` (marked // `#[ignore]`) exercises the failure path manually. @@ -154,6 +153,14 @@ static int test_driver_options_build_happy_path(void) { cfg.preferred_regions_len = 2; cosmos_operation_options_t operation_options = cosmos_operation_options_default(); operation_options.query_plan_mode = COSMOS_QUERY_PLAN_MODE_GATEWAY_ONLY; + operation_options.max_failover_retry_count = UINT32_MAX; + operation_options.max_session_retry_count = UINT32_MAX; + operation_options.throughput_bucket = UINT32_MAX; + operation_options.priority_level = 2; + operation_options.max_throttle_retry_count = 0; + operation_options.max_throttle_retry_wait_time_ms = 0; + operation_options.hedging_enabled = 1; + operation_options.availability_strategy = 1; cfg.operation_options = &operation_options; cosmos_driver_options_t *opts = NULL; diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/counted_strings.c b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/counted_strings.c index 471085d69b4..bf97c56e05a 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/counted_strings.c +++ b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/counted_strings.c @@ -8,9 +8,9 @@ #if UINTPTR_MAX == UINT64_MAX _Static_assert(sizeof(cosmos_string_view_t) == 16, "counted view layout"); _Static_assert(sizeof(cosmos_header_kv_t) == 32, "header pair layout"); -_Static_assert(sizeof(cosmos_runtime_options_t) == 64, "runtime layout"); -_Static_assert(sizeof(cosmos_operation_options_t) == 80, "operation options layout"); -_Static_assert(sizeof(cosmos_operation_request_t) == 224, "request layout"); +_Static_assert(sizeof(cosmos_runtime_options_t) == 72, "runtime layout"); +_Static_assert(sizeof(cosmos_operation_options_t) == 136, "operation options layout"); +_Static_assert(sizeof(cosmos_operation_request_t) == 232, "request layout"); _Static_assert(offsetof(cosmos_operation_request_t, item_id) == 32, "item view offset"); _Static_assert(offsetof(cosmos_operation_request_t, session_token) == 112, "session view offset"); _Static_assert(offsetof(cosmos_operation_request_t, patch_tracking_id) == 200, "tracking view offset"); diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/operation_construction.c b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/operation_construction.c index 19a9391f973..1899566f7fe 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/operation_construction.c +++ b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/operation_construction.c @@ -15,6 +15,16 @@ // 3. `out_pre_error` is optional (NULL is accepted). #include "test_common.h" +#include + +_Static_assert(sizeof(((cosmos_operation_options_t *)0)->max_failover_retry_count) == 8, + "ABI 0.2 carries the complete uint32 retry domain"); +#if UINTPTR_MAX == UINT64_MAX +_Static_assert(sizeof(cosmos_operation_options_t) == 136, "ABI 0.2 options layout"); +_Static_assert(sizeof(cosmos_operation_request_t) == 232, "ABI 0.2 request layout"); +_Static_assert(sizeof(cosmos_runtime_options_t) == 72, "ABI 0.2 runtime layout"); +_Static_assert(offsetof(cosmos_operation_request_t, options_snapshot) == 224, "snapshot offset"); +#endif // ───────────────────────────────────────────────────────────────────── // Section 1 — cosmos_operation_options_default sentinels @@ -24,6 +34,14 @@ static int test_options_default_is_all_unset(void) { int result = TEST_PASS; cosmos_operation_options_t opts = cosmos_operation_options_default(); + ASSERT(strcmp(cosmos_version(), AZURECOSMOSDRIVER_H_VERSION) == 0, + "loaded library matches generated header"); + ASSERT(strcmp(AZURECOSMOSDRIVER_H_VERSION, "0.2.0") == 0, "expected breaking ABI version"); + ASSERT(opts.throughput_bucket == -1 && opts.priority_level == 0, "throughput inherits"); + ASSERT(opts.max_throttle_retry_count == -1 && + opts.max_throttle_retry_wait_time_ms == -1, "throttling inherits"); + ASSERT(opts.hedging_enabled == 0 && opts.availability_strategy == 0 && + opts.hedge_threshold_ms == -1, "availability inherits"); ASSERT(COSMOS_READ_CONSISTENCY_STRATEGY_GLOBAL_STRONG == 4, "global strong ABI value (=%d)", @@ -46,9 +64,9 @@ static int test_options_default_is_all_unset(void) ASSERT(opts.session_capturing_disabled == 0, "session_capturing unset (=%d)", opts.session_capturing_disabled); ASSERT(opts.max_failover_retry_count < 0, - "max_failover unset (=%d)", opts.max_failover_retry_count); + "max_failover unset (=%lld)", (long long)opts.max_failover_retry_count); ASSERT(opts.max_session_retry_count < 0, - "max_session unset (=%d)", opts.max_session_retry_count); + "max_session unset (=%lld)", (long long)opts.max_session_retry_count); ASSERT(opts.end_to_end_timeout_ms < 0, "e2e timeout unset (=%lld)", (long long)opts.end_to_end_timeout_ms); ASSERT(opts.endpoint_unavailability_ttl_ms < 0, diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/runtime_lifecycle.c b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/runtime_lifecycle.c index 19ee99d78e1..88610216015 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/runtime_lifecycle.c +++ b/sdk/cosmos/azure_data_cosmos_driver_native/c_tests/runtime_lifecycle.c @@ -172,6 +172,49 @@ static int test_build_null_options_uses_defaults(void) return result; } +static int test_operation_defaults_and_atomic_updates(void) +{ + int result = TEST_PASS; + cosmos_runtime_t *runtime = NULL; + cosmos_error_t *error = NULL; + cosmos_operation_options_snapshot_t *snapshot = NULL; + int64_t timeout_ms = -2; + cosmos_operation_options_t defaults = cosmos_operation_options_default(); + defaults.end_to_end_timeout_ms = 2500; + cosmos_runtime_options_t runtime_options = cosmos_runtime_options_default(); + runtime_options.operation_options = &defaults; + REQUIRE(cosmos_runtime_build(&runtime_options, &runtime, &error) == COSMOS_STATUS_SUCCESS, + "runtime accepts operation defaults"); + defaults.end_to_end_timeout_ms = 5000; + REQUIRE(cosmos_operation_options_snapshot_create(runtime, NULL, NULL, &snapshot, &timeout_ms) == COSMOS_STATUS_SUCCESS, + "capture inherited defaults"); + ASSERT(timeout_ms == 2500, "construction copied inputs"); + cosmos_operation_options_snapshot_free(snapshot); + snapshot = NULL; + REQUIRE(cosmos_runtime_set_operation_options(runtime, &defaults) == COSMOS_STATUS_SUCCESS, + "atomic replacement succeeds"); + defaults.max_session_retry_count = (int64_t)UINT32_MAX + 1; + defaults.end_to_end_timeout_ms = 9000; + ASSERT(COSMOS_STATUS_SUB(cosmos_runtime_set_operation_options(runtime, &defaults)) == COSMOS_SUB_STATUS_CLIENT_FFI_INVALID_OPTION_VALUE, + "invalid update rejected"); + REQUIRE(cosmos_operation_options_snapshot_create(runtime, NULL, NULL, &snapshot, &timeout_ms) == COSMOS_STATUS_SUCCESS, + "capture after rejected update"); + ASSERT(timeout_ms == 5000, "invalid replacement preserves previous defaults"); + cosmos_operation_options_snapshot_free(snapshot); + snapshot = NULL; + REQUIRE(cosmos_runtime_set_operation_options(runtime, NULL) == COSMOS_STATUS_SUCCESS, + "NULL clears runtime defaults"); + REQUIRE(cosmos_operation_options_snapshot_create(runtime, NULL, NULL, &snapshot, &timeout_ms) == COSMOS_STATUS_SUCCESS, + "capture cleared defaults"); + ASSERT(timeout_ms == -1, "cleared latency inherits no timeout"); + +cleanup: + cosmos_operation_options_snapshot_free(snapshot); + cosmos_runtime_free(runtime); + cosmos_error_free(error); + return result; +} + TEST_SUITE_BEGIN("Runtime Construction & Lifecycle") TEST_REGISTER(options_default_all_unset) TEST_REGISTER(build_rejects_null_out_runtime) @@ -180,4 +223,5 @@ TEST_REGISTER(string_field_validation) TEST_REGISTER(cpu_refresh_interval_range) TEST_REGISTER(build_happy_path) TEST_REGISTER(build_null_options_uses_defaults) +TEST_REGISTER(operation_defaults_and_atomic_updates) TEST_SUITE_END("Runtime Construction & Lifecycle") diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h index 8b5a0f26459..7bf5c703a00 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h +++ b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h @@ -11,7 +11,7 @@ // Specifies the version of azurecosmosdriver this header file was generated from. // This should match the version of libazurecosmosdriver you are linking against. -#define AZURECOSMOSDRIVER_H_VERSION "0.1.0" +#define AZURECOSMOSDRIVER_H_VERSION "0.2.0" // Packed-status helpers (see cosmos_status_code_t). Emitted as macros so // they keep the SCREAMING_SNAKE_CASE spelling shared with the @@ -1030,6 +1030,11 @@ typedef struct cosmos_feed_range_t cosmos_feed_range_t; */ typedef struct cosmos_operation_handle_t cosmos_operation_handle_t; +/** + * Opaque admission snapshot. Owned by the host until freed; submits clone it. + */ +typedef struct cosmos_operation_options_snapshot_t cosmos_operation_options_snapshot_t; + /** * The C ABI handle for an immutable partition key (`cosmos_partition_key_t`). * @@ -1052,7 +1057,7 @@ typedef struct cosmos_partition_key_t cosmos_partition_key_t; * `block_on(...)` driver builder construction at FFI-call time and to * spawn the per-operation tasks that drive submits. * - `driver` — the underlying `azure_data_cosmos_driver` runtime that owns - * the per-account driver registry, container cache, account-metadata + * the account-metadata * cache, HTTP transport factory, and so on. Cloning the `Arc` is cheap * and is how the driver / account surfaces hand out handles. */ @@ -1548,11 +1553,11 @@ typedef struct cosmos_operation_options_t { /** * Max region-failover retries. `< 0` = unset. */ - int32_t max_failover_retry_count; + int64_t max_failover_retry_count; /** * Max session-consistency retries on 404/1002. `< 0` = unset. */ - int32_t max_session_retry_count; + int64_t max_session_retry_count; /** * End-to-end timeout (milliseconds). `< 0` = unset. */ @@ -1563,7 +1568,7 @@ typedef struct cosmos_operation_options_t { int64_t endpoint_unavailability_ttl_ms; /** * Excluded regions — array of counted UTF-8 region ids. - * NULL / `0` length = unset; non-NULL with `0` length is rejected. + * NULL / `0` length = unset; non-NULL with `0` length clears exclusions. */ const struct cosmos_string_view_t *excluded_regions; /** @@ -1630,6 +1635,36 @@ typedef struct cosmos_operation_options_t { * values to be rejected before materializing the enum. */ int32_t query_plan_mode; + /** + * Throughput bucket. `< 0` inherits; otherwise must fit `u32`. + */ + int64_t throughput_bucket; + /** + * Priority level: `0` inherits, `1` High, `2` Low. + */ + int32_t priority_level; + /** + * Throttling retry count. `< 0` inherits; otherwise must fit `u32`. + */ + int64_t max_throttle_retry_count; + /** + * Cumulative throttle retry wait per transport invocation, in milliseconds. + * `< 0` inherits; `0` explicitly disables waiting. + */ + int64_t max_throttle_retry_wait_time_ms; + /** + * Hedging master switch: `0` inherits, `1` false, `2` true. + */ + int8_t hedging_enabled; + /** + * Availability strategy: `0` inherits, `1` disabled, `2` hedging. + */ + int32_t availability_strategy; + /** + * Positive hedge threshold in milliseconds for strategy `2`. + * Must be negative (unset) for any other strategy. + */ + int64_t hedge_threshold_ms; } cosmos_operation_options_t; /** @@ -1774,6 +1809,11 @@ typedef struct cosmos_operation_request_t { * evict an entry earlier. `0` = use the driver default. */ uint32_t patch_tracking_retention_seconds; + /** + * Optional admission snapshot. Overrides the shared fields in `options`; + * `options.query_plan_mode` remains per submit. Borrowed only until submit returns. + */ + const struct cosmos_operation_options_snapshot_t *options_snapshot; } cosmos_operation_request_t; /** @@ -2072,6 +2112,10 @@ typedef struct cosmos_runtime_options_t { * `1000`–`60000`). `0` = unset. */ uint64_t cpu_refresh_interval_ms; + /** + * Runtime operation defaults, copied during construction. NULL inherits. + */ + const struct cosmos_operation_options_t *operation_options; } cosmos_runtime_options_t; #ifdef __cplusplus @@ -2571,9 +2615,8 @@ cosmos_status_code_t cosmos_diagnostics_to_json(const struct cosmos_diagnostics_ uintptr_t *out_len); /** - * Frees a driver handle. Drops the FFI-side `Arc` reference; the - * underlying driver remains alive in the runtime's cache until the - * owning `cosmos_runtime_t` is freed (spec section 4.4.1). NULL is a no-op. + * Frees a driver handle. In-flight operations retain their own references. + * Other drivers on the same runtime are unaffected. NULL is a no-op. */ void cosmos_driver_free(struct cosmos_driver_t *driver); @@ -2581,23 +2624,14 @@ void cosmos_driver_free(struct cosmos_driver_t *driver); * Synchronously gets or creates the driver for the supplied account. * * Bridges - * `CosmosDriverRuntime::get_or_create_driver` through the wrapper's + * `CosmosDriverRuntime::create_driver` through the wrapper's * own multi-threaded Tokio runtime via `block_on`. Suitable for * startup-time initialization; for runtime use prefer the async * `_submit` variant. * - * # Cache behavior (spec section 4.4.1) - * - * - The runtime caches drivers by endpoint URL. A second call with the - * same endpoint returns the cached driver and **silently ignores** - * `options`. - * - Two `AccountReference`s with the same endpoint but different - * credentials collide in the cache — first credential wins. - * - Cache eviction happens only when the owning `cosmos_runtime_t` is - * freed; freeing a `cosmos_driver_t` does not evict. - * - * The cache-hit advisory described in spec Section 4.4.1 is not emitted - * today — see the module-level `Cache-hit advisory` note for the rationale. + * Each call creates a fresh driver with its own account options. Cached container + * references can reuse credentials from another same-account driver on this runtime. + * No cache-hit advisory is emitted. * * # Parameters * @@ -2754,6 +2788,49 @@ void cosmos_feed_range_free(struct cosmos_feed_range_t *fr); */ struct cosmos_operation_options_t cosmos_operation_options_default(void); +/** + * Atomically replaces runtime operation defaults. NULL resets all defaults. + * + * All pointers must remain valid for this call. Input arrays and strings are + * copied; invalid options leave the previous defaults unchanged. + * + * # Safety + * + * The runtime and any non-NULL options and borrowed arrays must be live. + */ +cosmos_status_code_t cosmos_runtime_set_operation_options(const struct cosmos_runtime_t *runtime, + const struct cosmos_operation_options_t *options); + +/** + * Captures all configuration layers and starts the logical-operation budget. + * + * `client_options` and `request_options` may be NULL to inherit. Outputs are + * required. On success `out_timeout_ms` is the resolved Rust timeout (including + * its minimum clamp), or -1 when unset. A snapshot keeps its runtime generation + * through initialization, retries and patch stages; it must be submitted to a + * driver from this runtime. Invalid inputs do not modify outputs. + * + * # Safety + * + * Handles, option structs and their arrays must be valid for this call; output + * pointers must be writable. The host owns the returned snapshot. + */ +cosmos_status_code_t cosmos_operation_options_snapshot_create(const struct cosmos_runtime_t *runtime, + const struct cosmos_operation_options_t *client_options, + const struct cosmos_operation_options_t *request_options, + struct cosmos_operation_options_snapshot_t **out_snapshot, + int64_t *out_timeout_ms); + +/** + * Frees a snapshot. NULL is a no-op; already-submitted operations retain a copy. + * + * # Safety + * + * The pointer must be NULL or a live snapshot returned by snapshot creation, + * not concurrently borrowed by another call. + */ +void cosmos_operation_options_snapshot_free(struct cosmos_operation_options_snapshot_t *snapshot); + /** * Creates an immutable partition key from an inline component array in a * single call — the flat, standalone counterpart to the diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/cursor.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/cursor.rs index 078b04624a0..6680dc1c394 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/cursor.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/cursor.rs @@ -11,6 +11,7 @@ use crate::{ cursor_request::{build_cursor_request, CosmosCursorRequest}, driver::DriverHandle, error::{CosmosErrorCode, CosmosStatusCode, COSMOS_STATUS_SUCCESS}, + options_snapshot::OperationOptionsSnapshot, runtime::RuntimeContext, safety::MutexExt, string::CosmosStringView, @@ -78,6 +79,7 @@ struct CursorInner { driver: Arc, container: Option, options: OperationOptions, + snapshot: Option, state: Mutex, } @@ -413,6 +415,13 @@ pub extern "C" fn cosmos_cursor_open_submit( .ok_or(CosmosErrorCode::CosmosErrorCodeInvalidArgument)?; // SAFETY: request follows the versioned prefix and counted input contracts. let built = unsafe { build_cursor_request(request)? }; + if built + .snapshot + .as_ref() + .is_some_and(|snapshot| !snapshot.matches_driver(&driver.inner)) + { + return Err(CosmosErrorCode::CosmosErrorCodeInvalidArgument); + } cursor_queue.admit(None)?; Ok::<_, CosmosErrorCode>((queue, driver, built)) }; @@ -431,6 +440,7 @@ pub extern "C" fn cosmos_cursor_open_submit( let op = unsafe { Arc::clone(&(*handle).inner) }; let runtime = Arc::clone(queue.runtime()); runtime.tokio.spawn(async move { + let snapshot = built.snapshot.clone(); let work = async { let container = built.operation.container().cloned(); let plan = driver @@ -447,6 +457,7 @@ pub extern "C" fn cosmos_cursor_open_submit( driver: Arc::clone(&driver.inner), container, options: built.options, + snapshot: built.snapshot, state: Mutex::new(CursorState { plan: Some(plan), busy: false, @@ -455,7 +466,7 @@ pub extern "C" fn cosmos_cursor_open_submit( }), }))) }; - let result = run(work).await; + let result = run(OperationOptionsSnapshot::execute(snapshot, work)).await; if let Some(cursor_queue) = &queue.cursor { cursor_queue.publish(Delivery { result, @@ -531,11 +542,14 @@ fn submit_cursor( if cursor.state.lock_recover().exhausted { return Ok(ResultData::End); } - cursor - .driver - .execute_plan(plan, cursor.container.clone(), cursor.options.clone()) - .await - .map(|page| page.map_or(ResultData::End, |page| ResultData::Page(Box::new(page)))) + OperationOptionsSnapshot::execute( + cursor.snapshot.clone(), + cursor + .driver + .execute_plan(plan, cursor.container.clone(), cursor.options.clone()), + ) + .await + .map(|page| page.map_or(ResultData::End, |page| ResultData::Page(Box::new(page)))) }; let result = run(work).await; { diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/driver.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/driver.rs index d452a32db4f..d84747ac4a0 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/driver.rs @@ -11,16 +11,8 @@ //! wrapper's generic `tokio::spawn` → `cq_enqueue` plumbing lands once, //! with the operation submit pipeline. //! -//! ## Cache-hit advisory -//! -//! Spec section 4.4.1 describes an optional advisory for when the driver -//! returns a cached driver for an endpoint that already has an entry. The -//! merged `CosmosDriverRuntime::get_or_create_driver` API does not expose a -//! "was cached" signal, so detecting cache hits requires either a -//! driver-side enhancement (preferred) or wrapper-side cache shadowing -//! (hacky). The advisory is intentionally **not** implemented today -//! — `cosmos_driver_get_or_create_blocking` always returns `SUCCESS` on -//! a cached hit. This is tracked as a follow-up. +//! Every construction creates a fresh driver with its own credentials/options. +//! Runtime-owned container references retain their original credentials and are shared. //! use std::sync::Arc; @@ -86,9 +78,8 @@ impl DriverHandle { // FFI: lifecycle // ───────────────────────────────────────────────────────────────────────────── -/// Frees a driver handle. Drops the FFI-side `Arc` reference; the -/// underlying driver remains alive in the runtime's cache until the -/// owning `cosmos_runtime_t` is freed (spec section 4.4.1). NULL is a no-op. +/// Frees a driver handle. In-flight operations retain their own references. +/// Other drivers on the same runtime are unaffected. NULL is a no-op. #[no_mangle] pub extern "C" fn cosmos_driver_free(driver: *mut DriverHandle) { if driver.is_null() { @@ -105,23 +96,14 @@ pub extern "C" fn cosmos_driver_free(driver: *mut DriverHandle) { /// Synchronously gets or creates the driver for the supplied account. /// /// Bridges -/// `CosmosDriverRuntime::get_or_create_driver` through the wrapper's +/// `CosmosDriverRuntime::create_driver` through the wrapper's /// own multi-threaded Tokio runtime via `block_on`. Suitable for /// startup-time initialization; for runtime use prefer the async /// `_submit` variant. /// -/// # Cache behavior (spec section 4.4.1) -/// -/// - The runtime caches drivers by endpoint URL. A second call with the -/// same endpoint returns the cached driver and **silently ignores** -/// `options`. -/// - Two `AccountReference`s with the same endpoint but different -/// credentials collide in the cache — first credential wins. -/// - Cache eviction happens only when the owning `cosmos_runtime_t` is -/// freed; freeing a `cosmos_driver_t` does not evict. -/// -/// The cache-hit advisory described in spec Section 4.4.1 is not emitted -/// today — see the module-level `Cache-hit advisory` note for the rationale. +/// Each call creates a fresh driver with its own account options. Cached container +/// references can reuse credentials from another same-account driver on this runtime. +/// No cache-hit advisory is emitted. /// /// # Parameters /// diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs index 22cd4914778..8415d4aa4d7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs @@ -33,6 +33,7 @@ pub mod error; pub mod fault_injection; pub mod feed_range; pub mod op_request; +pub mod options_snapshot; pub mod partition_key; pub mod response_header; pub mod runtime; diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs index 60f670f37e7..2b8cd8fe8c3 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs @@ -41,8 +41,10 @@ use std::num::{NonZeroU16, NonZeroU32, NonZeroU8}; use azure_core::http::headers::{HeaderName, HeaderValue}; use azure_core::http::Etag; use azure_data_cosmos_driver::options::{ - BinaryEncodingOptions, ContentResponseOnWrite, EndToEndOperationLatencyPolicy, ExcludedRegions, - OperationOptions, PatchStrategy, QueryPlanMode, ReadConsistencyStrategy, Region, + AvailabilityStrategy, BinaryEncodingOptions, ContentResponseOnWrite, + EndToEndOperationLatencyPolicy, ExcludedRegions, HedgeThreshold, HedgingStrategy, + OperationOptions, PatchStrategy, PriorityLevel, QueryPlanMode, ReadConsistencyStrategy, Region, + ThrottlingRetryOptions, ThroughputControlOptions, }; use azure_data_cosmos_driver::{ models::{ @@ -58,6 +60,7 @@ use crate::container_ref::ContainerRefHandle; use crate::database_ref::DatabaseRefHandle; use crate::error::CosmosErrorCode; use crate::feed_range::FeedRangeHandle; +use crate::options_snapshot::OperationOptionsSnapshot; use crate::partition_key::{CosmosPartitionKeyComponent, PartitionKeyHandle}; // ───────────────────────────────────────────────────────────────────────────── @@ -69,13 +72,14 @@ use crate::partition_key::{CosmosPartitionKeyComponent, PartitionKeyHandle}; /// /// - **enum fields** (`*_strategy`, `content_response_on_write`): `0` = unset /// (inherit), any other value = the corresponding driver variant. -/// - **tri-state bools** (`session_capturing_disabled`): `0` = unset, -/// `1` = `false`, `2` = `true`. -/// - **i32 numeric fields** (retry counters): `< 0` = unset, -/// `>= 0` = the value. +/// - **tri-state bools** (`session_capturing_disabled`, `hedging_enabled`, +/// and binary-encoding flags): `0` = unset, `1` = `false`, `2` = `true`. +/// - **i64 numeric fields** (retry counters/bucket): `< 0` = unset, +/// `0..=u32::MAX` = the value; larger values are rejected. /// - **i64 duration fields** (`*_ms`): `< 0` = unset, `>= 0` = milliseconds. -/// - **string / array fields** (`excluded_regions`, `custom_headers`): -/// NULL / length `0` = unset. +/// - **array fields** (`excluded_regions`, `custom_headers`): NULL with length +/// `0` = unset (inherit); non-NULL with length `0` clears inherited values. +/// NULL with a nonzero length is invalid. /// /// It is a documentation marker only — the fields are plain integers / /// pointers so the struct stays `#[repr(C)]`. @@ -96,12 +100,14 @@ fn decode_tristate_bool(v: i8) -> Result, CosmosErrorCode> { } } -/// Decodes an `i32` numeric option (`< 0` = unset) into `Option`. -fn decode_opt_u32(v: i32) -> Option { +/// Decodes a full-width unsigned option with a signed unset sentinel. +fn decode_opt_u32(v: i64) -> Result, CosmosErrorCode> { if v < 0 { - None + Ok(None) } else { - Some(v as u32) + u32::try_from(v) + .map(Some) + .map_err(|_| CosmosErrorCode::CosmosErrorCodeInvalidOptionValue) } } @@ -415,15 +421,15 @@ pub struct CosmosOperationOptions { /// Disable automatic session token management. Tri-state bool. pub session_capturing_disabled: i8, /// Max region-failover retries. `< 0` = unset. - pub max_failover_retry_count: i32, + pub max_failover_retry_count: i64, /// Max session-consistency retries on 404/1002. `< 0` = unset. - pub max_session_retry_count: i32, + pub max_session_retry_count: i64, /// End-to-end timeout (milliseconds). `< 0` = unset. pub end_to_end_timeout_ms: i64, /// Endpoint unavailability TTL (milliseconds). `< 0` = unset. pub endpoint_unavailability_ttl_ms: i64, /// Excluded regions — array of counted UTF-8 region ids. - /// NULL / `0` length = unset; non-NULL with `0` length is rejected. + /// NULL / `0` length = unset; non-NULL with `0` length clears exclusions. pub excluded_regions: *const CosmosStringView, /// Number of entries in `excluded_regions`. pub excluded_regions_len: usize, @@ -478,6 +484,22 @@ pub struct CosmosOperationOptions { /// `0` (`Unset`) uses LocalPreferred. Raw `i32` storage allows invalid host /// values to be rejected before materializing the enum. pub query_plan_mode: i32, + /// Throughput bucket. `< 0` inherits; otherwise must fit `u32`. + pub throughput_bucket: i64, + /// Priority level: `0` inherits, `1` High, `2` Low. + pub priority_level: i32, + /// Throttling retry count. `< 0` inherits; otherwise must fit `u32`. + pub max_throttle_retry_count: i64, + /// Cumulative throttle retry wait per transport invocation, in milliseconds. + /// `< 0` inherits; `0` explicitly disables waiting. + pub max_throttle_retry_wait_time_ms: i64, + /// Hedging master switch: `0` inherits, `1` false, `2` true. + pub hedging_enabled: i8, + /// Availability strategy: `0` inherits, `1` disabled, `2` hedging. + pub availability_strategy: i32, + /// Positive hedge threshold in milliseconds for strategy `2`. + /// Must be negative (unset) for any other strategy. + pub hedge_threshold_ms: i64, } impl CosmosOperationOptions { @@ -506,8 +528,43 @@ impl CosmosOperationOptions { CosmosQueryPlanMode::from_i32(self.query_plan_mode)?; opts.session_capturing_disabled = decode_tristate_bool(self.session_capturing_disabled)?; - opts.max_failover_retry_count = decode_opt_u32(self.max_failover_retry_count); - opts.max_session_retry_count = decode_opt_u32(self.max_session_retry_count); + opts.max_failover_retry_count = decode_opt_u32(self.max_failover_retry_count)?; + opts.max_session_retry_count = decode_opt_u32(self.max_session_retry_count)?; + let mut throughput = ThroughputControlOptions::default(); + throughput.throughput_bucket = decode_opt_u32(self.throughput_bucket)?; + throughput.priority_level = match self.priority_level { + 0 => None, + 1 => Some(PriorityLevel::High), + 2 => Some(PriorityLevel::Low), + _ => return Err(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue), + }; + if throughput.throughput_bucket.is_some() || throughput.priority_level.is_some() { + opts.throughput_control = Some(throughput); + } + let mut throttling = ThrottlingRetryOptions::default(); + throttling.max_retry_count = decode_opt_u32(self.max_throttle_retry_count)?; + if self.max_throttle_retry_wait_time_ms >= 0 { + throttling.max_retry_wait_time = Some(std::time::Duration::from_millis( + self.max_throttle_retry_wait_time_ms as u64, + )); + } + if throttling.max_retry_count.is_some() || throttling.max_retry_wait_time.is_some() { + opts.throttling_retry_options = Some(throttling); + } + opts.hedging_enabled = decode_tristate_bool(self.hedging_enabled)?; + opts.availability_strategy = match (self.availability_strategy, self.hedge_threshold_ms) { + (0, threshold) if threshold < 0 => None, + (1, threshold) if threshold < 0 => Some(AvailabilityStrategy::Disabled), + (2, threshold) if threshold > 0 => { + let threshold = + HedgeThreshold::new(std::time::Duration::from_millis(threshold as u64)) + .ok_or(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue)?; + Some(AvailabilityStrategy::Hedging(HedgingStrategy::new( + threshold, + ))) + } + _ => return Err(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue), + }; if self.end_to_end_timeout_ms >= 0 { let dur = std::time::Duration::from_millis(self.end_to_end_timeout_ms as u64); @@ -554,8 +611,8 @@ impl CosmosOperationOptions { } /// Decodes a `(ptr, len)` region-id array into an [`ExcludedRegions`]. -/// NULL / `0` length yields `None`. A non-NULL pointer with `0` length is -/// rejected as a malformed input. +/// NULL / `0` length yields `None`. A non-NULL pointer with `0` length clears +/// inherited exclusions. /// /// # Safety /// @@ -570,9 +627,7 @@ unsafe fn decode_regions( return Ok(None); } if len == 0 { - // Non-NULL pointer with zero length is ambiguous; reject it so the - // host uses NULL to mean "unset" unambiguously. - return Err(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue); + return Ok(Some(ExcludedRegions::new())); } // SAFETY: caller contract above. let slice = unsafe { std::slice::from_raw_parts(regions, len) }; @@ -608,6 +663,13 @@ pub extern "C" fn cosmos_operation_options_default() -> CosmosOperationOptions { custom_headers_len: 0, binary_encoding_enabled: TRISTATE_UNSET, binary_encoding_request_text_response: TRISTATE_UNSET, + throughput_bucket: -1, + priority_level: 0, + max_throttle_retry_count: -1, + max_throttle_retry_wait_time_ms: -1, + hedging_enabled: TRISTATE_UNSET, + availability_strategy: 0, + hedge_threshold_ms: -1, } } @@ -854,6 +916,9 @@ pub struct CosmosOperationRequest { /// Age-based retention window in whole seconds. Capacity pressure can /// evict an entry earlier. `0` = use the driver default. pub patch_tracking_retention_seconds: u32, + /// Optional admission snapshot. Overrides the shared fields in `options`; + /// `options.query_plan_mode` remains per submit. Borrowed only until submit returns. + pub options_snapshot: *const OperationOptionsSnapshot, } // ───────────────────────────────────────────────────────────────────────────── @@ -862,6 +927,7 @@ pub struct CosmosOperationRequest { /// Holds the fully-built driver inputs ready to hand to a driver method. pub(crate) struct BuiltRequest { + pub(crate) snapshot: Option, pub(crate) operation: CosmosOperation, pub(crate) options: OperationOptions, pub(crate) patch_tracking_id: Option, @@ -918,7 +984,11 @@ pub(crate) unsafe fn build_request_with_operation( }?; let (operation, patch_tracking_id) = resolve_patch_tracking_id(operation); - let options = if req.options.is_null() { + // SAFETY: a non-NULL snapshot is a live handle borrowed for this call. + let snapshot = unsafe { req.options_snapshot.as_ref() }.cloned(); + let options = if let Some(snapshot) = &snapshot { + snapshot.options.clone() + } else if req.options.is_null() { OperationOptions::default() } else { // SAFETY: non-NULL checked; caller guarantees a valid struct. @@ -944,6 +1014,7 @@ pub(crate) unsafe fn build_request_with_operation( }; Ok(BuiltRequest { + snapshot, operation, options, patch_tracking_id, @@ -1416,7 +1487,8 @@ mod tests { fn operation_request_abi_layout_is_stable() { use std::mem::{offset_of, size_of}; - assert_eq!(size_of::(), 224); + assert_eq!(size_of::(), 232); + assert_eq!(offset_of!(CosmosOperationRequest, options_snapshot), 224); assert_eq!(offset_of!(CosmosOperationRequest, kind), 0); assert_eq!(offset_of!(CosmosOperationRequest, account), 8); assert_eq!(offset_of!(CosmosOperationRequest, database), 16); @@ -1532,11 +1604,73 @@ mod tests { #[test] fn opt_u32_treats_negative_as_unset() { - assert_eq!(decode_opt_u32(-1), None); - assert_eq!(decode_opt_u32(i32::MIN), None); - assert_eq!(decode_opt_u32(0), Some(0)); - assert_eq!(decode_opt_u32(7), Some(7)); - assert_eq!(decode_opt_u32(i32::MAX), Some(i32::MAX as u32)); + assert_eq!(decode_opt_u32(-1), Ok(None)); + assert_eq!(decode_opt_u32(i64::MIN), Ok(None)); + assert_eq!(decode_opt_u32(0), Ok(Some(0))); + assert_eq!(decode_opt_u32(7), Ok(Some(7))); + assert_eq!(decode_opt_u32(i64::from(u32::MAX)), Ok(Some(u32::MAX))); + assert_eq!( + decode_opt_u32(i64::from(u32::MAX) + 1), + Err(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue) + ); + } + + #[test] + fn flat_groups_preserve_zero_and_validate_discriminants() { + let mut options = cosmos_operation_options_default(); + options.throughput_bucket = 0; + options.max_throttle_retry_count = i64::from(u32::MAX); + options.max_throttle_retry_wait_time_ms = 0; + options.max_failover_retry_count = i64::from(u32::MAX); + options.max_session_retry_count = i64::from(u32::MAX); + options.hedging_enabled = 1; + options.availability_strategy = 2; + options.hedge_threshold_ms = 1; + // SAFETY: default option arrays are NULL/0. + let decoded = unsafe { options.to_driver() }.unwrap(); + assert_eq!( + decoded.throughput_control.unwrap().throughput_bucket, + Some(0) + ); + assert_eq!( + decoded + .throttling_retry_options + .as_ref() + .unwrap() + .max_retry_count, + Some(u32::MAX) + ); + assert_eq!( + decoded + .throttling_retry_options + .unwrap() + .max_retry_wait_time, + Some(std::time::Duration::ZERO) + ); + assert_eq!(decoded.max_failover_retry_count, Some(u32::MAX)); + assert_eq!(decoded.max_session_retry_count, Some(u32::MAX)); + assert_eq!(decoded.hedging_enabled, Some(false)); + for (kind, threshold) in [(0, 1), (1, 0), (2, -1), (2, 0), (3, -1)] { + options.availability_strategy = kind; + options.hedge_threshold_ms = threshold; + // SAFETY: arrays remain NULL/0. + assert_eq!( + unsafe { options.to_driver() }.unwrap_err(), + CosmosErrorCode::CosmosErrorCodeInvalidOptionValue + ); + } + } + + #[test] + fn explicit_empty_regions_clear_and_null_nonempty_is_rejected() { + let region = CosmosStringView::default(); + // SAFETY: zero length never dereferences this valid pointer. + assert_eq!( + unsafe { decode_regions(®ion, 0) }.unwrap(), + Some(ExcludedRegions::new()) + ); + // SAFETY: invalid NULL/1 is rejected before any dereference. + assert!(unsafe { decode_regions(std::ptr::null(), 1) }.is_err()); } #[test] @@ -1809,7 +1943,7 @@ mod tests { fn operation_options_abi_layout_is_stable() { use std::mem::{offset_of, size_of}; - assert_eq!(size_of::(), 80); + assert_eq!(size_of::(), 136); assert_eq!( offset_of!(CosmosOperationOptions, read_consistency_strategy), 0 @@ -1825,16 +1959,40 @@ mod tests { ); assert_eq!( offset_of!(CosmosOperationOptions, binary_encoding_enabled), - 72 + 80 ); assert_eq!( offset_of!( CosmosOperationOptions, binary_encoding_request_text_response ), - 73 + 81 + ); + assert_eq!(offset_of!(CosmosOperationOptions, query_plan_mode), 84); + assert_eq!( + offset_of!(CosmosOperationOptions, max_failover_retry_count), + 16 + ); + assert_eq!( + offset_of!(CosmosOperationOptions, max_session_retry_count), + 24 + ); + assert_eq!(offset_of!(CosmosOperationOptions, throughput_bucket), 88); + assert_eq!(offset_of!(CosmosOperationOptions, priority_level), 96); + assert_eq!( + offset_of!(CosmosOperationOptions, max_throttle_retry_count), + 104 + ); + assert_eq!( + offset_of!(CosmosOperationOptions, max_throttle_retry_wait_time_ms), + 112 + ); + assert_eq!(offset_of!(CosmosOperationOptions, hedging_enabled), 120); + assert_eq!( + offset_of!(CosmosOperationOptions, availability_strategy), + 124 ); - assert_eq!(offset_of!(CosmosOperationOptions, query_plan_mode), 76); + assert_eq!(offset_of!(CosmosOperationOptions, hedge_threshold_ms), 128); } #[test] diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot.rs new file mode 100644 index 00000000000..b932f724236 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot.rs @@ -0,0 +1,354 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Admission-time configuration and deadline snapshots for native hosts. + +#[cfg(test)] +mod submit_tests; + +use std::{ + future::Future, + sync::Arc, + time::{Duration, Instant}, +}; + +use azure_core::http::StatusCode; +use azure_data_cosmos_driver::{ + driver::{CosmosDriver, CosmosDriverRuntime}, + error::{status_codes::substatus::CLIENT_OPERATION_TIMEOUT, CosmosError as DriverError}, + models::CosmosStatus, + options::OperationOptions, +}; + +use crate::{ + error::{CosmosErrorCode, CosmosStatusCode}, + op_request::CosmosOperationOptions, + runtime::RuntimeContext, +}; + +/// Opaque admission snapshot. Owned by the host until freed; submits clone it. +#[derive(Clone)] +pub struct OperationOptionsSnapshot { + pub(crate) options: OperationOptions, + runtime: Arc, + started: Instant, + timeout: Option, +} + +impl OperationOptionsSnapshot { + pub(crate) fn matches_driver(&self, driver: &CosmosDriver) -> bool { + std::ptr::eq(self.runtime.as_ref(), driver.runtime()) + } + + pub(crate) async fn execute( + snapshot: Option, + work: impl Future>, + ) -> Result { + let Some(snapshot) = snapshot else { + return work.await; + }; + let Some(timeout) = snapshot.timeout else { + return work.await; + }; + let remaining = timeout.saturating_sub(snapshot.started.elapsed()); + if !remaining.is_zero() { + if let Ok(result) = tokio::time::timeout(remaining, work).await { + return result; + } + } + Err(DriverError::builder() + .with_status( + CosmosStatus::new(StatusCode::RequestTimeout) + .with_sub_status(CLIENT_OPERATION_TIMEOUT.value()), + ) + .with_message("end-to-end operation timeout exceeded after native admission") + .build()) + } +} + +/// Atomically replaces runtime operation defaults. NULL resets all defaults. +/// +/// All pointers must remain valid for this call. Input arrays and strings are +/// copied; invalid options leave the previous defaults unchanged. +/// +/// # Safety +/// +/// The runtime and any non-NULL options and borrowed arrays must be live. +#[no_mangle] +pub unsafe extern "C" fn cosmos_runtime_set_operation_options( + runtime: *const RuntimeContext, + options: *const CosmosOperationOptions, +) -> CosmosStatusCode { + let Some(runtime) = RuntimeContext::from_ptr(runtime) else { + return CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code(); + }; + // SAFETY: the host keeps options and their borrowed inputs alive for this call. + let options = match unsafe { decode_options(options) } { + Ok(options) => options, + Err(error) => return error.as_status_code(), + }; + runtime.driver.set_default_operation_options(options); + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() +} + +/// Captures all configuration layers and starts the logical-operation budget. +/// +/// `client_options` and `request_options` may be NULL to inherit. Outputs are +/// required. On success `out_timeout_ms` is the resolved Rust timeout (including +/// its minimum clamp), or -1 when unset. A snapshot keeps its runtime generation +/// through initialization, retries and patch stages; it must be submitted to a +/// driver from this runtime. Invalid inputs do not modify outputs. +/// +/// # Safety +/// +/// Handles, option structs and their arrays must be valid for this call; output +/// pointers must be writable. The host owns the returned snapshot. +#[no_mangle] +pub unsafe extern "C" fn cosmos_operation_options_snapshot_create( + runtime: *const RuntimeContext, + client_options: *const CosmosOperationOptions, + request_options: *const CosmosOperationOptions, + out_snapshot: *mut *mut OperationOptionsSnapshot, + out_timeout_ms: *mut i64, +) -> CosmosStatusCode { + let started = Instant::now(); + let Some(runtime) = RuntimeContext::from_ptr(runtime) else { + return CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code(); + }; + if out_snapshot.is_null() || out_timeout_ms.is_null() { + return CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code(); + } + // SAFETY: borrowed option inputs follow the documented host allocation contract. + let decoded = unsafe { + decode_options(client_options) + .and_then(|client| decode_options(request_options).map(|request| (client, request))) + }; + let (client, request) = match decoded { + Ok(value) => value, + Err(error) => return error.as_status_code(), + }; + let options = request.with_resolution_snapshot( + Arc::clone(runtime.driver.env_override_operation_options()), + Arc::clone(runtime.driver.env_operation_options()), + runtime.driver.default_operation_options(), + Arc::new(client), + ); + let timeout = options.resolution_snapshot_view().and_then(|view| { + view.end_to_end_latency_policy() + .map(|policy| policy.timeout()) + }); + let timeout_ms = match timeout { + Some(timeout) => match i64::try_from(timeout.as_millis()) { + Ok(timeout) => timeout, + Err(_) => return CosmosErrorCode::CosmosErrorCodeInvalidOptionValue.as_status_code(), + }, + None => -1, + }; + let snapshot = Box::new(OperationOptionsSnapshot { + options, + runtime: Arc::clone(&runtime.driver), + started, + timeout, + }); + // SAFETY: non-NULL outputs are writable according to the host contract. + unsafe { + *out_snapshot = Box::into_raw(snapshot); + *out_timeout_ms = timeout_ms; + } + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() +} + +/// Frees a snapshot. NULL is a no-op; already-submitted operations retain a copy. +/// +/// # Safety +/// +/// The pointer must be NULL or a live snapshot returned by snapshot creation, +/// not concurrently borrowed by another call. +#[no_mangle] +pub unsafe extern "C" fn cosmos_operation_options_snapshot_free( + snapshot: *mut OperationOptionsSnapshot, +) { + if !snapshot.is_null() { + // SAFETY: this call consumes the host's unique Box allocation. + drop(unsafe { Box::from_raw(snapshot) }); + } +} + +unsafe fn decode_options( + options: *const CosmosOperationOptions, +) -> Result { + // SAFETY: the caller guarantees validity of any non-NULL options and their inputs. + match unsafe { options.as_ref() } { + Some(options) => unsafe { options.to_driver() }, + None => Ok(OperationOptions::default()), + } +} + +#[cfg(test)] +mod tests { + use super::{ + cosmos_operation_options_snapshot_create, cosmos_operation_options_snapshot_free, + cosmos_runtime_set_operation_options, OperationOptionsSnapshot, + }; + use crate::{ + error::CosmosErrorCode, + op_request::{cosmos_operation_options_default, CosmosOperationOptions}, + runtime::{__test_only_create_default_runtime, cosmos_runtime_free, RuntimeContext}, + }; + use std::{ptr, time::Duration}; + + fn capture( + runtime: *const RuntimeContext, + client: &CosmosOperationOptions, + request: &CosmosOperationOptions, + ) -> (Box, i64) { + let mut snapshot = ptr::null_mut(); + let mut timeout = -2; + // SAFETY: all borrowed inputs/outputs live throughout capture; reclaim its unique allocation. + unsafe { + assert_eq!( + cosmos_operation_options_snapshot_create( + runtime, + client, + request, + &mut snapshot, + &mut timeout + ), + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() + ); + (Box::from_raw(snapshot), timeout) + } + } + + #[test] + fn updates_preserve_admission_layers_and_nested_inheritance() { + let runtime = __test_only_create_default_runtime(); + let mut defaults = cosmos_operation_options_default(); + defaults.throughput_bucket = i64::from(u32::MAX); + defaults.priority_level = 1; + defaults.max_throttle_retry_count = 9; + defaults.max_throttle_retry_wait_time_ms = 1234; + defaults.end_to_end_timeout_ms = 4000; + defaults.endpoint_unavailability_ttl_ms = 555; + defaults.binary_encoding_enabled = 1; + // SAFETY: runtime and defaults are live. + unsafe { + assert_eq!( + cosmos_runtime_set_operation_options(runtime, &defaults), + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() + ); + } + let mut client = cosmos_operation_options_default(); + client.priority_level = 2; + client.max_throttle_retry_count = 7; + let mut request = cosmos_operation_options_default(); + request.max_throttle_retry_wait_time_ms = 0; + request.binary_encoding_request_text_response = 2; + let (old, timeout) = capture(runtime, &client, &request); + assert_eq!(timeout, 4000); + // SAFETY: replacing all defaults with an empty group is supported. + unsafe { + assert_eq!( + cosmos_runtime_set_operation_options(runtime, ptr::null()), + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() + ); + } + let (new, timeout) = capture(runtime, &client, &request); + assert_eq!(timeout, -1); + let old_view = old.options.resolution_snapshot_view().unwrap(); + let new_view = new.options.resolution_snapshot_view().unwrap(); + assert_eq!( + old_view.throughput_control().throughput_bucket(), + Some(&u32::MAX) + ); + assert_eq!(new_view.throughput_control().throughput_bucket(), None); + assert_eq!( + old_view.throughput_control().priority_level(), + new_view.throughput_control().priority_level() + ); + assert_eq!( + old_view.throttling_retry_options().max_retry_count(), + Some(&7) + ); + assert_eq!( + old_view.throttling_retry_options().max_retry_wait_time(), + Some(&Duration::ZERO) + ); + assert_eq!( + old_view.endpoint_unavailability_ttl(), + Some(&Duration::from_millis(555)) + ); + assert_eq!(new_view.endpoint_unavailability_ttl(), None); + // Binary flags replace the entire group: request's unspecified enabled uses its own default. + assert!(old_view.binary_encoding().unwrap().enabled); + assert!(old_view.binary_encoding().unwrap().request_text_response); + cosmos_runtime_free(runtime); + } + + #[test] + fn invalid_update_is_atomic_and_sub_second_timeout_is_clamped() { + let runtime = __test_only_create_default_runtime(); + let mut defaults = cosmos_operation_options_default(); + defaults.end_to_end_timeout_ms = 20; + // SAFETY: runtime and input are live across both updates. + unsafe { + assert_eq!( + cosmos_runtime_set_operation_options(runtime, &defaults), + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() + ); + defaults.end_to_end_timeout_ms = 9000; + defaults.max_session_retry_count = i64::from(u32::MAX) + 1; + assert_eq!( + cosmos_runtime_set_operation_options(runtime, &defaults), + CosmosErrorCode::CosmosErrorCodeInvalidOptionValue.as_status_code() + ); + } + let (_, timeout) = capture( + runtime, + &cosmos_operation_options_default(), + &cosmos_operation_options_default(), + ); + assert_eq!(timeout, 1000); + cosmos_runtime_free(runtime); + } + + #[test] + fn deadline_includes_time_before_submit_and_never_polls_expired_work() { + let runtime = __test_only_create_default_runtime(); + let (mut snapshot, _) = capture( + runtime, + &cosmos_operation_options_default(), + &cosmos_operation_options_default(), + ); + snapshot.timeout = Some(Duration::from_secs(1)); + snapshot.started -= Duration::from_secs(2); + let runtime_ref = RuntimeContext::from_ptr(runtime).unwrap(); + let error = runtime_ref + .tokio + .block_on(OperationOptionsSnapshot::execute(Some(*snapshot), async { + panic!("expired work must not be polled"); + #[allow(unreachable_code)] + Ok::<(), azure_data_cosmos_driver::error::CosmosError>(()) + })) + .unwrap_err(); + assert!(error.status().is_timeout()); + cosmos_runtime_free(runtime); + } + + #[test] + fn snapshot_keeps_owned_inputs_after_host_free() { + let runtime = __test_only_create_default_runtime(); + let (snapshot, _) = capture( + runtime, + &cosmos_operation_options_default(), + &cosmos_operation_options_default(), + ); + let submitted = (*snapshot).clone(); + // SAFETY: transfer the Box back to its FFI deallocator exactly once. + unsafe { + cosmos_operation_options_snapshot_free(Box::into_raw(snapshot)); + } + cosmos_runtime_free(runtime); + assert!(submitted.options.resolution_snapshot_view().is_some()); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot/submit_tests.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot/submit_tests.rs new file mode 100644 index 00000000000..52faad6473f --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/options_snapshot/submit_tests.rs @@ -0,0 +1,448 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +use super::{ + cosmos_operation_options_snapshot_create, cosmos_operation_options_snapshot_free, + OperationOptionsSnapshot, +}; +use crate::{ + account_ref::{cosmos_account_ref_free, AccountRefHandle}, + completion::{ + cosmos_completion_queue_create, cosmos_completion_queue_free, + cosmos_completion_queue_free_completions, cosmos_completion_queue_shutdown, + cosmos_completion_queue_state, cosmos_completion_queue_wait, cosmos_operation_handle_free, + CompletionQueue, CosmosCompletion, CosmosCompletionOutcome, CosmosCompletionQueueState, + OperationHandle, + }, + cursor::{ + cosmos_cursor_completion_free, cosmos_cursor_completion_take_cursor, cosmos_cursor_free, + cosmos_cursor_next_submit, cosmos_cursor_open_submit, cosmos_cursor_queue_create, + cosmos_cursor_queue_wait, CosmosCursorCompletion, + }, + cursor_request::{cosmos_cursor_request_init, CosmosCursorRequest}, + driver::{cosmos_driver_free, DriverHandle}, + error::{CosmosErrorCode, CosmosStatusCode}, + op_request::{ + cosmos_operation_options_default, CosmosHeaderKv, CosmosOperationKind, + CosmosOperationOptions, CosmosOperationRequest, + }, + runtime::{cosmos_runtime_free, RuntimeContext}, + string::view, + submit::{cosmos_submit_operation, cosmos_submit_singleton_operation}, +}; +use async_trait::async_trait; +use azure_core::http::headers::{HeaderName, Headers}; +use azure_data_cosmos_driver::{ + driver::CosmosDriverRuntimeBuilder, + error::status_codes::substatus::CLIENT_OPERATION_TIMEOUT, + models::AccountReference, + options::{ConnectionPoolOptions, DriverOptions}, + test::{ + HttpClientConfig, HttpClientFactory, HttpRequest, HttpResponse, TransportClient, + TransportError, + }, +}; +use std::{ + mem::MaybeUninit, + ptr, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::sync::Notify; + +#[derive(Debug, Default)] +struct TransportState { + requests: Mutex>, + release: Notify, +} + +#[derive(Clone, Debug)] +struct SnapshotTransport(Arc); + +impl HttpClientFactory for SnapshotTransport { + fn build( + &self, + _: &ConnectionPoolOptions, + _: HttpClientConfig, + ) -> azure_data_cosmos_driver::error::Result> { + Ok(Arc::new(self.clone())) + } +} + +#[async_trait] +impl TransportClient for SnapshotTransport { + async fn send(&self, request: &HttpRequest) -> Result { + let body = if request.url.path() == "/" { + serde_json::to_vec(&serde_json::json!({ + "_self": "", "id": "test", "_rid": "test.documents.azure.com", + "media": "//media/", "addresses": "//addresses/", "_dbs": "//dbs/", + "writableLocations": [{"name": "East US", "databaseAccountEndpoint": "https://test-eastus.documents.azure.com/"}], + "readableLocations": [ + {"name": "East US", "databaseAccountEndpoint": "https://test-eastus.documents.azure.com/"}, + {"name": "West US", "databaseAccountEndpoint": "https://test-westus.documents.azure.com/"} + ], + "enableMultipleWriteLocations": false, + "userReplicationPolicy": {"minReplicaSetSize": 3, "maxReplicasetSize": 4}, + "userConsistencyPolicy": {"defaultConsistencyLevel": "Session"}, + "systemReplicationPolicy": {"minReplicaSetSize": 3, "maxReplicasetSize": 4}, + "readPolicy": {"primaryReadCoefficient": 1, "secondaryReadCoefficient": 1}, + "queryEngineConfiguration": "{}" + })).unwrap() + } else if request.url.path() == "/probe" { + Vec::new() + } else { + self.0.requests.lock().unwrap().push(request.clone()); + self.0.release.notified().await; + if request.url.path() == "/dbs" { + br#"{"Databases":[],"_count":0}"#.to_vec() + } else { + b"{}".to_vec() + } + }; + Ok(HttpResponse { + status: 200, + headers: Headers::new(), + body, + }) + } +} + +struct Fixture { + runtime: *mut RuntimeContext, + driver: *mut DriverHandle, + account: *mut AccountRefHandle, + queue: *mut CompletionQueue, + transport: Arc, +} + +impl Fixture { + fn new() -> Self { + let transport = Arc::new(TransportState::default()); + let builder = CosmosDriverRuntimeBuilder::new() + .with_mock_http_client_factory(Arc::new(SnapshotTransport(Arc::clone(&transport)))) + .with_connection_pool( + ConnectionPoolOptions::builder() + .with_is_http2_allowed(false) + .build() + .unwrap(), + ); + let runtime = RuntimeContext::new_with_builder(builder) + .ok() + .expect("mock runtime builds"); + let inner = RuntimeContext::from_ptr(runtime).unwrap(); + let account = AccountReference::with_master_key( + "https://test.documents.azure.com/".parse().unwrap(), + "dGVzdA==", + ); + let driver = inner + .tokio + .block_on( + inner + .driver + .create_driver(DriverOptions::builder(account.clone()).build()), + ) + .unwrap(); + Self { + runtime, + driver: DriverHandle::from_arc_into_raw(Arc::new(DriverHandle { inner: driver })), + account: Box::into_raw(Box::new(AccountRefHandle { inner: account })), + queue: cosmos_completion_queue_create(runtime, ptr::null()), + transport, + } + } + + fn snapshot(&self, options: &CosmosOperationOptions) -> *mut OperationOptionsSnapshot { + let mut snapshot = ptr::null_mut(); + let mut timeout = -2; + // SAFETY: fixture runtime, input arrays and both output slots are live. + let status = unsafe { + cosmos_operation_options_snapshot_create( + self.runtime, + ptr::null(), + options, + &mut snapshot, + &mut timeout, + ) + }; + assert_eq!( + status, + CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code() + ); + snapshot + } + + fn request( + &self, + snapshot: *const OperationOptionsSnapshot, + feed: bool, + ) -> CosmosOperationRequest { + // SAFETY: the flat request contains only integers, raw pointers and counted views. + let mut request: CosmosOperationRequest = unsafe { std::mem::zeroed() }; + request.account = self.account; + request.kind = if feed { + CosmosOperationKind::CosmosOperationKindReadAllDatabases + } else { + CosmosOperationKind::CosmosOperationKindReadOffer + } as i32; + request.resource_link = view(b"offer"); + request.max_item_count = -1; + request.options_snapshot = snapshot; + request + } + + fn completion(&self, wait_ms: u32) -> CosmosCompletion { + let mut slot = MaybeUninit::uninit(); + assert_eq!( + cosmos_completion_queue_wait(self.queue, slot.as_mut_ptr(), 1, wait_ms), + 1 + ); + // SAFETY: queue wait wrote one initialized completion. + unsafe { slot.assume_init() } + } +} + +impl Drop for Fixture { + fn drop(&mut self) { + cosmos_completion_queue_free(self.queue); + cosmos_driver_free(self.driver); + cosmos_account_ref_free(self.account); + cosmos_runtime_free(self.runtime); + } +} + +type Submit = extern "C" fn( + *const DriverHandle, + *const CosmosOperationRequest, + *mut CompletionQueue, + isize, + *mut CosmosStatusCode, +) -> *mut OperationHandle; + +fn submitter(feed: bool) -> Submit { + if feed { + cosmos_submit_operation + } else { + cosmos_submit_singleton_operation + } +} + +#[test] +fn submit_copies_snapshot_and_owned_collections_before_host_free() { + for feed in [false, true] { + let fixture = Fixture::new(); + let snapshot = { + let header_name = String::from("x-admission"); + let header_value = String::from("captured"); + let region = String::from("East US"); + let regions = [view(region.as_bytes())]; + let headers = [CosmosHeaderKv { + name: view(header_name.as_bytes()), + value: view(header_value.as_bytes()), + }]; + let mut options = cosmos_operation_options_default(); + options.hedging_enabled = 1; + options.excluded_regions = regions.as_ptr(); + options.excluded_regions_len = regions.len(); + options.custom_headers = headers.as_ptr(); + options.custom_headers_len = headers.len(); + fixture.snapshot(&options) + }; + let request = fixture.request(snapshot, feed); + let mut status = CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code(); + let operation = submitter(feed)(fixture.driver, &request, fixture.queue, 42, &mut status); + assert!(!operation.is_null()); + // SAFETY: submit must have cloned the snapshot before returning. + unsafe { + cosmos_operation_options_snapshot_free(snapshot); + } + fixture.transport.release.notify_one(); + let mut completion = fixture.completion(5000); + assert_eq!( + completion.outcome, + CosmosCompletionOutcome::CosmosCompletionOutcomeOk + ); + assert_eq!(completion.user_data, 42); + let requests = fixture.transport.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].url.host_str(), + Some("test-westus.documents.azure.com") + ); + assert_eq!( + requests[0] + .headers + .get_optional_str(&HeaderName::from_static("x-admission")), + Some("captured") + ); + drop(requests); + cosmos_completion_queue_free_completions(&mut completion, 1); + cosmos_operation_handle_free(operation); + } +} + +#[test] +fn submit_rejects_other_runtime_snapshot_before_queue_admission() { + let fixture = Fixture::new(); + let other = Fixture::new(); + let snapshot = other.snapshot(&cosmos_operation_options_default()); + for feed in [false, true] { + let request = fixture.request(snapshot, feed); + let mut status = CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code(); + assert!( + submitter(feed)(fixture.driver, &request, fixture.queue, 42, &mut status).is_null() + ); + assert_eq!( + status, + CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code() + ); + } + cosmos_completion_queue_shutdown(fixture.queue); + assert_eq!( + cosmos_completion_queue_state(fixture.queue), + CosmosCompletionQueueState::CosmosCompletionQueueStateDrained + ); + assert!(fixture.transport.requests.lock().unwrap().is_empty()); + // SAFETY: snapshot is owned by this test and neither rejected submit retained a borrow. + unsafe { + cosmos_operation_options_snapshot_free(snapshot); + } +} + +#[test] +fn submit_uses_remaining_admission_budget_instead_of_restarting_it() { + for feed in [false, true] { + let fixture = Fixture::new(); + let mut options = cosmos_operation_options_default(); + options.hedging_enabled = 1; + options.end_to_end_timeout_ms = 1000; + let snapshot = fixture.snapshot(&options); + // SAFETY: test owns the live snapshot exclusively, simulating time spent initializing. + unsafe { + (*snapshot).started -= Duration::from_millis(990); + } + + let request = fixture.request(snapshot, feed); + let mut status = CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code(); + let operation = submitter(feed)(fixture.driver, &request, fixture.queue, 42, &mut status); + assert!(!operation.is_null()); + // SAFETY: submit copied the handle; the transport intentionally remains blocked. + unsafe { + cosmos_operation_options_snapshot_free(snapshot); + } + let mut completion = fixture.completion(500); + assert_eq!( + completion.outcome, + CosmosCompletionOutcome::CosmosCompletionOutcomeError + ); + assert_eq!(completion.status.0 >> 16, 408); + assert_eq!( + completion.status.0 & 0xffff, + i32::from(CLIENT_OPERATION_TIMEOUT.value()), + ); + cosmos_completion_queue_free_completions(&mut completion, 1); + cosmos_operation_handle_free(operation); + } +} + +fn cursor_request( + fixture: &Fixture, + snapshot: *const OperationOptionsSnapshot, +) -> CosmosCursorRequest { + let mut request = MaybeUninit::uninit(); + cosmos_cursor_request_init(request.as_mut_ptr()); + // SAFETY: init writes the complete request. + let mut request = unsafe { request.assume_init() }; + request.operation = fixture.request(snapshot, true); + request +} + +fn cursor_completion(queue: *mut CompletionQueue) -> *mut CosmosCursorCompletion { + let mut completion = ptr::null_mut(); + let mut count = 0; + assert_eq!( + cosmos_cursor_queue_wait(queue, &mut completion, 1, 5000, &mut count).0, + 0 + ); + assert_eq!(count, 1); + completion +} + +#[test] +fn cursor_rejects_snapshot_from_other_runtime_before_admission() { + let fixture = Fixture::new(); + let other = Fixture::new(); + let queue = cosmos_cursor_queue_create(fixture.runtime, 1); + let snapshot = other.snapshot(&cosmos_operation_options_default()); + let request = cursor_request(&fixture, snapshot); + let mut status = CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code(); + let operation = cosmos_cursor_open_submit(fixture.driver, &request, queue, 42, &mut status); + assert!(operation.is_null()); + assert_eq!( + status, + CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code() + ); + cosmos_completion_queue_shutdown(queue); + assert_eq!( + cosmos_completion_queue_state(queue), + CosmosCompletionQueueState::CosmosCompletionQueueStateDrained + ); + assert!(fixture.transport.requests.lock().unwrap().is_empty()); + // SAFETY: the snapshot remains exclusively owned after rejected admission. + unsafe { + cosmos_operation_options_snapshot_free(snapshot); + } + cosmos_completion_queue_free(queue); +} + +#[test] +fn cursor_enforces_admission_deadline_at_open_and_between_pages() { + for expire_before_open in [true, false] { + let fixture = Fixture::new(); + let queue = cosmos_cursor_queue_create(fixture.runtime, 1); + let mut options = cosmos_operation_options_default(); + options.end_to_end_timeout_ms = 1000; + options.hedging_enabled = 1; + let snapshot = fixture.snapshot(&options); + if expire_before_open { + // SAFETY: the test exclusively owns the live snapshot before submission. + unsafe { + (*snapshot).started -= Duration::from_secs(2); + } + } + let request = cursor_request(&fixture, snapshot); + let mut status = CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code(); + let operation = cosmos_cursor_open_submit(fixture.driver, &request, queue, 42, &mut status); + assert!(!operation.is_null()); + // SAFETY: open must retain its own copy before returning. + unsafe { + cosmos_operation_options_snapshot_free(snapshot); + } + let opened = cursor_completion(queue); + cosmos_operation_handle_free(operation); + if expire_before_open { + // SAFETY: the queue returned an owned, initialized completion. + assert_eq!( + unsafe { (*opened).common.status.0 }, + (408 << 16) | i32::from(CLIENT_OPERATION_TIMEOUT.value()) + ); + } else { + let cursor = cosmos_cursor_completion_take_cursor(opened); + assert!(!cursor.is_null()); + std::thread::sleep(Duration::from_millis(1050)); + let next = cosmos_cursor_next_submit(cursor, 43, &mut status); + assert!(!next.is_null()); + let page = cursor_completion(queue); + // SAFETY: the queue returned an owned, initialized completion. + assert_eq!( + unsafe { (*page).common.status.0 }, + (408 << 16) | i32::from(CLIENT_OPERATION_TIMEOUT.value()) + ); + cosmos_cursor_completion_free(page); + cosmos_operation_handle_free(next); + cosmos_cursor_free(cursor); + } + assert!(fixture.transport.requests.lock().unwrap().is_empty()); + cosmos_cursor_completion_free(opened); + cosmos_completion_queue_free(queue); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime.rs index 31559ef7abe..738e9cd3197 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime.rs @@ -5,7 +5,7 @@ //! //! Owns both a Tokio multi-threaded [`Runtime`] (so the wrapper can drive //! `async fn` driver code from synchronous FFI entry points) and an -//! `Arc` (so cached drivers, container caches, and the +//! `Arc` (so transport resources and the //! account-metadata cache stay alive for the lifetime of the handle). //! //! The runtime pairs the wrapper-side Tokio runtime with the driver runtime @@ -35,7 +35,7 @@ use crate::runtime_builder::RuntimeBuildError; /// `block_on(...)` driver builder construction at FFI-call time and to /// spawn the per-operation tasks that drive submits. /// - `driver` — the underlying `azure_data_cosmos_driver` runtime that owns -/// the per-account driver registry, container cache, account-metadata +/// the account-metadata /// cache, HTTP transport factory, and so on. Cloning the `Arc` is cheap /// and is how the driver / account surfaces hand out handles. pub struct RuntimeContext { diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime_builder.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime_builder.rs index 27c2c7526ed..96abe0c2390 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime_builder.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/runtime_builder.rs @@ -10,7 +10,7 @@ //! [`cosmos_runtime_build`], which validates each set field, bridges the //! driver-side build through the wrapper's own Tokio runtime, and returns a //! fresh `cosmos_runtime_t *`. Complex nested config (`with_client_options` / -//! `with_connection_pool` / `with_operation_options` / +//! `with_connection_pool` / //! `with_fault_injection_rules`) is //! deliberately not surfaced yet — each requires its own flat options struct. //! @@ -21,6 +21,7 @@ use azure_data_cosmos_driver::driver::CosmosDriverRuntimeBuilder; use azure_data_cosmos_driver::options::{CorrelationId, UserAgentSuffix, WorkloadId}; use crate::error::{CosmosError, CosmosErrorCode, CosmosStatusCode}; +use crate::op_request::CosmosOperationOptions; use crate::runtime::RuntimeContext; // ───────────────────────────────────────────────────────────────────────────── @@ -148,6 +149,8 @@ pub struct CosmosRuntimeOptions { /// CPU/memory monitoring refresh interval in milliseconds (valid range /// `1000`–`60000`). `0` = unset. pub cpu_refresh_interval_ms: u64, + /// Runtime operation defaults, copied during construction. NULL inherits. + pub operation_options: *const CosmosOperationOptions, } impl CosmosRuntimeOptions { @@ -163,6 +166,11 @@ impl CosmosRuntimeOptions { &self, mut builder: CosmosDriverRuntimeBuilder, ) -> Result { + if !self.operation_options.is_null() { + // SAFETY: options and their borrowed arrays are valid throughout this call. + builder = builder + .with_default_operation_options(unsafe { (*self.operation_options).to_driver()? }); + } if self.workload_id != 0 { let Some(value) = WorkloadId::try_new(self.workload_id) else { return Err(CosmosErrorCode::CosmosErrorCodeInvalidOptionValue); @@ -225,6 +233,7 @@ pub extern "C" fn cosmos_runtime_options_default() -> CosmosRuntimeOptions { user_agent_suffix: CosmosStringView::default(), wrapping_sdk_identifier: CosmosStringView::default(), cpu_refresh_interval_ms: 0, + operation_options: std::ptr::null(), } } diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs index d36134d4f88..6e8a6d2164b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs @@ -38,6 +38,7 @@ use crate::driver::DriverHandle; use crate::driver_options::DriverOptionsHandle; use crate::error::{CosmosErrorCode, CosmosStatusCode}; use crate::op_request::{build_request, CosmosOperationRequest}; +use crate::options_snapshot::OperationOptionsSnapshot; use crate::runtime::RuntimeContext; /// Send-safe encoding of the opaque `user_data` cookie round-tripped @@ -361,6 +362,14 @@ fn submit_operation_with_builder( } }; + if built + .snapshot + .as_ref() + .is_some_and(|snapshot| !snapshot.matches_driver(&driver_arc)) + { + write_err(CosmosErrorCode::CosmosErrorCodeInvalidArgument); + return std::ptr::null_mut(); + } let (ctx, op_handle) = match pre_flight_spawn(queue, user_data) { Ok(pair) => pair, Err(code) => { @@ -371,6 +380,7 @@ fn submit_operation_with_builder( let runtime = Arc::clone(ctx.queue.runtime()); let crate::op_request::BuiltRequest { + snapshot, operation, options, patch_tracking_id, @@ -384,7 +394,7 @@ fn submit_operation_with_builder( spawn_oneshot( ctx, runtime, - async move { + OperationOptionsSnapshot::execute(snapshot, async move { // Plan with the inbound continuation, then execute a single // page. Mirrors `CosmosDriver::execute_operation` but threads // the continuation token through the planner and retains the @@ -424,7 +434,7 @@ fn submit_operation_with_builder( return Err(crate::cursor::legacy_representation_error()); } Ok((page, next)) - }, + }), |(page, next): (Option, Option)| SuccessKind::Response { response: page.map(Box::new), next_continuation: next, @@ -502,6 +512,14 @@ fn submit_singleton_operation_with_builder( } }; + if built + .snapshot + .as_ref() + .is_some_and(|snapshot| !snapshot.matches_driver(&driver_arc)) + { + write_err(CosmosErrorCode::CosmosErrorCodeInvalidArgument); + return std::ptr::null_mut(); + } let (ctx, op_handle) = match pre_flight_spawn(queue, user_data) { Ok(pair) => pair, Err(code) => { @@ -513,6 +531,7 @@ fn submit_singleton_operation_with_builder( let runtime = Arc::clone(ctx.queue.runtime()); // `continuation` is intentionally dropped: singletons do not paginate. let crate::op_request::BuiltRequest { + snapshot, operation, options, patch_tracking_id, @@ -525,11 +544,11 @@ fn submit_singleton_operation_with_builder( spawn_oneshot( ctx, runtime, - async move { + OperationOptionsSnapshot::execute(snapshot, async move { driver_arc .execute_singleton_operation(operation, options) .await - }, + }), |response: CosmosResponse| SuccessKind::Response { response: Some(Box::new(response)), next_continuation: None, diff --git a/sdk/cosmos/azure_data_cosmos_macros/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_macros/CHANGELOG.md index 9f40d0de6cc..b917d1a2268 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_macros/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Added `#[option(skip)]` for internal optional state excluded from configuration views and builder setters. ([#5366](https://github.com/Azure/azure-sdk-for-rust/pull/5366)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos_macros/README.md b/sdk/cosmos/azure_data_cosmos_macros/README.md index b55f9384876..fe40ca3e4eb 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/README.md +++ b/sdk/cosmos/azure_data_cosmos_macros/README.md @@ -19,6 +19,10 @@ and setters, view accessors, defaults, and environment-variable initializers. This lets a consuming crate feature-gate a configuration field without leaving generated references to it in builds where the feature is disabled. +Use `#[option(skip)]` on internal optional state that must not participate in +configuration resolution. It defaults to `None`, with no builder setter or view +accessor; it cannot be combined with other field-level option attributes. + See the Hierarchical Configuration Model specification for details. [Cosmos SDK project documentation]: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/cosmos/docs/README.md diff --git a/sdk/cosmos/azure_data_cosmos_macros/src/builder.rs b/sdk/cosmos/azure_data_cosmos_macros/src/builder.rs index 8812119be86..1708e4cbe2a 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/src/builder.rs +++ b/sdk/cosmos/azure_data_cosmos_macros/src/builder.rs @@ -22,20 +22,24 @@ pub fn generate_builder(input: &OptionsInput) -> Result { }); // Setter methods: each takes the inner type and wraps in Some. - let setters = input.fields.iter().map(|field| { - let cfg_attrs = &field.cfg_attrs; - let field_name = &field.ident; - let setter_name = format_ident!("with_{}", field_name); - let inner_type = &field.inner_type; - quote! { - #(#cfg_attrs)* - /// Sets this field on the builder. - #vis fn #setter_name(mut self, value: #inner_type) -> Self { - self.#field_name = Some(value); - self + let setters = input + .fields + .iter() + .filter(|field| !field.skip) + .map(|field| { + let cfg_attrs = &field.cfg_attrs; + let field_name = &field.ident; + let setter_name = format_ident!("with_{}", field_name); + let inner_type = &field.inner_type; + quote! { + #(#cfg_attrs)* + /// Sets this field on the builder. + #vis fn #setter_name(mut self, value: #inner_type) -> Self { + self.#field_name = Some(value); + self + } } - } - }); + }); // Build method: constructs the original struct. let build_fields = input.fields.iter().map(|field| { diff --git a/sdk/cosmos/azure_data_cosmos_macros/src/lib.rs b/sdk/cosmos/azure_data_cosmos_macros/src/lib.rs index 13dd5c280e0..787486c0cdd 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/src/lib.rs +++ b/sdk/cosmos/azure_data_cosmos_macros/src/lib.rs @@ -34,6 +34,8 @@ type Result = ::std::result::Result; /// /// # Field-Level Attributes /// +/// - `#[option(skip)]` — initializes internal `Option` state to `None` +/// without generating a builder setter or a layered accessor. /// - `#[option(env = "AZURE_COSMOS_...")]` — enables environment variable loading. /// - `#[option(env = "AZURE_COSMOS_...", overridable)]` — additionally recognizes a /// `{ENV}_OVERRIDE` kill-switch variable that takes precedence over **every** @@ -136,7 +138,8 @@ fn generate_default(input: &OptionsInput) -> Result { #[cfg(test)] mod tests { - use super::*; + use super::derive_cosmos_options_impl; + use syn::DeriveInput; #[test] fn conditional_field_gates_every_generated_member() { @@ -156,4 +159,38 @@ mod tests { "builder field, setter, build, new, view, env, override env, and Default must all be gated" ); } + + #[test] + fn skipped_state_has_no_configuration_accessors() { + let tokens = derive_cosmos_options_impl(syn::parse_quote! { + #[options(layers(runtime, account, operation))] + pub struct Options { + #[option(skip)] + state: Option, + #[option(env = "TEST_LIMIT")] + pub limit: Option, + } + }) + .unwrap() + .to_string(); + assert!(!tokens.contains("fn with_state")); + assert!(!tokens.contains("fn state")); + assert!(tokens.contains("state : None")); + assert!(tokens.contains("fn with_limit")); + } + + #[test] + fn skip_rejects_configuration_attributes() { + for attribute in [ + quote::quote!(#[option(skip, env = "TEST_LIMIT")]), + quote::quote!(#[option(skip, nested)]), + quote::quote!(#[option(skip, merge = "extend")]), + ] { + let input = syn::parse_quote! { + #[options(layers(runtime, operation))] + struct Options { #attribute state: Option } + }; + assert!(derive_cosmos_options_impl(input).is_err()); + } + } } diff --git a/sdk/cosmos/azure_data_cosmos_macros/src/parse.rs b/sdk/cosmos/azure_data_cosmos_macros/src/parse.rs index 8cd0fdb31e5..69d64c39bbd 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/src/parse.rs +++ b/sdk/cosmos/azure_data_cosmos_macros/src/parse.rs @@ -76,6 +76,8 @@ impl Layer { /// Parsed representation of a single field in the option struct. pub struct OptionField { + /// Internal state initialized to None, excluded from configuration resolution. + pub skip: bool, /// The field name. pub ident: Ident, /// Conditional compilation attributes to copy onto generated members. @@ -247,6 +249,7 @@ fn parse_fields(data: &DataStruct) -> Result> { })?; let ParsedOptionAttrs { + skip, env_var, merge, nested, @@ -255,6 +258,7 @@ fn parse_fields(data: &DataStruct) -> Result> { } = parse_option_attrs(&field.attrs)?; result.push(OptionField { + skip, ident, cfg_attrs: field .attrs @@ -277,6 +281,7 @@ fn parse_fields(data: &DataStruct) -> Result> { /// The parsed `#[option(...)]` field-level attributes for a single field. struct ParsedOptionAttrs { + skip: bool, env_var: Option, merge: Option, nested: bool, @@ -285,6 +290,7 @@ struct ParsedOptionAttrs { } fn parse_option_attrs(attrs: &[syn::Attribute]) -> Result { + let mut skip = false; let mut env_var = None; let mut merge = None; let mut nested = false; @@ -297,7 +303,10 @@ fn parse_option_attrs(attrs: &[syn::Attribute]) -> Result { } attr.parse_nested_meta(|meta| { - if meta.path.is_ident("env") { + if meta.path.is_ident("skip") { + skip = true; + Ok(()) + } else if meta.path.is_ident("env") { let value = meta.value()?; let lit: syn::LitStr = value.parse()?; env_var = Some(lit.value()); @@ -330,6 +339,12 @@ fn parse_option_attrs(attrs: &[syn::Attribute]) -> Result { } // Validate attribute combinations. + if skip && (env_var.is_some() || merge.is_some() || nested || overridable || parser.is_some()) { + return Err(Error::new( + Span::call_site(), + "`skip` cannot be combined with other option attributes", + )); + } if env_var.is_some() && merge.is_some() { return Err(Error::new( Span::call_site(), @@ -362,6 +377,7 @@ fn parse_option_attrs(attrs: &[syn::Attribute]) -> Result { } Ok(ParsedOptionAttrs { + skip, env_var, merge, nested, diff --git a/sdk/cosmos/azure_data_cosmos_macros/src/view.rs b/sdk/cosmos/azure_data_cosmos_macros/src/view.rs index 2bb8f183e75..b44cdd931df 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/src/view.rs +++ b/sdk/cosmos/azure_data_cosmos_macros/src/view.rs @@ -54,6 +54,7 @@ pub fn generate_view(input: &OptionsInput) -> Result { let accessors = input .fields .iter() + .filter(|field| !field.skip) .map(|field| { let cfg_attrs = &field.cfg_attrs; let accessor = generate_accessor(field, layers)?; diff --git a/sdk/cosmos/azure_data_cosmos_macros/tests/derive_cosmos_options.rs b/sdk/cosmos/azure_data_cosmos_macros/tests/derive_cosmos_options.rs index 6f57d6890c2..6a83d939d83 100644 --- a/sdk/cosmos/azure_data_cosmos_macros/tests/derive_cosmos_options.rs +++ b/sdk/cosmos/azure_data_cosmos_macros/tests/derive_cosmos_options.rs @@ -10,6 +10,13 @@ use std::sync::Arc; #[derive(CosmosOptions, Clone)] #[options(layers(runtime, account, operation))] pub struct ConditionalOptions { + #[cfg(any())] + #[option(skip)] + pub disabled_state: Option, + + #[option(skip)] + pub state: Option, + #[cfg(any())] #[option(env = "AZURE_COSMOS_DISABLED_FIELD", overridable)] pub disabled: Option, @@ -26,14 +33,18 @@ fn conditional_fields_gate_all_generated_members() { _ => panic!("unexpected environment variable: {key}"), }); assert_eq!(env.enabled, Some(3)); + assert!(env.state.is_none()); let override_options = ConditionalOptions::from_env_override_vars(|key| { assert_eq!(key, "AZURE_COSMOS_ENABLED_FIELD_OVERRIDE"); Ok("4".to_string()) }); assert_eq!(override_options.enabled, Some(4)); + assert!(override_options.state.is_none()); let runtime = Arc::new(ConditionalOptionsBuilder::new().with_enabled(5).build()); let operation = ConditionalOptions::default(); + assert!(runtime.state.is_none()); + assert!(operation.state.is_none()); assert!(operation.enabled.is_none()); let view = ConditionalOptionsView::new(Some(Arc::new(env)), Some(runtime), None, Some(&operation));