[Cosmos] Relocate `PartitionEndpointState` into the Runtime Layer

### Background and Motivation:

# Plan A — Relocate `PartitionEndpointState` into the runtime layer

## Scope confirmation

- **Selected scope**: Relocate type only. Move `PartitionEndpointState` (and its three sibling types) out of `driver/routing/` into the runtime-layer module. Each `CosmosDriver` continues to own its own instance via `LocationStateStore`. **Semantics unchanged, instance count unchanged, behavior diff = zero.**
- **Selected motivations**: (1) cross-account observability, (2) consolidate per-account memory / background tasks, (3) layering correction.

### Honest accounting of motivations vs. this plan

| Motivation | Delivered by Plan A? |
|---|---|
| Layering correction | ✅ Yes — the type now sits in the runtime-layer module, matching the conceptual layering. |
| Cross-account observability | ⚠️ **Enabled, not delivered.** Plan A is a prerequisite — a follow-up PR can add a runtime-level accessor that walks `driver_registry` and snapshots each store. That accessor is impossible today because the type sits inside `routing/` next to `LocationStateStore` and has no runtime-level home. |
| Consolidate background tasks / reduce per-account memory | ❌ **Not delivered.** Per-account failback loops still run via `BackgroundTaskManager` inside each `LocationStateStore`. Consolidating these into one runtime-level sweep loop requires re-keying the maps by `(AccountEndpoint, PartitionKeyRangeId)` and rewriting `expire_partition_overrides` to operate on a runtime-level registry — that's a separate, larger change. |

If both motivations 1 and 2 must land in the same PR, this plan must be re-scoped to include the runtime-level ownership migration (Plan B or C from the discussion). The plan below is Plan A only.

---

## Current state (discovery summary)

`PartitionEndpointState` is defined at `sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/partition_endpoint_state.rs` (`pub(crate)`). It is owned **per-account** by `LocationStateStore`, which is owned by `CosmosDriver`. Today's ownership chain:

```text
CosmosDriverRuntime (1 per process) ← runtime layer
└─ driver_registry: HashMap>
└─ CosmosDriver (N per runtime, 1 per account)
└─ Arc
├─ partitions: Atomic ← the type in question
└─ background_task_manager (per-account failback loop)
```

The state itself carries data that is **inherently per-account**:

- `circuit_breaker_overrides` / `failover_overrides` are keyed by `PartitionKeyRangeId`, which collides across accounts ("0", "1", … repeat per container).
- `per_partition_automatic_failover_enabled` is sourced from each account's properties.
- `per_partition_circuit_breaker_enabled` blends per-account properties with env defaults.
- `config: PartitionFailoverConfig` is resolved per-driver from layered options at driver init.
- The CAS swap site (`apply_partition`) and the failback loop (`start_failback_loop`) both assume a single per-account map.

**Types that move alongside `PartitionEndpointState`** (all defined in the same file, no external deps that would create cycles): `HealthStatus`, `PartitionFailoverEntry`, `PartitionFailoverConfig`.

**Types that stay in `routing/`** (moving them would create routing → runtime cycles): `PartitionKeyRangeId`, `CosmosEndpoint`, `LocationStateStore`, `LocationSnapshot`, all `routing_systems` pure functions.

---

## Target module layout

Today `runtime.rs` is a single file at `src/driver/runtime.rs`. Step 1 is to convert it into a directory module so the state types become its visible children:

```text
sdk/cosmos/azure_data_cosmos_driver/src/driver/
├── cosmos_driver.rs (unchanged path)
├── runtime/ ← NEW directory
│ ├── mod.rs ← was runtime.rs (CosmosDriverRuntime)
│ └── partition_endpoint_state.rs ← MOVED from driver/routing/
└── routing/
├── mod.rs (re-export removed for moved types)
├── location_state_store.rs (imports updated)
├── routing_systems.rs (imports updated)
├── partition_key_range_id.rs (unchanged — stays in routing/)
└── endpoint.rs (unchanged — stays in routing/)
```

**Why a directory, not appending types into `runtime.rs`**: keeps `runtime/mod.rs` focused on `CosmosDriverRuntime` itself; partition-state file stays its own ~250-line unit with its own test module.

**Why `PartitionKeyRangeId` and `CosmosEndpoint` do NOT move**: they're consumed by `LocationStateStore`, `routing_systems`, `operation_pipeline`, and the `PartitionKeyRangeCache` — all routing-layer code. Moving them would force `routing/` to depend on `runtime/`, which is the wrong direction (runtime is the lower layer the per-account driver builds on).

