feat(cosmos/workloads): add performance metrics collection for DR drill testing

[tvaron3]

Unified Workload with Performance Metrics for Cosmos DB DR Drill Testing

Summary

Consolidates 8 separate workload scripts into a single workload.py controlled entirely by environment variables, and adds a performance metrics
collection layer that reports latency, throughput, errors, and resource
utilization to a Cosmos DB results account — enabling Grafana dashboards and ADX-powered analysis for DR drill testing.

Motivation

Previously, testing different workload configurations (read-only, write-only, proxy, sync, etc.) required copying and modifying entire workload files.
Adding observability meant manually instrumenting each file. This PR solves
both problems:

1. One workload, many configurations — env vars control operations, proxy, sync, and client behavior
2. Built-in metrics — every operation is timed and reported automatically

Architecture

workload.py ├── WORKLOAD_OPERATIONS → read, write, query (comma-separated) ├── WORKLOAD_USE_PROXY → route through Envoy proxy ├── WORKLOAD_USE_SYNC → sync
vs async client ├── WORKLOAD_SKIP_CLOSE → simulate applications that
don't close the client └── PerfReporter (daemon thread, configurable interval) ├── HdrHistogram per operation (p50/p90/p99, O(1) record) ├── psutil for
CPU/memory metrics ├── Error status code + sub-status code tracking └──
Upserts PerfResult docs to results Cosmos DB account

New Files

File	Description
workload.py	Single unified workload — replaces all 8 operation-specific files
perf_stats.py	Per-operation latency tracking via HdrHistogram — O(1) record, O(1) percentile,
fixed 40KB memory per histogram
perf_reporter.py	Background daemon thread that upserts PerfResult documents to a results Cosmos DB account at a configurable interval
perf_config.py	All metrics configuration from environment variables

Modified Files

File	Change
workload_configs.py	All SDK configs now driven by environment variables (no manual editing)
workload_utils.py	Added timed operation wrappers using time.perf_counter_ns() with error status code capture

Deleted Files

r_workload.py, w_workload.py, r_proxy_workload.py, w_proxy_workload.py, r_w_q_workload.py, r_w_q_proxy_workload.py, r_w_q_sync_workload.py,
r_w_q_with_incorrect_client_workload.py, run_workloads.sh, dev.md

All replaced by workload.py with environment variable configuration. Infrastructure/orchestration scripts (run_workloads.sh, dev.md) moved to the
cosmos-sdk-copilot-toolkit repo.

Environment Variables

Workload behavior:

Variable	Default	Description
WORKLOAD_OPERATIONS	read,write,query	Comma-separated operations to run
WORKLOAD_USE_PROXY	false	Route through Envoy proxy
WORKLOAD_USE_SYNC	false	Use sync client instead of async
WORKLOAD_SKIP_CLOSE	false	Don't close client (simulate incorrect usage)

Cosmos DB client:

Variable	Default	Description
COSMOS_URI	(required)	Workload account endpoint
COSMOS_PREFERRED_LOCATIONS	""	Comma-separated preferred regions
COSMOS_CLIENT_EXCLUDED_LOCATIONS	""	Comma-separated excluded regions
COSMOS_USE_MULTIPLE_WRITABLE_LOCATIONS	false	Enable multi-region writes
AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER	false	Per-partition circuit breaker
COSMOS_CONCURRENT_REQUESTS	100	Concurrent operations per process
COSMOS_THROUGHPUT	1000000	Container provisioned throughput

Metrics reporting:

Variable	Default	Description
RESULTS_COSMOS_URI	""	Results account endpoint (enables metrics when set)
PERF_ENABLED	true	Toggle metrics collection
PERF_REPORT_INTERVAL	300	Seconds between metric reports

PerfResult Document Schema

Each document matches the Rust SDK perf tool schema, enabling
both SDKs to feed the same ADX → Grafana pipeline:

{
  "operation": "ReadItem",
  "count": 12345,
  "errors": 0,
  "min_ms": 1.2, "max_ms": 50.3, "mean_ms": 5.4,
  "p50_ms": 4.8, "p90_ms": 8.2, "p99_ms": 15.1,
  "cpu_percent": 45.2,
  "memory_bytes": 104857600,
  "system_cpu_percent": 62.1,
  "sdk_language": "python",
  "sdk_version": "4.15.0",
  "config_concurrency": 100,
  "config_ppcb_enabled": false
}

Usage Examples

# All operations, direct connection
WORKLOAD_OPERATIONS=read,write,query python3 workload.py

# Read-only via Envoy proxy
WORKLOAD_OPERATIONS=read WORKLOAD_USE_PROXY=true python3 workload.py

