DTX for Rust

[Fabian Meiswinkel (FabianMeiswinkel)]

Summary

This PR adds preview Distributed Transactions (DTX) support to the Rust Cosmos driver and SDK.

Distributed Transactions extend Cosmos DB atomicity beyond TransactionalBatch's single-container / single-partition-key boundary. With this change, Rust clients can prepare and execute multi-partition, multi-container transactions within a single Cosmos DB account through the service-side Distributed Transactions Coordinator (DTC).

The implementation is intentionally gated behind a disabled-by-default preview_dtx feature because the DTX service feature is still in preview/test-account rollout and is not production-ready.

Important References

• .NET SDK source of truth: Azure/azure-cosmos-dotnet-v3
• .NET DTX OpenSpec / proposal: Azure/azure-cosmos-dotnet-v3#5781
• DTX REST wire contract between SDK and service: Azure/azure-cosmos-dotnet-v3#6002
• .NET DTX API review: Azure/azure-cosmos-dotnet-v3#5877
• Response reordering / fail-closed handling: Azure/azure-cosmos-dotnet-v3#5974
• DTX 429 / RUBudgetExceeded retry handling: Azure/azure-cosmos-dotnet-v3#5989
• Patch filter predicate support: Azure/azure-cosmos-dotnet-v3#5995
• Malformed DTX session-token handling: Azure/azure-cosmos-dotnet-v3#5958
• Rust design spec with ADRs: sdk/cosmos/azure_data_cosmos_driver/docs/DISTRIBUTED_TRANSACTIONS_SPEC.md

What Changed

Driver (azure_data_cosmos_driver)

Adds preview DTX driver support behind preview_dtx:

• DTX wire models:
  • DistributedTransactionType
  • DistributedTransactionOperationKind
  • DistributedTransactionTarget
  • DistributedTransactionOperation
  • DistributedTransactionRequest
  • DistributedTransactionResponse
  • DistributedTransactionOperationResult
  • DistributedTransactionResultBody
• Request serialization for the DTC wire contract:
  • POST /operations/dtc
  • x-ms-cosmos-idempotency-token
  • x-ms-cosmos-operation-type
  • x-ms-cosmos-resource-type: DistributedTransactionBatch
  • JSON operations[] payload
• Response parsing:
  • reorder by operation index
  • fail closed on malformed successful coordinator responses
  • preserve coordinator isRetriable and diagnosticString
  • preserve raw coordinator headers and raw per-operation response payloads for diagnostics and parity testing
  • promote 207 MultiStatus according to the DTX contract
• Two-tier retry behavior:
  • outer loop for body-bearing coordinator responses using isRetriable
  • inner retry classification for bodyless coordinator/infrastructure failures
  • DTX-specific retry budgets matching the .NET implementation
  • bodyless 429 uses the shared throttle retry path; body-bearing DTX 429 is handled by the outer coordinator retry path
• Session handling:
  • pre-send session-token resolution for DTX operations
  • per-partition-key-range token resolution through the driver PKRange cache
  • parent-range fallback for freshly split children
  • response-side per-operation session-token merge
  • strict malformed-token rejection under Session consistency
• DTC sub-status constants:
  • DTC_COORDINATOR_RACE_CONFLICT
  • DTC_LEDGER_FAILURE
  • DTC_ACCOUNT_CONFIG_FAILURE
  • DTC_DISPATCH_FAILURE
  • DTC_OPERATION_ROLLED_BACK

SDK (azure_data_cosmos)

Adds preview public SDK transaction builders behind preview_dtx.

The SDK API now follows the same broad shape as TransactionalBatch: callers prepare a data-only transaction document and pass it to the account client for execution.

• DistributedWriteTransaction::new()
• DistributedReadTransaction::new()
• CosmosClient::commit_distributed_write(...)
• CosmosClient::execute_distributed_read(...)
• Builder operation methods target &ContainerClient; the SDK resolves the underlying driver ContainerReference internally.
• Same-account enforcement runs when the transaction is executed by CosmosClient.
• DistributedWriteTransaction
  • create_item
  • replace_item
  • upsert_item
  • delete_item
  • patch_item
• DistributedReadTransaction
  • read_item
• DistributedTransactionOperationOptions
  • per-operation session_token
  • per-operation ETag precondition
