1use std::collections::HashMap;
38use std::sync::Arc;
39use std::sync::atomic::{AtomicBool, Ordering};
40
41use crate::db::RuntimeConnection;
42use diesel_async::AsyncPgConnection;
43use diesel_async::pooled_connection::deadpool::Pool;
44
45pub use crate::config::SLOT_COUNT;
46use crate::config::{ConfigError, DatabaseConfig, ReplicaFallback};
47use crate::db::{DatabaseTopology, PoolError};
48use crate::error::AutumnError;
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
55pub struct ShardId(pub usize);
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
59pub struct SlotId(pub u16);
60
61#[derive(Debug, Clone, Copy)]
67pub enum ShardKey<'a> {
68 Int(i64),
70 Str(&'a str),
72 Bytes(&'a [u8]),
74}
75
76impl From<i64> for ShardKey<'_> {
77 fn from(key: i64) -> Self {
78 Self::Int(key)
79 }
80}
81
82impl From<i32> for ShardKey<'_> {
83 fn from(key: i32) -> Self {
84 Self::Int(i64::from(key))
85 }
86}
87
88impl<'a> From<&'a str> for ShardKey<'a> {
89 fn from(key: &'a str) -> Self {
90 Self::Str(key)
91 }
92}
93
94impl<'a> From<&'a String> for ShardKey<'a> {
95 fn from(key: &'a String) -> Self {
96 Self::Str(key)
97 }
98}
99
100impl<'a> From<&'a [u8]> for ShardKey<'a> {
101 fn from(key: &'a [u8]) -> Self {
102 Self::Bytes(key)
103 }
104}
105
106impl<'a> From<&'a [u8; 16]> for ShardKey<'a> {
107 fn from(key: &'a [u8; 16]) -> Self {
108 Self::Bytes(key)
109 }
110}
111
112const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
122const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
123
124fn fnv1a_64(bytes: &[u8]) -> u64 {
125 let mut hash = FNV_OFFSET_BASIS;
126 for byte in bytes {
127 hash ^= u64::from(*byte);
128 hash = hash.wrapping_mul(FNV_PRIME);
129 }
130 hash
131}
132
133const fn splitmix64(mut x: u64) -> u64 {
136 x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
137 x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
138 x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
139 x ^ (x >> 31)
140}
141
142#[must_use]
144fn key_hash64(key: ShardKey<'_>) -> u64 {
145 match key {
146 #[allow(clippy::cast_sign_loss)]
147 ShardKey::Int(value) => splitmix64(value as u64),
148 ShardKey::Str(value) => fnv1a_64(value.as_bytes()),
149 ShardKey::Bytes(value) => fnv1a_64(value),
150 }
151}
152
153#[must_use]
158pub fn slot_for_key(key: ShardKey<'_>) -> SlotId {
159 let hash = key_hash64(key);
160 #[allow(clippy::cast_possible_truncation)]
161 SlotId((hash % u64::from(SLOT_COUNT)) as u16)
162}
163
164pub trait ShardRouter: Send + Sync + 'static {
179 fn route<'a>(
181 &'a self,
182 key: ShardKey<'a>,
183 shards: &'a ShardSet,
184 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>>;
185}
186
187impl<R: ShardRouter + ?Sized> ShardRouter for Arc<R> {
196 fn route<'a>(
197 &'a self,
198 key: ShardKey<'a>,
199 shards: &'a ShardSet,
200 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
201 (**self).route(key, shards)
202 }
203}
204
205#[derive(Debug, Default, Clone, Copy)]
208pub struct HashShardRouter;
209
210impl ShardRouter for HashShardRouter {
211 fn route<'a>(
212 &'a self,
213 key: ShardKey<'a>,
214 shards: &'a ShardSet,
215 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
216 let slot = shards.slot_for_key(key);
217 Box::pin(std::future::ready(
218 shards
219 .inner
220 .slot_map
221 .get(usize::from(slot.0))
222 .map(|&idx| ShardId(idx))
223 .ok_or_else(|| {
224 AutumnError::service_unavailable_msg(format!(
225 "slot {} has no shard assigned (slot map inconsistent)",
226 slot.0
227 ))
228 }),
229 ))
230 }
231}
232
233pub struct DirectoryShardRouter {
262 control_pool: Pool<RuntimeConnection>,
263 fallback: Arc<dyn ShardRouter>,
264 cache: std::sync::RwLock<HashMap<String, DirectoryCacheEntry>>,
265 ttl: std::time::Duration,
266 statement_timeout_ms: u64,
272}
273
274pub const DEFAULT_DIRECTORY_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(30);
277
278#[cfg(feature = "db")]
287pub const SHARD_DIRECTORY_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
288 diesel_migrations::embed_migrations!("shard_directory_migrations");
289
290#[cfg(feature = "db")]
297pub const SHARD_MAP_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
298 diesel_migrations::embed_migrations!("shard_map_migrations");
299
300#[derive(Clone, Copy)]
301struct DirectoryCacheEntry {
302 shard: ShardId,
303 expires_at: std::time::Instant,
304}
305
306#[derive(diesel::QueryableByName)]
307struct ShardNameRow {
308 #[diesel(sql_type = diesel::sql_types::Text)]
309 shard_name: String,
310}
311
312const DIRECTORY_NOTIFY_CHANNEL: &str = "autumn_shard_directory";
317
318pub const DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL: std::time::Duration =
325 std::time::Duration::from_secs(5);
326
327impl std::fmt::Debug for DirectoryShardRouter {
328 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
329 f.debug_struct("DirectoryShardRouter")
330 .field("ttl", &self.ttl)
331 .field("cached_keys", &self.cache.read().map_or(0, |c| c.len()))
332 .finish_non_exhaustive()
333 }
334}
335
336impl DirectoryShardRouter {
337 #[must_use]
340 pub fn new(control_pool: Pool<RuntimeConnection>) -> Self {
341 Self::with_fallback(control_pool, Arc::new(HashShardRouter))
342 }
343
344 #[must_use]
347 pub fn with_fallback(
348 control_pool: Pool<RuntimeConnection>,
349 fallback: Arc<dyn ShardRouter>,
350 ) -> Self {
351 Self {
352 control_pool,
353 fallback,
354 cache: std::sync::RwLock::new(HashMap::new()),
355 ttl: DEFAULT_DIRECTORY_CACHE_TTL,
356 statement_timeout_ms: 0,
357 }
358 }
359
360 #[must_use]
365 pub const fn with_statement_timeout_ms(mut self, statement_timeout_ms: u64) -> Self {
366 self.statement_timeout_ms = statement_timeout_ms;
367 self
368 }
369
370 #[must_use]
372 pub const fn with_cache_ttl(mut self, ttl: std::time::Duration) -> Self {
373 self.ttl = ttl;
374 self
375 }
376
377 pub fn invalidate(&self, tenant_key: &str) {
381 if let Ok(mut cache) = self.cache.write() {
382 cache.remove(tenant_key);
383 }
384 }
385
386 pub fn invalidate_all(&self) {
388 if let Ok(mut cache) = self.cache.write() {
389 cache.clear();
390 }
391 }
392
393 #[must_use]
415 pub fn spawn_invalidation_listener(
416 router: Arc<Self>,
417 control_url: String,
418 sweep_interval: std::time::Duration,
419 ) -> tokio::task::JoinHandle<()> {
420 use diesel_async::{AsyncConnection as _, RunQueryDsl as _};
421 use futures::StreamExt as _;
422
423 tokio::spawn(async move {
424 loop {
425 let Ok(mut conn) = AsyncPgConnection::establish(&control_url).await else {
429 tokio::time::sleep(sweep_interval).await;
430 continue;
431 };
432 if diesel::sql_query(format!("LISTEN {DIRECTORY_NOTIFY_CHANNEL}"))
433 .execute(&mut conn)
434 .await
435 .is_err()
436 {
437 tokio::time::sleep(sweep_interval).await;
438 continue;
439 }
440 router.invalidate_all();
444
445 let mut notifications = std::pin::pin!(conn.notifications_stream());
451 loop {
452 match tokio::time::timeout(sweep_interval, notifications.next()).await {
453 Ok(Some(Ok(notification))) => router.invalidate(¬ification.payload),
454 Ok(Some(Err(_)) | None) => break,
455 Err(_elapsed) => router.sweep_expired(),
456 }
457 }
458 }
459 })
460 }
461
462 fn cache_get(&self, key: &str) -> Option<ShardId> {
463 let now = std::time::Instant::now();
464 {
465 let cache = self.cache.read().ok()?;
466 match cache.get(key) {
467 Some(entry) if entry.expires_at > now => return Some(entry.shard),
468 Some(_) => {}
473 None => return None,
474 }
475 }
476 let mut cache = self.cache.write().ok()?;
480 if cache.get(key).is_some_and(|entry| entry.expires_at <= now) {
481 cache.remove(key);
482 }
483 None
484 }
485
486 fn sweep_expired(&self) {
491 if let Ok(mut cache) = self.cache.write() {
492 let now = std::time::Instant::now();
493 cache.retain(|_, entry| entry.expires_at > now);
494 }
495 }
496
497 fn cache_put(&self, key: String, shard: ShardId) {
498 if let Ok(mut cache) = self.cache.write() {
499 cache.insert(
500 key,
501 DirectoryCacheEntry {
502 shard,
503 expires_at: std::time::Instant::now() + self.ttl,
504 },
505 );
506 }
507 }
508
509 async fn lookup_directory(
513 &self,
514 key: &str,
515 shards: &ShardSet,
516 ) -> Result<Option<ShardId>, AutumnError> {
517 use diesel::OptionalExtension as _;
518 use diesel_async::RunQueryDsl;
519
520 let mut conn = self.control_pool.get().await.map_err(|e| {
521 AutumnError::service_unavailable_msg(format!(
522 "DirectoryShardRouter could not acquire a control connection: {e}"
523 ))
524 })?;
525
526 diesel::sql_query(format!(
532 "SET statement_timeout = {}",
533 self.statement_timeout_ms
534 ))
535 .execute(&mut conn)
536 .await
537 .map_err(|e| {
538 AutumnError::service_unavailable_msg(format!(
539 "DirectoryShardRouter could not set statement_timeout: {e}"
540 ))
541 })?;
542
543 let row = diesel::sql_query(
544 "SELECT shard_name FROM _autumn_shard_directory WHERE tenant_key = $1",
545 )
546 .bind::<diesel::sql_types::Text, _>(key)
547 .get_result::<ShardNameRow>(&mut conn)
548 .await
549 .optional()
550 .map_err(|e| {
551 AutumnError::service_unavailable_msg(format!(
552 "DirectoryShardRouter directory lookup failed: {e}"
553 ))
554 })?;
555
556 let Some(row) = row else {
557 return Ok(None);
558 };
559
560 let shard = shards.by_name(&row.shard_name).ok_or_else(|| {
561 AutumnError::service_unavailable_msg(format!(
562 "shard directory pins tenant {key:?} to unknown shard {:?}",
563 row.shard_name
564 ))
565 })?;
566 Ok(Some(shard.id()))
567 }
568}
569
570impl ShardRouter for DirectoryShardRouter {
571 fn route<'a>(
572 &'a self,
573 key: ShardKey<'a>,
574 shards: &'a ShardSet,
575 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
576 Box::pin(async move {
577 let ShardKey::Str(key_str) = key else {
580 return self.fallback.route(key, shards).await;
581 };
582
583 if let Some(cached) = self.cache_get(key_str) {
584 return Ok(cached);
585 }
586
587 match self.lookup_directory(key_str, shards).await? {
596 Some(shard) => {
597 self.cache_put(key_str.to_owned(), shard);
598 Ok(shard)
599 }
600 None => self.fallback.route(key, shards).await,
601 }
602 })
603 }
604}
605
606#[derive(Debug)]
612pub(crate) struct ShardRuntime {
613 replica_fallback: ReplicaFallback,
614 replica_configured: bool,
615 connection_ready: AtomicBool,
616 migrations_ready: AtomicBool,
617 detail: std::sync::RwLock<Option<String>>,
618 migration_check: std::sync::RwLock<Option<(String, String)>>,
622 parity_checked_at: std::sync::Mutex<Option<std::time::Instant>>,
626}
627
628const PARITY_RECHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
635
636#[cfg_attr(not(test), allow(dead_code))]
640impl ShardRuntime {
641 fn new(replica_fallback: ReplicaFallback, replica_configured: bool) -> Self {
642 Self {
643 replica_fallback,
644 replica_configured,
645 connection_ready: AtomicBool::new(false),
646 migrations_ready: AtomicBool::new(true),
647 detail: std::sync::RwLock::new(
648 replica_configured.then(|| "replica has not passed a readiness check".to_owned()),
649 ),
650 migration_check: std::sync::RwLock::new(None),
651 parity_checked_at: std::sync::Mutex::new(None),
652 }
653 }
654
655 pub(crate) fn configure_migration_check(&self, primary_url: String, replica_url: String) {
656 *self
657 .migration_check
658 .write()
659 .expect("shard runtime lock poisoned") = Some((primary_url, replica_url));
660 }
661
662 fn migration_check(&self) -> Option<(String, String)> {
663 self.migration_check
664 .read()
665 .expect("shard runtime lock poisoned")
666 .clone()
667 }
668
669 pub(crate) fn parity_check_due(&self) -> bool {
672 let mut checked_at = self
673 .parity_checked_at
674 .lock()
675 .expect("shard runtime lock poisoned");
676 if checked_at.is_none_or(|at| at.elapsed() >= PARITY_RECHECK_INTERVAL) {
677 *checked_at = Some(std::time::Instant::now());
678 true
679 } else {
680 false
681 }
682 }
683
684 fn replica_ready(&self) -> bool {
685 self.connection_ready.load(Ordering::Relaxed)
686 && self.migrations_ready.load(Ordering::Relaxed)
687 }
688
689 fn refresh_detail(&self) {
690 if self.replica_ready() {
691 *self.detail.write().expect("shard runtime lock poisoned") = None;
692 }
693 }
694
695 pub(crate) fn mark_replica_connection_ready(&self) {
696 self.connection_ready.store(true, Ordering::Relaxed);
697 self.refresh_detail();
698 }
699
700 pub(crate) fn mark_replica_connection_unready(&self, detail: impl Into<String>) {
701 self.connection_ready.store(false, Ordering::Relaxed);
702 *self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
703 }
704
705 pub(crate) fn mark_replica_migrations_ready(&self) {
706 self.migrations_ready.store(true, Ordering::Relaxed);
707 self.refresh_detail();
708 }
709
710 pub(crate) fn mark_replica_migrations_unready(&self, detail: impl Into<String>) {
711 self.migrations_ready.store(false, Ordering::Relaxed);
712 *self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
713 }
714
715 pub(crate) fn detail(&self) -> Option<String> {
716 self.detail
717 .read()
718 .expect("shard runtime lock poisoned")
719 .clone()
720 }
721}
722
723#[derive(Clone)]
728pub struct Shard {
729 name: Arc<str>,
730 id: ShardId,
731 slots: Arc<[u16]>,
732 topology: DatabaseTopology,
733 runtime: Arc<ShardRuntime>,
734}
735
736impl Shard {
737 #[must_use]
739 pub fn name(&self) -> &str {
740 &self.name
741 }
742
743 #[must_use]
745 pub const fn id(&self) -> ShardId {
746 self.id
747 }
748
749 #[must_use]
751 pub fn slots(&self) -> &[u16] {
752 &self.slots
753 }
754
755 #[must_use]
757 pub const fn topology(&self) -> &DatabaseTopology {
758 &self.topology
759 }
760
761 #[must_use]
763 pub const fn primary_pool(&self) -> &Pool<RuntimeConnection> {
764 self.topology.primary()
765 }
766
767 #[must_use]
769 pub const fn replica_pool(&self) -> Option<&Pool<RuntimeConnection>> {
770 self.topology.replica()
771 }
772
773 #[must_use]
782 pub fn read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
783 self.read_pool_with_role().map(|(pool, _)| pool)
784 }
785
786 #[must_use]
803 pub fn read_route(&self) -> crate::repository::ReadRoute {
804 use crate::repository::ReadRoute;
805 if !self.runtime.replica_configured {
806 return ReadRoute::Primary;
807 }
808 self.read_pool().map_or(ReadRoute::Unavailable, |pool| {
809 ReadRoute::ReadPool(pool.clone())
810 })
811 }
812
813 pub(crate) fn replica_read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
827 if self.runtime.replica_configured && self.runtime.replica_ready() {
828 self.topology.replica()
829 } else {
830 None
831 }
832 }
833
834 pub(crate) fn read_pool_with_role(&self) -> Option<(&Pool<RuntimeConnection>, &'static str)> {
837 if !self.runtime.replica_configured {
838 return Some((self.topology.primary(), "primary"));
839 }
840 if self.runtime.replica_ready() {
841 return self.topology.replica().map(|pool| (pool, "replica"));
842 }
843 match self.runtime.replica_fallback {
844 ReplicaFallback::Primary => Some((self.topology.primary(), "primary")),
845 ReplicaFallback::FailReadiness => None,
846 }
847 }
848
849 #[cfg_attr(not(test), allow(dead_code))]
850 pub(crate) fn runtime(&self) -> &ShardRuntime {
851 &self.runtime
852 }
853}
854
855impl std::fmt::Debug for Shard {
856 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
857 f.debug_struct("Shard")
858 .field("name", &self.name)
859 .field("id", &self.id)
860 .field("slots", &self.slots)
861 .finish_non_exhaustive()
862 }
863}
864
865struct ShardSetInner {
866 shards: Vec<Shard>,
867 by_name: HashMap<String, usize>,
868 slot_map: Vec<usize>,
870 router: Arc<dyn ShardRouter>,
871}
872
873#[derive(Clone)]
879pub struct ShardSet {
880 inner: Arc<ShardSetInner>,
881}
882
883impl ShardSet {
884 #[must_use]
886 pub fn len(&self) -> usize {
887 self.inner.shards.len()
888 }
889
890 #[must_use]
892 pub fn is_empty(&self) -> bool {
893 self.inner.shards.is_empty()
894 }
895
896 #[must_use]
898 pub const fn slot_count(&self) -> u16 {
899 SLOT_COUNT
900 }
901
902 #[must_use]
904 pub fn get(&self, id: ShardId) -> Option<&Shard> {
905 self.inner.shards.get(id.0)
906 }
907
908 #[must_use]
910 pub fn by_name(&self, name: &str) -> Option<&Shard> {
911 self.inner
912 .by_name
913 .get(name)
914 .and_then(|&idx| self.inner.shards.get(idx))
915 }
916
917 pub fn iter(&self) -> impl Iterator<Item = &Shard> {
919 self.inner.shards.iter()
920 }
921
922 #[must_use]
925 pub fn slot_for_key<'k>(&self, key: impl Into<ShardKey<'k>>) -> SlotId {
926 slot_for_key(key.into())
927 }
928
929 #[must_use]
931 pub fn shard_for_slot(&self, slot: SlotId) -> Option<&Shard> {
932 self.inner
933 .slot_map
934 .get(usize::from(slot.0))
935 .and_then(|&idx| self.inner.shards.get(idx))
936 }
937
938 pub async fn route<'k>(&self, key: impl Into<ShardKey<'k>>) -> Result<&Shard, AutumnError> {
946 let key = key.into();
947 let id = self.inner.router.route(key, self).await?;
948 self.get(id).ok_or_else(|| {
949 AutumnError::service_unavailable_msg(format!(
950 "shard router returned out-of-range shard id {} (have {} shards)",
951 id.0,
952 self.len()
953 ))
954 })
955 }
956
957 #[must_use]
961 pub fn total_max_connections(&self) -> usize {
962 self.inner
963 .shards
964 .iter()
965 .map(|shard| {
966 shard.topology().primary().status().max_size
967 + shard
968 .topology()
969 .replica()
970 .map_or(0, |pool| pool.status().max_size)
971 })
972 .sum()
973 }
974
975 #[must_use]
982 pub fn owns_key<'k>(&self, shard_id: ShardId, key: impl Into<ShardKey<'k>>) -> bool {
983 let slot = self.slot_for_key(key);
984 self.shard_for_slot(slot)
985 .is_some_and(|s| s.id() == shard_id)
986 }
987
988 #[must_use]
992 pub fn slots_for_shard(&self, shard_id: ShardId) -> Option<&[u16]> {
993 self.inner.shards.get(shard_id.0).map(Shard::slots)
994 }
995
996 #[must_use]
1003 pub fn partition_by_shard<'k>(
1004 &self,
1005 keys: impl IntoIterator<Item = &'k str>,
1006 ) -> std::collections::HashMap<ShardId, Vec<&'k str>> {
1007 let mut map: std::collections::HashMap<ShardId, Vec<&'k str>> =
1008 std::collections::HashMap::new();
1009 for key in keys {
1010 let slot = self.slot_for_key(key);
1011 if let Some(shard) = self.shard_for_slot(slot) {
1012 map.entry(shard.id()).or_default().push(key);
1013 }
1014 }
1015 map
1016 }
1017
1018 #[doc(hidden)]
1045 pub async fn fan_out_shards<T, Fut, F>(&self, f: F) -> Result<Vec<T>, crate::AutumnError>
1046 where
1047 T: Send + 'static,
1048 Fut: std::future::Future<Output = Result<T, crate::AutumnError>> + Send + 'static,
1049 F: Fn(&Shard) -> Fut + Send + Sync,
1050 {
1051 use futures::StreamExt as _;
1052
1053 let mut slots: Vec<Option<T>> = (0..self.inner.shards.len()).map(|_| None).collect();
1056 let mut in_flight = futures::stream::FuturesUnordered::new();
1057
1058 for (idx, shard) in self.inner.shards.iter().enumerate() {
1059 if in_flight.len() >= FAN_OUT_CONCURRENCY
1060 && let Some((i, result)) = in_flight.next().await
1061 {
1062 slots[i] = Some(result?);
1063 }
1064 let fut = f(shard);
1065 in_flight.push(async move { (idx, fut.await) });
1066 }
1067 while let Some((i, result)) = in_flight.next().await {
1068 slots[i] = Some(result?);
1069 }
1070 Ok(slots
1073 .into_iter()
1074 .map(|slot| slot.expect("every shard produced a result"))
1075 .collect())
1076 }
1077}
1078
1079impl std::fmt::Debug for ShardSet {
1080 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1081 f.debug_struct("ShardSet")
1082 .field("shards", &self.inner.shards)
1083 .finish_non_exhaustive()
1084 }
1085}
1086
1087#[derive(Debug, thiserror::Error)]
1091#[non_exhaustive]
1092pub enum ShardSetBuildError {
1093 #[error("failed to build pool for shard {shard:?}: {source}")]
1095 Pool {
1096 shard: String,
1098 source: PoolError,
1100 },
1101 #[error(transparent)]
1103 Config(#[from] ConfigError),
1104 #[error("expected {expected} shard topologies, got {actual}")]
1106 TopologyCountMismatch {
1107 expected: usize,
1109 actual: usize,
1111 },
1112}
1113
1114pub fn create_shard_set(
1125 config: &DatabaseConfig,
1126 router: Arc<dyn ShardRouter>,
1127) -> Result<Option<ShardSet>, ShardSetBuildError> {
1128 if !config.has_shards() {
1129 return Ok(None);
1130 }
1131 let topologies = config
1132 .shards
1133 .iter()
1134 .map(|shard| {
1135 crate::db::create_shard_topology(shard, config).map_err(|source| {
1136 ShardSetBuildError::Pool {
1137 shard: shard.name.clone(),
1138 source,
1139 }
1140 })
1141 })
1142 .collect::<Result<Vec<_>, _>>()?;
1143 build_shard_set(config, topologies, router).map(Some)
1144}
1145
1146#[cfg(not(feature = "sqlite"))]
1169pub fn create_shard_set_transactional(
1170 config: &DatabaseConfig,
1171 router: Arc<dyn ShardRouter>,
1172) -> Result<Option<ShardSet>, ShardSetBuildError> {
1173 if !config.has_shards() {
1174 return Ok(None);
1175 }
1176
1177 let timeout = std::time::Duration::from_secs(config.connect_timeout_secs);
1178
1179 let topologies = config
1180 .shards
1181 .iter()
1182 .map(|shard| {
1183 let manager = diesel_async::pooled_connection::AsyncDieselConnectionManager::<
1184 diesel_async::AsyncPgConnection,
1185 >::new(&shard.primary_url);
1186 let pool = Pool::builder(manager)
1187 .max_size(1)
1188 .wait_timeout(Some(timeout))
1189 .create_timeout(Some(timeout))
1190 .runtime(deadpool::Runtime::Tokio1)
1191 .post_create(deadpool::managed::Hook::async_fn(
1192 |conn: &mut diesel_async::AsyncPgConnection, _| {
1193 Box::pin(async move {
1194 use diesel_async::AsyncConnection as _;
1195 use diesel_async::RunQueryDsl as _;
1196 conn.begin_test_transaction().await.map_err(|e| {
1197 deadpool::managed::HookError::Backend(
1198 diesel_async::pooled_connection::PoolError::QueryError(e),
1199 )
1200 })?;
1201 diesel::sql_query("SET autumn.test_transaction_started = 'true'")
1202 .execute(conn)
1203 .await
1204 .map_err(|e| {
1205 deadpool::managed::HookError::Backend(
1206 diesel_async::pooled_connection::PoolError::QueryError(e),
1207 )
1208 })?;
1209 Ok(())
1210 })
1211 },
1212 ))
1213 .build()
1214 .map_err(|source| ShardSetBuildError::Pool {
1215 shard: shard.name.clone(),
1216 source: crate::db::PoolError::Build(source),
1217 })?;
1218 Ok(crate::db::DatabaseTopology::primary_only(pool))
1219 })
1220 .collect::<Result<Vec<_>, ShardSetBuildError>>()?;
1221 build_shard_set(config, topologies, router).map(Some)
1222}
1223
1224pub fn build_shard_set(
1233 config: &DatabaseConfig,
1234 topologies: Vec<DatabaseTopology>,
1235 router: Arc<dyn ShardRouter>,
1236) -> Result<ShardSet, ShardSetBuildError> {
1237 if topologies.len() != config.shards.len() {
1238 return Err(ShardSetBuildError::TopologyCountMismatch {
1239 expected: config.shards.len(),
1240 actual: topologies.len(),
1241 });
1242 }
1243 let slot_map = config.resolved_slot_map()?;
1244
1245 let mut slots_per_shard: Vec<Vec<u16>> = vec![Vec::new(); config.shards.len()];
1246 for (slot, &owner) in slot_map.iter().enumerate() {
1247 #[allow(clippy::cast_possible_truncation)]
1248 slots_per_shard[owner].push(slot as u16);
1249 }
1250
1251 let shards: Vec<Shard> = config
1252 .shards
1253 .iter()
1254 .zip(topologies)
1255 .enumerate()
1256 .map(|(idx, (shard_config, topology))| {
1257 let replica_configured = topology.replica().is_some();
1258 Shard {
1259 name: Arc::from(shard_config.name.as_str()),
1260 id: ShardId(idx),
1261 slots: Arc::from(std::mem::take(&mut slots_per_shard[idx])),
1262 topology,
1263 runtime: Arc::new(ShardRuntime::new(
1264 shard_config.effective_replica_fallback(config),
1265 replica_configured,
1266 )),
1267 }
1268 })
1269 .collect();
1270 let mut by_name = HashMap::with_capacity(shards.len());
1275 for (idx, shard) in shards.iter().enumerate() {
1276 if by_name.insert(shard.name().to_owned(), idx).is_some() {
1277 return Err(ConfigError::Validation(format!(
1278 "database.shards: shard name {:?} is declared more than once; \
1279 shard names must be unique",
1280 shard.name()
1281 ))
1282 .into());
1283 }
1284 }
1285
1286 Ok(ShardSet {
1287 inner: Arc::new(ShardSetInner {
1288 shards,
1289 by_name,
1290 slot_map,
1291 router,
1292 }),
1293 })
1294}
1295
1296pub(crate) struct ShardHealthIndicator {
1313 shard: Shard,
1314}
1315
1316impl ShardHealthIndicator {
1317 pub(crate) const fn new(shard: Shard) -> Self {
1318 Self { shard }
1319 }
1320
1321 async fn refresh_replica_readiness(&self) {
1322 let Some(replica_pool) = self.shard.replica_pool() else {
1323 return;
1324 };
1325 match replica_pool.get().await {
1329 Ok(conn) => {
1330 drop(conn);
1331 self.shard.runtime().mark_replica_connection_ready();
1332 if self.shard.runtime().parity_check_due()
1333 && let Some((primary_url, replica_url)) = self.shard.runtime().migration_check()
1334 {
1335 let readiness = crate::migrate::check_replica_migration_readiness_blocking(
1336 primary_url,
1337 replica_url,
1338 )
1339 .await;
1340 if readiness.is_ready() {
1341 self.shard.runtime().mark_replica_migrations_ready();
1342 } else if let Some(detail) = readiness.detail() {
1343 self.shard.runtime().mark_replica_migrations_unready(detail);
1344 }
1345 }
1346 }
1347 Err(error) => self
1348 .shard
1349 .runtime()
1350 .mark_replica_connection_unready(format!("replica connection failed: {error}")),
1351 }
1352 }
1353}
1354
1355impl crate::actuator::HealthIndicator for ShardHealthIndicator {
1356 fn check(&self) -> futures::future::BoxFuture<'_, crate::actuator::HealthCheckOutput> {
1357 Box::pin(async move {
1358 self.refresh_replica_readiness().await;
1359
1360 let mut details = HashMap::new();
1361 let status = self.shard.primary_pool().status();
1362 details.insert("pool_size".to_owned(), serde_json::json!(status.max_size));
1363 details.insert(
1364 "active_connections".to_owned(),
1365 serde_json::json!((status.max_size as u64).saturating_sub(status.available as u64)),
1366 );
1367 details.insert(
1368 "idle_connections".to_owned(),
1369 serde_json::json!(status.available),
1370 );
1371 details.insert(
1372 "slots".to_owned(),
1373 serde_json::json!(self.shard.slots().len()),
1374 );
1375 if self.shard.replica_pool().is_some() {
1376 details.insert(
1377 "replica_ready".to_owned(),
1378 serde_json::json!(self.shard.runtime().replica_ready()),
1379 );
1380 if let Some(detail) = self.shard.runtime().detail() {
1381 details.insert("replica_detail".to_owned(), serde_json::json!(detail));
1382 }
1383 }
1384
1385 let primary_ok = match self.shard.primary_pool().get().await {
1395 Ok(conn) => {
1396 drop(conn);
1397 true
1398 }
1399 Err(error) => {
1400 details.insert(
1401 "primary_detail".to_owned(),
1402 serde_json::json!(format!("primary connection failed: {error}")),
1403 );
1404 false
1405 }
1406 };
1407 details.insert("primary_ready".to_owned(), serde_json::json!(primary_ok));
1408
1409 let output = if primary_ok && self.shard.read_pool().is_some() {
1413 crate::actuator::HealthCheckOutput::up()
1414 } else {
1415 crate::actuator::HealthCheckOutput::down()
1416 };
1417 output.with_details(details)
1418 })
1419 }
1420}
1421
1422pub(crate) fn register_shard_health_indicators(
1425 set: &ShardSet,
1426 registry: &crate::actuator::HealthIndicatorRegistry,
1427) {
1428 for shard in set.iter() {
1429 let name = format!("db:shard:{}", shard.name());
1430 if let Err(error) = registry.register(
1431 name,
1432 crate::actuator::IndicatorGroup::Readiness,
1433 Arc::new(ShardHealthIndicator::new(shard.clone())),
1434 ) {
1435 tracing::warn!("{error}");
1436 }
1437 }
1438}
1439
1440#[derive(Debug, Clone)]
1451pub struct ShardKeyOverride(pub String);
1452
1453pub struct Shards {
1476 set: ShardSet,
1477 ctx: crate::db::RequestDbContext,
1478}
1479
1480impl<S> axum::extract::FromRequestParts<S> for Shards
1481where
1482 S: crate::db::DbState + Send + Sync,
1483{
1484 type Rejection = AutumnError;
1485
1486 async fn from_request_parts(
1487 parts: &mut axum::http::request::Parts,
1488 state: &S,
1489 ) -> Result<Self, Self::Rejection> {
1490 let set = state.shards().cloned().ok_or_else(no_shards_configured)?;
1491 let ctx = crate::db::RequestDbContext::from_parts(parts, state);
1492 Ok(Self { set, ctx })
1493 }
1494}
1495
1496fn no_shards_configured() -> AutumnError {
1497 AutumnError::service_unavailable_msg(
1498 "No shards configured: declare [[database.shards]] in autumn.toml \
1499 (see docs/guide/sharding.md)",
1500 )
1501}
1502
1503fn cross_shard_seed(
1511 set: &ShardSet,
1512 ctx: &crate::db::RequestDbContext,
1513) -> Result<ShardRepositorySeed, AutumnError> {
1514 let shard = set.iter().next().ok_or_else(no_shards_configured)?;
1515 let mut seed =
1516 ShardRepositorySeed::from_ctx(shard.primary_pool(), ctx, shard.name(), shard.read_route());
1517 seed.route.clone_from(&ctx.route_key);
1518 Ok(seed)
1519}
1520
1521pub trait CrossShardRepository: Sized {
1527 #[doc(hidden)]
1530 fn __autumn_from_cross_shard(seed: ShardRepositorySeed, set: ShardSet) -> Self;
1531}
1532
1533pub struct CrossShard<R>(pub R);
1552
1553impl<R> std::ops::Deref for CrossShard<R> {
1554 type Target = R;
1555 fn deref(&self) -> &R {
1556 &self.0
1557 }
1558}
1559
1560impl<R> std::ops::DerefMut for CrossShard<R> {
1561 fn deref_mut(&mut self) -> &mut R {
1562 &mut self.0
1563 }
1564}
1565
1566impl<S, R> axum::extract::FromRequestParts<S> for CrossShard<R>
1567where
1568 S: crate::db::DbState + Send + Sync,
1569 R: CrossShardRepository,
1570{
1571 type Rejection = AutumnError;
1572
1573 async fn from_request_parts(
1574 parts: &mut axum::http::request::Parts,
1575 state: &S,
1576 ) -> Result<Self, Self::Rejection> {
1577 let shards =
1580 <Shards as axum::extract::FromRequestParts<S>>::from_request_parts(parts, state)
1581 .await?;
1582 let seed = cross_shard_seed(&shards.set, &shards.ctx)?;
1583 Ok(Self(R::__autumn_from_cross_shard(seed, shards.set)))
1584 }
1585}
1586
1587const FAN_OUT_CONCURRENCY: usize = 8;
1589
1590impl Shards {
1591 #[must_use]
1593 pub const fn set(&self) -> &ShardSet {
1594 &self.set
1595 }
1596
1597 pub fn iter(&self) -> impl Iterator<Item = &Shard> {
1599 self.set.iter()
1600 }
1601
1602 pub async fn db_for<'k>(
1609 &self,
1610 key: impl Into<ShardKey<'k>>,
1611 ) -> Result<crate::db::Db, AutumnError> {
1612 let shard = self.set.route(key).await?;
1613 self.checkout_primary(shard).await
1614 }
1615
1616 pub async fn read_for<'k>(
1626 &self,
1627 key: impl Into<ShardKey<'k>>,
1628 ) -> Result<crate::db::Db, AutumnError> {
1629 let shard = self.set.route(key).await?;
1630 let (pool, role) = shard.read_pool_with_role().ok_or_else(|| {
1631 AutumnError::service_unavailable_msg(format!(
1632 "shard {:?} replica is not ready and replica_fallback = \"fail_readiness\"",
1633 shard.name()
1634 ))
1635 })?;
1636 self.checkout(shard, pool, role).await
1637 }
1638
1639 pub async fn read_replica_for<'k>(
1656 &self,
1657 key: impl Into<ShardKey<'k>>,
1658 ) -> Result<crate::db::Db, AutumnError> {
1659 let shard = self.set.route(key).await?;
1660 let pool = shard.replica_read_pool().ok_or_else(|| {
1661 AutumnError::service_unavailable_msg(format!(
1662 "shard {:?} has no healthy replica; read_replica_for requires a \
1663 configured, ready replica (no primary fallback)",
1664 shard.name()
1665 ))
1666 })?;
1667 self.checkout(shard, pool, "replica").await
1668 }
1669
1670 pub async fn db_on(&self, shard_name: &str) -> Result<crate::db::Db, AutumnError> {
1678 let shard = self
1679 .set
1680 .by_name(shard_name)
1681 .ok_or_else(|| AutumnError::bad_request_msg(format!("unknown shard {shard_name:?}")))?;
1682 self.checkout_primary(shard).await
1683 }
1684
1685 pub async fn each_shard<T, Fut, F>(&self, f: F) -> Vec<(ShardId, Result<T, AutumnError>)>
1711 where
1712 T: Send,
1713 Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
1714 F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
1715 {
1716 use futures::StreamExt as _;
1723
1724 let mut results: Vec<Option<(ShardId, Result<T, AutumnError>)>> =
1725 std::iter::repeat_with(|| None)
1726 .take(self.set.len())
1727 .collect();
1728 let mut in_flight: futures::stream::FuturesUnordered<
1729 futures::future::BoxFuture<'_, (ShardId, Result<T, AutumnError>)>,
1730 > = futures::stream::FuturesUnordered::new();
1731
1732 for shard in self.set.iter() {
1733 if in_flight.len() >= FAN_OUT_CONCURRENCY
1734 && let Some((id, result)) = in_flight.next().await
1735 {
1736 results[id.0] = Some((id, result));
1737 }
1738 in_flight.push(Box::pin(self.run_on_shard(shard, &f)));
1739 }
1740 while let Some((id, result)) = in_flight.next().await {
1741 results[id.0] = Some((id, result));
1742 }
1743 results.into_iter().flatten().collect()
1744 }
1745
1746 async fn run_on_shard<T, Fut, F>(
1747 &self,
1748 shard: &Shard,
1749 f: &F,
1750 ) -> (ShardId, Result<T, AutumnError>)
1751 where
1752 T: Send,
1753 Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
1754 F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
1755 {
1756 let result = match self.checkout_primary(shard).await {
1757 Ok(db) => f(shard, db).await,
1758 Err(error) => Err(error),
1759 };
1760 (shard.id(), result)
1761 }
1762
1763 async fn checkout_primary(&self, shard: &Shard) -> Result<crate::db::Db, AutumnError> {
1764 self.checkout(shard, shard.primary_pool(), "primary").await
1765 }
1766
1767 async fn checkout(
1768 &self,
1769 shard: &Shard,
1770 pool: &Pool<RuntimeConnection>,
1771 role: &str,
1772 ) -> Result<crate::db::Db, AutumnError> {
1773 let ctx = self.ctx.clone();
1774 crate::db::Db::checkout(crate::db::DbCheckoutParams {
1775 pool,
1776 pool_name: &format!("shard:{}:{role}", shard.name()),
1777 shard: Some(shard.name()),
1778 statement_timeout: ctx.statement_timeout,
1779 route_key: ctx
1782 .route_key
1783 .map(|key| format!("{key} shard={}", shard.name())),
1784 metrics: ctx.metrics,
1785 slow_query_threshold: ctx.slow_query_threshold,
1786 interceptors: ctx.interceptors,
1787 })
1788 .await
1789 }
1790}
1791
1792#[doc(hidden)]
1805#[derive(Clone)]
1806pub struct ShardRepositorySeed {
1807 pub pool: Pool<RuntimeConnection>,
1808 pub statement_timeout_ms: u64,
1812 pub slow_query_threshold: std::time::Duration,
1813 pub route: Option<String>,
1816 pub read_route: crate::repository::ReadRoute,
1821}
1822
1823impl ShardRepositorySeed {
1824 pub(crate) fn from_ctx(
1825 pool: &Pool<RuntimeConnection>,
1826 ctx: &crate::db::RequestDbContext,
1827 shard_name: &str,
1828 read_route: crate::repository::ReadRoute,
1829 ) -> Self {
1830 const PG_TIMEOUT_MAX_MS: u64 = i32::MAX as u64;
1833 let statement_timeout_ms = ctx.statement_timeout.map_or(0, |d| {
1834 u64::try_from(d.as_millis().min(u128::from(PG_TIMEOUT_MAX_MS)))
1835 .unwrap_or(PG_TIMEOUT_MAX_MS)
1836 });
1837 Self {
1838 pool: pool.clone(),
1839 statement_timeout_ms,
1840 slow_query_threshold: ctx.slow_query_threshold,
1841 route: ctx
1842 .route_key
1843 .as_ref()
1844 .map(|key| format!("{key} shard={shard_name}")),
1845 read_route,
1846 }
1847 }
1848}
1849
1850#[must_use]
1859pub fn reshard_route_label(parent: Option<&str>, shard_name: &str) -> Option<String> {
1860 let parent = parent?;
1861 let base = parent.rsplit_once(" shard=").map_or(parent, |(key, _)| key);
1862 Some(format!("{base} shard={shard_name}"))
1863}
1864
1865pub struct ShardedDb {
1893 db: crate::db::Db,
1894 shard_name: Arc<str>,
1895 shard_id: ShardId,
1896 repo_seed: ShardRepositorySeed,
1897 shards: ShardSet,
1901}
1902
1903impl ShardedDb {
1904 #[must_use]
1906 pub fn shard(&self) -> &str {
1907 &self.shard_name
1908 }
1909
1910 #[must_use]
1912 pub const fn shard_id(&self) -> ShardId {
1913 self.shard_id
1914 }
1915
1916 #[must_use]
1918 pub const fn span(&self) -> &tracing::Span {
1919 self.db.span()
1920 }
1921
1922 pub async fn tx<'a, T, E, F>(&'a mut self, f: F) -> Result<T, AutumnError>
1930 where
1931 T: Send + 'a,
1932 E: From<diesel::result::Error> + Send + Sync + 'a,
1933 AutumnError: From<E>,
1934 F: for<'r> FnOnce(
1935 &'r mut crate::db::PooledConnection,
1936 ) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
1937 + Send
1938 + 'a,
1939 {
1940 self.db.tx(f).await
1941 }
1942
1943 pub async fn tx_with<'a, T, E, F>(
1952 &'a mut self,
1953 opts: crate::db::TxOptions,
1954 f: F,
1955 ) -> Result<T, AutumnError>
1956 where
1957 T: Send + 'a,
1958 E: From<diesel::result::Error> + Send + Sync + 'a,
1959 AutumnError: From<E>,
1960 F: for<'r> FnMut(
1961 &'r mut crate::db::RuntimeConnection,
1962 ) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
1963 + Send
1964 + 'a,
1965 {
1966 self.db.tx_with(opts, f).await
1967 }
1968
1969 pub const fn db_mut(&mut self) -> &mut crate::db::Db {
1972 &mut self.db
1973 }
1974
1975 #[doc(hidden)]
1978 #[must_use]
1979 pub const fn __autumn_repository_seed(&self) -> &ShardRepositorySeed {
1980 &self.repo_seed
1981 }
1982
1983 #[doc(hidden)]
1986 #[must_use]
1987 pub const fn __autumn_shard_set(&self) -> &ShardSet {
1988 &self.shards
1989 }
1990}
1991
1992impl std::ops::Deref for ShardedDb {
1993 type Target = RuntimeConnection;
1994 fn deref(&self) -> &Self::Target {
1995 &self.db
1996 }
1997}
1998
1999impl std::ops::DerefMut for ShardedDb {
2000 fn deref_mut(&mut self) -> &mut Self::Target {
2001 &mut self.db
2002 }
2003}
2004
2005impl AsMut<crate::db::Db> for ShardedDb {
2006 fn as_mut(&mut self) -> &mut crate::db::Db {
2007 &mut self.db
2008 }
2009}
2010
2011#[doc(hidden)]
2019pub async fn __autumn_resolve_repo_seed(
2020 parts: &mut axum::http::request::Parts,
2021 state: &crate::AppState,
2022) -> Result<(ShardRepositorySeed, ShardSet), AutumnError> {
2023 let shards = <Shards as axum::extract::FromRequestParts<crate::AppState>>::from_request_parts(
2024 parts, state,
2025 )
2026 .await?;
2027 let key = resolve_shard_key(parts, state).await?;
2028 let shard = shards.set.route(&key).await?;
2029 let shard_name = Arc::clone(&shard.name);
2030 let seed = ShardRepositorySeed::from_ctx(
2031 shard.primary_pool(),
2032 &shards.ctx,
2033 &shard_name,
2034 shard.read_route(),
2035 );
2036 let set = shards.set.clone();
2037 Ok((seed, set))
2038}
2039
2040impl axum::extract::FromRequestParts<crate::AppState> for ShardedDb {
2041 type Rejection = AutumnError;
2042
2043 async fn from_request_parts(
2044 parts: &mut axum::http::request::Parts,
2045 state: &crate::AppState,
2046 ) -> Result<Self, Self::Rejection> {
2047 let shards = Shards::from_request_parts(parts, state).await?;
2048 let key = resolve_shard_key(parts, state).await?;
2049
2050 let shard = shards.set.route(&key).await?;
2051 let shard_name = Arc::clone(&shard.name);
2052 let shard_id = shard.id();
2053 let repo_seed = ShardRepositorySeed::from_ctx(
2054 shard.primary_pool(),
2055 &shards.ctx,
2056 &shard_name,
2057 shard.read_route(),
2058 );
2059 let shard_set = shards.set.clone();
2060 let db = shards.checkout_primary(shard).await?;
2061 crate::read_your_writes::mark_write();
2062 Ok(Self {
2063 db,
2064 shard_name,
2065 shard_id,
2066 repo_seed,
2067 shards: shard_set,
2068 })
2069 }
2070}
2071
2072pub struct ShardedReadDb {
2098 db: crate::db::Db,
2099 shard_name: Arc<str>,
2100 shard_id: ShardId,
2101}
2102
2103impl ShardedReadDb {
2104 #[must_use]
2106 pub fn shard(&self) -> &str {
2107 &self.shard_name
2108 }
2109
2110 #[must_use]
2112 pub const fn shard_id(&self) -> ShardId {
2113 self.shard_id
2114 }
2115
2116 #[must_use]
2118 pub const fn span(&self) -> &tracing::Span {
2119 self.db.span()
2120 }
2121
2122 pub const fn db_mut(&mut self) -> &mut crate::db::Db {
2125 &mut self.db
2126 }
2127}
2128
2129impl std::ops::Deref for ShardedReadDb {
2130 type Target = RuntimeConnection;
2131 fn deref(&self) -> &Self::Target {
2132 &self.db
2133 }
2134}
2135
2136impl std::ops::DerefMut for ShardedReadDb {
2137 fn deref_mut(&mut self) -> &mut Self::Target {
2138 &mut self.db
2139 }
2140}
2141
2142impl AsMut<crate::db::Db> for ShardedReadDb {
2143 fn as_mut(&mut self) -> &mut crate::db::Db {
2144 &mut self.db
2145 }
2146}
2147
2148impl axum::extract::FromRequestParts<crate::AppState> for ShardedReadDb {
2149 type Rejection = AutumnError;
2150
2151 async fn from_request_parts(
2152 parts: &mut axum::http::request::Parts,
2153 state: &crate::AppState,
2154 ) -> Result<Self, Self::Rejection> {
2155 let shards = Shards::from_request_parts(parts, state).await?;
2156 let key = resolve_shard_key(parts, state).await?;
2157
2158 let shard = shards.set.route(&key).await?;
2159 let shard_name = Arc::clone(&shard.name);
2160 let shard_id = shard.id();
2161 let pool = shard.replica_read_pool().ok_or_else(|| {
2162 AutumnError::service_unavailable_msg(format!(
2163 "shard {:?} has no healthy replica; ShardedReadDb requires a \
2164 configured, ready replica (no primary fallback)",
2165 shard.name()
2166 ))
2167 })?;
2168 let db = shards.checkout(shard, pool, "replica").await?;
2169 Ok(Self {
2170 db,
2171 shard_name,
2172 shard_id,
2173 })
2174 }
2175}
2176
2177async fn resolve_shard_key(
2180 parts: &mut axum::http::request::Parts,
2181 state: &crate::AppState,
2182) -> Result<String, AutumnError> {
2183 if let Some(overridden) = parts.extensions.get::<ShardKeyOverride>() {
2184 return Ok(overridden.0.clone());
2185 }
2186 if let Ok(Some(tenant)) = crate::tenancy::CURRENT_TENANT.try_with(std::clone::Clone::clone) {
2187 return Ok(tenant);
2188 }
2189 let config = state
2190 .extension::<crate::config::AutumnConfig>()
2191 .ok_or_else(|| AutumnError::service_unavailable_msg("Config is not available"))?;
2192 crate::tenancy::extract_tenant_from_parts(parts, &config)
2193 .await
2194 .map_err(|error| {
2195 AutumnError::bad_request_msg(format!(
2196 "ShardedDb could not resolve a shard key: {error}. Enable [tenancy] so \
2197 the tenant id can route the request, or insert a ShardKeyOverride \
2198 request extension from middleware (see docs/guide/sharding.md)"
2199 ))
2200 })
2201}
2202
2203#[cfg(test)]
2204mod tests {
2205 use super::*;
2206 use crate::config::{ShardConfig, SlotSpec};
2207
2208 #[test]
2209 fn directory_invalidation_channel_and_interval_are_sane() {
2210 assert_eq!(DIRECTORY_NOTIFY_CHANNEL, "autumn_shard_directory");
2214 assert!(
2219 DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL < DEFAULT_DIRECTORY_CACHE_TTL,
2220 "sweep interval should beat the TTL"
2221 );
2222 }
2223
2224 fn shard_config(name: &str) -> ShardConfig {
2225 ShardConfig {
2226 name: name.to_owned(),
2227 primary_url: format!("postgres://localhost/{name}"),
2228 slots: None,
2229 replica_url: None,
2230 primary_pool_size: None,
2231 replica_pool_size: None,
2232 replica_fallback: None,
2233 }
2234 }
2235
2236 fn sharded_config(names: &[&str]) -> DatabaseConfig {
2237 DatabaseConfig {
2238 shards: names.iter().map(|name| shard_config(name)).collect(),
2239 ..Default::default()
2240 }
2241 }
2242
2243 fn shard_set(names: &[&str]) -> ShardSet {
2244 create_shard_set(&sharded_config(names), Arc::new(HashShardRouter))
2245 .expect("lazy pools should build")
2246 .expect("shards configured")
2247 }
2248
2249 #[test]
2256 fn golden_vector_str_keys() {
2257 let cases: &[(&str, u16)] = &[
2261 ("tenant-1", 12427),
2262 ("tenant-2", 12862),
2263 ("tenant-3", 13297),
2264 ("acme-corp", 11394),
2265 ("globex", 12846),
2266 ("initech", 11329),
2267 ("hooli", 3974),
2268 ("", 8997),
2269 ("a", 11404),
2270 ("00000000-0000-0000-0000-000000000001", 6206),
2271 ];
2272 for (key, expected_slot) in cases {
2273 assert_eq!(
2274 slot_for_key(ShardKey::Str(key)),
2275 SlotId(*expected_slot),
2276 "key {key:?} must keep routing to slot {expected_slot} forever",
2277 );
2278 }
2279 }
2280
2281 #[test]
2282 fn golden_vector_int_keys() {
2283 let cases: &[(i64, u16)] = &[
2287 (0, 3503),
2288 (1, 7361),
2289 (2, 5838),
2290 (42, 11925),
2291 (1_000_000, 1511),
2292 (-1, 11296),
2293 (i64::MAX, 7847),
2294 (i64::MIN, 13275),
2295 ];
2296 for (key, expected_slot) in cases {
2297 assert_eq!(
2298 slot_for_key(ShardKey::Int(*key)),
2299 SlotId(*expected_slot),
2300 "key {key} must keep routing to slot {expected_slot} forever",
2301 );
2302 }
2303 }
2304
2305 #[test]
2306 fn golden_vector_bytes_match_equivalent_str() {
2307 assert_eq!(
2309 slot_for_key(ShardKey::Bytes(b"tenant-1")),
2310 slot_for_key(ShardKey::Str("tenant-1")),
2311 );
2312 }
2313
2314 #[test]
2315 fn slots_stay_in_range_and_spread_roughly_uniformly() {
2316 let mut histogram = [0usize; 16];
2319 for i in 0..10_000i64 {
2320 let slot = slot_for_key(ShardKey::Int(i));
2321 assert!(slot.0 < SLOT_COUNT);
2322 histogram[usize::from(slot.0 / 1024)] += 1;
2323 }
2324 let expected = 10_000 / histogram.len();
2325 for (bucket, count) in histogram.iter().enumerate() {
2326 assert!(
2327 *count > expected / 2 && *count < expected * 2,
2328 "bucket {bucket} has {count} keys (expected ≈{expected})"
2329 );
2330 }
2331 }
2332
2333 #[tokio::test]
2336 async fn db_for_and_read_for_attempt_routed_checkouts() {
2337 let shards = shards_handle(&["alpha"]);
2340 let Err(error) = shards.db_for("tenant-1").await else {
2341 panic!("checkout must fail without a server");
2342 };
2343 assert!(!error.to_string().contains("Unknown shard"));
2344
2345 let Err(error) = shards.read_for("tenant-1").await else {
2347 panic!("checkout must fail without a server");
2348 };
2349 assert!(!error.to_string().contains("fail_readiness"));
2350 }
2351
2352 #[test]
2353 fn parity_recheck_is_throttled_per_window() {
2354 let set = shard_set(&["a"]);
2355 let runtime = set.get(ShardId(0)).expect("shard").runtime();
2356 runtime.configure_migration_check(
2357 "postgres://localhost/a".to_owned(),
2358 "postgres://localhost/a_ro".to_owned(),
2359 );
2360 assert!(runtime.migration_check().is_some());
2361
2362 assert!(runtime.parity_check_due(), "first check claims the window");
2363 assert!(
2364 !runtime.parity_check_due(),
2365 "checks within the window are suppressed"
2366 );
2367 }
2368
2369 #[tokio::test]
2370 async fn route_is_deterministic_and_respects_slot_map() {
2371 let mut config = sharded_config(&["a", "b"]);
2372 config.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2373 config.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
2374 let set = create_shard_set(&config, Arc::new(HashShardRouter))
2375 .expect("build")
2376 .expect("configured");
2377
2378 for key in ["k1", "k2", "k3", "k4", "k5"] {
2379 let slot = set.slot_for_key(key);
2380 let expected = if slot.0 >= 8192 { "b" } else { "a" };
2381 let routed = set.route(key).await.expect("route");
2382 assert_eq!(routed.name(), expected, "key {key:?} slot {}", slot.0);
2383 assert_eq!(set.route(key).await.expect("route").id(), routed.id());
2385 }
2386 }
2387
2388 #[tokio::test]
2389 async fn arc_shard_router_delegates_to_inner() {
2390 let config = sharded_config(&["a", "b"]);
2396 let set = create_shard_set(&config, Arc::new(Arc::new(HashShardRouter)))
2397 .expect("build")
2398 .expect("configured");
2399 let first = set.route("tenant-42").await.expect("route");
2400 let again = set.route("tenant-42").await.expect("route");
2401 assert_eq!(
2402 first.id(),
2403 again.id(),
2404 "Arc<R> routes deterministically through its inner router"
2405 );
2406 }
2407
2408 #[tokio::test]
2409 async fn moving_a_slot_in_config_moves_only_that_slot() {
2410 let mut before = sharded_config(&["a", "b"]);
2413 before.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2414 before.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
2415
2416 let mut after = sharded_config(&["a", "b", "c"]);
2417 after.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
2418 after.shards[1].slots = Some(vec![SlotSpec::Range("8192-12287".to_owned())]);
2419 after.shards[2].slots = Some(vec![SlotSpec::Range("12288-16383".to_owned())]);
2420
2421 let set_before = create_shard_set(&before, Arc::new(HashShardRouter))
2422 .expect("build")
2423 .expect("configured");
2424 let set_after = create_shard_set(&after, Arc::new(HashShardRouter))
2425 .expect("build")
2426 .expect("configured");
2427
2428 let mut moved = 0;
2429 for i in 0..200i64 {
2430 let slot = set_before.slot_for_key(i);
2431 assert_eq!(slot, set_after.slot_for_key(i), "key→slot never changes");
2432 let before_shard = set_before.route(i).await.expect("route");
2433 let after_shard = set_after.route(i).await.expect("route");
2434 if slot.0 >= 12288 {
2435 assert_eq!(before_shard.name(), "b");
2436 assert_eq!(after_shard.name(), "c");
2437 moved += 1;
2438 } else {
2439 assert_eq!(before_shard.name(), after_shard.name());
2440 }
2441 }
2442 assert!(moved > 0, "some keys must exercise the moved slot range");
2443 }
2444
2445 #[test]
2446 fn by_name_and_get_resolve_shards() {
2447 let set = shard_set(&["alpha", "beta"]);
2448 assert_eq!(set.len(), 2);
2449 assert_eq!(set.by_name("beta").expect("beta").id(), ShardId(1));
2450 assert_eq!(set.get(ShardId(0)).expect("alpha").name(), "alpha");
2451 assert!(set.by_name("gamma").is_none());
2452 assert!(set.get(ShardId(9)).is_none());
2453 let names: Vec<&str> = set.iter().map(Shard::name).collect();
2454 assert_eq!(names, ["alpha", "beta"]);
2455 }
2456
2457 #[test]
2458 fn auto_split_assigns_contiguous_slots() {
2459 let set = shard_set(&["a", "b"]);
2460 assert_eq!(set.slot_count(), SLOT_COUNT);
2461 assert_eq!(set.get(ShardId(0)).expect("a").slots().len(), 8192);
2462 assert_eq!(
2463 set.get(ShardId(1)).expect("b").slots(),
2464 (8192..16384).collect::<Vec<u16>>()
2465 );
2466 }
2467
2468 #[test]
2470 fn owns_key_agrees_with_route() {
2471 let set = shard_set(&["a", "b"]);
2475 assert!(
2476 set.owns_key(ShardId(0), "hooli"),
2477 "hooli (slot 3974) must be shard a"
2478 );
2479 assert!(
2480 !set.owns_key(ShardId(1), "hooli"),
2481 "hooli must not be shard b"
2482 );
2483 assert!(
2484 set.owns_key(ShardId(1), "a"),
2485 "key 'a' (slot 11404) must be shard b"
2486 );
2487 assert!(
2488 !set.owns_key(ShardId(0), "a"),
2489 "key 'a' must not be shard a"
2490 );
2491 }
2492
2493 #[test]
2494 fn slots_for_shard_returns_correct_slice() {
2495 let set = shard_set(&["a", "b"]);
2496 let a_slots = set.slots_for_shard(ShardId(0)).expect("shard a exists");
2497 let b_slots = set.slots_for_shard(ShardId(1)).expect("shard b exists");
2498 assert_eq!(a_slots.len(), 8192);
2499 assert_eq!(b_slots.len(), 8192);
2500 assert!(a_slots.iter().all(|&s| s < 8192));
2501 assert!(b_slots.iter().all(|&s| s >= 8192));
2502 assert!(set.slots_for_shard(ShardId(9)).is_none());
2503 }
2504
2505 #[test]
2506 fn partition_by_shard_groups_golden_keys() {
2507 let set = shard_set(&["a", "b"]);
2509 let keys = ["hooli", "a", "tenant-1"]; let map = set.partition_by_shard(keys.iter().copied());
2511 #[allow(clippy::similar_names)]
2512 let keys_on_a = map.get(&ShardId(0)).map_or(&[][..], Vec::as_slice);
2513 #[allow(clippy::similar_names)]
2514 let keys_on_b = map.get(&ShardId(1)).map_or(&[][..], Vec::as_slice);
2515 assert!(keys_on_a.contains(&"hooli"), "hooli must go to shard a");
2516 assert!(keys_on_b.contains(&"a"), "key 'a' must go to shard b");
2517 assert!(
2518 keys_on_b.contains(&"tenant-1"),
2519 "tenant-1 (slot 12427) must go to shard b"
2520 );
2521 assert_eq!(
2522 keys_on_a.len() + keys_on_b.len(),
2523 keys.len(),
2524 "no key dropped"
2525 );
2526 }
2527
2528 #[test]
2529 fn create_shard_set_returns_none_without_shards() {
2530 let config = DatabaseConfig::default();
2531 assert!(
2532 create_shard_set(&config, Arc::new(HashShardRouter))
2533 .expect("ok")
2534 .is_none()
2535 );
2536 }
2537
2538 #[test]
2539 fn build_shard_set_rejects_duplicate_names_without_config_validation() {
2540 let config = sharded_config(&["twin", "twin"]);
2544 let topologies = config
2545 .shards
2546 .iter()
2547 .map(|shard| crate::db::create_shard_topology(shard, &config).expect("lazy pools"))
2548 .collect();
2549
2550 let result = build_shard_set(&config, topologies, Arc::new(HashShardRouter));
2551
2552 let Err(ShardSetBuildError::Config(error)) = result else {
2553 panic!("duplicate shard names must be rejected, got {result:?}");
2554 };
2555 assert!(error.to_string().contains("twin"));
2556 }
2557
2558 #[test]
2559 fn build_shard_set_rejects_topology_count_mismatch() {
2560 let config = sharded_config(&["a", "b"]);
2561 let result = build_shard_set(&config, Vec::new(), Arc::new(HashShardRouter));
2562 assert!(matches!(
2563 result,
2564 Err(ShardSetBuildError::TopologyCountMismatch {
2565 expected: 2,
2566 actual: 0
2567 })
2568 ));
2569 }
2570
2571 fn shard_with_replica(fallback: ReplicaFallback) -> Shard {
2574 let mut config = sharded_config(&["a"]);
2575 config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
2576 config.shards[0].replica_fallback = Some(fallback);
2577 let set = create_shard_set(&config, Arc::new(HashShardRouter))
2578 .expect("build")
2579 .expect("configured");
2580 set.get(ShardId(0)).expect("shard").clone()
2581 }
2582
2583 #[test]
2584 fn read_pool_uses_primary_when_no_replica() {
2585 let set = shard_set(&["a"]);
2586 let shard = set.get(ShardId(0)).expect("shard");
2587 assert!(shard.read_pool().is_some());
2588 assert!(shard.replica_pool().is_none());
2589 }
2590
2591 #[test]
2592 fn read_pool_requires_readiness_check_before_replica_traffic() {
2593 let shard = shard_with_replica(ReplicaFallback::Primary);
2594 assert!(shard.read_pool().is_some());
2596 assert!(shard.runtime().detail().is_some());
2597
2598 shard.runtime().mark_replica_connection_ready();
2599 assert!(shard.runtime().replica_ready());
2600 assert!(shard.read_pool().is_some());
2601 assert!(shard.runtime().detail().is_none());
2602 }
2603
2604 #[test]
2605 fn read_pool_fails_closed_under_fail_readiness() {
2606 let shard = shard_with_replica(ReplicaFallback::FailReadiness);
2607 assert!(
2608 shard.read_pool().is_none(),
2609 "unchecked replica fails closed"
2610 );
2611
2612 shard.runtime().mark_replica_connection_ready();
2613 assert!(shard.read_pool().is_some());
2614
2615 shard
2616 .runtime()
2617 .mark_replica_migrations_unready("replica lags primary");
2618 assert!(shard.read_pool().is_none());
2619 assert!(shard.runtime().detail().expect("detail").contains("lags"));
2620 }
2621
2622 const PRIMARY_SIZE: usize = 7;
2625 const REPLICA_SIZE: usize = 3;
2626
2627 fn shard_with_sized_replica(fallback: ReplicaFallback) -> Shard {
2630 let mut config = sharded_config(&["a"]);
2631 config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
2632 config.shards[0].replica_fallback = Some(fallback);
2633 config.shards[0].primary_pool_size = Some(PRIMARY_SIZE);
2634 config.shards[0].replica_pool_size = Some(REPLICA_SIZE);
2635 let set = create_shard_set(&config, Arc::new(HashShardRouter))
2636 .expect("build")
2637 .expect("configured");
2638 set.get(ShardId(0)).expect("shard").clone()
2639 }
2640
2641 fn read_pool_size(route: &crate::repository::ReadRoute) -> Option<usize> {
2644 match route {
2645 crate::repository::ReadRoute::ReadPool(pool) => Some(pool.status().max_size),
2646 crate::repository::ReadRoute::Primary | crate::repository::ReadRoute::Unavailable => {
2647 None
2648 }
2649 }
2650 }
2651
2652 #[test]
2653 fn read_route_is_primary_without_replica() {
2654 let set = shard_set(&["a"]);
2655 let shard = set.get(ShardId(0)).expect("shard");
2656 assert!(
2657 matches!(shard.read_route(), crate::repository::ReadRoute::Primary),
2658 "a shard with no replica must keep reads on the primary"
2659 );
2660 }
2661
2662 #[test]
2663 fn read_route_targets_replica_when_ready() {
2664 let shard = shard_with_sized_replica(ReplicaFallback::Primary);
2665 shard.runtime().mark_replica_connection_ready();
2666 assert!(shard.runtime().replica_ready());
2667 assert_eq!(
2668 read_pool_size(&shard.read_route()),
2669 Some(REPLICA_SIZE),
2670 "a ready replica must route reads to the replica pool"
2671 );
2672 }
2673
2674 #[test]
2675 fn read_route_falls_back_to_primary_when_unready_and_policy_allows() {
2676 let shard = shard_with_sized_replica(ReplicaFallback::Primary);
2678 assert_eq!(
2679 read_pool_size(&shard.read_route()),
2680 Some(PRIMARY_SIZE),
2681 "primary fallback must route reads to the primary pool"
2682 );
2683 }
2684
2685 #[test]
2686 fn read_route_is_unavailable_when_unready_and_fallback_forbidden() {
2687 let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
2688 assert!(
2689 matches!(
2690 shard.read_route(),
2691 crate::repository::ReadRoute::Unavailable
2692 ),
2693 "fail_readiness must not silently fall back to the primary"
2694 );
2695 }
2696
2697 #[test]
2698 fn repository_seed_snapshots_the_shard_read_route() {
2699 let shard = shard_with_sized_replica(ReplicaFallback::Primary);
2700 shard.runtime().mark_replica_connection_ready();
2701 let ctx = crate::db::RequestDbContext {
2702 statement_timeout: None,
2703 route_key: Some("GET /notes".to_owned()),
2704 metrics: None,
2705 slow_query_threshold: std::time::Duration::from_millis(500),
2706 interceptors: Vec::new(),
2707 };
2708 let seed = ShardRepositorySeed::from_ctx(
2709 shard.primary_pool(),
2710 &ctx,
2711 shard.name(),
2712 shard.read_route(),
2713 );
2714 assert_eq!(
2715 read_pool_size(&seed.read_route),
2716 Some(REPLICA_SIZE),
2717 "the seed must carry the shard's read route for from_shard"
2718 );
2719 }
2720
2721 fn shards_handle(names: &[&str]) -> Shards {
2724 Shards {
2725 set: shard_set(names),
2726 ctx: crate::db::RequestDbContext {
2727 statement_timeout: None,
2728 route_key: Some("GET /test".to_owned()),
2729 metrics: None,
2730 slow_query_threshold: std::time::Duration::from_millis(500),
2731 interceptors: Vec::new(),
2732 },
2733 }
2734 }
2735
2736 #[tokio::test]
2737 async fn db_on_rejects_unknown_shard_names() {
2738 let shards = shards_handle(&["alpha"]);
2739 let Err(error) = shards.db_on("beta").await else {
2740 panic!("unknown shard name must be rejected");
2741 };
2742 assert!(error.to_string().contains("beta"));
2743 }
2744
2745 #[tokio::test]
2746 async fn read_for_fails_closed_without_checkout_under_fail_readiness() {
2747 let mut config = sharded_config(&["a"]);
2748 config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
2749 config.shards[0].replica_fallback = Some(ReplicaFallback::FailReadiness);
2750 let shards = Shards {
2751 set: create_shard_set(&config, Arc::new(HashShardRouter))
2752 .expect("build")
2753 .expect("configured"),
2754 ctx: crate::db::RequestDbContext {
2755 statement_timeout: None,
2756 route_key: None,
2757 metrics: None,
2758 slow_query_threshold: std::time::Duration::from_millis(500),
2759 interceptors: Vec::new(),
2760 },
2761 };
2762
2763 let Err(error) = shards.read_for("tenant-1").await else {
2766 panic!("unready replica under fail_readiness must be rejected");
2767 };
2768 assert!(error.to_string().contains("fail_readiness"));
2769 }
2770
2771 #[test]
2772 fn shards_exposes_set_and_iter() {
2773 let shards = shards_handle(&["alpha", "beta"]);
2774 assert_eq!(shards.set().len(), 2);
2775 let names: Vec<&str> = shards.iter().map(Shard::name).collect();
2776 assert_eq!(names, ["alpha", "beta"]);
2777 }
2778
2779 #[tokio::test]
2780 async fn route_rejects_out_of_range_router_results() {
2781 struct BadRouter;
2782 impl ShardRouter for BadRouter {
2783 fn route<'a>(
2784 &'a self,
2785 _key: ShardKey<'a>,
2786 _shards: &'a ShardSet,
2787 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
2788 Box::pin(std::future::ready(Ok(ShardId(99))))
2789 }
2790 }
2791
2792 let set = create_shard_set(&sharded_config(&["a"]), Arc::new(BadRouter))
2793 .expect("build")
2794 .expect("configured");
2795 let error = set.route("k").await.expect_err("out of range");
2796 assert!(error.to_string().contains("out-of-range"));
2797 }
2798
2799 #[test]
2800 fn shard_key_from_impls_route_consistently() {
2801 assert_eq!(
2803 slot_for_key(ShardKey::from(42i32)),
2804 slot_for_key(ShardKey::from(42i64)),
2805 );
2806 let owned = "tenant-1".to_owned();
2808 let bytes: [u8; 16] = *b"0123456789abcdef";
2809 assert_eq!(
2810 slot_for_key(ShardKey::from(&owned)),
2811 slot_for_key(ShardKey::from("tenant-1")),
2812 );
2813 assert_eq!(
2814 slot_for_key(ShardKey::from(&bytes)),
2815 slot_for_key(ShardKey::from(&b"0123456789abcdef"[..])),
2816 );
2817 }
2818
2819 #[test]
2820 fn build_errors_and_debug_render_usefully() {
2821 let error = ShardSetBuildError::TopologyCountMismatch {
2822 expected: 2,
2823 actual: 0,
2824 };
2825 assert!(error.to_string().contains("expected 2"));
2826
2827 let set = shard_set(&["alpha"]);
2828 let debug = format!("{set:?}");
2829 assert!(
2830 debug.contains("alpha"),
2831 "ShardSet Debug names shards: {debug}"
2832 );
2833 let shard_debug = format!("{:?}", set.get(ShardId(0)).expect("shard"));
2834 assert!(shard_debug.contains("alpha"));
2835 assert_eq!(
2836 set.shard_for_slot(SlotId(0)).expect("owner").name(),
2837 "alpha"
2838 );
2839 }
2840
2841 fn shard_with_unreachable_replica(fallback: ReplicaFallback) -> Shard {
2844 let mut config = sharded_config(&["a"]);
2845 config.connect_timeout_secs = 1;
2847 config.shards[0].replica_url = Some("postgres://localhost:1/a_ro".to_owned());
2848 config.shards[0].replica_fallback = Some(fallback);
2849 let set = create_shard_set(&config, Arc::new(HashShardRouter))
2850 .expect("build")
2851 .expect("configured");
2852 set.get(ShardId(0)).expect("shard").clone()
2853 }
2854
2855 #[tokio::test]
2856 async fn shard_indicator_gates_readiness_for_fail_readiness_replica() {
2857 use crate::actuator::HealthIndicator as _;
2858
2859 let shard = shard_with_unreachable_replica(ReplicaFallback::FailReadiness);
2860 let indicator = ShardHealthIndicator::new(shard);
2861 let output = indicator.check().await;
2862
2863 assert!(
2864 !output.status.is_healthy(),
2865 "unreachable replica under fail_readiness must report Down"
2866 );
2867 assert_eq!(output.details["replica_ready"], serde_json::json!(false));
2868 assert!(output.details.contains_key("replica_detail"));
2869 }
2870
2871 #[tokio::test]
2872 async fn shard_indicator_reports_down_when_primary_unreachable() {
2873 use crate::actuator::HealthIndicator as _;
2874
2875 let shard = shard_with_unreachable_replica(ReplicaFallback::Primary);
2883 let indicator = ShardHealthIndicator::new(shard);
2884 let output = indicator.check().await;
2885
2886 assert!(
2887 !output.status.is_healthy(),
2888 "unreachable primary must report Down even under primary fallback"
2889 );
2890 assert_eq!(output.details["primary_ready"], serde_json::json!(false));
2891 assert!(output.details.contains_key("primary_detail"));
2892 }
2893
2894 #[tokio::test]
2895 async fn register_shard_health_indicators_names_components() {
2896 let set = shard_set(&["alpha", "beta"]);
2897 let registry = crate::actuator::HealthIndicatorRegistry::new();
2898
2899 register_shard_health_indicators(&set, ®istry);
2900 register_shard_health_indicators(&set, ®istry);
2902
2903 let results = registry.run_all().await;
2904 let mut names: Vec<&str> = results
2908 .iter()
2909 .map(|r| r.name.as_str())
2910 .filter(|name| name.starts_with("db:shard:"))
2911 .collect();
2912 names.sort_unstable();
2913 assert_eq!(names, ["db:shard:alpha", "db:shard:beta"]);
2914 assert!(
2915 results
2916 .iter()
2917 .filter(|r| r.name.starts_with("db:shard:"))
2918 .all(|r| matches!(r.group, crate::actuator::IndicatorGroup::Readiness)),
2919 "shard indicators gate readiness"
2920 );
2921 }
2922
2923 #[test]
2924 fn total_max_connections_sums_every_pool() {
2925 let mut config = sharded_config(&["a", "b"]);
2926 config.pool_size = 7;
2927 config.shards[1].replica_url = Some("postgres://localhost/b_ro".to_owned());
2928 config.shards[1].replica_pool_size = Some(3);
2929 let set = create_shard_set(&config, Arc::new(HashShardRouter))
2930 .expect("build")
2931 .expect("configured");
2932 assert_eq!(set.total_max_connections(), 17);
2934 }
2935
2936 #[test]
2939 fn repo_seed_from_ctx_preserves_statement_timeout() {
2940 let set = shard_set(&["shard0"]);
2941 let shard = set.get(ShardId(0)).expect("shard");
2942 let ctx = crate::db::RequestDbContext {
2943 statement_timeout: Some(std::time::Duration::from_secs(3)),
2944 route_key: Some("GET /test".to_owned()),
2945 metrics: None,
2946 slow_query_threshold: std::time::Duration::from_millis(200),
2947 interceptors: Vec::new(),
2948 };
2949 let seed =
2950 ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
2951 assert_eq!(seed.statement_timeout_ms, 3_000, "timeout preserved as ms");
2952 assert_eq!(
2953 seed.slow_query_threshold,
2954 std::time::Duration::from_millis(200),
2955 "slow threshold preserved"
2956 );
2957 assert_eq!(
2958 seed.route.as_deref(),
2959 Some("GET /test shard=shard0"),
2960 "route tagged with shard name"
2961 );
2962 }
2963
2964 #[test]
2965 fn reshard_route_label_retags_with_target_shard() {
2966 assert_eq!(
2969 reshard_route_label(Some("GET /admin shard=shard0"), "shard2").as_deref(),
2970 Some("GET /admin shard=shard2"),
2971 );
2972 assert_eq!(reshard_route_label(None, "shard2"), None);
2974 assert_eq!(
2976 reshard_route_label(Some("GET /admin"), "shard2").as_deref(),
2977 Some("GET /admin shard=shard2"),
2978 );
2979 }
2980
2981 #[test]
2982 fn cross_shard_wrapper_derefs_to_inner() {
2983 let mut w = CrossShard(7i32);
2984 assert_eq!(*w, 7); *w = 9; assert_eq!(w.0, 9);
2987 }
2988
2989 #[test]
2990 fn cross_shard_seed_is_tenant_free_and_untagged() {
2991 let set = shard_set(&["shard0", "shard1"]);
2994 let ctx = crate::db::RequestDbContext {
2995 statement_timeout: Some(std::time::Duration::from_millis(1500)),
2996 route_key: Some("GET /admin".to_owned()),
2997 metrics: None,
2998 slow_query_threshold: std::time::Duration::from_millis(250),
2999 interceptors: Vec::new(),
3000 };
3001 let seed = cross_shard_seed(&set, &ctx).expect("seed");
3002 assert_eq!(seed.route.as_deref(), Some("GET /admin"));
3006 assert!(!seed.route.as_deref().unwrap().contains("shard="));
3007 assert_eq!(seed.statement_timeout_ms, 1500);
3008 assert_eq!(
3009 seed.slow_query_threshold,
3010 std::time::Duration::from_millis(250)
3011 );
3012 }
3013
3014 #[test]
3015 fn repo_seed_none_timeout_maps_to_zero() {
3016 let set = shard_set(&["shard0"]);
3017 let shard = set.get(ShardId(0)).expect("shard");
3018 let ctx = crate::db::RequestDbContext {
3019 statement_timeout: None,
3020 route_key: None,
3021 metrics: None,
3022 slow_query_threshold: std::time::Duration::from_millis(500),
3023 interceptors: Vec::new(),
3024 };
3025 let seed =
3026 ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
3027 assert_eq!(seed.statement_timeout_ms, 0, "None timeout maps to 0");
3028 assert!(seed.route.is_none(), "None route_key propagates as None");
3029 }
3030
3031 #[test]
3032 fn repo_seed_timeout_capped_at_i32_max() {
3033 let set = shard_set(&["shard0"]);
3034 let shard = set.get(ShardId(0)).expect("shard");
3035 let ctx = crate::db::RequestDbContext {
3036 statement_timeout: Some(std::time::Duration::from_secs(u64::MAX / 1_000)),
3037 route_key: None,
3038 metrics: None,
3039 slow_query_threshold: std::time::Duration::from_millis(500),
3040 interceptors: Vec::new(),
3041 };
3042 let seed =
3043 ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
3044 assert_eq!(
3045 seed.statement_timeout_ms,
3046 i32::MAX as u64,
3047 "timeout capped at i32::MAX ms"
3048 );
3049 }
3050
3051 #[test]
3054 fn replica_read_pool_is_none_without_replica() {
3055 let set = shard_set(&["a"]);
3056 let shard = set.get(ShardId(0)).expect("shard");
3057 assert!(
3058 shard.replica_read_pool().is_none(),
3059 "no replica configured → replica_read_pool must be None"
3060 );
3061 }
3062
3063 #[test]
3064 fn replica_read_pool_is_none_when_unready_even_under_primary_fallback() {
3065 let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3068 assert!(
3069 shard.replica_read_pool().is_none(),
3070 "unready replica under primary fallback must still return None for replica_read_pool"
3071 );
3072 }
3073
3074 #[test]
3075 fn replica_read_pool_is_none_when_unready_under_fail_readiness() {
3076 let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
3077 assert!(
3078 shard.replica_read_pool().is_none(),
3079 "unready replica under fail_readiness must return None"
3080 );
3081 }
3082
3083 #[test]
3084 fn replica_read_pool_targets_replica_when_ready() {
3085 let shard = shard_with_sized_replica(ReplicaFallback::Primary);
3086 shard.runtime().mark_replica_connection_ready();
3087 assert!(shard.runtime().replica_ready());
3088 assert_eq!(
3089 shard.replica_read_pool().map(|p| p.status().max_size),
3090 Some(REPLICA_SIZE),
3091 "a ready replica must be returned by replica_read_pool"
3092 );
3093 }
3094
3095 #[tokio::test]
3096 async fn read_replica_for_fails_when_no_replica_configured() {
3097 let shards = shards_handle(&["a"]);
3098 let Err(error) = shards.read_replica_for("tenant-1").await else {
3099 panic!("no replica configured must be rejected");
3100 };
3101 let msg = error.to_string();
3104 assert!(
3105 msg.contains("replica"),
3106 "error must name the missing replica: {msg}"
3107 );
3108 assert!(
3109 !msg.contains("fail_readiness"),
3110 "error must not mention fallback policy: {msg}"
3111 );
3112 }
3113
3114 #[tokio::test]
3115 async fn read_replica_for_fails_when_replica_unready_under_primary_fallback() {
3116 let mut config = sharded_config(&["a"]);
3118 config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
3119 config.shards[0].replica_fallback = Some(ReplicaFallback::Primary);
3120 let shards = Shards {
3121 set: create_shard_set(&config, Arc::new(HashShardRouter))
3122 .expect("build")
3123 .expect("configured"),
3124 ctx: crate::db::RequestDbContext {
3125 statement_timeout: None,
3126 route_key: None,
3127 metrics: None,
3128 slow_query_threshold: std::time::Duration::from_millis(500),
3129 interceptors: Vec::new(),
3130 },
3131 };
3132 let Err(error) = shards.read_replica_for("tenant-1").await else {
3133 panic!("unready replica must be rejected even under primary fallback");
3134 };
3135 assert!(
3136 error.to_string().contains("replica"),
3137 "must name the replica: {error}"
3138 );
3139 }
3140}