# Write-only with metrics reporting
WORKLOAD_OPERATIONS=write \
  RESULTS_COSMOS_URI=https://my-results.documents.azure.com:443/ \
  python3 workload.py

# Simulate unclosed client
WORKLOAD_SKIP_CLOSE=true python3 workload.py

# Scale to 5 processes
for i in $(seq 5); do
  WORKLOAD_OPERATIONS=read,write,query nohup python3 workload.py &
done

[Copilot]

Pull request overview

This PR modernizes the Cosmos DB DR drill workload harness by consolidating multiple per-scenario scripts into a single environment-variable-driven workload runner, and adds an optional performance metrics collection/reporting layer that can upsert results to a separate Cosmos DB account for dashboarding and analysis.

Changes:

• Consolidates prior workload entrypoints into a single workload.py controlled by environment variables.
• Adds performance statistics collection (per-operation latency + errors) and a background reporter that periodically upserts PerfResult documents.
• Refactors workload configuration and operation helpers to be environment-variable-driven and to support timed operation wrappers.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
sdk/cosmos/azure-cosmos/tests/workloads/workload.py	New unified workload entrypoint (sync/async + proxy + perf hooks).
sdk/cosmos/azure-cosmos/tests/workloads/workload_utils.py	Adds timed operation wrappers and Cosmos error status/substatus extraction for perf stats.
sdk/cosmos/azure-cosmos/tests/workloads/workload_configs.py	Moves workload/client configuration to environment variables with parsing/validation.
sdk/cosmos/azure-cosmos/tests/workloads/perf_stats.py	New thread-safe per-operation latency histograms + error aggregation (HdrHistogram-backed).
sdk/cosmos/azure-cosmos/tests/workloads/perf_reporter.py	New daemon-thread reporter to drain stats and upsert PerfResult/Error docs to a results Cosmos account.
sdk/cosmos/azure-cosmos/tests/workloads/perf_config.py	New perf reporter configuration builder from environment variables (includes git SHA + defaults).
sdk/cosmos/azure-cosmos/cspell.json	Adds workload/perf-specific dictionary exceptions for spell checking.
sdk/cosmos/azure-cosmos/tests/workloads/w_workload.py	Removes legacy write-only workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/w_proxy_workload.py	Removes legacy write-via-proxy workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/r_workload.py	Removes legacy read/query workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload.py	Removes legacy mixed workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_with_incorrect_client_workload.py	Removes legacy “incorrect client usage” script (replaced by WORKLOAD_SKIP_CLOSE).
sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_sync_workload.py	Removes legacy sync mixed workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_proxy_workload.py	Removes legacy mixed workload + proxy script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/r_proxy_workload.py	Removes legacy read/query-via-proxy workload script (replaced by workload.py).
sdk/cosmos/azure-cosmos/tests/workloads/setup_env.sh	Removes local environment bootstrap script (moved out of repo per PR description).
sdk/cosmos/azure-cosmos/tests/workloads/run_workloads.sh	Removes legacy multi-script launcher (superseded by env-var-driven workload).
sdk/cosmos/azure-cosmos/tests/workloads/dev.md	Removes legacy local runbook doc (moved out of repo per PR description).

[kushagraThapar]

maybe put a cap to the errors we record? or first in first out type of queue, so that we only keep the last few errors?

[tvaron3]

Addressed in 8c53514. Replaced unbounded list with collections.deque(maxlen=2000) to keep the last 2000 errors. In practice, errors at this level should be minimal since the SDK handles retries internally — what surfaces here are non-retryable failures like 404s.

[kushagraThapar]

I like how you did in the async version -> do_query(ri=random_item)

[tvaron3]

Good eye — fixed in 8c53514. Updated sync version to use def _do_query(ri=random_item) matching the async pattern.

[kushagraThapar]

curious if we can be more streamlined in generation of pk - like using something related to workload id or something else which is more stable and deterministic rather than using uuid now. Since we are really scaling here, using UUID will just create too many cross pk queries, etc.

[tvaron3]

We only use change feed on the results container (→ ADX via Cosmos DB data connection), not cross-partition queries directly. The random partition_key distributes writes evenly across physical partitions. Could follow up with workload_id-based partitioning if we add direct querying on the results container in the future.

[kushagraThapar]

so every flush call upserts one item at a time? won't this be a perf bottleneck?

[tvaron3]

At the default 5-minute report interval (PERF_REPORT_INTERVAL=300), each flush upserts ~3-4 items (one per operation type + errors). That's <0.5s of work on a background daemon thread every 5 minutes — not a bottleneck. The Python sync SDK also doesn't have a bulk/batch API. If we reduce the interval significantly, we could revisit.