---

## What moves and what stays

| Item | Current path | After Plan A | Notes |
|---|---|---|---|
| `PartitionEndpointState` | `crate::driver::routing::partition_endpoint_state::PartitionEndpointState` | `crate::driver::runtime::partition_endpoint_state::PartitionEndpointState` | Same visibility (`pub(crate)`), same fields |
| `HealthStatus` | same module as above | same module as above | Moves with the type |
| `PartitionFailoverEntry` | same module | same module | Moves with the type |
| `PartitionFailoverConfig` (incl. `from_options`) | same module | same module | Continues to depend on `OperationOptionsView`, no cycle |
| `LocationStateStore` | `crate::driver::routing::location_state_store` | **unchanged** | Per-account owner stays in routing; imports `PartitionEndpointState` from the new path |
| `LocationSnapshot` | same as above | **unchanged** | Holds `Arc` — just an import-path change |
| `routing_systems::mark_partition_unavailable` / `expire_partition_overrides` / `remove_probe_succeeded_entry` | `crate::driver::routing::routing_systems` | **unchanged** | Pure functions stay co-located with `LocationStateStore`; imports updated |
| `PartitionKeyRangeId` | `crate::driver::routing::partition_key_range_id` | **unchanged** | Routing-layer identity type |
| `CosmosEndpoint` | `crate::driver::routing::endpoint` | **unchanged** | Routing-layer identity type |
| `__internal_testing` re-exports | `src/testing.rs` | Update import paths | Re-exported names unchanged so benchmarks don't break |

---

## Step-by-step execution sequence

Each step compiles and passes tests independently so the change can be split or bisected.

### Step 1 — Convert `runtime.rs` to a directory module (mechanical, no semantic change)

```powershell
git mv sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime.rs `
sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime/mod.rs
```

**Validate**: `cargo build -p azure_data_cosmos_driver`. No imports change because the module path stays `crate::driver::runtime`.

### Step 2 — Move the file

```powershell
git mv sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/partition_endpoint_state.rs `
sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime/partition_endpoint_state.rs
```

- Add `pub(crate) mod partition_endpoint_state;` to `src/driver/runtime/mod.rs`.
- Remove the `mod partition_endpoint_state;` (and any `pub(crate) use`) from `src/driver/routing/mod.rs`.

**Expected**: build breaks on every import — that's Step 3.

### Step 3 — Update all `use` paths

Mechanical find-and-replace across these files:

| File | Change |
|---|---|
| `src/driver/routing/location_state_store.rs` | `use super::partition_endpoint_state::…` → `use crate::driver::runtime::partition_endpoint_state::…` |
| `src/driver/routing/routing_systems.rs` | Same path update; also update the `#[cfg(test)] mod tests` imports |
| `src/driver/pipeline/operation_pipeline.rs` | Split the existing `use crate::driver::routing::{partition_endpoint_state::HealthStatus, …}` into two `use` statements: one for `crate::driver::routing::…` (for `PartitionKeyRangeId`, `remove_probe_succeeded_entry`) and one for `crate::driver::runtime::partition_endpoint_state::HealthStatus` |
| `src/driver/cosmos_driver.rs` | `PartitionFailoverConfig::from_options(&init_view)` call site — update import path |
| `src/testing.rs` | `pub use crate::driver::routing::partition_endpoint_state::{…}` → `pub use crate::driver::runtime::partition_endpoint_state::{…}` (keep the exported name list identical: `HealthStatus, PartitionEndpointState, PartitionFailoverConfig, PartitionFailoverEntry`) |

**Validate**: `cargo build -p azure_data_cosmos_driver --all-features`.

### Step 4 — Confirm no `routing/` → `runtime/` cycle

- The runtime module already exists at `crate::driver::runtime`. After the move, `crate::driver::routing` imports from `crate::driver::runtime` (one direction).
- Verify `crate::driver::runtime` has no `use crate::driver::routing::…`. Today `runtime.rs` doesn't reference `routing/` at all; the move preserves that property because `PartitionEndpointState`'s only non-`std` dependency was `OperationOptionsView` (already in `crate::options`), which doesn't touch routing.

### Step 5 — Run all validation gates

Per repo `AGENTS.md`:

```powershell
cargo fmt -p azure_data_cosmos_driver
cargo build -p azure_data_cosmos_driver --all-features
cargo clippy -p azure_data_cosmos_driver --all-features --all-targets
cargo doc -p azure_data_cosmos_driver --no-deps --all-features
cargo test -p azure_data_cosmos_driver --all-features
```

Also build the benchmarks crate to validate the `__internal_testing` re-export path:

```powershell
cargo build -p azure_data_cosmos_benchmarks --release --features dhat-heap --example ppcb_state_dhat
cargo build -p azure_data_cosmos_benchmarks --release --features dhat-heap --example ppcb_state_real_dhat
```

---

## Test impact

- **New tests**: none.
- **Deleted tests**: none.

Every existing `#[cfg(test)] mod tests` block stays in its current file:

- `partition_endpoint_state.rs` moves with its own tests intact.
- `routing_systems.rs` and `location_state_store.rs` tests stay put and just update import paths.

The 15+ unit tests catalogued in discovery (`mark_partition_unavailable_*`, `expire_partition_overrides_*`, `apply_partition_keeps_installed_pointer_live_until_store_drop`, `apply_partition_respects_ppaf_failover_*`, etc.) continue to pass unchanged once imports compile.

The `_test_canary` field on `PartitionEndpointState` stays exactly as-is — the regression test in `location_state_store.rs` still works because the use-after-free invariant is unrelated to module location.

---

## Documentation updates

| File | What to update |
|---|---|
| `azure_data_cosmos_driver/docs/PARTITION_LEVEL_FAILOVER_SPEC.md` | Every `crate::driver::routing::partition_endpoint_state::*` reference → `crate::driver::runtime::partition_endpoint_state::*`. Update the "Component design" section to note that the state lives in the runtime-layer module but is instantiated per-account by `LocationStateStore`. |
| `azure_data_cosmos_driver/docs/TRANSPORT_PIPELINE_SPEC.md` | Path mentions, if any. |
| `azure_data_cosmos_driver/docs/PARTITION_KEY_RANGE_CACHE_SPEC.md` | Path mentions, if any. |
| `azure_data_cosmos_benchmarks/docs/PPCB_MEMORY_ANALYSIS.md` | Path mentions in §§ 15, 16, 17, 18; the `azure_data_cosmos_driver::testing` re-export name is unchanged so callers in examples don't change. |
| `azure_data_cosmos_driver/CHANGELOG.md` | Add an unreleased entry under "Other Changes" per `cosmos.changelog.instructions.md`: "Moved `PartitionEndpointState` (and `HealthStatus` / `PartitionFailoverEntry` / `PartitionFailoverConfig`) from `driver::routing` to the runtime layer at `driver::runtime`. Internal refactor — no public-API impact." |

---

## Public-API / compatibility surface

- **`azure_data_cosmos_driver` public API**: unchanged. `PartitionEndpointState` is `pub(crate)`; the move doesn't cross the public boundary.
- **`azure_data_cosmos` SDK**: unchanged. Discovery confirmed it has no references to `PartitionEndpointState`.
- **`__internal_testing` feature consumers** (benchmarks, and any external crate building with that feature): re-export *names* via `azure_data_cosmos_driver::testing::*` are unchanged. Only the upstream definition path changed, which is opaque to consumers.
- **Wire format / serialization**: N/A — type is never serialized.

---

## Change footprint

- **File moves**: 2 (`partition_endpoint_state.rs` + `runtime.rs` → `runtime/mod.rs`).
- **Files touched for import updates**: ~6 (`location_state_store`, `routing_systems`, `operation_pipeline`, `cosmos_driver`, `testing`, `routing/mod`, `runtime/mod`).
- **Doc files touched**: ~3–4 spec files + 1 changelog + 1 benchmarks doc.
- **Net behavioral diff**: zero. Verified by the existing 20+ unit tests passing without modification beyond import paths.

---

## Follow-up PRs (out of scope for Plan A, enabled by it)

1. **Cross-account observability accessor** — Add `pub(crate) fn iter_partition_states(&self) -> impl Iterator)>` to `CosmosDriverRuntime`. Walks `driver_registry`, calls each driver's `LocationStateStore::snapshot()`, yields the partition state. Enables diagnostics aggregation and process-wide metrics.
2. **Consolidated failback sweep** (Plan B/C territory) — Move `BackgroundTaskManager` ownership of the failback loop to `CosmosDriverRuntime`. Replace N per-account loops with one runtime-level sweep that iterates the driver registry. Requires re-keying or per-account dispatch; significantly larger blast radius.