1#![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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
66pub struct ShardId(pub usize);
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
70pub struct SlotId(pub u16);
71
72#[derive(Debug, Clone, Copy)]
78pub enum ShardKey<'a> {
79 Int(i64),
81 Str(&'a str),
83 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
123const 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
144const 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#[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#[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#[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
199pub trait ShardRouter: Send + Sync + 'static {
214 fn route<'a>(
216 &'a self,
217 key: ShardKey<'a>,
218 shards: &'a ShardSet,
219 ) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>>;
220}
221
222impl<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#[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
268pub 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: u64,
307}
308
309pub const DEFAULT_DIRECTORY_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(30);
312
313#[cfg(feature = "db")]
322pub const SHARD_DIRECTORY_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
323 diesel_migrations::embed_migrations!("shard_directory_migrations");
324
325#[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#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
352const DIRECTORY_NOTIFY_CHANNEL: &str = "autumn_shard_directory";
353
354pub 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 #[must_use]
376 pub fn new(control_pool: Pool<RuntimeConnection>) -> Self {
377 Self::with_fallback(control_pool, Arc::new(HashShardRouter))
378 }
379
380 #[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 #[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 #[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 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 pub fn invalidate_all(&self) {
424 if let Ok(mut cache) = self.cache.write() {
425 cache.clear();
426 }
427 }
428
429 #[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 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 router.invalidate_all();
481
482 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(¬ification.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 Some(_) => {}
510 None => return None,
511 }
512 }
513 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 #[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 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 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 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 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 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#[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 migration_check: std::sync::RwLock<Option<(String, String)>>,
662 parity_checked_at: std::sync::Mutex<Option<std::time::Instant>>,
666}
667
668const PARITY_RECHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
675
676#[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 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#[derive(Clone)]
770pub struct Shard {
771 name: Arc<str>,
772 id: ShardId,
773 slots: Arc<[u16]>,
774 topology: DatabaseTopology,
775 runtime: Arc<ShardRuntime>,
776 metric_label: Arc<str>,
780 #[cfg(feature = "sqlite")]
782 fleet_db: Option<crate::db::fleet::FleetDatabase>,
783}
784
785pub type ShardLease = Arc<dyn std::any::Any + Send + Sync>;
791
792pub const FLEET_METRIC_LABEL: &str = "fleet";
794
795impl Shard {
796 #[must_use]
799 pub fn name(&self) -> &str {
800 &self.name
801 }
802
803 #[must_use]
807 pub fn metric_label(&self) -> &str {
808 &self.metric_label
809 }
810
811 #[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 #[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 #[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 #[must_use]
859 pub const fn id(&self) -> ShardId {
860 self.id
861 }
862
863 #[must_use]
865 pub fn slots(&self) -> &[u16] {
866 &self.slots
867 }
868
869 #[must_use]
871 pub const fn topology(&self) -> &DatabaseTopology {
872 &self.topology
873 }
874
875 #[must_use]
877 pub const fn primary_pool(&self) -> &Pool<RuntimeConnection> {
878 self.topology.primary()
879 }
880
881 #[must_use]
883 pub const fn replica_pool(&self) -> Option<&Pool<RuntimeConnection>> {
884 self.topology.replica()
885 }
886
887 #[must_use]
896 pub fn read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
897 self.read_pool_with_role().map(|(pool, _)| pool)
898 }
899
900 #[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 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 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: Vec<usize>,
984 router: Arc<dyn ShardRouter>,
985 #[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#[derive(Clone)]
1003pub struct ShardSet {
1004 inner: Arc<ShardSetInner>,
1005}
1006
1007impl ShardSet {
1008 #[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 #[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 #[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 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 #[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 #[must_use]
1116 pub fn len(&self) -> usize {
1117 self.inner.shards.len()
1118 }
1119
1120 #[must_use]
1122 pub fn is_empty(&self) -> bool {
1123 self.inner.shards.is_empty()
1124 }
1125
1126 #[must_use]
1128 pub const fn slot_count(&self) -> u16 {
1129 SLOT_COUNT
1130 }
1131
1132 #[must_use]
1134 pub fn get(&self, id: ShardId) -> Option<&Shard> {
1135 self.inner.shards.get(id.0)
1136 }
1137
1138 #[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 pub fn iter(&self) -> impl Iterator<Item = &Shard> {
1149 self.inner.shards.iter()
1150 }
1151
1152 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 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 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#[derive(Debug, thiserror::Error)]
1352#[non_exhaustive]
1353pub enum ShardSetBuildError {
1354 #[error("failed to build pool for shard {shard:?}: {source}")]
1356 Pool {
1357 shard: String,
1359 source: PoolError,
1361 },
1362 #[error(transparent)]
1364 Config(#[from] ConfigError),
1365 #[error("expected {expected} shard topologies, got {actual}")]
1367 TopologyCountMismatch {
1368 expected: usize,
1370 actual: usize,
1372 },
1373}
1374
1375pub 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#[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
1485pub 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 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
1562pub(crate) struct ShardHealthIndicator {
1579 shard: Shard,
1580 primary_ping: crate::db_ping::DbPingCheck,
1583 replica_ping: crate::db_ping::DbPingCheck,
1584 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 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 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
1662const 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 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 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 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 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
1737pub(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#[derive(Debug, Clone)]
1777pub struct ShardKeyOverride(pub String);
1778
1779pub 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
1829fn cross_shard_seed(
1837 set: &ShardSet,
1838 ctx: &crate::db::RequestDbContext,
1839) -> Result<ShardRepositorySeed, AutumnError> {
1840 #[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
1860pub trait CrossShardRepository: Sized {
1866 #[doc(hidden)]
1869 fn __autumn_from_cross_shard(seed: ShardRepositorySeed, set: ShardSet) -> Self;
1870}
1871
1872pub 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 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
1926const FAN_OUT_CONCURRENCY: usize = 8;
1928
1929impl Shards {
1930 #[must_use]
1932 pub const fn set(&self) -> &ShardSet {
1933 &self.set
1934 }
1935
1936 pub fn iter(&self) -> impl Iterator<Item = &Shard> {
1938 self.set.iter()
1939 }
1940
1941 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 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 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 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 #[cfg(feature = "sqlite")]
2024 #[must_use]
2025 pub fn fleet(&self) -> Option<&crate::db::fleet::DatabaseFleet> {
2026 self.set.fleet()
2027 }
2028
2029 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 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 #[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 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 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 #[cfg(all(feature = "reporting", not(feature = "sqlite")))]
2179 capture_gap: None,
2180 clock: ctx.clock,
2181 })
2182 .await
2183 }
2184}
2185
2186#[doc(hidden)]
2199#[derive(Clone)]
2200pub struct ShardRepositorySeed {
2201 pub pool: Pool<RuntimeConnection>,
2202 pub statement_timeout_ms: u64,
2206 pub slow_query_threshold: std::time::Duration,
2207 pub route: Option<String>,
2210 pub read_route: crate::repository::ReadRoute,
2215 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 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#[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
2263pub struct ShardedDb {
2291 db: crate::db::Db,
2292 shard_name: Arc<str>,
2293 shard_id: ShardId,
2294 repo_seed: ShardRepositorySeed,
2295 shards: ShardSet,
2299}
2300
2301impl ShardedDb {
2302 #[must_use]
2304 pub fn shard(&self) -> &str {
2305 &self.shard_name
2306 }
2307
2308 #[must_use]
2310 pub const fn shard_id(&self) -> ShardId {
2311 self.shard_id
2312 }
2313
2314 #[must_use]
2316 pub const fn span(&self) -> &tracing::Span {
2317 self.db.span()
2318 }
2319
2320 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 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 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 pub const fn db_mut(&mut self) -> &mut crate::db::Db {
2391 &mut self.db
2392 }
2393
2394 #[doc(hidden)]
2397 #[must_use]
2398 pub const fn __autumn_repository_seed(&self) -> &ShardRepositorySeed {
2399 &self.repo_seed
2400 }
2401
2402 #[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#[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 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
2494pub struct ShardedReadDb {
2520 db: crate::db::Db,
2521 shard_name: Arc<str>,
2522 shard_id: ShardId,
2523}
2524
2525impl ShardedReadDb {
2526 #[must_use]
2528 pub fn shard(&self) -> &str {
2529 &self.shard_name
2530 }
2531
2532 #[must_use]
2534 pub const fn shard_id(&self) -> ShardId {
2535 self.shard_id
2536 }
2537
2538 #[must_use]
2540 pub const fn span(&self) -> &tracing::Span {
2541 self.db.span()
2542 }
2543
2544 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
2599async 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 #[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 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 assert_eq!(DIRECTORY_NOTIFY_CHANNEL, "autumn_shard_directory");
2665 assert!(
2670 DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL < DEFAULT_DIRECTORY_CACHE_TTL,
2671 "sweep interval should beat the TTL"
2672 );
2673 }
2674
2675 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 #[test]
2743 fn golden_vector_str_keys() {
2744 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 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 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 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 #[tokio::test]
2823 async fn db_for_and_read_for_attempt_routed_checkouts() {
2824 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 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 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 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 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 #[test]
2967 fn owns_key_agrees_with_route() {
2968 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 let set = shard_set(&["a", "b"]);
3006 let keys = ["hooli", "a", "tenant-1"]; 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 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 #[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 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 #[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 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 #[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 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 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 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 assert_eq!(
3334 slot_for_key(ShardKey::from(42i32)),
3335 slot_for_key(ShardKey::from(42i64)),
3336 );
3337 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 #[cfg(not(feature = "sqlite"))]
3375 fn shard_with_unreachable_replica(fallback: ReplicaFallback) -> Shard {
3376 let mut config = sharded_config(&["a"]);
3377 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 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 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 #[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 ®istry,
3492 crate::health_cache::DEFAULT_PING_TIMEOUT,
3493 );
3494 register_shard_health_indicators(
3496 &set,
3497 ®istry,
3498 crate::health_cache::DEFAULT_PING_TIMEOUT,
3499 );
3500
3501 let results = registry.run_all().await;
3502 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 assert_eq!(set.total_max_connections(), 17);
3533 }
3534
3535 #[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 assert_eq!(set.total_max_connections(), 14);
3548 }
3549
3550 #[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 assert_eq!(
3585 reshard_route_label(Some("GET /admin shard=shard0"), "shard2").as_deref(),
3586 Some("GET /admin shard=shard2"),
3587 );
3588 assert_eq!(reshard_route_label(None, "shard2"), None);
3590 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); *w = 9; assert_eq!(w.0, 9);
3603 }
3604
3605 #[test]
3606 fn cross_shard_seed_is_tenant_free_and_untagged() {
3607 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 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 #[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 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 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 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}