[aayush3011]

🟡 Recommendation — Resource Waste in Multi-Client Mode

In multi-client mode (WORKLOAD_NUM_CLIENTS > 1), each async task creates its own Stats, PerfReporter, daemon thread, and sync CosmosClient. Three issues:

1. Fragmented workload IDs: Each Stats calls get_perf_config(), which generates a fresh uuid.uuid4() for workload_id if PERF_WORKLOAD_ID isn't set. Results from the same process get different IDs — making correlation impossible in Grafana/ADX.
2. Fragmented metrics: N separate summaries per flush interval, each with 1/Nth of the data. P99 from 1/5th of operations is statistically meaningless.
3. Resource waste: N daemon threads + N sync CosmosClient instances (each with its own HTTP connection pool) just for writing metrics.

Suggestion: Create one shared Stats + PerfReporter at the run_multi_client_async level and pass it into each task. Stats is already thread-safe (uses locks).

[tvaron3]

Addressed in 314405e. run_multi_client_async now creates a single shared Stats + PerfReporter and passes them to each run_workload_async task. Single workload_id, consolidated metrics, one daemon thread. Individual tasks only create their own reporter if none is passed (single-client mode).

[aayush3011]

🟡 Recommendation — Unbounded Memory Growth

_errors accumulates every error as a dict (with full traceback.format_exc() strings of ~2-5KB each) with no size cap. Between 5-minute drain intervals, a sustained failure scenario could spike memory.

This tool is designed for DR drill testing — the exact scenario where error storms (region failovers) are expected. With 100 concurrent requests failing at ~12 errors/sec × 300 seconds = ~3,600 errors × ~5KB each ≈ 18MB per drain interval.

The aggregate error count is already tracked separately in _error_counts, so no count information is lost.

Suggestion: Add a maximum error buffer size (e.g., 1000). Keep _error_counts for accurate totals, cap the detailed list.

[tvaron3]

Addressed in earlier commit (8c53514). Replaced with collections.deque(maxlen=2000) — keeps last 2000 errors (~10MB max). Errors at this level should be minimal due to SDK internal retries; what surfaces are non-retryable failures. _error_counts still tracks accurate totals regardless of the cap.

[aayush3011]

🟡 Recommendation — Partition Key Anti-Pattern

Both id and partition_key are set to independent random UUIDs for every document. This creates a new logical partition for every single result document.

This is a known Cosmos DB anti-pattern:

• Querying "all results for workload X" or "all results from this host" requires expensive cross-partition fan-out queries.
• Results from the same test run are scattered across thousands of partitions.
• For ADX/Grafana dashboards doing time-series queries, this means high RU cost and latency.

Suggestion: Use workload_id as the partition key value. All results from one workload run land in one partition → efficient point reads and range scans for dashboards.

[tvaron3]

We only use change feed on the results container (→ ADX via Cosmos DB data connection), not cross-partition queries directly. The random partition_key distributes writes evenly. Could follow up with workload_id-based partitioning if direct querying on the results container is needed in the future.

[aayush3011]

🟡 Recommendation — Data-Loss API Trap

drain_summaries() and drain_errors() both delegate to drain_all(), which atomically clears all histograms, error counts, and error details. Calling drain_summaries() silently and permanently discards all accumulated errors (and vice versa).

The method names suggest "drain only summaries" but the implementation drains (and discards) everything. A future caller might reasonably write:

summaries = stats.drain_summaries()
errors = stats.drain_errors()  # Returns empty! Data was already cleared.

Suggestion: Either (a) remove these methods — force callers to use drain_all(), or (b) add a prominent docstring warning about the destructive behavior.

[tvaron3]

Addressed in 314405e. Removed drain_summaries() and drain_errors() — callers now use drain_all() only, which is the only method used by PerfReporter._flush().

[aayush3011]

🟡 Recommendation — Silent 10x Cost Increase

Default THROUGHPUT changed from 100000 (100K RU/s) to 1000000 (1M RU/s) — a 10x increase.

The untouched initial-setup.py imports THROUGHPUT and uses it in ThroughputProperties(THROUGHPUT) when creating containers. Anyone running initial-setup.py without explicitly setting COSMOS_THROUGHPUT will provision 10x the previous default (~$80/hr vs ~$8/hr).

Suggestion: If intentional for CI pipelines, add a comment explaining the increase and document it as a breaking change. If not, revert to 100000.

[tvaron3]

Addressed in 314405e. Reverted default to 100K RU/s. Added comment that DR drills should set COSMOS_THROUGHPUT=1000000 explicitly.

[aayush3011]