• DistributedTransactionPatchOperationOptions
  • per-operation session_token
  • per-operation ETag precondition
  • patch filter_predicate
• DistributedTransactionResponse
  • status
  • is_success_status_code
  • is_completed_status_code
  • len
  • is_empty
  • operation_result
  • headers
  • diagnostic_string
  • idempotency_token
  • is_retriable
  • error_message
  • diagnostics
  • activity_id
  • request_charge
  • retry_after_ms
• DistributedTransactionOperationResult
  • index
  • status_code
  • sub_status_code
  • is_success_status_code
  • is_completed_status_code
  • etag
  • session_token
  • partition_key_range_id
  • request_charge
  • resource<T>()

DistributedWriteTransaction and DistributedReadTransaction are intentionally data-only builders. They do not carry a client and do not execute themselves. Execution happens through CosmosClient::commit_distributed_write(...) or CosmosClient::execute_distributed_read(...), which consumes the prepared transaction and validates that every operation targets the same Cosmos DB account as the executing client.

In-Memory Emulator Support

DTX is not available in the normal emulator or emulator vnext today, and live DTX testing requires special test accounts. To enable local iteration, this PR adds DTX support to the in-memory emulator:

• POST /operations/dtc
• write transaction prepare / commit / rollback behavior
• read transaction snapshot result rewriting
• multi-container DTX sanity coverage
• DTX Patch support, including filter predicate evaluation
• prepared-operation rollback surfaced as 453 / 5415
• read snapshot failure rewrite to 424 FailedDependency
• no-partial-commit checks
• 304 NotModified read transaction handling

Public API Example

use azure_data_cosmos::{
    CosmosClient,
    DistributedReadTransaction,
    DistributedWriteTransaction,
    PartitionKey,
};
use azure_data_cosmos::clients::{
    ContainerClient,
    DistributedTransactionOperationOptions,
    DistributedTransactionPatchOperationOptions,
};
use azure_data_cosmos::models::{PatchInstructions, PatchOperation};
use azure_data_cosmos::options::Precondition;

#[derive(serde::Serialize, serde::Deserialize)]
struct Account {
    id: String,
    pk: String,
    balance: i64,
}

#[derive(serde::Serialize)]
struct AuditEvent {
    id: String,
    pk: String,
    from: String,
    to: String,
    amount: i64,
}

async fn transfer_between_accounts(
    client: &CosmosClient,
    accounts: &ContainerClient,
    audit: &ContainerClient,
    from: Account,
    to: Account,
    from_etag: azure_core::http::Etag,
    to_etag: azure_core::http::Etag,
) -> azure_data_cosmos::Result<()> {
    let amount = 100;

    let updated_from = Account {
        balance: from.balance - amount,
        ..from
    };

    let updated_to = Account {
        balance: to.balance + amount,
        ..to
    };

    let write_transaction = DistributedWriteTransaction::new()
        .replace_item(
            accounts,
            PartitionKey::from(updated_from.pk.clone()),
            updated_from.id.clone(),
            &updated_from,
            Some(
                DistributedTransactionOperationOptions::default()
                    .with_precondition(Precondition::if_match(from_etag)),
            ),
        )?
        .replace_item(
            accounts,
            PartitionKey::from(updated_to.pk.clone()),
            updated_to.id.clone(),
            &updated_to,
            Some(
                DistributedTransactionOperationOptions::default()
                    .with_precondition(Precondition::if_match(to_etag)),
            ),
        )?
        .create_item(
            audit,
            PartitionKey::from("transfers"),
            "transfer-0001",
            AuditEvent {
                id: "transfer-0001".to_owned(),
                pk: "transfers".to_owned(),
                from: updated_from.id.clone(),
                to: updated_to.id.clone(),
                amount,
            },
            None,
        )?;

    let write_response = client
        .commit_distributed_write(write_transaction)
        .await?;

    if write_response.is_completed_status_code() {
        tracing::info!(
            idempotency_token = %write_response.idempotency_token(),
            activity_id = ?write_response.activity_id(),
            request_charge = ?write_response.request_charge(),
            "distributed transaction committed"
        );

        return Ok(());
    }

    for index in 0..write_response.len() {
        if let Some(result) = write_response.operation_result(index) {
            if result.status_code().as_u16() == 424 {
                continue;
            }

            tracing::warn!(
                operation_index = result.index(),
                status = ?result.status_code(),
                sub_status = ?result.sub_status_code(),
                pk_range = ?result.partition_key_range_id(),
                "distributed transaction operation failed"
            );
        }
    }

    if write_response.is_retriable() {
        tracing::warn!(
            idempotency_token = %write_response.idempotency_token(),
            diagnostic = ?write_response.diagnostic_string(),
            "distributed transaction ended retriable after SDK retry budget; reconcile before retrying"
        );
    }

    Err(azure_data_cosmos::Error::message(format!(
        "distributed transaction failed: status={:?}, diagnostic={:?}, error={:?}",
        write_response.status(),
        write_response.diagnostic_string(),
        write_response.error_message(),
    )))
}

