Skip to main content

autumn_web/
sharding.rs

1//! Horizontal database sharding.
2//!
3//! Autumn routes sharded data in two steps: a routing key (typically the
4//! tenant id) hashes onto a fixed set of [`SLOT_COUNT`] (16384) **logical
5//! slots** — the same constant Redis Cluster and Valkey use — and each
6//! slot maps to one physical shard via the `[[database.shards]]`
7//! configuration. The key→slot hash is a permanent contract — it is
8//! deterministic across processes, replicas, and Autumn versions — while
9//! the slot→shard map is plain configuration. Resharding therefore means
10//! moving whole slots between shards and flipping the map, never
11//! rehashing keys.
12//!
13//! Each shard is a full [`DatabaseTopology`] (primary + optional read
14//! replica), so the primary/replica story composes with sharding.
15//!
16//! Most framework state (jobs, scheduler locks, sessions, feature flags)
17//! lives on the control topology configured by
18//! `database.primary_url`/`database.url`. The outbox, commit hooks, version
19//! history and derivation state are on each shard. With
20//! `jobs.postgres.shard_local = true`, jobs are on each shard too (ADR 0018).
21//!
22//! # Example
23//!
24//! ```toml
25//! [database]
26//! primary_url = "postgres://db-control/app"
27//!
28//! [[database.shards]]
29//! name = "shard0"
30//! primary_url = "postgres://db-shard0/app"
31//! slots = ["0-8191"]
32//!
33//! [[database.shards]]
34//! name = "shard1"
35//! primary_url = "postgres://db-shard1/app"
36//! slots = ["8192-16383"]
37//! ```
38
39// autumn-determinism-gate: production code in this module must read time and
40// mint identifiers through the framework's injected seams (ClockSource /
41// Entropy), never `Instant::now()` / `Utc::now()` / `SystemTime::now()` /
42// `Uuid::new_v4()` directly. See CONTRIBUTING.md "Determinism seam gate"
43// (issue #1797). Justify exceptions with
44// #[allow(clippy::disallowed_methods, reason = "…")] at the narrowest scope.
45#![cfg_attr(not(test), deny(clippy::disallowed_methods))]
46
47use std::collections::HashMap;
48use std::sync::Arc;
49use std::sync::atomic::{AtomicBool, Ordering};
50
51use crate::db::RuntimeConnection;
52#[cfg(feature = "postgres")]
53use diesel_async::AsyncPgConnection;
54use diesel_async::pooled_connection::deadpool::Pool;
55
56pub use crate::config::SLOT_COUNT;
57use crate::config::{ConfigError, DatabaseConfig, ReplicaFallback};
58use crate::db::{DatabaseTopology, PoolError};
59use crate::error::AutumnError;
60
61/// Index of a physical shard within the configured shard set.
62///
63/// Stable only for a given configuration; use [`Shard::name`] for
64/// identity that survives configuration edits.
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
66pub struct ShardId(pub usize);
67
68/// A logical routing slot in <code>0..[SLOT_COUNT]</code>.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
70pub struct SlotId(pub u16);
71
72/// Borrowed routing key.
73///
74/// Tenant ids in Autumn are strings, so [`ShardKey::Str`] is the common
75/// variant; `Int` and `Bytes` cover numeric primary keys and UUIDs
76/// (`ShardKey::from(uuid.as_bytes())`).
77#[derive(Debug, Clone, Copy)]
78pub enum ShardKey<'a> {
79    /// Numeric key (e.g. a `BIGINT` primary key).
80    Int(i64),
81    /// Textual key (e.g. a tenant id).
82    Str(&'a str),
83    /// Raw bytes (e.g. a UUID).
84    Bytes(&'a [u8]),
85}
86
87impl From<i64> for ShardKey<'_> {
88    fn from(key: i64) -> Self {
89        Self::Int(key)
90    }
91}
92
93impl From<i32> for ShardKey<'_> {
94    fn from(key: i32) -> Self {
95        Self::Int(i64::from(key))
96    }
97}
98
99impl<'a> From<&'a str> for ShardKey<'a> {
100    fn from(key: &'a str) -> Self {
101        Self::Str(key)
102    }
103}
104
105impl<'a> From<&'a String> for ShardKey<'a> {
106    fn from(key: &'a String) -> Self {
107        Self::Str(key)
108    }
109}
110
111impl<'a> From<&'a [u8]> for ShardKey<'a> {
112    fn from(key: &'a [u8]) -> Self {
113        Self::Bytes(key)
114    }
115}
116
117impl<'a> From<&'a [u8; 16]> for ShardKey<'a> {
118    fn from(key: &'a [u8; 16]) -> Self {
119        Self::Bytes(key)
120    }
121}
122
123// ── Deterministic key hashing ────────────────────────────────────────────────
124//
125// The key→slot function is a PERMANENT CONTRACT: every process, replica,
126// and future Autumn version must route the same key to the same slot, or
127// data written by one replica becomes invisible to another. That rules out
128// std's SipHash (randomly keyed per process). FNV-1a and splitmix64 are
129// fixed, well-known functions; the golden-vector tests below pin their
130// output forever.
131
132const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
133const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
134
135fn fnv1a_64(bytes: &[u8]) -> u64 {
136    let mut hash = FNV_OFFSET_BASIS;
137    for byte in bytes {
138        hash ^= u64::from(*byte);
139        hash = hash.wrapping_mul(FNV_PRIME);
140    }
141    hash
142}
143
144/// splitmix64 finalizer — mixes integer keys so that sequential ids
145/// spread uniformly across slots.
146const fn splitmix64(mut x: u64) -> u64 {
147    x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
148    x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
149    x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
150    x ^ (x >> 31)
151}
152
153/// Deterministic 64-bit hash of a routing key.
154#[must_use]
155fn key_hash64(key: ShardKey<'_>) -> u64 {
156    match key {
157        #[allow(clippy::cast_sign_loss)]
158        ShardKey::Int(value) => splitmix64(value as u64),
159        ShardKey::Str(value) => fnv1a_64(value.as_bytes()),
160        ShardKey::Bytes(value) => fnv1a_64(value),
161    }
162}
163
164/// The [`ShardId`] of a fleet database. A slot database is its slot. In a
165/// tenant fleet many tenants share a slot, so the id is a 64-bit hash of the
166/// tenant id instead: two tenant databases never share an id, and the same
167/// tenant always gets the same one. A fleet database's id indexes nothing
168/// ([`ShardSet::get`] is for configured shards); it only tells databases apart.
169/// It is a hash, so distinctness is probabilistic: negligible risk with a
170/// 64-bit `usize`, real past tens of thousands of tenants with a 32-bit one.
171/// [`Shard::name`] is the identity to key results by.
172#[cfg(any(feature = "sqlite", test))]
173#[must_use]
174pub(crate) fn fleet_shard_id(key: &crate::fleet_layout::FleetDbKey) -> ShardId {
175    match key {
176        crate::fleet_layout::FleetDbKey::Slot(slot) => ShardId(usize::from(*slot)),
177        crate::fleet_layout::FleetDbKey::Tenant { id, .. } =>
178        {
179            #[allow(
180                clippy::cast_possible_truncation,
181                reason = "an id, not an index: a 32-bit target keeps 32 well-mixed bits"
182            )]
183            ShardId(fnv1a_64(id.as_str().as_bytes()) as usize)
184        }
185    }
186}
187
188/// Map a routing key onto a logical slot in <code>0..[SLOT_COUNT]</code>.
189///
190/// This function is deterministic across processes and versions; see the
191/// module docs.
192#[must_use]
193pub fn slot_for_key(key: ShardKey<'_>) -> SlotId {
194    let hash = key_hash64(key);
195    #[allow(clippy::cast_possible_truncation)]
196    SlotId((hash % u64::from(SLOT_COUNT)) as u16)
197}
198
199// ── Router ───────────────────────────────────────────────────────────────────
200
201/// Pluggable shard routing strategy.
202///
203/// The default [`HashShardRouter`] hashes the key onto a logical slot and
204/// resolves the slot's owner from configuration. Implement this trait for
205/// directory/lookup routing (e.g. a control-plane table mapping tenants to
206/// shards, with hot "whale" tenants pinned to dedicated shards) and
207/// install it with
208/// [`AppBuilder::with_shard_router`](crate::app::AppBuilder::with_shard_router).
209///
210/// Routing is async so directory routers can consult a cache or the
211/// control database. Custom routers can still compose with the hash via
212/// [`ShardSet::slot_for_key`] and [`ShardSet::shard_for_slot`].
213pub trait ShardRouter: Send + Sync + 'static {
214    /// Resolve the shard that owns `key`.
215    fn route<'a>(
216        &'a self,
217        key: ShardKey<'a>,
218        shards: &'a ShardSet,
219    ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>>;
220}
221
222/// `Arc<R>` routes through its inner router. This lets a caller build a single
223/// `Arc<DirectoryShardRouter>`, share one clone with
224/// [`DirectoryShardRouter::spawn_invalidation_listener`] (which takes an
225/// `Arc<Self>`) and install another clone via `AppBuilder::with_shard_router` —
226/// both then read and invalidate the **same** cache. Without it the
227/// manually-installed router
228/// and the listener would hold separate caches, so directory re-pins would stay
229/// stale until the TTL despite the documented manual-listener path.
230impl<R: ShardRouter + ?Sized> ShardRouter for Arc<R> {
231    fn route<'a>(
232        &'a self,
233        key: ShardKey<'a>,
234        shards: &'a ShardSet,
235    ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
236        (**self).route(key, shards)
237    }
238}
239
240/// Default router: key → logical slot (deterministic hash) → shard
241/// (configured slot map).
242#[derive(Debug, Default, Clone, Copy)]
243pub struct HashShardRouter;
244
245impl ShardRouter for HashShardRouter {
246    fn route<'a>(
247        &'a self,
248        key: ShardKey<'a>,
249        shards: &'a ShardSet,
250    ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
251        let slot = shards.slot_for_key(key);
252        Box::pin(std::future::ready(
253            shards
254                .inner
255                .slot_map
256                .get(usize::from(slot.0))
257                .map(|&idx| ShardId(idx))
258                .ok_or_else(|| {
259                    AutumnError::service_unavailable_msg(format!(
260                        "slot {} has no shard assigned (slot map inconsistent)",
261                        slot.0
262                    ))
263                }),
264        ))
265    }
266}
267
268/// A [`ShardRouter`] that consults an explicit `_autumn_shard_directory`
269/// table on the control database, falling back to the hash router for any
270/// tenant without a directory row.
271///
272/// This is the routing half of "move a tenant to a specific shard": a row in
273/// `_autumn_shard_directory(tenant_key, shard_name)` pins that tenant to a
274/// named shard regardless of where the slot hash would place it. Tenants with
275/// no row route by [`HashShardRouter`], so the directory only needs entries
276/// for relocated/"whale" tenants.
277///
278/// Directory **hits** (a real pin) are cached for
279/// [`DEFAULT_DIRECTORY_CACHE_TTL`] so steady-state routing of pinned tenants
280/// issues no control-DB query. **Misses are not cached** — an unpinned tenant
281/// re-reads the directory on every route. This keeps the move workflow safe:
282/// once an operator inserts a directory row, no other process can keep routing
283/// that tenant to its old hash shard from a stale cached miss (there is no
284/// cross-process invalidation), so `move-slot --confirm` won't delete rows that
285/// late writes landed on the source. After changing a directory row, call
286/// [`invalidate`](Self::invalidate) for that key so the next route re-reads it.
287/// (NOTIFY-based cross-process invalidation is a planned follow-up; today the
288/// TTL bounds hit staleness and `invalidate` clears the local entry
289/// immediately.)
290///
291/// Install with
292/// [`AppBuilder::with_directory_shard_router`](crate::app::AppBuilder::with_directory_shard_router).
293///
294/// Only string keys are looked up in the directory (tenants are strings);
295/// numeric/byte keys route straight through the fallback.
296pub struct DirectoryShardRouter {
297    control_pool: Pool<RuntimeConnection>,
298    fallback: Arc<dyn ShardRouter>,
299    cache: std::sync::RwLock<HashMap<String, DirectoryCacheEntry>>,
300    ttl: std::time::Duration,
301    /// `statement_timeout` (ms) applied to the control-plane directory lookup so
302    /// a stuck control query / lock on `_autumn_shard_directory` fails within
303    /// the configured timeout instead of hanging every tenant-routed request.
304    /// `0` disables it. The router checks out a raw pooled connection (no
305    /// request context), so the timeout is set explicitly here.
306    statement_timeout_ms: u64,
307}
308
309/// Default time a resolved tenant→shard mapping is cached before re-reading
310/// the directory table.
311pub const DEFAULT_DIRECTORY_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(30);
312
313/// The `_autumn_shard_directory` table migration as a standalone embedded set.
314///
315/// Embedded separately so the app can auto-create the table at startup when
316/// directory routing is enabled (the `migrations/` copy is applied by
317/// `autumn migrate` for the control plane; this mirror lets auto-migrate
318/// deployments create the table without a manual migrate). The migration is
319/// `CREATE TABLE IF NOT EXISTS`, so applying it from either set is idempotent.
320/// Keep both copies in sync.
321#[cfg(feature = "db")]
322pub const SHARD_DIRECTORY_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
323    diesel_migrations::embed_migrations!("shard_directory_migrations");
324
325/// The `_autumn_shard_map` table migration as a standalone embedded set.
326///
327/// Embedded separately so the boot-time shard-map guard can auto-create its
328/// control table at startup (the `migrations/` copy is applied by
329/// `autumn migrate`). The migration is `CREATE TABLE IF NOT EXISTS`, so
330/// applying it from either set is idempotent. Keep both copies in sync.
331#[cfg(feature = "db")]
332pub const SHARD_MAP_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
333    diesel_migrations::embed_migrations!("shard_map_migrations");
334
335#[derive(Clone, Copy)]
336struct DirectoryCacheEntry {
337    shard: ShardId,
338    expires_at: std::time::Instant,
339}
340
341#[derive(diesel::QueryableByName)]
342struct ShardNameRow {
343    #[diesel(sql_type = diesel::sql_types::Text)]
344    shard_name: String,
345}
346
347/// Postgres `LISTEN`/`NOTIFY` channel the directory trigger fires on. The
348/// invalidation listener subscribes to it; the trigger
349/// (`autumn_notify_shard_directory_change`, in the shard-directory migration)
350/// must `pg_notify` the same channel. Keep the two in sync.
351#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
352const DIRECTORY_NOTIFY_CHANNEL: &str = "autumn_shard_directory";
353
354/// How often the invalidation listener wakes while idle to sweep expired cache
355/// entries and notice a dropped LISTEN connection.
356///
357/// Invalidation delivery itself is event-driven — a `NOTIFY` delivered at
358/// commit — so this only bounds idle housekeeping; kept well under
359/// [`DEFAULT_DIRECTORY_CACHE_TTL`].
360pub const DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL: std::time::Duration =
361    std::time::Duration::from_secs(5);
362
363impl std::fmt::Debug for DirectoryShardRouter {
364    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
365        f.debug_struct("DirectoryShardRouter")
366            .field("ttl", &self.ttl)
367            .field("cached_keys", &self.cache.read().map_or(0, |c| c.len()))
368            .finish_non_exhaustive()
369    }
370}
371
372impl DirectoryShardRouter {
373    /// Build a directory router over the given control pool, falling back to
374    /// [`HashShardRouter`] and using [`DEFAULT_DIRECTORY_CACHE_TTL`].
375    #[must_use]
376    pub fn new(control_pool: Pool<RuntimeConnection>) -> Self {
377        Self::with_fallback(control_pool, Arc::new(HashShardRouter))
378    }
379
380    /// Build a directory router with an explicit fallback router and the
381    /// default cache TTL.
382    #[must_use]
383    pub fn with_fallback(
384        control_pool: Pool<RuntimeConnection>,
385        fallback: Arc<dyn ShardRouter>,
386    ) -> Self {
387        Self {
388            control_pool,
389            fallback,
390            cache: std::sync::RwLock::new(HashMap::new()),
391            ttl: DEFAULT_DIRECTORY_CACHE_TTL,
392            statement_timeout_ms: 0,
393        }
394    }
395
396    /// Bound the control-plane directory lookup with `statement_timeout`
397    /// (milliseconds); `0` disables it. Typically the app's configured database
398    /// statement timeout, so a stuck control query fails fast instead of hanging
399    /// tenant routing.
400    #[must_use]
401    pub const fn with_statement_timeout_ms(mut self, statement_timeout_ms: u64) -> Self {
402        self.statement_timeout_ms = statement_timeout_ms;
403        self
404    }
405
406    /// Override the cache TTL.
407    #[must_use]
408    pub const fn with_cache_ttl(mut self, ttl: std::time::Duration) -> Self {
409        self.ttl = ttl;
410        self
411    }
412
413    /// Drop the cached mapping for `tenant_key`, forcing the next route to
414    /// re-read the directory. Call this after inserting, updating, or deleting
415    /// that tenant's directory row.
416    pub fn invalidate(&self, tenant_key: &str) {
417        if let Ok(mut cache) = self.cache.write() {
418            cache.remove(tenant_key);
419        }
420    }
421
422    /// Drop every cached mapping.
423    pub fn invalidate_all(&self) {
424        if let Ok(mut cache) = self.cache.write() {
425            cache.clear();
426        }
427    }
428
429    /// Spawn a background task that `LISTEN`s on the control DB's
430    /// `autumn_shard_directory` notification channel and invalidates this
431    /// router's cached pin whenever a tenant's directory row changes.
432    ///
433    /// This covers writes made on other replicas or directly via operator SQL
434    /// (the channel is fired by a trigger, not app code). Without it the cache
435    /// only refreshes when the TTL expires; with it a re-pin during a slot move
436    /// is picked up the moment it commits.
437    ///
438    /// Postgres delivers `NOTIFY` at **commit** (never before), so the
439    /// invalidation arrives exactly when the new mapping becomes visible: a
440    /// slow-committing re-pin cannot be skipped the way a timestamp-cursor poll
441    /// could. The cache TTL stays the backstop for any window where the LISTEN
442    /// connection is down.
443    ///
444    /// `control_url` is the control database URL backing this router's control
445    /// pool. Must be called from within a Tokio runtime; the returned handle can
446    /// be detached, and the task runs for the life of the process. The framework
447    /// spawns this automatically when directory routing is enabled via the
448    /// built-in path. `sweep_interval` only bounds idle housekeeping (see
449    /// [`DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL`]).
450    #[cfg(feature = "postgres")]
451    #[must_use]
452    pub fn spawn_invalidation_listener(
453        router: Arc<Self>,
454        control_url: String,
455        sweep_interval: std::time::Duration,
456    ) -> tokio::task::JoinHandle<()> {
457        use diesel_async::{AsyncConnection as _, RunQueryDsl as _};
458        use futures::StreamExt as _;
459
460        tokio::spawn(async move {
461            loop {
462                // (Re)connect and subscribe. On any failure back off for one
463                // sweep interval and retry; the cache TTL backstops staleness
464                // while we're disconnected.
465                let Ok(mut conn) = AsyncPgConnection::establish(&control_url).await else {
466                    tokio::time::sleep(sweep_interval).await;
467                    continue;
468                };
469                if diesel::sql_query(format!("LISTEN {DIRECTORY_NOTIFY_CHANNEL}"))
470                    .execute(&mut conn)
471                    .await
472                    .is_err()
473                {
474                    tokio::time::sleep(sweep_interval).await;
475                    continue;
476                }
477                // A re-pin may have committed between losing the previous
478                // connection and (re)subscribing; those NOTIFYs are gone, so drop
479                // the whole cache and let it repopulate lazily from the directory.
480                router.invalidate_all();
481
482                // Drain notifications until the stream errors or ends, then fall
483                // through to the outer loop and reconnect. Each idle
484                // `sweep_interval` we reclaim expired entries that were never
485                // looked up again (lazy eviction in `cache_get` only fires on
486                // re-observation).
487                let mut notifications = std::pin::pin!(conn.notifications_stream());
488                loop {
489                    match tokio::time::timeout(sweep_interval, notifications.next()).await {
490                        Ok(Some(Ok(notification))) => router.invalidate(&notification.payload),
491                        Ok(Some(Err(_)) | None) => break,
492                        Err(_elapsed) => router.sweep_expired(),
493                    }
494                }
495            }
496        })
497    }
498
499    fn cache_get(&self, key: &str) -> Option<ShardId> {
500        let now = crate::time::ambient_instant();
501        {
502            let cache = self.cache.read().ok()?;
503            match cache.get(key) {
504                Some(entry) if entry.expires_at > now => return Some(entry.shard),
505                // Miss, or present-but-expired: fall through. `None` is returned
506                // either way; an expired entry is additionally evicted below so a
507                // long-running process doesn't retain every pinned tenant it has
508                // ever looked up (the TTL bounds staleness, not memory).
509                Some(_) => {}
510                None => return None,
511            }
512        }
513        // Evict the expired entry under the write lock. Re-check expiry (against
514        // the same `now`) so we don't drop a fresh entry written by `cache_put`
515        // between releasing the read lock and taking the write lock.
516        let mut cache = self.cache.write().ok()?;
517        if cache.get(key).is_some_and(|entry| entry.expires_at <= now) {
518            cache.remove(key);
519        }
520        None
521    }
522
523    /// Drop every expired entry from the cache. Lazy eviction in `cache_get`
524    /// only reclaims keys that are looked up again; this bounds memory for
525    /// pinned tenants that are never re-observed. Called periodically by the
526    /// invalidation listener.
527    #[cfg_attr(not(feature = "postgres"), allow(dead_code))]
528    fn sweep_expired(&self) {
529        if let Ok(mut cache) = self.cache.write() {
530            let now = crate::time::ambient_instant();
531            cache.retain(|_, entry| entry.expires_at > now);
532        }
533    }
534
535    fn cache_put(&self, key: String, shard: ShardId) {
536        if let Ok(mut cache) = self.cache.write() {
537            cache.insert(
538                key,
539                DirectoryCacheEntry {
540                    shard,
541                    expires_at: crate::time::ambient_instant() + self.ttl,
542                },
543            );
544        }
545    }
546
547    /// Look up a tenant key in the directory table. Returns the resolved
548    /// `ShardId` on a directory hit, or `None` when the tenant has no row
549    /// (the caller then falls back to the hash router).
550    async fn lookup_directory(
551        &self,
552        key: &str,
553        shards: &ShardSet,
554    ) -> Result<Option<ShardId>, AutumnError> {
555        use diesel::OptionalExtension as _;
556        use diesel_async::RunQueryDsl;
557
558        // Staging fault injection (#3071). Inert outside a fault scope.
559        crate::fault_injection::__database_fault().await?;
560        let mut conn = self.control_pool.get().await.map_err(|e| {
561            AutumnError::service_unavailable_msg(format!(
562                "DirectoryShardRouter could not acquire a control connection: {e}"
563            ))
564        })?;
565
566        // Bound the lookup so a stuck control query / lock doesn't hang routing.
567        // Always issued (even for 0 = disabled) because the raw pooled checkout
568        // can return a connection carrying a shorter route-specific timeout set
569        // by a prior `Db`/repository checkout; mirror the normal checkout path,
570        // which always sets `statement_timeout`.
571        diesel::sql_query(format!(
572            "SET statement_timeout = {}",
573            self.statement_timeout_ms
574        ))
575        .execute(&mut conn)
576        .await
577        .map_err(|e| {
578            AutumnError::service_unavailable_msg(format!(
579                "DirectoryShardRouter could not set statement_timeout: {e}"
580            ))
581        })?;
582
583        let row = diesel::sql_query(
584            "SELECT shard_name FROM _autumn_shard_directory WHERE tenant_key = $1",
585        )
586        .bind::<diesel::sql_types::Text, _>(key)
587        .get_result::<ShardNameRow>(&mut conn)
588        .await
589        .optional()
590        .map_err(|e| {
591            AutumnError::service_unavailable_msg(format!(
592                "DirectoryShardRouter directory lookup failed: {e}"
593            ))
594        })?;
595
596        let Some(row) = row else {
597            return Ok(None);
598        };
599
600        let shard = shards.by_name(&row.shard_name).ok_or_else(|| {
601            AutumnError::service_unavailable_msg(format!(
602                "shard directory pins tenant {key:?} to unknown shard {:?}",
603                row.shard_name
604            ))
605        })?;
606        Ok(Some(shard.id()))
607    }
608}
609
610impl ShardRouter for DirectoryShardRouter {
611    fn route<'a>(
612        &'a self,
613        key: ShardKey<'a>,
614        shards: &'a ShardSet,
615    ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
616        Box::pin(async move {
617            // Only string keys participate in the directory (tenants are
618            // strings); numeric/byte keys route straight through the fallback.
619            let ShardKey::Str(key_str) = key else {
620                return self.fallback.route(key, shards).await;
621            };
622
623            if let Some(cached) = self.cache_get(key_str) {
624                return Ok(cached);
625            }
626
627            // Only cache real directory hits. A miss routes through the hash
628            // fallback WITHOUT caching: during a tenant move the operator
629            // inserts a directory row, and a cached miss on another process
630            // (e.g. a replica) would keep routing that tenant to its old hash
631            // shard until the TTL expired — there is no cross-process
632            // invalidation — and `move-slot --confirm` could then delete rows
633            // those stale writes had landed on the source. Re-querying unpinned
634            // tenants each route is the safe default.
635            match self.lookup_directory(key_str, shards).await? {
636                Some(shard) => {
637                    self.cache_put(key_str.to_owned(), shard);
638                    Ok(shard)
639                }
640                None => self.fallback.route(key, shards).await,
641            }
642        })
643    }
644}
645
646// ── Shard runtime state ──────────────────────────────────────────────────────
647
648/// Mutable per-shard replica readiness, updated by the per-shard health
649/// indicator on readiness probes (mirrors the control replica's
650/// [`ProbeState`](crate::probe::ProbeState) lifecycle).
651#[derive(Debug)]
652pub(crate) struct ShardRuntime {
653    replica_fallback: ReplicaFallback,
654    replica_configured: bool,
655    connection_ready: AtomicBool,
656    migrations_ready: AtomicBool,
657    detail: std::sync::RwLock<Option<String>>,
658    /// `(primary_url, replica_url)` for re-running the migration parity
659    /// check from the per-shard health indicator. `None` when the app
660    /// registered no migrations.
661    migration_check: std::sync::RwLock<Option<(String, String)>>,
662    /// When the parity comparison last ran, for throttling: unlike the
663    /// pooled connectivity check, parity opens fresh synchronous
664    /// connections to both roles, so it must not run on every probe.
665    parity_checked_at: std::sync::Mutex<Option<std::time::Instant>>,
666}
667
668/// Minimum interval between migration parity re-checks per shard.
669///
670/// Readiness probes fire every few seconds per replica; the parity check
671/// opens fresh synchronous connections to the shard's primary *and*
672/// replica, so running it per probe per shard would exhaust Postgres
673/// connection limits as shard counts grow.
674const PARITY_RECHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
675
676// Mutators are driven by startup migration parity checks and the
677// per-shard health indicators; some are exercised only by tests until
678// the health wiring lands.
679#[cfg_attr(not(test), allow(dead_code))]
680impl ShardRuntime {
681    fn new(replica_fallback: ReplicaFallback, replica_configured: bool) -> Self {
682        Self {
683            replica_fallback,
684            replica_configured,
685            connection_ready: AtomicBool::new(false),
686            migrations_ready: AtomicBool::new(true),
687            detail: std::sync::RwLock::new(
688                replica_configured.then(|| "replica has not passed a readiness check".to_owned()),
689            ),
690            migration_check: std::sync::RwLock::new(None),
691            parity_checked_at: std::sync::Mutex::new(None),
692        }
693    }
694
695    pub(crate) fn configure_migration_check(&self, primary_url: String, replica_url: String) {
696        *self
697            .migration_check
698            .write()
699            .expect("shard runtime lock poisoned") = Some((primary_url, replica_url));
700    }
701
702    fn migration_check(&self) -> Option<(String, String)> {
703        self.migration_check
704            .read()
705            .expect("shard runtime lock poisoned")
706            .clone()
707    }
708
709    /// Whether the throttle window has elapsed; claims the slot when it
710    /// has, so concurrent probes run at most one parity check per window.
711    pub(crate) fn parity_check_due(&self) -> bool {
712        let mut checked_at = self
713            .parity_checked_at
714            .lock()
715            .expect("shard runtime lock poisoned");
716        let now = crate::time::ambient_instant();
717        if checked_at.is_none_or(|at| now.saturating_duration_since(at) >= PARITY_RECHECK_INTERVAL)
718        {
719            *checked_at = Some(now);
720            true
721        } else {
722            false
723        }
724    }
725
726    fn replica_ready(&self) -> bool {
727        self.connection_ready.load(Ordering::Relaxed)
728            && self.migrations_ready.load(Ordering::Relaxed)
729    }
730
731    fn refresh_detail(&self) {
732        if self.replica_ready() {
733            *self.detail.write().expect("shard runtime lock poisoned") = None;
734        }
735    }
736
737    pub(crate) fn mark_replica_connection_ready(&self) {
738        self.connection_ready.store(true, Ordering::Relaxed);
739        self.refresh_detail();
740    }
741
742    pub(crate) fn mark_replica_connection_unready(&self, detail: impl Into<String>) {
743        self.connection_ready.store(false, Ordering::Relaxed);
744        *self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
745    }
746
747    pub(crate) fn mark_replica_migrations_ready(&self) {
748        self.migrations_ready.store(true, Ordering::Relaxed);
749        self.refresh_detail();
750    }
751
752    pub(crate) fn mark_replica_migrations_unready(&self, detail: impl Into<String>) {
753        self.migrations_ready.store(false, Ordering::Relaxed);
754        *self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
755    }
756
757    pub(crate) fn detail(&self) -> Option<String> {
758        self.detail
759            .read()
760            .expect("shard runtime lock poisoned")
761            .clone()
762    }
763}
764
765// ── Shard / ShardSet ─────────────────────────────────────────────────────────
766
767/// One physical shard: a named [`DatabaseTopology`] plus its slot
768/// assignment and runtime replica state.
769#[derive(Clone)]
770pub struct Shard {
771    name: Arc<str>,
772    id: ShardId,
773    slots: Arc<[u16]>,
774    topology: DatabaseTopology,
775    runtime: Arc<ShardRuntime>,
776    /// The `shard=` value on metrics and route labels: the name of a
777    /// configured shard, `fleet` for a fleet database (whose name is
778    /// unbounded in a tenant fleet).
779    metric_label: Arc<str>,
780    /// The fleet database behind this shard, for a fleet-backed set.
781    #[cfg(feature = "sqlite")]
782    fleet_db: Option<crate::db::fleet::FleetDatabase>,
783}
784
785/// A handle that keeps a fleet database open while it is held (ADR 0019).
786///
787/// The fleet never closes a database while any lease on it is alive, so a
788/// connection acquired lazily — a generated repository's first query, a
789/// fan-out future — never finds its pool closed. Opaque on purpose.
790pub type ShardLease = Arc<dyn std::any::Any + Send + Sync>;
791
792/// The metric label of every fleet database. See [`Shard::metric_label`].
793pub const FLEET_METRIC_LABEL: &str = "fleet";
794
795impl Shard {
796    /// Stable shard name: the configured name, or `tenant:<id>` /
797    /// `slot:<nnnnn>` for a fleet database (ADR 0019).
798    #[must_use]
799    pub fn name(&self) -> &str {
800        &self.name
801    }
802
803    /// The bounded label metrics and route keys carry for this shard: its
804    /// name for a configured shard, [`FLEET_METRIC_LABEL`] for a fleet
805    /// database. Spans and logs carry [`name`](Self::name).
806    #[must_use]
807    pub fn metric_label(&self) -> &str {
808        &self.metric_label
809    }
810
811    /// A lease that keeps this shard's database open while held: `Some` for
812    /// a fleet database, `None` for a configured shard (always open).
813    #[must_use]
814    #[allow(
815        clippy::missing_const_for_fn,
816        clippy::unused_self,
817        reason = "reads the fleet handle under the sqlite feature"
818    )]
819    pub fn lease(&self) -> Option<ShardLease> {
820        #[cfg(feature = "sqlite")]
821        {
822            self.fleet_db
823                .as_ref()
824                .map(crate::db::fleet::FleetDatabase::lease)
825        }
826        #[cfg(not(feature = "sqlite"))]
827        {
828            None
829        }
830    }
831
832    /// The fleet database behind this shard, when the set is fleet-backed.
833    #[cfg(feature = "sqlite")]
834    #[must_use]
835    pub const fn fleet_database(&self) -> Option<&crate::db::fleet::FleetDatabase> {
836        self.fleet_db.as_ref()
837    }
838
839    /// A shard over one open fleet database.
840    #[cfg(feature = "sqlite")]
841    pub(crate) fn for_fleet_db(db: crate::db::fleet::FleetDatabase) -> Self {
842        static RUNTIME: std::sync::LazyLock<Arc<ShardRuntime>> = std::sync::LazyLock::new(|| {
843            Arc::new(ShardRuntime::new(ReplicaFallback::Primary, false))
844        });
845        let slot = db.key().routing_slot();
846        Self {
847            name: Arc::from(db.key().name()),
848            id: fleet_shard_id(db.key()),
849            slots: Arc::from([slot]),
850            topology: DatabaseTopology::primary_only(db.pool().clone()),
851            runtime: Arc::clone(&RUNTIME),
852            metric_label: Arc::from(FLEET_METRIC_LABEL),
853            fleet_db: Some(db),
854        }
855    }
856
857    /// Position of this shard in the configured set.
858    #[must_use]
859    pub const fn id(&self) -> ShardId {
860        self.id
861    }
862
863    /// Logical slots owned by this shard, in ascending order.
864    #[must_use]
865    pub fn slots(&self) -> &[u16] {
866        &self.slots
867    }
868
869    /// This shard's primary/replica pool topology.
870    #[must_use]
871    pub const fn topology(&self) -> &DatabaseTopology {
872        &self.topology
873    }
874
875    /// This shard's primary/write pool.
876    #[must_use]
877    pub const fn primary_pool(&self) -> &Pool<RuntimeConnection> {
878        self.topology.primary()
879    }
880
881    /// This shard's replica pool, when configured.
882    #[must_use]
883    pub const fn replica_pool(&self) -> Option<&Pool<RuntimeConnection>> {
884        self.topology.replica()
885    }
886
887    /// Pool for read-only work, honoring this shard's `replica_fallback`
888    /// and runtime replica readiness (mirrors
889    /// [`AppState::read_pool`](crate::AppState::read_pool)):
890    ///
891    /// - no replica configured → the primary pool;
892    /// - replica configured and ready → the replica pool;
893    /// - replica unready, fallback `primary` → the primary pool;
894    /// - replica unready, fallback `fail_readiness` → `None`.
895    #[must_use]
896    pub fn read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
897        self.read_pool_with_role().map(|(pool, _)| pool)
898    }
899
900    /// Snapshot this shard's read-routing decision as a
901    /// [`ReadRoute`](crate::repository::ReadRoute), the per-shard analogue of
902    /// [`ReadRoute::from_state`](crate::repository::ReadRoute::from_state).
903    ///
904    /// [`ShardedDb`] captures this at extraction time so a generated
905    /// `#[repository]` built with `from_shard` routes its read-only methods
906    /// to the shard's replica automatically — mirroring [`read_pool`] and
907    /// honoring the shard's `replica_fallback` policy and replica readiness:
908    ///
909    /// - no replica configured → [`Primary`](crate::repository::ReadRoute::Primary);
910    /// - replica ready → [`ReadPool`](crate::repository::ReadRoute::ReadPool) over the replica;
911    /// - replica unready, fallback `primary` → `ReadPool` over the primary;
912    /// - replica unready, fallback `fail_readiness` →
913    ///   [`Unavailable`](crate::repository::ReadRoute::Unavailable).
914    ///
915    /// [`read_pool`]: Self::read_pool
916    #[must_use]
917    pub fn read_route(&self) -> crate::repository::ReadRoute {
918        use crate::repository::ReadRoute;
919        if !self.runtime.replica_configured {
920            return ReadRoute::Primary;
921        }
922        self.read_pool().map_or(ReadRoute::Unavailable, |pool| {
923            ReadRoute::ReadPool(pool.clone())
924        })
925    }
926
927    /// The shard's replica pool for **explicit replica-only** reads.
928    ///
929    /// Returns the replica pool only when a replica is configured *and* has
930    /// passed its readiness checks. Never returns the primary pool — this is
931    /// the `replica_fallback`-independent counterpart to [`read_pool`]:
932    ///
933    /// - no replica configured → `None`;
934    /// - replica configured but unready (regardless of `replica_fallback`) → `None`;
935    /// - replica configured and ready → `Some(replica_pool)`.
936    ///
937    /// Backs [`ShardedReadDb`], which always requires a healthy replica.
938    ///
939    /// [`read_pool`]: Self::read_pool
940    pub(crate) fn replica_read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
941        if self.runtime.replica_configured && self.runtime.replica_ready() {
942            self.topology.replica()
943        } else {
944            None
945        }
946    }
947
948    /// [`read_pool`](Self::read_pool) plus the role label of the returned
949    /// pool, for interceptor/metric naming.
950    pub(crate) fn read_pool_with_role(&self) -> Option<(&Pool<RuntimeConnection>, &'static str)> {
951        if !self.runtime.replica_configured {
952            return Some((self.topology.primary(), "primary"));
953        }
954        if self.runtime.replica_ready() {
955            return self.topology.replica().map(|pool| (pool, "replica"));
956        }
957        match self.runtime.replica_fallback {
958            ReplicaFallback::Primary => Some((self.topology.primary(), "primary")),
959            ReplicaFallback::FailReadiness => None,
960        }
961    }
962
963    #[cfg_attr(not(test), allow(dead_code))]
964    pub(crate) fn runtime(&self) -> &ShardRuntime {
965        &self.runtime
966    }
967}
968
969impl std::fmt::Debug for Shard {
970    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
971        f.debug_struct("Shard")
972            .field("name", &self.name)
973            .field("id", &self.id)
974            .field("slots", &self.slots)
975            .finish_non_exhaustive()
976    }
977}
978
979struct ShardSetInner {
980    shards: Vec<Shard>,
981    by_name: HashMap<String, usize>,
982    /// `slot_map[slot]` is the index into `shards` of the slot's owner.
983    slot_map: Vec<usize>,
984    router: Arc<dyn ShardRouter>,
985    /// The fleet behind a fleet-backed set (ADR 0019), with the control pool
986    /// that seeds tenant-free cross-shard repositories.
987    #[cfg(feature = "sqlite")]
988    fleet: Option<FleetBacking>,
989}
990
991#[cfg(feature = "sqlite")]
992struct FleetBacking {
993    fleet: crate::db::fleet::DatabaseFleet,
994    control_pool: Pool<RuntimeConnection>,
995}
996
997/// The configured set of shards plus the routing strategy.
998///
999/// Cheap to clone (a single `Arc`). Available from
1000/// [`AppState::shards`](crate::AppState::shards) and through the
1001/// [`Shards`] extractor.
1002#[derive(Clone)]
1003pub struct ShardSet {
1004    inner: Arc<ShardSetInner>,
1005}
1006
1007impl ShardSet {
1008    /// A set whose shards are the databases of a [`DatabaseFleet`]
1009    /// (ADR 0019): every key resolves to its tenant's (or slot's) own
1010    /// `SQLite` file, opened on first use.
1011    ///
1012    /// `control_pool` seeds tenant-free [`CrossShard`] repositories before
1013    /// they fan out. A fleet-backed set has no configured shards:
1014    /// [`len`](Self::len), [`iter`](Self::iter), [`get`](Self::get) and the
1015    /// slot-map helpers see none, and [`route`](Self::route) refuses —
1016    /// resolve keys with [`resolve`](Self::resolve).
1017    ///
1018    /// [`DatabaseFleet`]: crate::db::fleet::DatabaseFleet
1019    #[cfg(feature = "sqlite")]
1020    #[must_use]
1021    pub fn from_fleet(
1022        fleet: crate::db::fleet::DatabaseFleet,
1023        control_pool: Pool<RuntimeConnection>,
1024    ) -> Self {
1025        Self {
1026            inner: Arc::new(ShardSetInner {
1027                shards: Vec::new(),
1028                by_name: HashMap::new(),
1029                slot_map: Vec::new(),
1030                router: Arc::new(HashShardRouter),
1031                fleet: Some(FleetBacking {
1032                    fleet,
1033                    control_pool,
1034                }),
1035            }),
1036        }
1037    }
1038
1039    /// The fleet behind a fleet-backed set.
1040    #[cfg(feature = "sqlite")]
1041    #[must_use]
1042    pub fn fleet(&self) -> Option<&crate::db::fleet::DatabaseFleet> {
1043        self.inner.fleet.as_ref().map(|backing| &backing.fleet)
1044    }
1045
1046    /// Whether this set is backed by a fleet rather than configured shards.
1047    #[must_use]
1048    #[allow(
1049        clippy::missing_const_for_fn,
1050        reason = "reads through the Arc under the sqlite feature, which is not const"
1051    )]
1052    pub fn is_fleet(&self) -> bool {
1053        #[cfg(feature = "sqlite")]
1054        {
1055            self.inner.fleet.is_some()
1056        }
1057        #[cfg(not(feature = "sqlite"))]
1058        {
1059            false
1060        }
1061    }
1062
1063    /// Resolve the shard that owns `key`, as an owned (cheap, `Arc`-backed)
1064    /// [`Shard`].
1065    ///
1066    /// For configured shards this is [`route`](Self::route). For a
1067    /// fleet-backed set it opens the key's database (creating it when the
1068    /// fleet creates on demand). Every framework extractor resolves through
1069    /// here, so both kinds of set serve [`ShardedDb`], [`Shards`] and
1070    /// `#[repository(sharded)]`.
1071    ///
1072    /// # Errors
1073    ///
1074    /// The router's error, or the fleet's: `400` for a tenant id that cannot
1075    /// name a file, `404` for an unprovisioned tenant, `503` for pending
1076    /// migrations.
1077    pub async fn resolve<'k>(&self, key: impl Into<ShardKey<'k>>) -> Result<Shard, AutumnError> {
1078        let key = key.into();
1079        #[cfg(feature = "sqlite")]
1080        if let Some(backing) = &self.inner.fleet {
1081            let db_key = backing.fleet.key_for_shard_key(key)?;
1082            let db = backing.fleet.open(&db_key).await?;
1083            return Ok(Shard::for_fleet_db(db));
1084        }
1085        self.route(key).await.cloned()
1086    }
1087
1088    /// Resolve a shard by name: a configured shard's name, or a fleet
1089    /// database's `tenant:<id>` / `slot:<n>` name (opened, never created).
1090    ///
1091    /// # Errors
1092    ///
1093    /// `400` for an unknown or malformed name, or the fleet's open error.
1094    #[allow(
1095        clippy::unused_async,
1096        reason = "async for the fleet arm, which opens the database (sqlite feature)"
1097    )]
1098    pub async fn resolve_name(&self, name: &str) -> Result<Shard, AutumnError> {
1099        #[cfg(feature = "sqlite")]
1100        if let Some(backing) = &self.inner.fleet {
1101            let key = backing.fleet.key_for_name(name)?;
1102            if !backing.fleet.exists(&key) {
1103                return Err(crate::db::fleet::FleetError::NotFound { name: key.name() }.into());
1104            }
1105            let db = backing.fleet.open(&key).await?;
1106            return Ok(Shard::for_fleet_db(db));
1107        }
1108        self.by_name(name)
1109            .cloned()
1110            .ok_or_else(|| AutumnError::bad_request_msg(format!("unknown shard {name:?}")))
1111    }
1112
1113    /// Number of configured shards (`0` for a fleet-backed set, whose
1114    /// databases are opened on demand; see [`is_fleet`](Self::is_fleet)).
1115    #[must_use]
1116    pub fn len(&self) -> usize {
1117        self.inner.shards.len()
1118    }
1119
1120    /// Whether the set contains no shards.
1121    #[must_use]
1122    pub fn is_empty(&self) -> bool {
1123        self.inner.shards.is_empty()
1124    }
1125
1126    /// Number of logical slots — the fixed [`SLOT_COUNT`] (16384).
1127    #[must_use]
1128    pub const fn slot_count(&self) -> u16 {
1129        SLOT_COUNT
1130    }
1131
1132    /// Shard by positional id.
1133    #[must_use]
1134    pub fn get(&self, id: ShardId) -> Option<&Shard> {
1135        self.inner.shards.get(id.0)
1136    }
1137
1138    /// Shard by configured name.
1139    #[must_use]
1140    pub fn by_name(&self, name: &str) -> Option<&Shard> {
1141        self.inner
1142            .by_name
1143            .get(name)
1144            .and_then(|&idx| self.inner.shards.get(idx))
1145    }
1146
1147    /// Iterate shards in declaration order.
1148    pub fn iter(&self) -> impl Iterator<Item = &Shard> {
1149        self.inner.shards.iter()
1150    }
1151
1152    /// Map a routing key onto its logical slot (deterministic hash; see
1153    /// the module docs for the permanence guarantee).
1154    #[must_use]
1155    pub fn slot_for_key<'k>(&self, key: impl Into<ShardKey<'k>>) -> SlotId {
1156        slot_for_key(key.into())
1157    }
1158
1159    /// Owner of a logical slot per the configured slot map.
1160    #[must_use]
1161    pub fn shard_for_slot(&self, slot: SlotId) -> Option<&Shard> {
1162        self.inner
1163            .slot_map
1164            .get(usize::from(slot.0))
1165            .and_then(|&idx| self.inner.shards.get(idx))
1166    }
1167
1168    /// Resolve the shard that owns `key` via the installed
1169    /// [`ShardRouter`].
1170    ///
1171    /// # Errors
1172    ///
1173    /// Returns the router's error, or an internal error if the router
1174    /// produced an out-of-range [`ShardId`].
1175    pub async fn route<'k>(&self, key: impl Into<ShardKey<'k>>) -> Result<&Shard, AutumnError> {
1176        let key = key.into();
1177        if self.is_fleet() {
1178            return Err(AutumnError::internal_server_error_msg(
1179                "ShardSet::route cannot borrow a fleet database, which is opened on demand; \
1180                 call ShardSet::resolve instead",
1181            ));
1182        }
1183        let id = self.inner.router.route(key, self).await?;
1184        self.get(id).ok_or_else(|| {
1185            AutumnError::service_unavailable_msg(format!(
1186                "shard router returned out-of-range shard id {} (have {} shards)",
1187                id.0,
1188                self.len()
1189            ))
1190        })
1191    }
1192
1193    /// Total configured `max_size` across every pool in the set
1194    /// (primaries plus replicas). Logged at startup so N-shard
1195    /// deployments notice multiplied connection counts.
1196    #[must_use]
1197    pub fn total_max_connections(&self) -> usize {
1198        #[cfg(feature = "sqlite")]
1199        if let Some(backing) = &self.inner.fleet {
1200            return backing
1201                .fleet
1202                .max_open()
1203                .saturating_mul(backing.fleet.pool_size());
1204        }
1205        self.inner
1206            .shards
1207            .iter()
1208            .map(|shard| {
1209                shard.topology().primary().status().max_size
1210                    + shard
1211                        .topology()
1212                        .replica()
1213                        .map_or(0, |pool| pool.status().max_size)
1214            })
1215            .sum()
1216    }
1217
1218    /// Whether `key` is owned by the shard at the given index in declaration order.
1219    ///
1220    /// Uses the hash-based slot assignment, **not** the installed router (which
1221    /// may override routing for individual tenants via a directory). Use this for
1222    /// tooling / slot-move scripts where you need to verify ownership without
1223    /// issuing an async router call.
1224    #[must_use]
1225    pub fn owns_key<'k>(&self, shard_id: ShardId, key: impl Into<ShardKey<'k>>) -> bool {
1226        let slot = self.slot_for_key(key);
1227        self.shard_for_slot(slot)
1228            .is_some_and(|s| s.id() == shard_id)
1229    }
1230
1231    /// All logical slots assigned to the shard at index `shard_id`.
1232    ///
1233    /// Returns `None` when the id is out of range.
1234    #[must_use]
1235    pub fn slots_for_shard(&self, shard_id: ShardId) -> Option<&[u16]> {
1236        self.inner.shards.get(shard_id.0).map(Shard::slots)
1237    }
1238
1239    /// Partition string `keys` by their owning shard based on hash-slot assignment.
1240    ///
1241    /// Keys are grouped in declaration order; the returned map may have fewer
1242    /// entries than `self.len()` when some shards own none of the given keys.
1243    /// Useful for slot-move tooling that needs to issue `WHERE tenant_id = ANY($1)`
1244    /// per destination shard.
1245    #[must_use]
1246    pub fn partition_by_shard<'k>(
1247        &self,
1248        keys: impl IntoIterator<Item = &'k str>,
1249    ) -> std::collections::HashMap<ShardId, Vec<&'k str>> {
1250        let mut map: std::collections::HashMap<ShardId, Vec<&'k str>> =
1251            std::collections::HashMap::new();
1252        for key in keys {
1253            let slot = self.slot_for_key(key);
1254            if let Some(shard) = self.shard_for_slot(slot) {
1255                map.entry(shard.id()).or_default().push(key);
1256            }
1257        }
1258        map
1259    }
1260
1261    /// Fan out a closure over every shard concurrently, collecting one result
1262    /// per shard.  Fails the whole call if **any** shard errors.
1263    ///
1264    /// Intended for cross-shard read fan-out from `across_tenants()` reads on
1265    /// `#[repository(tenant_scoped, sharded)]` repositories.  The closure
1266    /// receives each [`Shard`] so it can build a sub-repo that honors that
1267    /// shard's read routing (replica/primary/fail-closed) and the parent
1268    /// request context; the sub-repo must set `__autumn_shards = None` so
1269    /// recursion is impossible.
1270    ///
1271    /// The closure is invoked synchronously per shard (the `&Shard` borrow ends
1272    /// when it returns the owned, `'static` future), so the futures can run
1273    /// concurrently without borrowing the [`ShardSet`].
1274    ///
1275    /// Concurrency is bounded at [`FAN_OUT_CONCURRENCY`] so a cross-tenant admin
1276    /// read does not check out a connection from every shard at once (which
1277    /// could spike load or exhaust connection limits on large fleets), matching
1278    /// the public [`Shards::each_shard`] pipeline. Results are returned in shard
1279    /// **declaration order** (not completion order), so order-dependent merges
1280    /// such as `search`'s per-shard ranking concatenation are deterministic.
1281    /// Fails the whole call on the first shard error.
1282    ///
1283    /// This is a framework-internal primitive used by generated repository
1284    /// code.  It is `pub` so that downstream crates can call it from
1285    /// `#[repository]`-generated `impl` blocks, but it is not part of the
1286    /// stable public API.
1287    #[doc(hidden)]
1288    pub async fn fan_out_shards<T, Fut, F>(&self, f: F) -> Result<Vec<T>, crate::AutumnError>
1289    where
1290        T: Send + 'static,
1291        Fut: std::future::Future<Output = Result<T, crate::AutumnError>> + Send + 'static,
1292        F: Fn(&Shard) -> Fut + Send + Sync,
1293    {
1294        use futures::StreamExt as _;
1295
1296        // A fleet fans out over every database on disk, in key order, opening
1297        // each (never creating one).
1298        #[cfg(feature = "sqlite")]
1299        if let Some(backing) = &self.inner.fleet {
1300            let f = &f;
1301            return backing
1302                .fleet
1303                .each(FAN_OUT_CONCURRENCY, |db| {
1304                    let shard = Shard::for_fleet_db(db);
1305                    f(&shard)
1306                })
1307                .await
1308                .map_err(AutumnError::from)?
1309                .into_iter()
1310                .map(|(_, result)| result)
1311                .collect();
1312        }
1313
1314        // Results are placed by shard index so declaration order is preserved
1315        // even though `FuturesUnordered` yields them in completion order.
1316        let mut slots: Vec<Option<T>> = (0..self.inner.shards.len()).map(|_| None).collect();
1317        let mut in_flight = futures::stream::FuturesUnordered::new();
1318
1319        for (idx, shard) in self.inner.shards.iter().enumerate() {
1320            if in_flight.len() >= FAN_OUT_CONCURRENCY
1321                && let Some((i, result)) = in_flight.next().await
1322            {
1323                slots[i] = Some(result?);
1324            }
1325            let fut = f(shard);
1326            in_flight.push(async move { (idx, fut.await) });
1327        }
1328        while let Some((i, result)) = in_flight.next().await {
1329            slots[i] = Some(result?);
1330        }
1331        // Every shard pushed exactly one future and all were drained above, so
1332        // on the success path every slot is filled.
1333        Ok(slots
1334            .into_iter()
1335            .map(|slot| slot.expect("every shard produced a result"))
1336            .collect())
1337    }
1338}
1339
1340impl std::fmt::Debug for ShardSet {
1341    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1342        f.debug_struct("ShardSet")
1343            .field("shards", &self.inner.shards)
1344            .finish_non_exhaustive()
1345    }
1346}
1347
1348// ── Construction ─────────────────────────────────────────────────────────────
1349
1350/// Error building a [`ShardSet`] from configuration.
1351#[derive(Debug, thiserror::Error)]
1352#[non_exhaustive]
1353pub enum ShardSetBuildError {
1354    /// A shard's connection pool could not be constructed.
1355    #[error("failed to build pool for shard {shard:?}: {source}")]
1356    Pool {
1357        /// Name of the failing shard.
1358        shard: String,
1359        /// Underlying pool construction error.
1360        source: PoolError,
1361    },
1362    /// The slot map could not be resolved from configuration.
1363    #[error(transparent)]
1364    Config(#[from] ConfigError),
1365    /// A custom provider returned the wrong number of shard topologies.
1366    #[error("expected {expected} shard topologies, got {actual}")]
1367    TopologyCountMismatch {
1368        /// Number of configured shards.
1369        expected: usize,
1370        /// Number of topologies supplied.
1371        actual: usize,
1372    },
1373}
1374
1375/// Build a [`ShardSet`] from configuration using the default deadpool
1376/// factory for every shard topology.
1377///
1378/// Returns `Ok(None)` when no `[[database.shards]]` entries are
1379/// configured.
1380///
1381/// # Errors
1382///
1383/// Returns [`ShardSetBuildError`] when a pool cannot be constructed or
1384/// the slot map is invalid.
1385pub fn create_shard_set(
1386    config: &DatabaseConfig,
1387    router: Arc<dyn ShardRouter>,
1388) -> Result<Option<ShardSet>, ShardSetBuildError> {
1389    if !config.has_shards() {
1390        return Ok(None);
1391    }
1392    let topologies = config
1393        .shards
1394        .iter()
1395        .map(|shard| {
1396            crate::db::create_shard_topology(shard, config).map_err(|source| {
1397                ShardSetBuildError::Pool {
1398                    shard: shard.name.clone(),
1399                    source,
1400                }
1401            })
1402        })
1403        .collect::<Result<Vec<_>, _>>()?;
1404    build_shard_set(config, topologies, router).map(Some)
1405}
1406
1407/// Build a [`ShardSet`] where every shard primary pool uses `max_size(1)` and
1408/// wraps each connection in a test transaction that is rolled back when the
1409/// connection is returned.
1410///
1411/// This mirrors the transactional control-pool logic in `TestApp` so that
1412/// shard repositories in integration tests see rolled-back state between test
1413/// runs.
1414///
1415/// **Deadlock caveat:** with `max_size(1)` a handler that checks out the same
1416/// shard connection twice in a single request will deadlock (same as the
1417/// control pool).  Use a separate non-transactional shard set when a test
1418/// requires concurrent shard checkouts.
1419///
1420/// # Errors
1421///
1422/// Returns [`ShardSetBuildError`] when no shards are configured, any pool
1423/// cannot be built, or the slot map is invalid.
1424///
1425/// Postgres-only: each shard pool is established with a `begin_test_transaction`
1426/// rollback hook (Postgres transactional test isolation), so this helper is not
1427/// compiled under the `sqlite` feature — its sole caller, the Postgres
1428/// transactional `TestApp` harness (`crate::test`), is likewise Postgres-only.
1429#[cfg(not(feature = "sqlite"))]
1430pub fn create_shard_set_transactional(
1431    config: &DatabaseConfig,
1432    router: Arc<dyn ShardRouter>,
1433) -> Result<Option<ShardSet>, ShardSetBuildError> {
1434    if !config.has_shards() {
1435        return Ok(None);
1436    }
1437
1438    let timeout = std::time::Duration::from_secs(config.connect_timeout_secs);
1439
1440    let topologies = config
1441        .shards
1442        .iter()
1443        .map(|shard| {
1444            let manager = diesel_async::pooled_connection::AsyncDieselConnectionManager::<
1445                diesel_async::AsyncPgConnection,
1446            >::new(&shard.primary_url);
1447            let pool = Pool::builder(manager)
1448                .max_size(1)
1449                .wait_timeout(Some(timeout))
1450                .create_timeout(Some(timeout))
1451                .runtime(deadpool::Runtime::Tokio1)
1452                .post_create(deadpool::managed::Hook::async_fn(
1453                    |conn: &mut diesel_async::AsyncPgConnection, _| {
1454                        Box::pin(async move {
1455                            use diesel_async::AsyncConnection as _;
1456                            use diesel_async::RunQueryDsl as _;
1457                            conn.begin_test_transaction().await.map_err(|e| {
1458                                deadpool::managed::HookError::Backend(
1459                                    diesel_async::pooled_connection::PoolError::QueryError(e),
1460                                )
1461                            })?;
1462                            diesel::sql_query("SET autumn.test_transaction_started = 'true'")
1463                                .execute(conn)
1464                                .await
1465                                .map_err(|e| {
1466                                    deadpool::managed::HookError::Backend(
1467                                        diesel_async::pooled_connection::PoolError::QueryError(e),
1468                                    )
1469                                })?;
1470                            Ok(())
1471                        })
1472                    },
1473                ))
1474                .build()
1475                .map_err(|source| ShardSetBuildError::Pool {
1476                    shard: shard.name.clone(),
1477                    source: crate::db::PoolError::Build(source),
1478                })?;
1479            Ok(crate::db::DatabaseTopology::primary_only(pool))
1480        })
1481        .collect::<Result<Vec<_>, ShardSetBuildError>>()?;
1482    build_shard_set(config, topologies, router).map(Some)
1483}
1484
1485/// Assemble a [`ShardSet`] from pre-built topologies (one per configured
1486/// shard, in declaration order). Used by custom
1487/// [`DatabasePoolProvider`](crate::db::DatabasePoolProvider)s and tests.
1488///
1489/// # Errors
1490///
1491/// Returns [`ShardSetBuildError`] when the topology count does not match
1492/// the configuration or the slot map is invalid.
1493pub fn build_shard_set(
1494    config: &DatabaseConfig,
1495    topologies: Vec<DatabaseTopology>,
1496    router: Arc<dyn ShardRouter>,
1497) -> Result<ShardSet, ShardSetBuildError> {
1498    if topologies.len() != config.shards.len() {
1499        return Err(ShardSetBuildError::TopologyCountMismatch {
1500            expected: config.shards.len(),
1501            actual: topologies.len(),
1502        });
1503    }
1504    let slot_map = config.resolved_slot_map()?;
1505
1506    let mut slots_per_shard: Vec<Vec<u16>> = vec![Vec::new(); config.shards.len()];
1507    for (slot, &owner) in slot_map.iter().enumerate() {
1508        #[allow(clippy::cast_possible_truncation)]
1509        slots_per_shard[owner].push(slot as u16);
1510    }
1511
1512    let shards: Vec<Shard> = config
1513        .shards
1514        .iter()
1515        .zip(topologies)
1516        .enumerate()
1517        .map(|(idx, (shard_config, topology))| {
1518            let replica_configured = topology.replica().is_some();
1519            Shard {
1520                name: Arc::from(shard_config.name.as_str()),
1521                id: ShardId(idx),
1522                slots: Arc::from(std::mem::take(&mut slots_per_shard[idx])),
1523                topology,
1524                runtime: Arc::new(ShardRuntime::new(
1525                    shard_config.effective_replica_fallback(config),
1526                    replica_configured,
1527                )),
1528                metric_label: Arc::from(shard_config.name.as_str()),
1529                #[cfg(feature = "sqlite")]
1530                fleet_db: None,
1531            }
1532        })
1533        .collect();
1534    // `AutumnConfig::validate()` already rejects duplicate names, but this
1535    // builder is public and reachable with unvalidated configs (custom
1536    // loaders, direct callers); a silently-shadowed map would make one of
1537    // the duplicates unaddressable via by_name/db_on and health components.
1538    let mut by_name = HashMap::with_capacity(shards.len());
1539    for (idx, shard) in shards.iter().enumerate() {
1540        if by_name.insert(shard.name().to_owned(), idx).is_some() {
1541            return Err(ConfigError::Validation(format!(
1542                "database.shards: shard name {:?} is declared more than once; \
1543                 shard names must be unique",
1544                shard.name()
1545            ))
1546            .into());
1547        }
1548    }
1549
1550    Ok(ShardSet {
1551        inner: Arc::new(ShardSetInner {
1552            shards,
1553            by_name,
1554            slot_map,
1555            router,
1556            #[cfg(feature = "sqlite")]
1557            fleet: None,
1558        }),
1559    })
1560}
1561
1562// ── Health ───────────────────────────────────────────────────────────────────
1563
1564/// Framework health indicator registered per shard as `db:shard:<name>`.
1565///
1566/// Mirrors the control topology's lifecycle: on every readiness probe it
1567/// live-checks primary and replica connectivity and re-runs the migration
1568/// parity comparison, feeding the shard's runtime state (which gates
1569/// [`Shard::read_pool`]). A shard whose primary is unreachable reports `Down`
1570/// even when a replica can still serve reads, since writes and primary reads
1571/// would fail.
1572///
1573/// Reports `Down` — gating `/ready` — when the shard primary is unreachable,
1574/// or when the shard's replica is unready **and** its `replica_fallback` is
1575/// `fail_readiness`. A `primary`-fallback shard with a reachable primary
1576/// degrades to primary reads and stays `Up` with the replica state in its
1577/// details.
1578pub(crate) struct ShardHealthIndicator {
1579    shard: Shard,
1580    /// Pings on dedicated connections, so a busy shard pool is not a failed
1581    /// shard (#3059). No own cache: the indicator registry caches.
1582    primary_ping: crate::db_ping::DbPingCheck,
1583    replica_ping: crate::db_ping::DbPingCheck,
1584    /// Time limit of one ping (`health.ping_timeout_ms`).
1585    ping_timeout: std::time::Duration,
1586}
1587
1588impl ShardHealthIndicator {
1589    pub(crate) fn new(shard: Shard, ping_timeout: std::time::Duration) -> Self {
1590        let primary_ping = crate::db_ping::DbPingCheck::new("primary");
1591        let replica_ping = crate::db_ping::DbPingCheck::new("replica");
1592        for ping in [&primary_ping, &replica_ping] {
1593            ping.configure(std::time::Duration::ZERO, ping_timeout);
1594        }
1595        Self {
1596            shard,
1597            primary_ping,
1598            replica_ping,
1599            ping_timeout,
1600        }
1601    }
1602
1603    async fn refresh_replica_readiness(&self) {
1604        let Some(replica_pool) = self.shard.replica_pool() else {
1605            return;
1606        };
1607        let status = self.replica_ping.check(replica_pool).await;
1608        if !status.up {
1609            self.shard.runtime().mark_replica_connection_unready(
1610                status
1611                    .error
1612                    .unwrap_or_else(|| "replica ping failed".to_owned()),
1613            );
1614            return;
1615        }
1616        self.shard.runtime().mark_replica_connection_ready();
1617        // The parity comparison opens new connections to both roles, so it
1618        // is throttled.
1619        if self.shard.runtime().parity_check_due()
1620            && let Some((primary_url, replica_url)) = self.shard.runtime().migration_check()
1621        {
1622            self.record_parity(
1623                crate::migrate::check_replica_migration_readiness_blocking(
1624                    primary_url,
1625                    replica_url,
1626                ),
1627                PARITY_TIMEOUT,
1628            )
1629            .await;
1630        }
1631    }
1632
1633    /// Run the parity comparison with its own time limit, and record the
1634    /// result. A comparison that does not finish in time marks the replica
1635    /// migrations unready (fail closed): the throttle already claimed this
1636    /// window, so the old result must not stay in use.
1637    async fn record_parity(
1638        &self,
1639        check: impl std::future::Future<Output = crate::migrate::ReplicaMigrationReadiness>,
1640        timeout: std::time::Duration,
1641    ) {
1642        match tokio::time::timeout(timeout, check).await {
1643            Ok(readiness) if readiness.is_ready() => {
1644                self.shard.runtime().mark_replica_migrations_ready();
1645            }
1646            Ok(readiness) => {
1647                if let Some(detail) = readiness.detail() {
1648                    self.shard.runtime().mark_replica_migrations_unready(detail);
1649                }
1650            }
1651            Err(_elapsed) => self
1652                .shard
1653                .runtime()
1654                .mark_replica_migrations_unready(format!(
1655                    "replica migration parity check timed out after {} ms",
1656                    timeout.as_millis()
1657                )),
1658        }
1659    }
1660}
1661
1662/// Time limit for the throttled migration parity comparison. It opens new
1663/// connections to both roles, so it gets more time than a ping.
1664const PARITY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
1665
1666impl crate::actuator::HealthIndicator for ShardHealthIndicator {
1667    fn check(&self) -> futures::future::BoxFuture<'_, crate::actuator::HealthCheckOutput> {
1668        Box::pin(async move {
1669            // Ping both roles at the same time: the wall time is the slower
1670            // ping, not the sum.
1671            let ((), primary) = tokio::join!(
1672                self.refresh_replica_readiness(),
1673                self.primary_ping.check(self.shard.primary_pool())
1674            );
1675
1676            let mut details = HashMap::new();
1677            let status = self.shard.primary_pool().status();
1678            details.insert("pool_size".to_owned(), serde_json::json!(status.max_size));
1679            details.insert(
1680                "active_connections".to_owned(),
1681                serde_json::json!((status.max_size as u64).saturating_sub(status.available as u64)),
1682            );
1683            details.insert(
1684                "idle_connections".to_owned(),
1685                serde_json::json!(status.available),
1686            );
1687            details.insert(
1688                "slots".to_owned(),
1689                serde_json::json!(self.shard.slots().len()),
1690            );
1691            if self.shard.replica_pool().is_some() {
1692                details.insert(
1693                    "replica_ready".to_owned(),
1694                    serde_json::json!(self.shard.runtime().replica_ready()),
1695                );
1696                if let Some(detail) = self.shard.runtime().detail() {
1697                    details.insert("replica_detail".to_owned(), serde_json::json!(detail));
1698                }
1699            }
1700
1701            // Live-check the shard primary. `read_pool()` alone is not enough:
1702            // a primary-only shard's `read_pool()` always returns the primary
1703            // pool (so it is `Some` even when the primary is down), and a
1704            // replicated shard's `read_pool()` can be `Some` via a healthy
1705            // replica while the primary is unreachable — yet all shard writes
1706            // and primary reads would fail at request time. Probe the primary
1707            // (like the replica connectivity check above) and gate `/ready` on
1708            // it so load balancers stop routing to an instance that cannot
1709            // reach a shard primary.
1710            let primary_ok = primary.up;
1711            if let Some(error) = primary.error {
1712                details.insert("primary_detail".to_owned(), serde_json::json!(error));
1713            }
1714            details.insert("primary_ready".to_owned(), serde_json::json!(primary_ok));
1715
1716            // `read_pool()` is `None` exactly when the replica is unready under
1717            // `fail_readiness`. Report `Up` only when the primary is reachable
1718            // *and* a read pool is available; either failing gates `/ready`.
1719            let output = if primary_ok && self.shard.read_pool().is_some() {
1720                crate::actuator::HealthCheckOutput::up()
1721            } else {
1722                crate::actuator::HealthCheckOutput::down()
1723            };
1724            output.with_details(details)
1725        })
1726    }
1727
1728    /// The ping time limit, plus the parity time limit, plus a margin. A
1729    /// late ping or parity check then reports its own result, not the
1730    /// registry's `UNKNOWN` timeout result.
1731    fn timeout_ms(&self) -> u64 {
1732        let limit = self.ping_timeout + PARITY_TIMEOUT + std::time::Duration::from_millis(500);
1733        u64::try_from(limit.as_millis()).unwrap_or(u64::MAX)
1734    }
1735}
1736
1737/// Register one `db:shard:<name>` readiness indicator per configured
1738/// shard onto `registry`. Called once at startup by `build_state`.
1739pub(crate) fn register_shard_health_indicators(
1740    set: &ShardSet,
1741    registry: &crate::actuator::HealthIndicatorRegistry,
1742    ping_timeout: std::time::Duration,
1743) {
1744    #[cfg(feature = "sqlite")]
1745    if let Some(fleet) = set.fleet()
1746        && let Err(error) = registry.register(
1747            "db:fleet",
1748            crate::actuator::IndicatorGroup::Readiness,
1749            Arc::new(crate::db::fleet::FleetHealthIndicator::new(fleet.clone())),
1750        )
1751    {
1752        tracing::warn!("{error}");
1753    }
1754    for shard in set.iter() {
1755        let name = format!("db:shard:{}", shard.name());
1756        if let Err(error) = registry.register(
1757            name,
1758            crate::actuator::IndicatorGroup::Readiness,
1759            Arc::new(ShardHealthIndicator::new(shard.clone(), ping_timeout)),
1760        ) {
1761            tracing::warn!("{error}");
1762        }
1763    }
1764}
1765
1766// ── Extractors ───────────────────────────────────────────────────────────────
1767
1768/// Request-extension escape hatch for [`ShardedDb`] key resolution.
1769///
1770/// Insert this from middleware (or tests) to route a request to a
1771/// specific shard key, bypassing tenant extraction:
1772///
1773/// ```rust,ignore
1774/// request.extensions_mut().insert(ShardKeyOverride("tenant-42".to_owned()));
1775/// ```
1776#[derive(Debug, Clone)]
1777pub struct ShardKeyOverride(pub String);
1778
1779/// Explicit shard access extractor.
1780///
1781/// Extract once, then route per call. Captures the request's database
1782/// context (route-level statement timeout, metrics key, interceptors) at
1783/// extraction so every checkout carries the same instrumentation as the
1784/// plain [`Db`](crate::db::Db) extractor.
1785///
1786/// Rejects with `503 Service Unavailable` when no `[[database.shards]]`
1787/// are configured.
1788///
1789/// # Examples
1790///
1791/// ```rust,no_run
1792/// use autumn_web::prelude::*;
1793///
1794/// #[get("/users/{user_id}/bookmarks")]
1795/// async fn list(shards: Shards, Path(user_id): Path<i64>) -> AutumnResult<&'static str> {
1796///     let mut db = shards.db_for(user_id).await?;
1797///     // run Diesel queries against the owning shard's primary
1798///     Ok("ok")
1799/// }
1800/// ```
1801pub struct Shards {
1802    set: ShardSet,
1803    ctx: crate::db::RequestDbContext,
1804}
1805
1806impl<S> axum::extract::FromRequestParts<S> for Shards
1807where
1808    S: crate::db::DbState + Send + Sync,
1809{
1810    type Rejection = AutumnError;
1811
1812    async fn from_request_parts(
1813        parts: &mut axum::http::request::Parts,
1814        state: &S,
1815    ) -> Result<Self, Self::Rejection> {
1816        let set = state.shards().cloned().ok_or_else(no_shards_configured)?;
1817        let ctx = crate::db::RequestDbContext::from_parts(parts, state);
1818        Ok(Self { set, ctx })
1819    }
1820}
1821
1822fn no_shards_configured() -> AutumnError {
1823    AutumnError::service_unavailable_msg(
1824        "No shards configured: declare [[database.shards]] in autumn.toml \
1825         (see docs/guide/sharding.md)",
1826    )
1827}
1828
1829/// Build a tenant-free repository seed for cross-shard admin reads.
1830///
1831/// Unlike [`__autumn_resolve_repo_seed`], this resolves no tenant key. It seeds
1832/// from the first configured shard (its primary pool and read route) so the
1833/// pre-fan-out connection the trait methods acquire succeeds, then strips the
1834/// shard tag from the route label — the fan-out re-tags each per-shard query
1835/// with the shard actually executing it (see [`reshard_route_label`]).
1836fn cross_shard_seed(
1837    set: &ShardSet,
1838    ctx: &crate::db::RequestDbContext,
1839) -> Result<ShardRepositorySeed, AutumnError> {
1840    // A fleet has no first shard; the control pool serves the pre-fan-out
1841    // connection, and the fan-out re-tags each query with `shard=fleet`.
1842    #[cfg(feature = "sqlite")]
1843    if let Some(backing) = &set.inner.fleet {
1844        let mut seed = ShardRepositorySeed::from_ctx(
1845            &backing.control_pool,
1846            ctx,
1847            FLEET_METRIC_LABEL,
1848            crate::repository::ReadRoute::Primary,
1849        );
1850        seed.route = ctx.route_key.as_ref().map(ToString::to_string);
1851        return Ok(seed);
1852    }
1853    let shard = set.iter().next().ok_or_else(no_shards_configured)?;
1854    let mut seed =
1855        ShardRepositorySeed::from_ctx(shard.primary_pool(), ctx, shard.name(), shard.read_route());
1856    seed.route = ctx.route_key.as_ref().map(ToString::to_string);
1857    Ok(seed)
1858}
1859
1860/// Marks a repository built for tenant-free cross-shard reads.
1861///
1862/// Implemented by the `#[repository(tenant_scoped, sharded)]` macro and used by
1863/// [`CrossShard`] to construct the repository from a [`ShardSet`] without
1864/// resolving a tenant. Not intended to be implemented by hand.
1865pub trait CrossShardRepository: Sized {
1866    /// Construct the repository in `across_tenants()` mode from a tenant-free
1867    /// seed and the full shard set.
1868    #[doc(hidden)]
1869    fn __autumn_from_cross_shard(seed: ShardRepositorySeed, set: ShardSet) -> Self;
1870}
1871
1872/// Axum extractor for tenant-free cross-shard reads on a
1873/// `#[repository(tenant_scoped, sharded)]` repository.
1874///
1875/// Cross-tenant admin endpoints normally have no tenant header or task-local, so
1876/// the standard repository extractor — which resolves a tenant to route to a
1877/// single shard — rejects them during extraction. `CrossShard<R>` instead loads
1878/// the full [`ShardSet`] without a tenant and yields a repository already in
1879/// `across_tenants()` mode: reads fan out across every
1880/// shard, while writes are rejected (cross-shard writes are unsupported).
1881///
1882/// ```ignore
1883/// async fn admin_list(
1884///     CrossShard(repo): CrossShard<PgBookmarkRepository>,
1885/// ) -> AutumnResult<Json<Vec<Bookmark>>> {
1886///     // fans out across all shards
1887///     Ok(Json(repo.find_all().await?))
1888/// }
1889/// ```
1890pub struct CrossShard<R>(pub R);
1891
1892impl<R> std::ops::Deref for CrossShard<R> {
1893    type Target = R;
1894    fn deref(&self) -> &R {
1895        &self.0
1896    }
1897}
1898
1899impl<R> std::ops::DerefMut for CrossShard<R> {
1900    fn deref_mut(&mut self) -> &mut R {
1901        &mut self.0
1902    }
1903}
1904
1905impl<S, R> axum::extract::FromRequestParts<S> for CrossShard<R>
1906where
1907    S: crate::db::DbState + Send + Sync,
1908    R: CrossShardRepository,
1909{
1910    type Rejection = AutumnError;
1911
1912    async fn from_request_parts(
1913        parts: &mut axum::http::request::Parts,
1914        state: &S,
1915    ) -> Result<Self, Self::Rejection> {
1916        // Load the shard set without resolving a tenant (the whole point), then
1917        // seed from it and build the repo in across_tenants() fan-out mode.
1918        let shards =
1919            <Shards as axum::extract::FromRequestParts<S>>::from_request_parts(parts, state)
1920                .await?;
1921        let seed = cross_shard_seed(&shards.set, &shards.ctx)?;
1922        Ok(Self(R::__autumn_from_cross_shard(seed, shards.set)))
1923    }
1924}
1925
1926/// How many shards `each_shard` queries concurrently.
1927const FAN_OUT_CONCURRENCY: usize = 8;
1928
1929impl Shards {
1930    /// The underlying [`ShardSet`].
1931    #[must_use]
1932    pub const fn set(&self) -> &ShardSet {
1933        &self.set
1934    }
1935
1936    /// Iterate shards in declaration order.
1937    pub fn iter(&self) -> impl Iterator<Item = &Shard> {
1938        self.set.iter()
1939    }
1940
1941    /// Check out a connection to the **primary** of the shard that owns
1942    /// `key`.
1943    ///
1944    /// # Errors
1945    ///
1946    /// Returns the router's error or a checkout failure.
1947    pub async fn db_for<'k>(
1948        &self,
1949        key: impl Into<ShardKey<'k>>,
1950    ) -> Result<crate::db::Db, AutumnError> {
1951        let shard = self.set.resolve(key).await?;
1952        self.checkout_primary(&shard).await
1953    }
1954
1955    /// Check out a **read** connection to the shard that owns `key`,
1956    /// honoring the shard's replica topology, readiness, and
1957    /// `replica_fallback`.
1958    ///
1959    /// # Errors
1960    ///
1961    /// Returns the router's error, a checkout failure, or
1962    /// `503 Service Unavailable` when the shard's replica is unready and
1963    /// its fallback is `fail_readiness`.
1964    pub async fn read_for<'k>(
1965        &self,
1966        key: impl Into<ShardKey<'k>>,
1967    ) -> Result<crate::db::Db, AutumnError> {
1968        let shard = self.set.resolve(key).await?;
1969        let (pool, role) = shard.read_pool_with_role().ok_or_else(|| {
1970            AutumnError::service_unavailable_msg(format!(
1971                "shard {:?} replica is not ready and replica_fallback = \"fail_readiness\"",
1972                shard.name()
1973            ))
1974        })?;
1975        self.checkout(&shard, pool, role).await
1976    }
1977
1978    /// Check out a **replica-only** connection to the shard that owns `key`.
1979    ///
1980    /// Unlike [`read_for`], this method ignores the shard's `replica_fallback`
1981    /// policy and **never** falls back to the primary. It returns `503 Service
1982    /// Unavailable` whenever a healthy, ready replica is unavailable — whether
1983    /// no replica is configured, or the replica has not yet passed its
1984    /// readiness checks. Use this for analytics/reporting paths that must
1985    /// guarantee replica-only semantics.
1986    ///
1987    /// # Errors
1988    ///
1989    /// Returns the router's error, a checkout failure, or
1990    /// `503 Service Unavailable` when no healthy replica is available for the
1991    /// resolved shard.
1992    ///
1993    /// [`read_for`]: Self::read_for
1994    pub async fn read_replica_for<'k>(
1995        &self,
1996        key: impl Into<ShardKey<'k>>,
1997    ) -> Result<crate::db::Db, AutumnError> {
1998        let shard = self.set.resolve(key).await?;
1999        let pool = shard.replica_read_pool().ok_or_else(|| {
2000            AutumnError::service_unavailable_msg(format!(
2001                "shard {:?} has no healthy replica; read_replica_for requires a \
2002                 configured, ready replica (no primary fallback)",
2003                shard.name()
2004            ))
2005        })?;
2006        self.checkout(&shard, pool, "replica").await
2007    }
2008
2009    /// Check out a connection to a shard's primary **by name** —
2010    /// intended for admin/operational paths, not request routing.
2011    ///
2012    /// # Errors
2013    ///
2014    /// Returns a bad-request error for an unknown name, or a checkout
2015    /// failure.
2016    pub async fn db_on(&self, shard_name: &str) -> Result<crate::db::Db, AutumnError> {
2017        let shard = self.set.resolve_name(shard_name).await?;
2018        self.checkout_primary(&shard).await
2019    }
2020
2021    /// The fleet behind a fleet-backed set (ADR 0019): provision, delete,
2022    /// back up and list tenant databases from a handler.
2023    #[cfg(feature = "sqlite")]
2024    #[must_use]
2025    pub fn fleet(&self) -> Option<&crate::db::fleet::DatabaseFleet> {
2026        self.set.fleet()
2027    }
2028
2029    /// Run `f` against the primary of **every** shard, concurrently
2030    /// (bounded), collecting per-shard results in declaration order.
2031    ///
2032    /// Failures are collected rather than short-circuited so aggregate/
2033    /// admin endpoints can report partial outages. Fan-out latency is
2034    /// roughly the slowest shard, not the sum — but remember that
2035    /// scatter/gather amplifies tail latency: the more shards, the more
2036    /// likely one is slow.
2037    ///
2038    /// There are **no cross-shard transactions**: each closure invocation
2039    /// commits or fails independently, and concurrent writers mean the
2040    /// collected results can observe torn aggregates.
2041    ///
2042    /// The returned future cannot borrow the `&Shard` argument — copy
2043    /// what you need (e.g. `shard.name().to_owned()`) before the
2044    /// `async move` block:
2045    ///
2046    /// ```rust,ignore
2047    /// let counts = shards
2048    ///     .each_shard(|shard, mut db| {
2049    ///         let name = shard.name().to_owned();
2050    ///         async move { /* query with db, label with name */ Ok(0i64) }
2051    ///     })
2052    ///     .await;
2053    /// ```
2054    pub async fn each_shard<T, Fut, F>(&self, f: F) -> Vec<(ShardId, Result<T, AutumnError>)>
2055    where
2056        T: Send,
2057        Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
2058        F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
2059    {
2060        // FuturesUnordered keeps the pipeline full at FAN_OUT_CONCURRENCY
2061        // (no head-of-line blocking on a slow shard); results are placed
2062        // by ShardId so declaration order is preserved. Futures come from
2063        // a named async fn rather than a closure returning an async block,
2064        // which would trip rustc #89976 when the handler future is checked
2065        // for Send.
2066        use futures::StreamExt as _;
2067
2068        #[cfg(feature = "sqlite")]
2069        if self.set.is_fleet() {
2070            return self.each_fleet_database(&f).await;
2071        }
2072
2073        let mut results: Vec<Option<(ShardId, Result<T, AutumnError>)>> =
2074            std::iter::repeat_with(|| None)
2075                .take(self.set.len())
2076                .collect();
2077        let mut in_flight: futures::stream::FuturesUnordered<
2078            futures::future::BoxFuture<'_, (ShardId, Result<T, AutumnError>)>,
2079        > = futures::stream::FuturesUnordered::new();
2080
2081        for shard in self.set.iter() {
2082            if in_flight.len() >= FAN_OUT_CONCURRENCY
2083                && let Some((id, result)) = in_flight.next().await
2084            {
2085                results[id.0] = Some((id, result));
2086            }
2087            in_flight.push(Box::pin(self.run_on_shard(shard, &f)));
2088        }
2089        while let Some((id, result)) = in_flight.next().await {
2090            results[id.0] = Some((id, result));
2091        }
2092        results.into_iter().flatten().collect()
2093    }
2094
2095    /// [`each_shard`](Self::each_shard) over a fleet: every database on
2096    /// disk, in key order, each opened (never created). A fleet that cannot
2097    /// be listed yields one error entry with `ShardId(usize::MAX)`.
2098    #[cfg(feature = "sqlite")]
2099    async fn each_fleet_database<T, Fut, F>(&self, f: &F) -> Vec<(ShardId, Result<T, AutumnError>)>
2100    where
2101        T: Send,
2102        Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
2103        F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
2104    {
2105        use futures::StreamExt as _;
2106        let Some(fleet) = self.set.fleet() else {
2107            return Vec::new();
2108        };
2109        let keys = match fleet.list().await {
2110            // A file of a slot another host owns is a stale copy (issue #3233).
2111            Ok(keys) => keys
2112                .into_iter()
2113                .filter(|key| fleet.owns(key))
2114                .collect::<Vec<_>>(),
2115            Err(error) => return vec![(ShardId(usize::MAX), Err(error.into()))],
2116        };
2117        futures::stream::iter(keys)
2118            .map(
2119                |key| -> futures::future::BoxFuture<'_, (ShardId, Result<T, AutumnError>)> {
2120                    Box::pin(async move {
2121                        match self.set.resolve_name(&key.name()).await {
2122                            Ok(shard) => self.run_on_shard(&shard, f).await,
2123                            Err(error) => (fleet_shard_id(&key), Err(error)),
2124                        }
2125                    })
2126                },
2127            )
2128            .buffered(FAN_OUT_CONCURRENCY)
2129            .collect()
2130            .await
2131    }
2132
2133    async fn run_on_shard<T, Fut, F>(
2134        &self,
2135        shard: &Shard,
2136        f: &F,
2137    ) -> (ShardId, Result<T, AutumnError>)
2138    where
2139        T: Send,
2140        Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
2141        F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
2142    {
2143        let result = match self.checkout_primary(shard).await {
2144            Ok(db) => f(shard, db).await,
2145            Err(error) => Err(error),
2146        };
2147        (shard.id(), result)
2148    }
2149
2150    async fn checkout_primary(&self, shard: &Shard) -> Result<crate::db::Db, AutumnError> {
2151        self.checkout(shard, shard.primary_pool(), "primary").await
2152    }
2153
2154    async fn checkout(
2155        &self,
2156        shard: &Shard,
2157        pool: &Pool<RuntimeConnection>,
2158        role: &str,
2159    ) -> Result<crate::db::Db, AutumnError> {
2160        let ctx = self.ctx.clone();
2161        crate::db::Db::checkout(crate::db::DbCheckoutParams {
2162            pool,
2163            pool_name: &format!("shard:{}:{role}", shard.metric_label()),
2164            shard: Some(shard.name()),
2165            statement_timeout: ctx.statement_timeout,
2166            idle_in_transaction_timeout: ctx.idle_in_transaction_timeout,
2167            // Tag the route metric with the shard so per-shard latency
2168            // separates in /actuator/metrics.
2169            route_key: ctx.route_key.map(|key| {
2170                crate::db::RouteKey::Composed(format!("{key} shard={}", shard.metric_label()))
2171            }),
2172            metrics: ctx.metrics,
2173            slow_query_threshold: ctx.slow_query_threshold,
2174            interceptors: ctx.interceptors,
2175            // Shard traffic is not recorded in this slice, and the shard gap
2176            // is noted separately by `note_shard_capture_gap` on this same
2177            // checkout path — the control topology's gap does not apply here.
2178            #[cfg(all(feature = "reporting", not(feature = "sqlite")))]
2179            capture_gap: None,
2180            clock: ctx.clock,
2181        })
2182        .await
2183    }
2184}
2185
2186/// Instrumentation seed for building a `#[repository]` over a shard.
2187///
2188/// Carries the shard's primary pool plus the three request-derived
2189/// observability values captured by [`ShardedDb`] at extraction time.
2190/// Generated repositories read this via `__autumn_repository_seed()` when
2191/// their `from_shard` constructor is called, so they apply the same
2192/// statement timeout, slow-query threshold, and route label as the
2193/// [`Shards`] extractor does when checking out a [`Db`](crate::db::Db).
2194///
2195/// This type is sealed behind `#[doc(hidden)]`; it is part of the
2196/// framework's internal ABI for generated code and must not be considered
2197/// a stable public API.
2198#[doc(hidden)]
2199#[derive(Clone)]
2200pub struct ShardRepositorySeed {
2201    pub pool: Pool<RuntimeConnection>,
2202    /// Statement timeout in milliseconds (`0` = no limit, matching the
2203    /// Postgres `statement_timeout = 0` convention).  Capped at
2204    /// `i32::MAX` ms to match the Postgres signed-integer constraint.
2205    pub statement_timeout_ms: u64,
2206    pub slow_query_threshold: std::time::Duration,
2207    /// Shard-tagged route label (e.g. `"GET /bookmarks shard=shard0"`),
2208    /// or `None` when no `MatchedPath` was present in the request.
2209    pub route: Option<String>,
2210    /// The shard's read-routing decision, snapshotted at extraction time so
2211    /// `from_shard` repositories send read-only methods to the shard's
2212    /// replica when one is healthy (issue #1274). Built via
2213    /// [`Shard::read_route`].
2214    pub read_route: crate::repository::ReadRoute,
2215    /// Keeps a fleet database open while a repository built from this seed
2216    /// lives (see [`Shard::lease`]). `None` for configured shards.
2217    pub lease: Option<ShardLease>,
2218}
2219
2220impl ShardRepositorySeed {
2221    pub(crate) fn from_ctx(
2222        pool: &Pool<RuntimeConnection>,
2223        ctx: &crate::db::RequestDbContext,
2224        shard_name: &str,
2225        read_route: crate::repository::ReadRoute,
2226    ) -> Self {
2227        // Postgres `statement_timeout` is a signed 32-bit integer (ms); cap
2228        // to `i32::MAX` so the cast back to a `u64` field is always lossless.
2229        const PG_TIMEOUT_MAX_MS: u64 = i32::MAX as u64;
2230        let statement_timeout_ms = ctx.statement_timeout.map_or(0, |d| {
2231            u64::try_from(d.as_millis().min(u128::from(PG_TIMEOUT_MAX_MS)))
2232                .unwrap_or(PG_TIMEOUT_MAX_MS)
2233        });
2234        Self {
2235            pool: pool.clone(),
2236            statement_timeout_ms,
2237            slow_query_threshold: ctx.slow_query_threshold,
2238            route: ctx
2239                .route_key
2240                .as_ref()
2241                .map(|key| format!("{key} shard={shard_name}")),
2242            read_route,
2243            lease: None,
2244        }
2245    }
2246}
2247
2248/// Re-tag a fan-out sub-repo's route label with the shard executing the query.
2249///
2250/// Keeps per-shard DB metrics and slow-query logs attributed to the shard that
2251/// actually runs the query rather than the originally-routed shard. The parent
2252/// label is `"<key> shard=<orig>"` (see `ShardRepositorySeed::from_ctx`); this
2253/// swaps the `shard=` tag for `shard_name` while preserving the base route key.
2254/// Returns `None` when the parent had no label (no `MatchedPath`), so unlabelled
2255/// repos stay unlabelled.
2256#[must_use]
2257pub fn reshard_route_label(parent: Option<&str>, shard_name: &str) -> Option<String> {
2258    let parent = parent?;
2259    let base = parent.rsplit_once(" shard=").map_or(parent, |(key, _)| key);
2260    Some(format!("{base} shard={shard_name}"))
2261}
2262
2263/// Tenant-routed shard connection extractor.
2264///
2265/// Resolves the routing key automatically and checks out a connection to
2266/// the owning shard's primary. Key resolution order:
2267///
2268/// 1. a [`ShardKeyOverride`] request extension (middleware/test escape
2269///    hatch),
2270/// 2. the tenant id established by the tenancy middleware
2271///    ([`tenancy::CURRENT_TENANT`](crate::tenancy::CURRENT_TENANT)),
2272/// 3. direct tenant extraction from the request per the `[tenancy]`
2273///    configuration.
2274///
2275/// Dereferences to `AsyncPgConnection` exactly like
2276/// [`Db`](crate::db::Db), and exposes [`tx`](Self::tx) with the same
2277/// transaction semantics.
2278///
2279/// # Examples
2280///
2281/// ```rust,no_run
2282/// use autumn_web::prelude::*;
2283///
2284/// #[get("/bookmarks")]
2285/// async fn list(mut db: ShardedDb) -> AutumnResult<String> {
2286///     // queries run on the tenant's shard
2287///     Ok(format!("served from shard {}", db.shard()))
2288/// }
2289/// ```
2290pub struct ShardedDb {
2291    db: crate::db::Db,
2292    shard_name: Arc<str>,
2293    shard_id: ShardId,
2294    repo_seed: ShardRepositorySeed,
2295    // The full shard set, so `Repo::from_shard(&db).across_tenants()` can fan
2296    // out across shards exactly like the generated extractor path (cheap to
2297    // clone — `ShardSet` is `Arc`-backed).
2298    shards: ShardSet,
2299}
2300
2301impl ShardedDb {
2302    /// Name of the shard this connection belongs to.
2303    #[must_use]
2304    pub fn shard(&self) -> &str {
2305        &self.shard_name
2306    }
2307
2308    /// Id of the shard this connection belongs to.
2309    #[must_use]
2310    pub const fn shard_id(&self) -> ShardId {
2311        self.shard_id
2312    }
2313
2314    /// Connection-scoped tracing span (see [`Db::span`](crate::db::Db::span)).
2315    #[must_use]
2316    pub const fn span(&self) -> &tracing::Span {
2317        self.db.span()
2318    }
2319
2320    /// Run an async closure inside a transaction **on this shard**.
2321    /// Same semantics as [`Db::tx`](crate::db::Db::tx); the transaction
2322    /// never spans shards.
2323    ///
2324    /// # Errors
2325    ///
2326    /// See [`Db::tx`](crate::db::Db::tx).
2327    pub async fn tx<'a, T, E, F>(&'a mut self, f: F) -> Result<T, AutumnError>
2328    where
2329        T: Send + 'a,
2330        E: From<diesel::result::Error> + Send + Sync + 'a,
2331        AutumnError: From<E>,
2332        F: for<'r> FnOnce(
2333                &'r mut crate::db::PooledConnection,
2334            ) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
2335            + Send
2336            + 'a,
2337    {
2338        self.db.tx(f).await
2339    }
2340
2341    /// Run an async closure inside a `BEGIN IMMEDIATE` transaction **on this
2342    /// shard**. Same semantics as [`Db::tx_immediate`](crate::db::Db::tx_immediate);
2343    /// the transaction never spans shards.
2344    ///
2345    /// # Errors
2346    ///
2347    /// See [`Db::tx_immediate`](crate::db::Db::tx_immediate).
2348    pub async fn tx_immediate<'a, T, E, F>(&'a mut self, f: F) -> Result<T, AutumnError>
2349    where
2350        T: Send + 'a,
2351        E: From<diesel::result::Error> + Send + Sync + 'a,
2352        AutumnError: From<E>,
2353        F: for<'r> FnOnce(
2354                &'r mut crate::db::RuntimeConnection,
2355            ) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
2356            + Send
2357            + 'a,
2358    {
2359        self.db.tx_immediate(f).await
2360    }
2361
2362    /// Run an async closure inside a transaction **on this shard** with explicit
2363    /// [`TxOptions`](crate::db::TxOptions) (isolation level + retry). Same
2364    /// semantics as [`Db::tx_with`](crate::db::Db::tx_with); the transaction
2365    /// never spans shards.
2366    ///
2367    /// # Errors
2368    ///
2369    /// See [`Db::tx_with`](crate::db::Db::tx_with).
2370    pub async fn tx_with<'a, T, E, F>(
2371        &'a mut self,
2372        opts: crate::db::TxOptions,
2373        f: F,
2374    ) -> Result<T, AutumnError>
2375    where
2376        T: Send + 'a,
2377        E: From<diesel::result::Error> + Send + Sync + 'a,
2378        AutumnError: From<E>,
2379        F: for<'r> FnMut(
2380                &'r mut crate::db::RuntimeConnection,
2381            ) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
2382            + Send
2383            + 'a,
2384    {
2385        self.db.tx_with(opts, f).await
2386    }
2387
2388    /// Borrow the underlying [`Db`](crate::db::Db) (e.g. to pass to
2389    /// helpers written against the unsharded extractor).
2390    pub const fn db_mut(&mut self) -> &mut crate::db::Db {
2391        &mut self.db
2392    }
2393
2394    /// Instrumentation seed for `from_shard` on generated repositories.
2395    /// Internal ABI; not a stable public API.
2396    #[doc(hidden)]
2397    #[must_use]
2398    pub const fn __autumn_repository_seed(&self) -> &ShardRepositorySeed {
2399        &self.repo_seed
2400    }
2401
2402    /// The full shard set, so `from_shard`-built repositories can fan out under
2403    /// `across_tenants()`. Internal ABI; not a stable public API.
2404    #[doc(hidden)]
2405    #[must_use]
2406    pub const fn __autumn_shard_set(&self) -> &ShardSet {
2407        &self.shards
2408    }
2409}
2410
2411impl std::ops::Deref for ShardedDb {
2412    type Target = RuntimeConnection;
2413    fn deref(&self) -> &Self::Target {
2414        &self.db
2415    }
2416}
2417
2418impl std::ops::DerefMut for ShardedDb {
2419    fn deref_mut(&mut self) -> &mut Self::Target {
2420        &mut self.db
2421    }
2422}
2423
2424impl AsMut<crate::db::Db> for ShardedDb {
2425    fn as_mut(&mut self) -> &mut crate::db::Db {
2426        &mut self.db
2427    }
2428}
2429
2430/// Internal ABI for generated `#[repository(sharded)]` extractors.
2431///
2432/// Resolves the tenant→shard routing from a request, builds the
2433/// [`ShardRepositorySeed`] that carries the shard's pool and observability
2434/// context, and returns a cheap clone of the [`ShardSet`] for cross-shard
2435/// fan-out. Unlike [`ShardedDb::from_request_parts`] it does **not** check
2436/// out a connection, so generated repositories can acquire their own lazily.
2437#[doc(hidden)]
2438pub async fn __autumn_resolve_repo_seed(
2439    parts: &mut axum::http::request::Parts,
2440    state: &crate::AppState,
2441) -> Result<(ShardRepositorySeed, ShardSet), AutumnError> {
2442    let shards = <Shards as axum::extract::FromRequestParts<crate::AppState>>::from_request_parts(
2443        parts, state,
2444    )
2445    .await?;
2446    let key = resolve_shard_key(parts, state).await?;
2447    let shard = shards.set.resolve(&key).await?;
2448    let mut seed = ShardRepositorySeed::from_ctx(
2449        shard.primary_pool(),
2450        &shards.ctx,
2451        shard.metric_label(),
2452        shard.read_route(),
2453    );
2454    // A generated repository acquires its connection lazily, possibly long
2455    // after this; the lease keeps a fleet database from closing meanwhile.
2456    seed.lease = shard.lease();
2457    let set = shards.set.clone();
2458    Ok((seed, set))
2459}
2460
2461impl axum::extract::FromRequestParts<crate::AppState> for ShardedDb {
2462    type Rejection = AutumnError;
2463
2464    async fn from_request_parts(
2465        parts: &mut axum::http::request::Parts,
2466        state: &crate::AppState,
2467    ) -> Result<Self, Self::Rejection> {
2468        let shards = Shards::from_request_parts(parts, state).await?;
2469        let key = resolve_shard_key(parts, state).await?;
2470
2471        let shard = shards.set.resolve(&key).await?;
2472        let shard_name = Arc::clone(&shard.name);
2473        let shard_id = shard.id();
2474        let mut repo_seed = ShardRepositorySeed::from_ctx(
2475            shard.primary_pool(),
2476            &shards.ctx,
2477            shard.metric_label(),
2478            shard.read_route(),
2479        );
2480        repo_seed.lease = shard.lease();
2481        let shard_set = shards.set.clone();
2482        let db = shards.checkout_primary(&shard).await?;
2483        crate::read_your_writes::mark_write();
2484        Ok(Self {
2485            db,
2486            shard_name,
2487            shard_id,
2488            repo_seed,
2489            shards: shard_set,
2490        })
2491    }
2492}
2493
2494/// Explicit replica-only shard connection extractor.
2495///
2496/// Resolves the routing key exactly like [`ShardedDb`] and checks out a
2497/// connection to the **replica** of the owning shard. Unlike [`ShardedDb`]'s
2498/// transparent read-routing (which follows the shard's `replica_fallback`
2499/// policy), `ShardedReadDb` **always** requires a healthy replica and returns
2500/// `503 Service Unavailable` immediately if one is not available — it never
2501/// silently falls back to the primary.
2502///
2503/// Use this extractor for analytics, reporting, or admin scatter-gather
2504/// handlers where replica-only semantics must be guaranteed:
2505///
2506/// ```rust,no_run
2507/// use autumn_web::prelude::*;
2508///
2509/// #[get("/analytics")]
2510/// async fn analytics(db: ShardedReadDb) -> impl IntoResponse {
2511///     // guaranteed replica connection; 503 if none is configured or healthy
2512///     "ok"
2513/// }
2514/// ```
2515///
2516/// Pairs with the transparent default routing provided by [`ShardedDb`]:
2517/// that is the opt-out-free default; `ShardedReadDb` is the explicit
2518/// replica-only override (see issue #1275).
2519pub struct ShardedReadDb {
2520    db: crate::db::Db,
2521    shard_name: Arc<str>,
2522    shard_id: ShardId,
2523}
2524
2525impl ShardedReadDb {
2526    /// Name of the shard this connection belongs to.
2527    #[must_use]
2528    pub fn shard(&self) -> &str {
2529        &self.shard_name
2530    }
2531
2532    /// Id of the shard this connection belongs to.
2533    #[must_use]
2534    pub const fn shard_id(&self) -> ShardId {
2535        self.shard_id
2536    }
2537
2538    /// Connection-scoped tracing span (see [`Db::span`](crate::db::Db::span)).
2539    #[must_use]
2540    pub const fn span(&self) -> &tracing::Span {
2541        self.db.span()
2542    }
2543
2544    /// Borrow the underlying [`Db`](crate::db::Db) (e.g. to pass to
2545    /// helpers written against the unsharded extractor).
2546    pub const fn db_mut(&mut self) -> &mut crate::db::Db {
2547        &mut self.db
2548    }
2549}
2550
2551impl std::ops::Deref for ShardedReadDb {
2552    type Target = RuntimeConnection;
2553    fn deref(&self) -> &Self::Target {
2554        &self.db
2555    }
2556}
2557
2558impl std::ops::DerefMut for ShardedReadDb {
2559    fn deref_mut(&mut self) -> &mut Self::Target {
2560        &mut self.db
2561    }
2562}
2563
2564impl AsMut<crate::db::Db> for ShardedReadDb {
2565    fn as_mut(&mut self) -> &mut crate::db::Db {
2566        &mut self.db
2567    }
2568}
2569
2570impl axum::extract::FromRequestParts<crate::AppState> for ShardedReadDb {
2571    type Rejection = AutumnError;
2572
2573    async fn from_request_parts(
2574        parts: &mut axum::http::request::Parts,
2575        state: &crate::AppState,
2576    ) -> Result<Self, Self::Rejection> {
2577        let shards = Shards::from_request_parts(parts, state).await?;
2578        let key = resolve_shard_key(parts, state).await?;
2579
2580        let shard = shards.set.resolve(&key).await?;
2581        let shard_name = Arc::clone(&shard.name);
2582        let shard_id = shard.id();
2583        let pool = shard.replica_read_pool().ok_or_else(|| {
2584            AutumnError::service_unavailable_msg(format!(
2585                "shard {:?} has no healthy replica; ShardedReadDb requires a \
2586                 configured, ready replica (no primary fallback)",
2587                shard.name()
2588            ))
2589        })?;
2590        let db = shards.checkout(&shard, pool, "replica").await?;
2591        Ok(Self {
2592            db,
2593            shard_name,
2594            shard_id,
2595        })
2596    }
2597}
2598
2599/// Resolve the routing key for [`ShardedDb`]; see its docs for the
2600/// resolution order.
2601async fn resolve_shard_key(
2602    parts: &mut axum::http::request::Parts,
2603    state: &crate::AppState,
2604) -> Result<String, AutumnError> {
2605    if let Some(overridden) = parts.extensions.get::<ShardKeyOverride>() {
2606        return Ok(overridden.0.clone());
2607    }
2608    if let Ok(Some(tenant)) = crate::tenancy::CURRENT_TENANT.try_with(std::clone::Clone::clone) {
2609        return Ok(tenant);
2610    }
2611    let config = state
2612        .extension::<crate::config::AutumnConfig>()
2613        .ok_or_else(|| AutumnError::service_unavailable_msg("Config is not available"))?;
2614    crate::tenancy::extract_tenant_from_parts(parts, &config)
2615        .await
2616        .map_err(|error| {
2617            AutumnError::bad_request_msg(format!(
2618                "ShardedDb could not resolve a shard key: {error}. Enable [tenancy] so \
2619                 the tenant id can route the request, or insert a ShardKeyOverride \
2620                 request extension from middleware (see docs/guide/sharding.md)"
2621            ))
2622        })
2623}
2624
2625#[cfg(test)]
2626mod tests {
2627    use super::*;
2628
2629    /// Tenant databases that share a routing slot are still distinct
2630    /// databases: their `ShardId`s must differ, or callers aggregating by id
2631    /// would merge two tenants' results. A slot database keeps its slot.
2632    #[test]
2633    fn fleet_shard_ids_are_unique_per_database() {
2634        use crate::fleet_layout::{FleetDbKey, TenantDbId};
2635        let tenant = |id: &str| FleetDbKey::Tenant {
2636            id: TenantDbId::parse(id).unwrap(),
2637            slot: slot_for_key(ShardKey::Str(id)).0,
2638        };
2639        // Find two tenant ids that land on the same slot.
2640        let mut seen = std::collections::HashMap::new();
2641        let (a, b) = (0..100_000)
2642            .map(|n| format!("t{n}"))
2643            .find_map(|id| {
2644                let slot = slot_for_key(ShardKey::Str(&id)).0;
2645                seen.insert(slot, id.clone()).map(|earlier| (earlier, id))
2646            })
2647            .expect("16384 slots collide well within 100k ids");
2648        assert_eq!(tenant(&a).routing_slot(), tenant(&b).routing_slot());
2649        assert_ne!(fleet_shard_id(&tenant(&a)), fleet_shard_id(&tenant(&b)));
2650        assert_eq!(
2651            fleet_shard_id(&tenant(&a)),
2652            fleet_shard_id(&tenant(&a)),
2653            "stable"
2654        );
2655        assert_eq!(fleet_shard_id(&FleetDbKey::Slot(42)), ShardId(42));
2656    }
2657    use crate::config::{ShardConfig, SlotSpec};
2658
2659    #[test]
2660    fn directory_invalidation_channel_and_interval_are_sane() {
2661        // The listener LISTENs on the same channel the migration's trigger
2662        // fires via `pg_notify`. If these drift, invalidations are never
2663        // delivered.
2664        assert_eq!(DIRECTORY_NOTIFY_CHANNEL, "autumn_shard_directory");
2665        // The idle sweep interval must be shorter than the cache TTL so a
2666        // never-re-observed expired entry is reclaimed before the TTL would
2667        // have done so anyway (delivery of an actual invalidation is immediate,
2668        // independent of this interval).
2669        assert!(
2670            DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL < DEFAULT_DIRECTORY_CACHE_TTL,
2671            "sweep interval should beat the TTL"
2672        );
2673    }
2674
2675    // Backend-parametric: `build_sqlite_pool` refuses a `postgres://` target,
2676    // so an inline spelling here made every shard test panic at fixture
2677    // construction under `--features sqlite`. The routing, slot-map, seed and
2678    // health-registration assertions below are backend-independent and now run
2679    // on both. Pools are lazy, so nothing connects unless a test checks out.
2680    //
2681    // WHICH LAYER THIS EXERCISES. `create_shard_set` is public API and applies
2682    // no backend screen, which is what lets these tests run on either backend.
2683    // A shard topology is nonetheless Postgres-only in PRODUCTION:
2684    // `sqlite_sharding_unsupported_guard` (app.rs) refuses any SQLite shard
2685    // primary at boot, and `database_backend_consistency` refuses `shards`
2686    // beside a SQLite primary at config time. So these tests say "the routing
2687    // and pool machinery is backend-independent", not "SQLite sharding is
2688    // supported" — and if a screen is ever added to `create_shard_set` itself,
2689    // they are meant to fail loudly and be gated then.
2690    fn shard_config(name: &str) -> ShardConfig {
2691        ShardConfig {
2692            name: name.to_owned(),
2693            primary_url: crate::test_urls::primary(name),
2694            slots: None,
2695            replica_url: None,
2696            primary_pool_size: None,
2697            replica_pool_size: None,
2698            replica_fallback: None,
2699        }
2700    }
2701
2702    fn sharded_config(names: &[&str]) -> DatabaseConfig {
2703        DatabaseConfig {
2704            shards: names.iter().map(|name| shard_config(name)).collect(),
2705            ..Default::default()
2706        }
2707    }
2708
2709    fn shard_set(names: &[&str]) -> ShardSet {
2710        create_shard_set(&sharded_config(names), Arc::new(HashShardRouter))
2711            .expect("lazy pools should build")
2712            .expect("shards configured")
2713    }
2714
2715    #[tokio::test]
2716    async fn an_injected_database_fault_fails_the_directory_lookup() {
2717        use crate::fault_injection::{FaultKind, FaultRule, FaultTarget, with_faults};
2718
2719        let config = DatabaseConfig {
2720            primary_url: Some(crate::test_urls::primary("control")),
2721            ..Default::default()
2722        };
2723        let control_pool = crate::db::create_pool(&config)
2724            .expect("a lazy pool builds")
2725            .expect("a URL is set");
2726        let router = DirectoryShardRouter::new(control_pool);
2727        let shards = shard_set(&["a"]);
2728        let rules = [FaultRule::new(FaultTarget::Database, FaultKind::Error, 1.0)];
2729        let Err(error) = with_faults(&rules, router.lookup_directory("tenant", &shards)).await
2730        else {
2731            panic!("the injected fault must fail the control checkout");
2732        };
2733        assert!(error.to_string().contains("fault injection"), "{error}");
2734    }
2735
2736    // ── key→slot golden vectors ─────────────────────────────────────────
2737    //
2738    // These values are a PERMANENT CONTRACT. If one of these assertions
2739    // fails, the change re-routes every existing sharded deployment's
2740    // keys — do not update the expected values; fix the hash instead.
2741
2742    #[test]
2743    fn golden_vector_str_keys() {
2744        // Expected slots computed independently (Python reference
2745        // implementation of FNV-1a 64 mod 16384) when the contract was
2746        // established.
2747        let cases: &[(&str, u16)] = &[
2748            ("tenant-1", 12427),
2749            ("tenant-2", 12862),
2750            ("tenant-3", 13297),
2751            ("acme-corp", 11394),
2752            ("globex", 12846),
2753            ("initech", 11329),
2754            ("hooli", 3974),
2755            ("", 8997),
2756            ("a", 11404),
2757            ("00000000-0000-0000-0000-000000000001", 6206),
2758        ];
2759        for (key, expected_slot) in cases {
2760            assert_eq!(
2761                slot_for_key(ShardKey::Str(key)),
2762                SlotId(*expected_slot),
2763                "key {key:?} must keep routing to slot {expected_slot} forever",
2764            );
2765        }
2766    }
2767
2768    #[test]
2769    fn golden_vector_int_keys() {
2770        // Expected slots computed independently (Python reference
2771        // implementation of splitmix64 mod 16384) when the contract was
2772        // established.
2773        let cases: &[(i64, u16)] = &[
2774            (0, 3503),
2775            (1, 7361),
2776            (2, 5838),
2777            (42, 11925),
2778            (1_000_000, 1511),
2779            (-1, 11296),
2780            (i64::MAX, 7847),
2781            (i64::MIN, 13275),
2782        ];
2783        for (key, expected_slot) in cases {
2784            assert_eq!(
2785                slot_for_key(ShardKey::Int(*key)),
2786                SlotId(*expected_slot),
2787                "key {key} must keep routing to slot {expected_slot} forever",
2788            );
2789        }
2790    }
2791
2792    #[test]
2793    fn golden_vector_bytes_match_equivalent_str() {
2794        // Str and Bytes share FNV-1a, so identical bytes route identically.
2795        assert_eq!(
2796            slot_for_key(ShardKey::Bytes(b"tenant-1")),
2797            slot_for_key(ShardKey::Str("tenant-1")),
2798        );
2799    }
2800
2801    #[test]
2802    fn slots_stay_in_range_and_spread_roughly_uniformly() {
2803        // 10k keys over 16384 slots is too sparse for per-slot bounds, so
2804        // check uniformity over 16 contiguous buckets of 1024 slots each.
2805        let mut histogram = [0usize; 16];
2806        for i in 0..10_000i64 {
2807            let slot = slot_for_key(ShardKey::Int(i));
2808            assert!(slot.0 < SLOT_COUNT);
2809            histogram[usize::from(slot.0 / 1024)] += 1;
2810        }
2811        let expected = 10_000 / histogram.len();
2812        for (bucket, count) in histogram.iter().enumerate() {
2813            assert!(
2814                *count > expected / 2 && *count < expected * 2,
2815                "bucket {bucket} has {count} keys (expected ≈{expected})"
2816            );
2817        }
2818    }
2819
2820    // ── ShardSet behavior ───────────────────────────────────────────────
2821
2822    #[tokio::test]
2823    async fn db_for_and_read_for_attempt_routed_checkouts() {
2824        // Nothing can be reached, so both calls must surface checkout
2825        // failures (not routing errors) after resolving the shard. The target
2826        // is spelled per backend: the default fixture's SQLite target is a
2827        // live in-memory database, which would make both checkouts SUCCEED and
2828        // the assertions vacuous.
2829        let mut config = sharded_config(&["alpha"]);
2830        config.connect_timeout_secs = 1;
2831        config.shards[0].primary_url = crate::test_urls::unreachable("alpha");
2832        let shards = shards_handle_from(
2833            create_shard_set(&config, Arc::new(HashShardRouter))
2834                .expect("build")
2835                .expect("configured"),
2836        );
2837        let Err(error) = shards.db_for("tenant-1").await else {
2838            panic!("checkout must fail without a server");
2839        };
2840        assert!(!error.to_string().contains("Unknown shard"));
2841
2842        // Without a replica, reads route to the primary role.
2843        let Err(error) = shards.read_for("tenant-1").await else {
2844            panic!("checkout must fail without a server");
2845        };
2846        assert!(!error.to_string().contains("fail_readiness"));
2847    }
2848
2849    #[test]
2850    fn parity_recheck_is_throttled_per_window() {
2851        let set = shard_set(&["a"]);
2852        let runtime = set.get(ShardId(0)).expect("shard").runtime();
2853        runtime.configure_migration_check(
2854            "postgres://localhost/a".to_owned(),
2855            "postgres://localhost/a_ro".to_owned(),
2856        );
2857        assert!(runtime.migration_check().is_some());
2858
2859        assert!(runtime.parity_check_due(), "first check claims the window");
2860        assert!(
2861            !runtime.parity_check_due(),
2862            "checks within the window are suppressed"
2863        );
2864    }
2865
2866    #[tokio::test]
2867    async fn route_is_deterministic_and_respects_slot_map() {
2868        let mut config = sharded_config(&["a", "b"]);
2869        config.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2870        config.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
2871        let set = create_shard_set(&config, Arc::new(HashShardRouter))
2872            .expect("build")
2873            .expect("configured");
2874
2875        for key in ["k1", "k2", "k3", "k4", "k5"] {
2876            let slot = set.slot_for_key(key);
2877            let expected = if slot.0 >= 8192 { "b" } else { "a" };
2878            let routed = set.route(key).await.expect("route");
2879            assert_eq!(routed.name(), expected, "key {key:?} slot {}", slot.0);
2880            // Same key always lands on the same shard.
2881            assert_eq!(set.route(key).await.expect("route").id(), routed.id());
2882        }
2883    }
2884
2885    #[tokio::test]
2886    async fn arc_shard_router_delegates_to_inner() {
2887        // `Arc<R>: ShardRouter` is what lets a custom DirectoryShardRouter be
2888        // shared between routing (`with_shard_router`) and its invalidation
2889        // listener (`spawn_invalidation_listener`, which needs `Arc<Self>`).
2890        // Install one through a ShardSet and confirm routing flows to the inner
2891        // router rather than failing the trait bound.
2892        let config = sharded_config(&["a", "b"]);
2893        let set = create_shard_set(&config, Arc::new(Arc::new(HashShardRouter)))
2894            .expect("build")
2895            .expect("configured");
2896        let first = set.route("tenant-42").await.expect("route");
2897        let again = set.route("tenant-42").await.expect("route");
2898        assert_eq!(
2899            first.id(),
2900            again.id(),
2901            "Arc<R> routes deterministically through its inner router"
2902        );
2903    }
2904
2905    #[tokio::test]
2906    async fn moving_a_slot_in_config_moves_only_that_slot() {
2907        // "Reshard" by reassigning slots 12288-16383 from shard b to a new
2908        // shard c: keys in slots 0-12287 must not move.
2909        let mut before = sharded_config(&["a", "b"]);
2910        before.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2911        before.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
2912
2913        let mut after = sharded_config(&["a", "b", "c"]);
2914        after.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2915        after.shards[1].slots = Some(vec![SlotSpec::Range("8192-12287".to_owned())]);
2916        after.shards[2].slots = Some(vec![SlotSpec::Range("12288-16383".to_owned())]);
2917
2918        let set_before = create_shard_set(&before, Arc::new(HashShardRouter))
2919            .expect("build")
2920            .expect("configured");
2921        let set_after = create_shard_set(&after, Arc::new(HashShardRouter))
2922            .expect("build")
2923            .expect("configured");
2924
2925        let mut moved = 0;
2926        for i in 0..200i64 {
2927            let slot = set_before.slot_for_key(i);
2928            assert_eq!(slot, set_after.slot_for_key(i), "key→slot never changes");
2929            let before_shard = set_before.route(i).await.expect("route");
2930            let after_shard = set_after.route(i).await.expect("route");
2931            if slot.0 >= 12288 {
2932                assert_eq!(before_shard.name(), "b");
2933                assert_eq!(after_shard.name(), "c");
2934                moved += 1;
2935            } else {
2936                assert_eq!(before_shard.name(), after_shard.name());
2937            }
2938        }
2939        assert!(moved > 0, "some keys must exercise the moved slot range");
2940    }
2941
2942    #[test]
2943    fn by_name_and_get_resolve_shards() {
2944        let set = shard_set(&["alpha", "beta"]);
2945        assert_eq!(set.len(), 2);
2946        assert_eq!(set.by_name("beta").expect("beta").id(), ShardId(1));
2947        assert_eq!(set.get(ShardId(0)).expect("alpha").name(), "alpha");
2948        assert!(set.by_name("gamma").is_none());
2949        assert!(set.get(ShardId(9)).is_none());
2950        let names: Vec<&str> = set.iter().map(Shard::name).collect();
2951        assert_eq!(names, ["alpha", "beta"]);
2952    }
2953
2954    #[test]
2955    fn auto_split_assigns_contiguous_slots() {
2956        let set = shard_set(&["a", "b"]);
2957        assert_eq!(set.slot_count(), SLOT_COUNT);
2958        assert_eq!(set.get(ShardId(0)).expect("a").slots().len(), 8192);
2959        assert_eq!(
2960            set.get(ShardId(1)).expect("b").slots(),
2961            (8192..16384).collect::<Vec<u16>>()
2962        );
2963    }
2964
2965    // §3 slot-move helpers
2966    #[test]
2967    fn owns_key_agrees_with_route() {
2968        // shard a owns slots 0-8191, shard b owns 8192-16383.
2969        // "hooli" → slot 3974 → shard a (ShardId(0)).
2970        // "a"     → slot 11404 → shard b (ShardId(1)).
2971        let set = shard_set(&["a", "b"]);
2972        assert!(
2973            set.owns_key(ShardId(0), "hooli"),
2974            "hooli (slot 3974) must be shard a"
2975        );
2976        assert!(
2977            !set.owns_key(ShardId(1), "hooli"),
2978            "hooli must not be shard b"
2979        );
2980        assert!(
2981            set.owns_key(ShardId(1), "a"),
2982            "key 'a' (slot 11404) must be shard b"
2983        );
2984        assert!(
2985            !set.owns_key(ShardId(0), "a"),
2986            "key 'a' must not be shard a"
2987        );
2988    }
2989
2990    #[test]
2991    fn slots_for_shard_returns_correct_slice() {
2992        let set = shard_set(&["a", "b"]);
2993        let a_slots = set.slots_for_shard(ShardId(0)).expect("shard a exists");
2994        let b_slots = set.slots_for_shard(ShardId(1)).expect("shard b exists");
2995        assert_eq!(a_slots.len(), 8192);
2996        assert_eq!(b_slots.len(), 8192);
2997        assert!(a_slots.iter().all(|&s| s < 8192));
2998        assert!(b_slots.iter().all(|&s| s >= 8192));
2999        assert!(set.slots_for_shard(ShardId(9)).is_none());
3000    }
3001
3002    #[test]
3003    fn partition_by_shard_groups_golden_keys() {
3004        // Using the golden-vector keys: "hooli"→3974 (shard a), "a"→11404 (shard b).
3005        let set = shard_set(&["a", "b"]);
3006        let keys = ["hooli", "a", "tenant-1"]; // tenant-1 → 12427 → shard b
3007        let map = set.partition_by_shard(keys.iter().copied());
3008        #[allow(clippy::similar_names)]
3009        let keys_on_a = map.get(&ShardId(0)).map_or(&[][..], Vec::as_slice);
3010        #[allow(clippy::similar_names)]
3011        let keys_on_b = map.get(&ShardId(1)).map_or(&[][..], Vec::as_slice);
3012        assert!(keys_on_a.contains(&"hooli"), "hooli must go to shard a");
3013        assert!(keys_on_b.contains(&"a"), "key 'a' must go to shard b");
3014        assert!(
3015            keys_on_b.contains(&"tenant-1"),
3016            "tenant-1 (slot 12427) must go to shard b"
3017        );
3018        assert_eq!(
3019            keys_on_a.len() + keys_on_b.len(),
3020            keys.len(),
3021            "no key dropped"
3022        );
3023    }
3024
3025    #[test]
3026    fn create_shard_set_returns_none_without_shards() {
3027        let config = DatabaseConfig::default();
3028        assert!(
3029            create_shard_set(&config, Arc::new(HashShardRouter))
3030                .expect("ok")
3031                .is_none()
3032        );
3033    }
3034
3035    #[test]
3036    fn build_shard_set_rejects_duplicate_names_without_config_validation() {
3037        // The builder is public: configs that bypassed
3038        // AutumnConfig::validate() must still not produce a shadowed
3039        // by_name map.
3040        let config = sharded_config(&["twin", "twin"]);
3041        let topologies = config
3042            .shards
3043            .iter()
3044            .map(|shard| crate::db::create_shard_topology(shard, &config).expect("lazy pools"))
3045            .collect();
3046
3047        let result = build_shard_set(&config, topologies, Arc::new(HashShardRouter));
3048
3049        let Err(ShardSetBuildError::Config(error)) = result else {
3050            panic!("duplicate shard names must be rejected, got {result:?}");
3051        };
3052        assert!(error.to_string().contains("twin"));
3053    }
3054
3055    #[test]
3056    fn build_shard_set_rejects_topology_count_mismatch() {
3057        let config = sharded_config(&["a", "b"]);
3058        let result = build_shard_set(&config, Vec::new(), Arc::new(HashShardRouter));
3059        assert!(matches!(
3060            result,
3061            Err(ShardSetBuildError::TopologyCountMismatch {
3062                expected: 2,
3063                actual: 0
3064            })
3065        ));
3066    }
3067
3068    // ── read_pool / replica fallback semantics ──────────────────────────
3069
3070    // ── Postgres-only fixtures ───────────────────────────────────────────
3071    //
3072    // Everything from here that carries a `replica_url` is `#[cfg(not(feature
3073    // = "sqlite"))]`. A SQLite read replica is refused at both layers —
3074    // `database_backend_consistency` rejects any `replica_url` beside a SQLite
3075    // primary, and `reject_unusable_sqlite_replica` rejects one that is
3076    // in-memory or names a different file — so there is no SQLite spelling to
3077    // swap in. Gated rather than parametrised, so no test asserts a
3078    // configuration production refuses; the same treatment `db::`'s two
3079    // replica tests already got. The primary-only fixtures above run on both
3080    // backends.
3081    #[cfg(not(feature = "sqlite"))]
3082    fn shard_with_replica(fallback: ReplicaFallback) -> Shard {
3083        let mut config = sharded_config(&["a"]);
3084        config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
3085        config.shards[0].replica_fallback = Some(fallback);
3086        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3087            .expect("build")
3088            .expect("configured");
3089        set.get(ShardId(0)).expect("shard").clone()
3090    }
3091
3092    #[test]
3093    fn read_pool_uses_primary_when_no_replica() {
3094        let set = shard_set(&["a"]);
3095        let shard = set.get(ShardId(0)).expect("shard");
3096        assert!(shard.read_pool().is_some());
3097        assert!(shard.replica_pool().is_none());
3098    }
3099
3100    #[cfg(not(feature = "sqlite"))]
3101    #[test]
3102    fn read_pool_requires_readiness_check_before_replica_traffic() {
3103        let shard = shard_with_replica(ReplicaFallback::Primary);
3104        // Unchecked replica: fallback policy routes reads to the primary.
3105        assert!(shard.read_pool().is_some());
3106        assert!(shard.runtime().detail().is_some());
3107
3108        shard.runtime().mark_replica_connection_ready();
3109        assert!(shard.runtime().replica_ready());
3110        assert!(shard.read_pool().is_some());
3111        assert!(shard.runtime().detail().is_none());
3112    }
3113
3114    #[cfg(not(feature = "sqlite"))]
3115    #[test]
3116    fn read_pool_fails_closed_under_fail_readiness() {
3117        let shard = shard_with_replica(ReplicaFallback::FailReadiness);
3118        assert!(
3119            shard.read_pool().is_none(),
3120            "unchecked replica fails closed"
3121        );
3122
3123        shard.runtime().mark_replica_connection_ready();
3124        assert!(shard.read_pool().is_some());
3125
3126        shard
3127            .runtime()
3128            .mark_replica_migrations_unready("replica lags primary");
3129        assert!(shard.read_pool().is_none());
3130        assert!(shard.runtime().detail().expect("detail").contains("lags"));
3131    }
3132
3133    // ── read_route: per-shard ReadRoute snapshot (issue #1274) ───────────
3134
3135    // Sizes only the Postgres-only replica fixtures below assert on.
3136    #[cfg(not(feature = "sqlite"))]
3137    const PRIMARY_SIZE: usize = 7;
3138    #[cfg(not(feature = "sqlite"))]
3139    const REPLICA_SIZE: usize = 3;
3140
3141    #[cfg(not(feature = "sqlite"))]
3142    /// A one-shard set whose primary and replica pools have *distinct*
3143    /// `max_size` so `read_route()` reveals which pool it selected.
3144    fn shard_with_sized_replica(fallback: ReplicaFallback) -> Shard {
3145        let mut config = sharded_config(&["a"]);
3146        config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
3147        config.shards[0].replica_fallback = Some(fallback);
3148        config.shards[0].primary_pool_size = Some(PRIMARY_SIZE);
3149        config.shards[0].replica_pool_size = Some(REPLICA_SIZE);
3150        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3151            .expect("build")
3152            .expect("configured");
3153        set.get(ShardId(0)).expect("shard").clone()
3154    }
3155
3156    /// `max_size` of the pool a `ReadPool` route would acquire from, or
3157    /// `None` for the `Primary` / `Unavailable` variants.
3158    #[cfg(not(feature = "sqlite"))]
3159    fn read_pool_size(route: &crate::repository::ReadRoute) -> Option<usize> {
3160        match route {
3161            crate::repository::ReadRoute::ReadPool(pool) => Some(pool.status().max_size),
3162            crate::repository::ReadRoute::Primary | crate::repository::ReadRoute::Unavailable => {
3163                None
3164            }
3165        }
3166    }
3167
3168    #[test]
3169    fn read_route_is_primary_without_replica() {
3170        let set = shard_set(&["a"]);
3171        let shard = set.get(ShardId(0)).expect("shard");
3172        assert!(
3173            matches!(shard.read_route(), crate::repository::ReadRoute::Primary),
3174            "a shard with no replica must keep reads on the primary"
3175        );
3176    }
3177
3178    #[cfg(not(feature = "sqlite"))]
3179    #[test]
3180    fn read_route_targets_replica_when_ready() {
3181        let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3182        shard.runtime().mark_replica_connection_ready();
3183        assert!(shard.runtime().replica_ready());
3184        assert_eq!(
3185            read_pool_size(&shard.read_route()),
3186            Some(REPLICA_SIZE),
3187            "a ready replica must route reads to the replica pool"
3188        );
3189    }
3190
3191    #[cfg(not(feature = "sqlite"))]
3192    #[test]
3193    fn read_route_falls_back_to_primary_when_unready_and_policy_allows() {
3194        // Replica configured but never checked → fallback policy applies.
3195        let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3196        assert_eq!(
3197            read_pool_size(&shard.read_route()),
3198            Some(PRIMARY_SIZE),
3199            "primary fallback must route reads to the primary pool"
3200        );
3201    }
3202
3203    #[cfg(not(feature = "sqlite"))]
3204    #[test]
3205    fn read_route_is_unavailable_when_unready_and_fallback_forbidden() {
3206        let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
3207        assert!(
3208            matches!(
3209                shard.read_route(),
3210                crate::repository::ReadRoute::Unavailable
3211            ),
3212            "fail_readiness must not silently fall back to the primary"
3213        );
3214    }
3215
3216    #[cfg(not(feature = "sqlite"))]
3217    #[test]
3218    fn repository_seed_snapshots_the_shard_read_route() {
3219        let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3220        shard.runtime().mark_replica_connection_ready();
3221        let ctx = crate::db::RequestDbContext {
3222            idle_in_transaction_timeout: None,
3223            statement_timeout: None,
3224            route_key: Some(crate::db::RouteKey::Composed("GET /notes".to_owned())),
3225            metrics: None,
3226            slow_query_threshold: std::time::Duration::from_millis(500),
3227            interceptors: Vec::new(),
3228            clock: std::sync::Arc::new(crate::time::SystemClock),
3229        };
3230        let seed = ShardRepositorySeed::from_ctx(
3231            shard.primary_pool(),
3232            &ctx,
3233            shard.name(),
3234            shard.read_route(),
3235        );
3236        assert_eq!(
3237            read_pool_size(&seed.read_route),
3238            Some(REPLICA_SIZE),
3239            "the seed must carry the shard's read route for from_shard"
3240        );
3241    }
3242
3243    // ── Shards routing surface ──────────────────────────────────────────
3244
3245    fn shards_handle(names: &[&str]) -> Shards {
3246        shards_handle_from(shard_set(names))
3247    }
3248
3249    fn shards_handle_from(set: ShardSet) -> Shards {
3250        Shards {
3251            set,
3252            ctx: crate::db::RequestDbContext {
3253                idle_in_transaction_timeout: None,
3254                statement_timeout: None,
3255                route_key: Some(crate::db::RouteKey::Composed("GET /test".to_owned())),
3256                metrics: None,
3257                slow_query_threshold: std::time::Duration::from_millis(500),
3258                interceptors: Vec::new(),
3259                clock: std::sync::Arc::new(crate::time::SystemClock),
3260            },
3261        }
3262    }
3263
3264    #[tokio::test]
3265    async fn db_on_rejects_unknown_shard_names() {
3266        let shards = shards_handle(&["alpha"]);
3267        let Err(error) = shards.db_on("beta").await else {
3268            panic!("unknown shard name must be rejected");
3269        };
3270        assert!(error.to_string().contains("beta"));
3271    }
3272
3273    #[cfg(not(feature = "sqlite"))]
3274    #[tokio::test]
3275    async fn read_for_fails_closed_without_checkout_under_fail_readiness() {
3276        let mut config = sharded_config(&["a"]);
3277        config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
3278        config.shards[0].replica_fallback = Some(ReplicaFallback::FailReadiness);
3279        let shards = Shards {
3280            set: create_shard_set(&config, Arc::new(HashShardRouter))
3281                .expect("build")
3282                .expect("configured"),
3283            ctx: crate::db::RequestDbContext {
3284                idle_in_transaction_timeout: None,
3285                statement_timeout: None,
3286                route_key: None,
3287                metrics: None,
3288                slow_query_threshold: std::time::Duration::from_millis(500),
3289                interceptors: Vec::new(),
3290                clock: std::sync::Arc::new(crate::time::SystemClock),
3291            },
3292        };
3293
3294        // The replica has not passed a readiness check, so the rejection
3295        // must be the fallback-policy error, not a connection failure.
3296        let Err(error) = shards.read_for("tenant-1").await else {
3297            panic!("unready replica under fail_readiness must be rejected");
3298        };
3299        assert!(error.to_string().contains("fail_readiness"));
3300    }
3301
3302    #[test]
3303    fn shards_exposes_set_and_iter() {
3304        let shards = shards_handle(&["alpha", "beta"]);
3305        assert_eq!(shards.set().len(), 2);
3306        let names: Vec<&str> = shards.iter().map(Shard::name).collect();
3307        assert_eq!(names, ["alpha", "beta"]);
3308    }
3309
3310    #[tokio::test]
3311    async fn route_rejects_out_of_range_router_results() {
3312        struct BadRouter;
3313        impl ShardRouter for BadRouter {
3314            fn route<'a>(
3315                &'a self,
3316                _key: ShardKey<'a>,
3317                _shards: &'a ShardSet,
3318            ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
3319                Box::pin(std::future::ready(Ok(ShardId(99))))
3320            }
3321        }
3322
3323        let set = create_shard_set(&sharded_config(&["a"]), Arc::new(BadRouter))
3324            .expect("build")
3325            .expect("configured");
3326        let error = set.route("k").await.expect_err("out of range");
3327        assert!(error.to_string().contains("out-of-range"));
3328    }
3329
3330    #[test]
3331    fn shard_key_from_impls_route_consistently() {
3332        // i32 widens to the same slot as the equivalent i64.
3333        assert_eq!(
3334            slot_for_key(ShardKey::from(42i32)),
3335            slot_for_key(ShardKey::from(42i64)),
3336        );
3337        // Owned strings, str slices, and byte arrays agree.
3338        let owned = "tenant-1".to_owned();
3339        let bytes: [u8; 16] = *b"0123456789abcdef";
3340        assert_eq!(
3341            slot_for_key(ShardKey::from(&owned)),
3342            slot_for_key(ShardKey::from("tenant-1")),
3343        );
3344        assert_eq!(
3345            slot_for_key(ShardKey::from(&bytes)),
3346            slot_for_key(ShardKey::from(&b"0123456789abcdef"[..])),
3347        );
3348    }
3349
3350    #[test]
3351    fn build_errors_and_debug_render_usefully() {
3352        let error = ShardSetBuildError::TopologyCountMismatch {
3353            expected: 2,
3354            actual: 0,
3355        };
3356        assert!(error.to_string().contains("expected 2"));
3357
3358        let set = shard_set(&["alpha"]);
3359        let debug = format!("{set:?}");
3360        assert!(
3361            debug.contains("alpha"),
3362            "ShardSet Debug names shards: {debug}"
3363        );
3364        let shard_debug = format!("{:?}", set.get(ShardId(0)).expect("shard"));
3365        assert!(shard_debug.contains("alpha"));
3366        assert_eq!(
3367            set.shard_for_slot(SlotId(0)).expect("owner").name(),
3368            "alpha"
3369        );
3370    }
3371
3372    // ── per-shard health indicator ──────────────────────────────────────
3373
3374    #[cfg(not(feature = "sqlite"))]
3375    fn shard_with_unreachable_replica(fallback: ReplicaFallback) -> Shard {
3376        let mut config = sharded_config(&["a"]);
3377        // Nothing listens on these URLs; keep the failing checks fast.
3378        config.connect_timeout_secs = 1;
3379        config.shards[0].replica_url = Some("postgres://localhost:1/a_ro".to_owned());
3380        config.shards[0].replica_fallback = Some(fallback);
3381        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3382            .expect("build")
3383            .expect("configured");
3384        set.get(ShardId(0)).expect("shard").clone()
3385    }
3386
3387    #[cfg(not(feature = "sqlite"))]
3388    #[tokio::test]
3389    async fn parity_check_timeout_fails_closed() {
3390        let shard = shard_with_unreachable_replica(ReplicaFallback::FailReadiness);
3391        let indicator = ShardHealthIndicator::new(shard, std::time::Duration::from_secs(2));
3392
3393        indicator
3394            .record_parity(std::future::pending(), std::time::Duration::from_millis(10))
3395            .await;
3396
3397        let detail = indicator.shard.runtime().detail().unwrap_or_default();
3398        assert!(detail.contains("parity check timed out"), "{detail}");
3399    }
3400
3401    #[cfg(not(feature = "sqlite"))]
3402    #[tokio::test]
3403    async fn shard_indicator_uses_the_configured_ping_timeout() {
3404        use crate::actuator::HealthIndicator as _;
3405
3406        let shard = shard_with_unreachable_replica(ReplicaFallback::FailReadiness);
3407        let indicator = ShardHealthIndicator::new(shard, std::time::Duration::from_millis(700));
3408
3409        // 700 ms ping + 5 s parity + 500 ms margin.
3410        assert_eq!(indicator.timeout_ms(), 6_200);
3411    }
3412
3413    #[cfg(not(feature = "sqlite"))]
3414    #[tokio::test]
3415    async fn shard_indicator_gates_readiness_for_fail_readiness_replica() {
3416        use crate::actuator::HealthIndicator as _;
3417
3418        let shard = shard_with_unreachable_replica(ReplicaFallback::FailReadiness);
3419        let indicator = ShardHealthIndicator::new(shard, crate::health_cache::DEFAULT_PING_TIMEOUT);
3420        let output = indicator.check().await;
3421
3422        assert!(
3423            !output.status.is_healthy(),
3424            "unreachable replica under fail_readiness must report Down"
3425        );
3426        assert_eq!(output.details["replica_ready"], serde_json::json!(false));
3427        assert!(output.details.contains_key("replica_detail"));
3428    }
3429
3430    #[cfg(not(feature = "sqlite"))]
3431    #[tokio::test]
3432    async fn shard_indicator_reports_down_when_primary_unreachable() {
3433        use crate::actuator::HealthIndicator as _;
3434
3435        // `ReplicaFallback::Primary` would normally let a dead replica degrade
3436        // to primary reads and stay Up — but here the primary is also
3437        // unreachable, so the primary connectivity gate must force Down: an
3438        // instance that cannot reach the shard primary fails all writes and
3439        // primary reads, so `/ready` must not stay green. (The healthy-primary
3440        // + dead-replica fallback path needs a live primary and is exercised by
3441        // the `read_pool`/`read_route` fallback tests above, not the indicator.)
3442        let shard = shard_with_unreachable_replica(ReplicaFallback::Primary);
3443        let indicator = ShardHealthIndicator::new(shard, crate::health_cache::DEFAULT_PING_TIMEOUT);
3444        let output = indicator.check().await;
3445
3446        assert!(
3447            !output.status.is_healthy(),
3448            "unreachable primary must report Down even under primary fallback"
3449        );
3450        assert_eq!(output.details["primary_ready"], serde_json::json!(false));
3451        assert!(output.details.contains_key("primary_detail"));
3452    }
3453
3454    // The primary-connectivity gate, on both backends. The two indicator tests
3455    // above reach it through the unreachable-REPLICA fixture, which SQLite has
3456    // no spelling for; this one needs no replica, and "unreachable" is spelled
3457    // per backend (nothing listens on TCP port 1 / `sqlite3_open` cannot create
3458    // a file in a directory that does not exist).
3459    #[tokio::test]
3460    async fn shard_indicator_reports_down_when_primary_cannot_be_reached() {
3461        use crate::actuator::HealthIndicator as _;
3462
3463        let mut config = sharded_config(&["a"]);
3464        config.connect_timeout_secs = 1;
3465        config.shards[0].primary_url = crate::test_urls::unreachable("a");
3466        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3467            .expect("build")
3468            .expect("configured");
3469        let shard = set.get(ShardId(0)).expect("shard").clone();
3470
3471        let output = ShardHealthIndicator::new(shard, crate::health_cache::DEFAULT_PING_TIMEOUT)
3472            .check()
3473            .await;
3474
3475        assert!(
3476            !output.status.is_healthy(),
3477            "an unreachable primary fails all writes and primary reads, so \
3478             /ready must not stay green"
3479        );
3480        assert_eq!(output.details["primary_ready"], serde_json::json!(false));
3481        assert!(output.details.contains_key("primary_detail"));
3482    }
3483
3484    #[tokio::test]
3485    async fn register_shard_health_indicators_names_components() {
3486        let set = shard_set(&["alpha", "beta"]);
3487        let registry = crate::actuator::HealthIndicatorRegistry::new();
3488
3489        register_shard_health_indicators(
3490            &set,
3491            &registry,
3492            crate::health_cache::DEFAULT_PING_TIMEOUT,
3493        );
3494        // Re-registration is ignored with a warning rather than panicking.
3495        register_shard_health_indicators(
3496            &set,
3497            &registry,
3498            crate::health_cache::DEFAULT_PING_TIMEOUT,
3499        );
3500
3501        let results = registry.run_all().await;
3502        // run_all also appends process-global results (e.g. circuit
3503        // breakers created by concurrently-running tests), so assert on
3504        // the shard components only.
3505        let mut names: Vec<&str> = results
3506            .iter()
3507            .map(|r| r.name.as_str())
3508            .filter(|name| name.starts_with("db:shard:"))
3509            .collect();
3510        names.sort_unstable();
3511        assert_eq!(names, ["db:shard:alpha", "db:shard:beta"]);
3512        assert!(
3513            results
3514                .iter()
3515                .filter(|r| r.name.starts_with("db:shard:"))
3516                .all(|r| matches!(r.group, crate::actuator::IndicatorGroup::Readiness)),
3517            "shard indicators gate readiness"
3518        );
3519    }
3520
3521    #[cfg(not(feature = "sqlite"))]
3522    #[test]
3523    fn total_max_connections_sums_every_pool() {
3524        let mut config = sharded_config(&["a", "b"]);
3525        config.pool_size = 7;
3526        config.shards[1].replica_url = Some("postgres://localhost/b_ro".to_owned());
3527        config.shards[1].replica_pool_size = Some(3);
3528        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3529            .expect("build")
3530            .expect("configured");
3531        // a primary (7) + b primary (7) + b replica (3).
3532        assert_eq!(set.total_max_connections(), 17);
3533    }
3534
3535    // The replica-free half of the sum, on both backends. The fixture above
3536    // needs a replica to make the total interesting and so is Postgres-only;
3537    // `cache=shared` is exempt from the single-slot clamp, so a configured
3538    // `pool_size` reaches a SQLite pool too and the arithmetic still holds.
3539    #[test]
3540    fn total_max_connections_sums_every_primary_pool() {
3541        let mut config = sharded_config(&["a", "b"]);
3542        config.pool_size = 7;
3543        let set = create_shard_set(&config, Arc::new(HashShardRouter))
3544            .expect("build")
3545            .expect("configured");
3546        // a primary (7) + b primary (7).
3547        assert_eq!(set.total_max_connections(), 14);
3548    }
3549
3550    // ── ShardRepositorySeed (#1273) ─────────────────────────────────────
3551
3552    #[test]
3553    fn repo_seed_from_ctx_preserves_statement_timeout() {
3554        let set = shard_set(&["shard0"]);
3555        let shard = set.get(ShardId(0)).expect("shard");
3556        let ctx = crate::db::RequestDbContext {
3557            idle_in_transaction_timeout: None,
3558            statement_timeout: Some(std::time::Duration::from_secs(3)),
3559            route_key: Some(crate::db::RouteKey::Composed("GET /test".to_owned())),
3560            metrics: None,
3561            slow_query_threshold: std::time::Duration::from_millis(200),
3562            interceptors: Vec::new(),
3563            clock: std::sync::Arc::new(crate::time::SystemClock),
3564        };
3565        let seed =
3566            ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
3567        assert_eq!(seed.statement_timeout_ms, 3_000, "timeout preserved as ms");
3568        assert_eq!(
3569            seed.slow_query_threshold,
3570            std::time::Duration::from_millis(200),
3571            "slow threshold preserved"
3572        );
3573        assert_eq!(
3574            seed.route.as_deref(),
3575            Some("GET /test shard=shard0"),
3576            "route tagged with shard name"
3577        );
3578    }
3579
3580    #[test]
3581    fn reshard_route_label_retags_with_target_shard() {
3582        // Fan-out sub-repo on shard2 must report under shard2, not the
3583        // originally-routed shard0, so per-shard metrics stay accurate.
3584        assert_eq!(
3585            reshard_route_label(Some("GET /admin shard=shard0"), "shard2").as_deref(),
3586            Some("GET /admin shard=shard2"),
3587        );
3588        // No parent label (no MatchedPath) stays unlabelled.
3589        assert_eq!(reshard_route_label(None, "shard2"), None);
3590        // A label without a shard tag still gets tagged for the target shard.
3591        assert_eq!(
3592            reshard_route_label(Some("GET /admin"), "shard2").as_deref(),
3593            Some("GET /admin shard=shard2"),
3594        );
3595    }
3596
3597    #[test]
3598    fn cross_shard_wrapper_derefs_to_inner() {
3599        let mut w = CrossShard(7i32);
3600        assert_eq!(*w, 7); // Deref
3601        *w = 9; // DerefMut
3602        assert_eq!(w.0, 9);
3603    }
3604
3605    #[test]
3606    fn cross_shard_seed_is_tenant_free_and_untagged() {
3607        // No tenant is resolved; the seed is built straight from the set so an
3608        // admin CrossShard<R> extractor can construct the repo without a header.
3609        let set = shard_set(&["shard0", "shard1"]);
3610        let ctx = crate::db::RequestDbContext {
3611            idle_in_transaction_timeout: None,
3612            statement_timeout: Some(std::time::Duration::from_millis(1500)),
3613            route_key: Some(crate::db::RouteKey::Composed("GET /admin".to_owned())),
3614            metrics: None,
3615            slow_query_threshold: std::time::Duration::from_millis(250),
3616            interceptors: Vec::new(),
3617            clock: std::sync::Arc::new(crate::time::SystemClock),
3618        };
3619        let seed = cross_shard_seed(&set, &ctx).expect("seed");
3620        // The route carries only the base key — the fan-out re-tags it per
3621        // executing shard via reshard_route_label, so it must NOT be pre-tagged
3622        // with the seed shard.
3623        assert_eq!(seed.route.as_deref(), Some("GET /admin"));
3624        assert!(!seed.route.as_deref().unwrap().contains("shard="));
3625        assert_eq!(seed.statement_timeout_ms, 1500);
3626        assert_eq!(
3627            seed.slow_query_threshold,
3628            std::time::Duration::from_millis(250)
3629        );
3630    }
3631
3632    #[test]
3633    fn repo_seed_none_timeout_maps_to_zero() {
3634        let set = shard_set(&["shard0"]);
3635        let shard = set.get(ShardId(0)).expect("shard");
3636        let ctx = crate::db::RequestDbContext {
3637            idle_in_transaction_timeout: None,
3638            statement_timeout: None,
3639            route_key: None,
3640            metrics: None,
3641            slow_query_threshold: std::time::Duration::from_millis(500),
3642            interceptors: Vec::new(),
3643            clock: std::sync::Arc::new(crate::time::SystemClock),
3644        };
3645        let seed =
3646            ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
3647        assert_eq!(seed.statement_timeout_ms, 0, "None timeout maps to 0");
3648        assert!(seed.route.is_none(), "None route_key propagates as None");
3649    }
3650
3651    #[test]
3652    fn repo_seed_timeout_capped_at_i32_max() {
3653        let set = shard_set(&["shard0"]);
3654        let shard = set.get(ShardId(0)).expect("shard");
3655        let ctx = crate::db::RequestDbContext {
3656            idle_in_transaction_timeout: None,
3657            statement_timeout: Some(std::time::Duration::from_secs(u64::MAX / 1_000)),
3658            route_key: None,
3659            metrics: None,
3660            slow_query_threshold: std::time::Duration::from_millis(500),
3661            interceptors: Vec::new(),
3662            clock: std::sync::Arc::new(crate::time::SystemClock),
3663        };
3664        let seed =
3665            ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
3666        assert_eq!(
3667            seed.statement_timeout_ms,
3668            i32::MAX as u64,
3669            "timeout capped at i32::MAX ms"
3670        );
3671    }
3672
3673    // ── ShardedReadDb / replica_read_pool (issue #1275) ─────────────────
3674
3675    #[test]
3676    fn replica_read_pool_is_none_without_replica() {
3677        let set = shard_set(&["a"]);
3678        let shard = set.get(ShardId(0)).expect("shard");
3679        assert!(
3680            shard.replica_read_pool().is_none(),
3681            "no replica configured → replica_read_pool must be None"
3682        );
3683    }
3684
3685    #[cfg(not(feature = "sqlite"))]
3686    #[test]
3687    fn replica_read_pool_is_none_when_unready_even_under_primary_fallback() {
3688        // The key difference from read_pool(): even with ReplicaFallback::Primary,
3689        // replica_read_pool never falls back to the primary — returns None.
3690        let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3691        assert!(
3692            shard.replica_read_pool().is_none(),
3693            "unready replica under primary fallback must still return None for replica_read_pool"
3694        );
3695    }
3696
3697    #[cfg(not(feature = "sqlite"))]
3698    #[test]
3699    fn replica_read_pool_is_none_when_unready_under_fail_readiness() {
3700        let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
3701        assert!(
3702            shard.replica_read_pool().is_none(),
3703            "unready replica under fail_readiness must return None"
3704        );
3705    }
3706
3707    #[cfg(not(feature = "sqlite"))]
3708    #[test]
3709    fn replica_read_pool_targets_replica_when_ready() {
3710        let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3711        shard.runtime().mark_replica_connection_ready();
3712        assert!(shard.runtime().replica_ready());
3713        assert_eq!(
3714            shard.replica_read_pool().map(|p| p.status().max_size),
3715            Some(REPLICA_SIZE),
3716            "a ready replica must be returned by replica_read_pool"
3717        );
3718    }
3719
3720    #[tokio::test]
3721    async fn read_replica_for_fails_when_no_replica_configured() {
3722        let shards = shards_handle(&["a"]);
3723        let Err(error) = shards.read_replica_for("tenant-1").await else {
3724            panic!("no replica configured must be rejected");
3725        };
3726        // Error must mention replica (not checkout failure) and must not
3727        // mention fail_readiness (this path is policy-independent).
3728        let msg = error.to_string();
3729        assert!(
3730            msg.contains("replica"),
3731            "error must name the missing replica: {msg}"
3732        );
3733        assert!(
3734            !msg.contains("fail_readiness"),
3735            "error must not mention fallback policy: {msg}"
3736        );
3737    }
3738
3739    #[cfg(not(feature = "sqlite"))]
3740    #[tokio::test]
3741    async fn read_replica_for_fails_when_replica_unready_under_primary_fallback() {
3742        // Unlike read_for, read_replica_for must NOT fall back to the primary.
3743        let mut config = sharded_config(&["a"]);
3744        config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
3745        config.shards[0].replica_fallback = Some(ReplicaFallback::Primary);
3746        let shards = Shards {
3747            set: create_shard_set(&config, Arc::new(HashShardRouter))
3748                .expect("build")
3749                .expect("configured"),
3750            ctx: crate::db::RequestDbContext {
3751                idle_in_transaction_timeout: None,
3752                statement_timeout: None,
3753                route_key: None,
3754                metrics: None,
3755                slow_query_threshold: std::time::Duration::from_millis(500),
3756                interceptors: Vec::new(),
3757                clock: std::sync::Arc::new(crate::time::SystemClock),
3758            },
3759        };
3760        let Err(error) = shards.read_replica_for("tenant-1").await else {
3761            panic!("unready replica must be rejected even under primary fallback");
3762        };
3763        assert!(
3764            error.to_string().contains("replica"),
3765            "must name the replica: {error}"
3766        );
3767    }
3768}