🟡 Recommendation — Crash on Invalid Input

Five config values use bare int() without error handling: CONCURRENT_REQUESTS, CONCURRENT_QUERIES, WORKLOAD_NUM_CLIENTS, NUMBER_OF_LOGICAL_PARTITIONS, and THROUGHPUT.

Meanwhile, perf_config.py in the same PR introduces _safe_int() for exactly this purpose.

A typo like COSMOS_CONCURRENT_REQUESTS=1O0 (letter O) crashes at import time with an unhelpful ValueError.

Suggestion: Use the _safe_int pattern from perf_config.py for consistency, or wrap with try/except that identifies which env var is invalid.

[tvaron3]

Addressed in 314405e. All int env vars in workload_configs.py now use _safe_int() from perf_config.py with fallback defaults.

[aayush3011]

🟡 Recommendation — Configuration Inconsistency

Eight os.environ.get() / _safe_int_env() calls inside _flush() re-read configuration from the environment on every flush interval (every 5 minutes). But workload_configs.py captures these same values once at import time.

Two issues:

1. Semantic inconsistency: If env vars were modified after startup (unusual but possible), the workload uses old values while metrics report new ones.
2. Which is the source of truth? workload_configs.py reads at import; _flush() re-reads every interval.

Suggestion: Capture these values once in __init__() or start() and store as instance attributes. Or accept a config dict from the caller.

[tvaron3]

The re-reading in _flush() is intentional — it allows config changes at runtime without restarting workloads (useful during DR drills when excluded regions change). Will refactor workload_configs.py to also re-read on access in a follow-up for full consistency.

[aayush3011]

🟡 Recommendation — Inconsistent Error Handling

The sync path does not protect reporter.stop() in the finally block, but the async path (line ~102) wraps it in try/except.

If the workload throws an exception AND reporter.stop() also throws (e.g., network error flushing final metrics), the original exception is masked by the cleanup exception.

This matters because DR drills are the scenario most likely to have both workload errors AND reporter errors simultaneously.

Suggestion: Apply the same try/except pattern as the async path. Consider logging rather than silently swallowing.

[tvaron3]

Addressed in 314405e. Added try/except around reporter.stop() in the sync path, matching the async pattern.

[aayush3011]

🟢 Suggestion — Misplaced Side Effect

A function named create_logger() shouldn't set global environment variables that control SDK behavior. The circuit breaker config isn't applied until a logger is created, which is invisible to callers.

This was flagged in a past PR review (#40843) as a misplaced concern.

Suggestion: Move the env var setting to workload_configs.py (at import time, alongside the other config), or to the __main__ block in workload.py.

[tvaron3]

Addressed in 314405e. Moved os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] setting from create_logger() to workload_configs.py alongside the other config.

[aayush3011]

🟢 Suggestion — Duplicated Client Constructor

Both branches construct the client identically — only the __aenter__() call differs. If constructor args are modified, both branches must be updated.

Suggestion:

client = AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, **client_kwargs)
if not WORKLOAD_SKIP_CLOSE:
    await client.__aenter__()

[tvaron3]

Addressed in 314405e. Consolidated to a single constructor call with conditional __aenter__().

[tvaron3]

Addressed in 314405e. Consolidated to a single constructor call with conditional __aenter__().

[aayush3011]

🟢 Suggestion — Inconsistent Field Naming

Every field in the PerfResult document uses snake_case except TIMESTAMP which is ALL_CAPS. This creates inconsistency in Kusto/ADX queries.

Suggestion: Rename to timestamp for consistency, or add a comment if the ALL_CAPS convention is required for Rust SDK schema compatibility.

[tvaron3]

Keeping TIMESTAMP in ALL_CAPS — it's required for Rust SDK PerfResults schema compatibility (both SDKs write to the same ADX table). Added a comment explaining this.

[tvaron3]

Keeping TIMESTAMP in ALL_CAPS — it's required for Rust SDK PerfResults schema compatibility (both SDKs write to the same ADX table). Added a comment explaining this.

[aayush3011]

🟢 Suggestion — Duplicated Safe-Int Helper

This PR introduces two nearly identical safe-int-parsing helpers:

• perf_config.py: _safe_int(value, default) — parses a value
• perf_reporter.py: _safe_int_env(name, default) — reads env var then parses

Bug fixes in one won't propagate to the other.

Suggestion: Consolidate into a single helper in perf_config.py and import it here.

[tvaron3]

Addressed in 314405e. Consolidated _safe_int_env into perf_config.py and imported it in perf_reporter.py.

[tvaron3]

Addressed in 314405e. Consolidated _safe_int_env into perf_config.py and imported it in perf_reporter.py.