async fn conditional_patch_example(
    client: &CosmosClient,
    container: &ContainerClient,
) -> azure_data_cosmos::Result<()> {
    let patch = PatchInstructions::from(vec![
        PatchOperation::set("/status", serde_json::json!("completed")),
    ]);

    let write_transaction = DistributedWriteTransaction::new()
        .patch_item(
            container,
            PartitionKey::from("order-123"),
            "order-123",
            patch,
            Some(
                DistributedTransactionPatchOperationOptions::default()
                    .with_filter_predicate("from c where c.status = 'pending'"),
            ),
        )?;

    let response = client
        .commit_distributed_write(write_transaction)
        .await?;

    if !response.is_completed_status_code() {
        for index in 0..response.len() {
            if let Some(result) = response.operation_result(index) {
                if result.status_code().as_u16() != 424 {
                    tracing::warn!(
                        index = result.index(),
                        status = ?result.status_code(),
                        sub_status = ?result.sub_status_code(),
                        "patch DTX operation failed"
                    );
                }
            }
        }
    }

    Ok(())
}

async fn distributed_read_example<T>(
    client: &CosmosClient,
    container_a: &ContainerClient,
    container_b: &ContainerClient,
) -> azure_data_cosmos::Result<(Option<T>, Option<T>)>
where
    T: serde::de::DeserializeOwned,
{
    let read_transaction = DistributedReadTransaction::new()
        .read_item(
            container_a,
            PartitionKey::from("tenant-1"),
            "item-1",
            None,
        )
        .read_item(
            container_b,
            PartitionKey::from("tenant-2"),
            "item-2",
            None,
        );

    let response = client
        .execute_distributed_read(read_transaction)
        .await?;

    let first = response
        .operation_result(0)
        .map(|result| result.resource::<T>())
        .transpose()?
        .flatten();

    let second = response
        .operation_result(1)
        .map(|result| result.resource::<T>())
        .transpose()?
        .flatten();

    Ok((first, second))
}

[Copilot]

Pull request overview

Adds a preview (“disabled by default”) Distributed Transactions (DTX) capability to the Cosmos Rust driver (azure_data_cosmos_driver) and exposes a corresponding gated builder-based API in the public SDK (azure_data_cosmos), plus emulator support, tests, and a design/spec document.

Changes:

• Introduces preview_dtx-gated driver wire models, request/response handling, retries, and session-token behavior for POST /operations/dtc.
• Adds preview_dtx-gated SDK transaction builders and response accessors, and updates the in-memory emulator to handle /operations/dtc.
• Adds integration tests (emulator + optional live comparison), documentation/spec, changelog entries, and cSpell terms.

Show a summary per file

File	Description
sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs	Gates new DTX live-comparison test module behind preview_dtx.
sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/dtx_live_comparison.rs	Adds dual-backend emulator vs live DTX parity test (ignored by default).
sdk/cosmos/azure_data_cosmos/src/clients/mod.rs	Re-exports DTX SDK types and wires new module behind preview_dtx.
sdk/cosmos/azure_data_cosmos/src/clients/distributed_transaction.rs	Adds preview SDK DTX builders and response accessors.
sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs	Adds preview factory methods for distributed read/write transactions.
sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs	Exposes resolved ContainerReference under preview_dtx.
sdk/cosmos/azure_data_cosmos/CHANGELOG.md	Documents the new preview DTX SDK surface.
sdk/cosmos/azure_data_cosmos/Cargo.toml	Adds uuid dependency + preview_dtx feature wiring to driver.
sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs	Adds DTX emulator test module behind preview_dtx.
sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distributed_transaction.rs	Adds emulator integration tests for /operations/dtc.
sdk/cosmos/azure_data_cosmos_driver/src/models/vector_session_token.rs	Updates vector session token parsing/merging to support -1 sentinel.
sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs	Exposes DTX models + adds new ResourceType/OperationType variants behind feature.
sdk/cosmos/azure_data_cosmos_driver/src/models/distributed_transaction.rs	Adds DTX wire models, serialization, parsing, status promotion, and unit tests.
sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_resource_reference.rs	Adds resource-path/signing handling for DTX coordinator endpoint.
sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs	Adds CosmosOperation::distributed_transaction and idempotency/read-only tests.
sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/store.rs	Adds counter rollback helper for DTX abort semantics.
sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs	Adds /operations/dtc handler implementing DTX read/write semantics.
sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs	Routes POST /operations/dtc and parses it into an emulator operation.
sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs	Adds DTX-related substatus codes behind preview_dtx.
sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/transport_pipeline.rs	Adjusts 429 retry eligibility for body-bearing DTX responses; adds DTX detection.
sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/mod.rs	Marks DTX resource type as dataplane for pipeline selection.
sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/session_manager.rs	Adds DTX session-token resolve/merge logic behind preview_dtx.
sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/session_container.rs	Adds strict session-token merge helper behind preview_dtx.
sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs	Adds DTX-specific retry classification for bodyless coordinator/infra failures.
sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs	Adds handling for new OperationAction::DtxRetry.
sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/components.rs	Adds DTX retry counters/budgets and new pipeline action behind feature.
sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs	Adds execute_distributed_transaction outer retry loop and request header stamping.
sdk/cosmos/azure_data_cosmos_driver/docs/DISTRIBUTED_TRANSACTIONS_SPEC.md	Adds a detailed DTX design/spec/ADR document.
sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md	Documents new preview DTX driver functionality.
sdk/cosmos/azure_data_cosmos_driver/Cargo.toml	Adds preview_dtx feature flag.
sdk/cosmos/.cspell.json	Adds DTX-related terms to cSpell dictionary.

Review details

• Files reviewed: 31/31 changed files
• Comments generated: 2
• Review effort level: Low

[Fabian Meiswinkel (FabianMeiswinkel)]

/azp run rust - cosmos - weekly

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Ashley Stanton-Nurse (analogrelay)]

My only naming nit is in an internal API, so it's non-blocking.

[Fabian Meiswinkel (FabianMeiswinkel)]

/azp run rust - cosmos - weekly

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Fabian Meiswinkel (FabianMeiswinkel)]

/azp run rust - cosmos - weekly

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Fabian Meiswinkel (FabianMeiswinkel)]

Test failure in weekly is unrelated

So the failure is at initial GetDatabaseAccount metadata fetch to the live account, before the actual query/fault-injection behavior under test. The log also shows many other live operations succeeded before these, which makes this look like live-account/network flakiness or account endpoint reachability pressure, not a DTX semantic regression.

A few relevant points:

The failing tests are not DTX tests.
The failing target is [azure_data_cosmos --test emulator](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html), but in this live leg it runs because [test-resources.bicep](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html) defaults [testCategory = 'emulator'](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html); that naming is misleading here.
PR 4702 does touch [cosmos_driver.rs](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html), [driver/transport/mod.rs](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html), and [transport_pipeline.rs](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html), so there is some broad metadata/transport overlap, but the observed failure is a 5s TCP/connect timeout on account metadata fetch, not a changed request classification, DTX route, retry classification, or session behavior.
I reverted my speculative local [cosmos_proxy.rs](vscode-file://vscode-app/c:/Users/fabianm/AppData/Local/Programs/Microsoft%20VS%20Code/4fe60c8b1c/resources/app/out/vs/code/electron-browser/workbench/workbench.html) edit. Current working tree is clean.
My read: rerun the failed live leg first. If it repeats, then the next useful investigation is around metadata fetch timeout/retry behavior under live CI load, not DTX specifically.