use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::db::RuntimeConnection;
use diesel_async::AsyncPgConnection;
use diesel_async::pooled_connection::deadpool::Pool;
pub use crate::config::SLOT_COUNT;
use crate::config::{ConfigError, DatabaseConfig, ReplicaFallback};
use crate::db::{DatabaseTopology, PoolError};
use crate::error::AutumnError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ShardId(pub usize);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SlotId(pub u16);
#[derive(Debug, Clone, Copy)]
pub enum ShardKey<'a> {
Int(i64),
Str(&'a str),
Bytes(&'a [u8]),
}
impl From<i64> for ShardKey<'_> {
fn from(key: i64) -> Self {
Self::Int(key)
}
}
impl From<i32> for ShardKey<'_> {
fn from(key: i32) -> Self {
Self::Int(i64::from(key))
}
}
impl<'a> From<&'a str> for ShardKey<'a> {
fn from(key: &'a str) -> Self {
Self::Str(key)
}
}
impl<'a> From<&'a String> for ShardKey<'a> {
fn from(key: &'a String) -> Self {
Self::Str(key)
}
}
impl<'a> From<&'a [u8]> for ShardKey<'a> {
fn from(key: &'a [u8]) -> Self {
Self::Bytes(key)
}
}
impl<'a> From<&'a [u8; 16]> for ShardKey<'a> {
fn from(key: &'a [u8; 16]) -> Self {
Self::Bytes(key)
}
}
const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
fn fnv1a_64(bytes: &[u8]) -> u64 {
let mut hash = FNV_OFFSET_BASIS;
for byte in bytes {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
const fn splitmix64(mut x: u64) -> u64 {
x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
x ^ (x >> 31)
}
#[must_use]
fn key_hash64(key: ShardKey<'_>) -> u64 {
match key {
#[allow(clippy::cast_sign_loss)]
ShardKey::Int(value) => splitmix64(value as u64),
ShardKey::Str(value) => fnv1a_64(value.as_bytes()),
ShardKey::Bytes(value) => fnv1a_64(value),
}
}
#[must_use]
pub fn slot_for_key(key: ShardKey<'_>) -> SlotId {
let hash = key_hash64(key);
#[allow(clippy::cast_possible_truncation)]
SlotId((hash % u64::from(SLOT_COUNT)) as u16)
}
pub trait ShardRouter: Send + Sync + 'static {
fn route<'a>(
&'a self,
key: ShardKey<'a>,
shards: &'a ShardSet,
) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>>;
}
impl<R: ShardRouter + ?Sized> ShardRouter for Arc<R> {
fn route<'a>(
&'a self,
key: ShardKey<'a>,
shards: &'a ShardSet,
) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
(**self).route(key, shards)
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct HashShardRouter;
impl ShardRouter for HashShardRouter {
fn route<'a>(
&'a self,
key: ShardKey<'a>,
shards: &'a ShardSet,
) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
let slot = shards.slot_for_key(key);
Box::pin(std::future::ready(
shards
.inner
.slot_map
.get(usize::from(slot.0))
.map(|&idx| ShardId(idx))
.ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"slot {} has no shard assigned (slot map inconsistent)",
slot.0
))
}),
))
}
}
pub struct DirectoryShardRouter {
control_pool: Pool<RuntimeConnection>,
fallback: Arc<dyn ShardRouter>,
cache: std::sync::RwLock<HashMap<String, DirectoryCacheEntry>>,
ttl: std::time::Duration,
statement_timeout_ms: u64,
}
pub const DEFAULT_DIRECTORY_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(30);
#[cfg(feature = "db")]
pub const SHARD_DIRECTORY_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
diesel_migrations::embed_migrations!("shard_directory_migrations");
#[cfg(feature = "db")]
pub const SHARD_MAP_MIGRATIONS: diesel_migrations::EmbeddedMigrations =
diesel_migrations::embed_migrations!("shard_map_migrations");
#[derive(Clone, Copy)]
struct DirectoryCacheEntry {
shard: ShardId,
expires_at: std::time::Instant,
}
#[derive(diesel::QueryableByName)]
struct ShardNameRow {
#[diesel(sql_type = diesel::sql_types::Text)]
shard_name: String,
}
const DIRECTORY_NOTIFY_CHANNEL: &str = "autumn_shard_directory";
pub const DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL: std::time::Duration =
std::time::Duration::from_secs(5);
impl std::fmt::Debug for DirectoryShardRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DirectoryShardRouter")
.field("ttl", &self.ttl)
.field("cached_keys", &self.cache.read().map_or(0, |c| c.len()))
.finish_non_exhaustive()
}
}
impl DirectoryShardRouter {
#[must_use]
pub fn new(control_pool: Pool<RuntimeConnection>) -> Self {
Self::with_fallback(control_pool, Arc::new(HashShardRouter))
}
#[must_use]
pub fn with_fallback(
control_pool: Pool<RuntimeConnection>,
fallback: Arc<dyn ShardRouter>,
) -> Self {
Self {
control_pool,
fallback,
cache: std::sync::RwLock::new(HashMap::new()),
ttl: DEFAULT_DIRECTORY_CACHE_TTL,
statement_timeout_ms: 0,
}
}
#[must_use]
pub const fn with_statement_timeout_ms(mut self, statement_timeout_ms: u64) -> Self {
self.statement_timeout_ms = statement_timeout_ms;
self
}
#[must_use]
pub const fn with_cache_ttl(mut self, ttl: std::time::Duration) -> Self {
self.ttl = ttl;
self
}
pub fn invalidate(&self, tenant_key: &str) {
if let Ok(mut cache) = self.cache.write() {
cache.remove(tenant_key);
}
}
pub fn invalidate_all(&self) {
if let Ok(mut cache) = self.cache.write() {
cache.clear();
}
}
#[must_use]
pub fn spawn_invalidation_listener(
router: Arc<Self>,
control_url: String,
sweep_interval: std::time::Duration,
) -> tokio::task::JoinHandle<()> {
use diesel_async::{AsyncConnection as _, RunQueryDsl as _};
use futures::StreamExt as _;
tokio::spawn(async move {
loop {
let Ok(mut conn) = AsyncPgConnection::establish(&control_url).await else {
tokio::time::sleep(sweep_interval).await;
continue;
};
if diesel::sql_query(format!("LISTEN {DIRECTORY_NOTIFY_CHANNEL}"))
.execute(&mut conn)
.await
.is_err()
{
tokio::time::sleep(sweep_interval).await;
continue;
}
router.invalidate_all();
let mut notifications = std::pin::pin!(conn.notifications_stream());
loop {
match tokio::time::timeout(sweep_interval, notifications.next()).await {
Ok(Some(Ok(notification))) => router.invalidate(¬ification.payload),
Ok(Some(Err(_)) | None) => break,
Err(_elapsed) => router.sweep_expired(),
}
}
}
})
}
fn cache_get(&self, key: &str) -> Option<ShardId> {
let now = std::time::Instant::now();
{
let cache = self.cache.read().ok()?;
match cache.get(key) {
Some(entry) if entry.expires_at > now => return Some(entry.shard),
Some(_) => {}
None => return None,
}
}
let mut cache = self.cache.write().ok()?;
if cache.get(key).is_some_and(|entry| entry.expires_at <= now) {
cache.remove(key);
}
None
}
fn sweep_expired(&self) {
if let Ok(mut cache) = self.cache.write() {
let now = std::time::Instant::now();
cache.retain(|_, entry| entry.expires_at > now);
}
}
fn cache_put(&self, key: String, shard: ShardId) {
if let Ok(mut cache) = self.cache.write() {
cache.insert(
key,
DirectoryCacheEntry {
shard,
expires_at: std::time::Instant::now() + self.ttl,
},
);
}
}
async fn lookup_directory(
&self,
key: &str,
shards: &ShardSet,
) -> Result<Option<ShardId>, AutumnError> {
use diesel::OptionalExtension as _;
use diesel_async::RunQueryDsl;
let mut conn = self.control_pool.get().await.map_err(|e| {
AutumnError::service_unavailable_msg(format!(
"DirectoryShardRouter could not acquire a control connection: {e}"
))
})?;
diesel::sql_query(format!(
"SET statement_timeout = {}",
self.statement_timeout_ms
))
.execute(&mut conn)
.await
.map_err(|e| {
AutumnError::service_unavailable_msg(format!(
"DirectoryShardRouter could not set statement_timeout: {e}"
))
})?;
let row = diesel::sql_query(
"SELECT shard_name FROM _autumn_shard_directory WHERE tenant_key = $1",
)
.bind::<diesel::sql_types::Text, _>(key)
.get_result::<ShardNameRow>(&mut conn)
.await
.optional()
.map_err(|e| {
AutumnError::service_unavailable_msg(format!(
"DirectoryShardRouter directory lookup failed: {e}"
))
})?;
let Some(row) = row else {
return Ok(None);
};
let shard = shards.by_name(&row.shard_name).ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"shard directory pins tenant {key:?} to unknown shard {:?}",
row.shard_name
))
})?;
Ok(Some(shard.id()))
}
}
impl ShardRouter for DirectoryShardRouter {
fn route<'a>(
&'a self,
key: ShardKey<'a>,
shards: &'a ShardSet,
) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
Box::pin(async move {
let ShardKey::Str(key_str) = key else {
return self.fallback.route(key, shards).await;
};
if let Some(cached) = self.cache_get(key_str) {
return Ok(cached);
}
match self.lookup_directory(key_str, shards).await? {
Some(shard) => {
self.cache_put(key_str.to_owned(), shard);
Ok(shard)
}
None => self.fallback.route(key, shards).await,
}
})
}
}
#[derive(Debug)]
pub(crate) struct ShardRuntime {
replica_fallback: ReplicaFallback,
replica_configured: bool,
connection_ready: AtomicBool,
migrations_ready: AtomicBool,
detail: std::sync::RwLock<Option<String>>,
migration_check: std::sync::RwLock<Option<(String, String)>>,
parity_checked_at: std::sync::Mutex<Option<std::time::Instant>>,
}
const PARITY_RECHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
#[cfg_attr(not(test), allow(dead_code))]
impl ShardRuntime {
fn new(replica_fallback: ReplicaFallback, replica_configured: bool) -> Self {
Self {
replica_fallback,
replica_configured,
connection_ready: AtomicBool::new(false),
migrations_ready: AtomicBool::new(true),
detail: std::sync::RwLock::new(
replica_configured.then(|| "replica has not passed a readiness check".to_owned()),
),
migration_check: std::sync::RwLock::new(None),
parity_checked_at: std::sync::Mutex::new(None),
}
}
pub(crate) fn configure_migration_check(&self, primary_url: String, replica_url: String) {
*self
.migration_check
.write()
.expect("shard runtime lock poisoned") = Some((primary_url, replica_url));
}
fn migration_check(&self) -> Option<(String, String)> {
self.migration_check
.read()
.expect("shard runtime lock poisoned")
.clone()
}
pub(crate) fn parity_check_due(&self) -> bool {
let mut checked_at = self
.parity_checked_at
.lock()
.expect("shard runtime lock poisoned");
if checked_at.is_none_or(|at| at.elapsed() >= PARITY_RECHECK_INTERVAL) {
*checked_at = Some(std::time::Instant::now());
true
} else {
false
}
}
fn replica_ready(&self) -> bool {
self.connection_ready.load(Ordering::Relaxed)
&& self.migrations_ready.load(Ordering::Relaxed)
}
fn refresh_detail(&self) {
if self.replica_ready() {
*self.detail.write().expect("shard runtime lock poisoned") = None;
}
}
pub(crate) fn mark_replica_connection_ready(&self) {
self.connection_ready.store(true, Ordering::Relaxed);
self.refresh_detail();
}
pub(crate) fn mark_replica_connection_unready(&self, detail: impl Into<String>) {
self.connection_ready.store(false, Ordering::Relaxed);
*self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
}
pub(crate) fn mark_replica_migrations_ready(&self) {
self.migrations_ready.store(true, Ordering::Relaxed);
self.refresh_detail();
}
pub(crate) fn mark_replica_migrations_unready(&self, detail: impl Into<String>) {
self.migrations_ready.store(false, Ordering::Relaxed);
*self.detail.write().expect("shard runtime lock poisoned") = Some(detail.into());
}
pub(crate) fn detail(&self) -> Option<String> {
self.detail
.read()
.expect("shard runtime lock poisoned")
.clone()
}
}
#[derive(Clone)]
pub struct Shard {
name: Arc<str>,
id: ShardId,
slots: Arc<[u16]>,
topology: DatabaseTopology,
runtime: Arc<ShardRuntime>,
}
impl Shard {
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub const fn id(&self) -> ShardId {
self.id
}
#[must_use]
pub fn slots(&self) -> &[u16] {
&self.slots
}
#[must_use]
pub const fn topology(&self) -> &DatabaseTopology {
&self.topology
}
#[must_use]
pub const fn primary_pool(&self) -> &Pool<RuntimeConnection> {
self.topology.primary()
}
#[must_use]
pub const fn replica_pool(&self) -> Option<&Pool<RuntimeConnection>> {
self.topology.replica()
}
#[must_use]
pub fn read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
self.read_pool_with_role().map(|(pool, _)| pool)
}
#[must_use]
pub fn read_route(&self) -> crate::repository::ReadRoute {
use crate::repository::ReadRoute;
if !self.runtime.replica_configured {
return ReadRoute::Primary;
}
self.read_pool().map_or(ReadRoute::Unavailable, |pool| {
ReadRoute::ReadPool(pool.clone())
})
}
pub(crate) fn replica_read_pool(&self) -> Option<&Pool<RuntimeConnection>> {
if self.runtime.replica_configured && self.runtime.replica_ready() {
self.topology.replica()
} else {
None
}
}
pub(crate) fn read_pool_with_role(&self) -> Option<(&Pool<RuntimeConnection>, &'static str)> {
if !self.runtime.replica_configured {
return Some((self.topology.primary(), "primary"));
}
if self.runtime.replica_ready() {
return self.topology.replica().map(|pool| (pool, "replica"));
}
match self.runtime.replica_fallback {
ReplicaFallback::Primary => Some((self.topology.primary(), "primary")),
ReplicaFallback::FailReadiness => None,
}
}
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) fn runtime(&self) -> &ShardRuntime {
&self.runtime
}
}
impl std::fmt::Debug for Shard {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Shard")
.field("name", &self.name)
.field("id", &self.id)
.field("slots", &self.slots)
.finish_non_exhaustive()
}
}
struct ShardSetInner {
shards: Vec<Shard>,
by_name: HashMap<String, usize>,
slot_map: Vec<usize>,
router: Arc<dyn ShardRouter>,
}
#[derive(Clone)]
pub struct ShardSet {
inner: Arc<ShardSetInner>,
}
impl ShardSet {
#[must_use]
pub fn len(&self) -> usize {
self.inner.shards.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.inner.shards.is_empty()
}
#[must_use]
pub const fn slot_count(&self) -> u16 {
SLOT_COUNT
}
#[must_use]
pub fn get(&self, id: ShardId) -> Option<&Shard> {
self.inner.shards.get(id.0)
}
#[must_use]
pub fn by_name(&self, name: &str) -> Option<&Shard> {
self.inner
.by_name
.get(name)
.and_then(|&idx| self.inner.shards.get(idx))
}
pub fn iter(&self) -> impl Iterator<Item = &Shard> {
self.inner.shards.iter()
}
#[must_use]
pub fn slot_for_key<'k>(&self, key: impl Into<ShardKey<'k>>) -> SlotId {
slot_for_key(key.into())
}
#[must_use]
pub fn shard_for_slot(&self, slot: SlotId) -> Option<&Shard> {
self.inner
.slot_map
.get(usize::from(slot.0))
.and_then(|&idx| self.inner.shards.get(idx))
}
pub async fn route<'k>(&self, key: impl Into<ShardKey<'k>>) -> Result<&Shard, AutumnError> {
let key = key.into();
let id = self.inner.router.route(key, self).await?;
self.get(id).ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"shard router returned out-of-range shard id {} (have {} shards)",
id.0,
self.len()
))
})
}
#[must_use]
pub fn total_max_connections(&self) -> usize {
self.inner
.shards
.iter()
.map(|shard| {
shard.topology().primary().status().max_size
+ shard
.topology()
.replica()
.map_or(0, |pool| pool.status().max_size)
})
.sum()
}
#[must_use]
pub fn owns_key<'k>(&self, shard_id: ShardId, key: impl Into<ShardKey<'k>>) -> bool {
let slot = self.slot_for_key(key);
self.shard_for_slot(slot)
.is_some_and(|s| s.id() == shard_id)
}
#[must_use]
pub fn slots_for_shard(&self, shard_id: ShardId) -> Option<&[u16]> {
self.inner.shards.get(shard_id.0).map(Shard::slots)
}
#[must_use]
pub fn partition_by_shard<'k>(
&self,
keys: impl IntoIterator<Item = &'k str>,
) -> std::collections::HashMap<ShardId, Vec<&'k str>> {
let mut map: std::collections::HashMap<ShardId, Vec<&'k str>> =
std::collections::HashMap::new();
for key in keys {
let slot = self.slot_for_key(key);
if let Some(shard) = self.shard_for_slot(slot) {
map.entry(shard.id()).or_default().push(key);
}
}
map
}
#[doc(hidden)]
pub async fn fan_out_shards<T, Fut, F>(&self, f: F) -> Result<Vec<T>, crate::AutumnError>
where
T: Send + 'static,
Fut: std::future::Future<Output = Result<T, crate::AutumnError>> + Send + 'static,
F: Fn(&Shard) -> Fut + Send + Sync,
{
use futures::StreamExt as _;
let mut slots: Vec<Option<T>> = (0..self.inner.shards.len()).map(|_| None).collect();
let mut in_flight = futures::stream::FuturesUnordered::new();
for (idx, shard) in self.inner.shards.iter().enumerate() {
if in_flight.len() >= FAN_OUT_CONCURRENCY
&& let Some((i, result)) = in_flight.next().await
{
slots[i] = Some(result?);
}
let fut = f(shard);
in_flight.push(async move { (idx, fut.await) });
}
while let Some((i, result)) = in_flight.next().await {
slots[i] = Some(result?);
}
Ok(slots
.into_iter()
.map(|slot| slot.expect("every shard produced a result"))
.collect())
}
}
impl std::fmt::Debug for ShardSet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ShardSet")
.field("shards", &self.inner.shards)
.finish_non_exhaustive()
}
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ShardSetBuildError {
#[error("failed to build pool for shard {shard:?}: {source}")]
Pool {
shard: String,
source: PoolError,
},
#[error(transparent)]
Config(#[from] ConfigError),
#[error("expected {expected} shard topologies, got {actual}")]
TopologyCountMismatch {
expected: usize,
actual: usize,
},
}
pub fn create_shard_set(
config: &DatabaseConfig,
router: Arc<dyn ShardRouter>,
) -> Result<Option<ShardSet>, ShardSetBuildError> {
if !config.has_shards() {
return Ok(None);
}
let topologies = config
.shards
.iter()
.map(|shard| {
crate::db::create_shard_topology(shard, config).map_err(|source| {
ShardSetBuildError::Pool {
shard: shard.name.clone(),
source,
}
})
})
.collect::<Result<Vec<_>, _>>()?;
build_shard_set(config, topologies, router).map(Some)
}
#[cfg(not(feature = "sqlite"))]
pub fn create_shard_set_transactional(
config: &DatabaseConfig,
router: Arc<dyn ShardRouter>,
) -> Result<Option<ShardSet>, ShardSetBuildError> {
if !config.has_shards() {
return Ok(None);
}
let timeout = std::time::Duration::from_secs(config.connect_timeout_secs);
let topologies = config
.shards
.iter()
.map(|shard| {
let manager = diesel_async::pooled_connection::AsyncDieselConnectionManager::<
diesel_async::AsyncPgConnection,
>::new(&shard.primary_url);
let pool = Pool::builder(manager)
.max_size(1)
.wait_timeout(Some(timeout))
.create_timeout(Some(timeout))
.runtime(deadpool::Runtime::Tokio1)
.post_create(deadpool::managed::Hook::async_fn(
|conn: &mut diesel_async::AsyncPgConnection, _| {
Box::pin(async move {
use diesel_async::AsyncConnection as _;
use diesel_async::RunQueryDsl as _;
conn.begin_test_transaction().await.map_err(|e| {
deadpool::managed::HookError::Backend(
diesel_async::pooled_connection::PoolError::QueryError(e),
)
})?;
diesel::sql_query("SET autumn.test_transaction_started = 'true'")
.execute(conn)
.await
.map_err(|e| {
deadpool::managed::HookError::Backend(
diesel_async::pooled_connection::PoolError::QueryError(e),
)
})?;
Ok(())
})
},
))
.build()
.map_err(|source| ShardSetBuildError::Pool {
shard: shard.name.clone(),
source: crate::db::PoolError::Build(source),
})?;
Ok(crate::db::DatabaseTopology::primary_only(pool))
})
.collect::<Result<Vec<_>, ShardSetBuildError>>()?;
build_shard_set(config, topologies, router).map(Some)
}
pub fn build_shard_set(
config: &DatabaseConfig,
topologies: Vec<DatabaseTopology>,
router: Arc<dyn ShardRouter>,
) -> Result<ShardSet, ShardSetBuildError> {
if topologies.len() != config.shards.len() {
return Err(ShardSetBuildError::TopologyCountMismatch {
expected: config.shards.len(),
actual: topologies.len(),
});
}
let slot_map = config.resolved_slot_map()?;
let mut slots_per_shard: Vec<Vec<u16>> = vec![Vec::new(); config.shards.len()];
for (slot, &owner) in slot_map.iter().enumerate() {
#[allow(clippy::cast_possible_truncation)]
slots_per_shard[owner].push(slot as u16);
}
let shards: Vec<Shard> = config
.shards
.iter()
.zip(topologies)
.enumerate()
.map(|(idx, (shard_config, topology))| {
let replica_configured = topology.replica().is_some();
Shard {
name: Arc::from(shard_config.name.as_str()),
id: ShardId(idx),
slots: Arc::from(std::mem::take(&mut slots_per_shard[idx])),
topology,
runtime: Arc::new(ShardRuntime::new(
shard_config.effective_replica_fallback(config),
replica_configured,
)),
}
})
.collect();
let mut by_name = HashMap::with_capacity(shards.len());
for (idx, shard) in shards.iter().enumerate() {
if by_name.insert(shard.name().to_owned(), idx).is_some() {
return Err(ConfigError::Validation(format!(
"database.shards: shard name {:?} is declared more than once; \
shard names must be unique",
shard.name()
))
.into());
}
}
Ok(ShardSet {
inner: Arc::new(ShardSetInner {
shards,
by_name,
slot_map,
router,
}),
})
}
pub(crate) struct ShardHealthIndicator {
shard: Shard,
}
impl ShardHealthIndicator {
pub(crate) const fn new(shard: Shard) -> Self {
Self { shard }
}
async fn refresh_replica_readiness(&self) {
let Some(replica_pool) = self.shard.replica_pool() else {
return;
};
match replica_pool.get().await {
Ok(conn) => {
drop(conn);
self.shard.runtime().mark_replica_connection_ready();
if self.shard.runtime().parity_check_due()
&& let Some((primary_url, replica_url)) = self.shard.runtime().migration_check()
{
let readiness = crate::migrate::check_replica_migration_readiness_blocking(
primary_url,
replica_url,
)
.await;
if readiness.is_ready() {
self.shard.runtime().mark_replica_migrations_ready();
} else if let Some(detail) = readiness.detail() {
self.shard.runtime().mark_replica_migrations_unready(detail);
}
}
}
Err(error) => self
.shard
.runtime()
.mark_replica_connection_unready(format!("replica connection failed: {error}")),
}
}
}
impl crate::actuator::HealthIndicator for ShardHealthIndicator {
fn check(&self) -> futures::future::BoxFuture<'_, crate::actuator::HealthCheckOutput> {
Box::pin(async move {
self.refresh_replica_readiness().await;
let mut details = HashMap::new();
let status = self.shard.primary_pool().status();
details.insert("pool_size".to_owned(), serde_json::json!(status.max_size));
details.insert(
"active_connections".to_owned(),
serde_json::json!((status.max_size as u64).saturating_sub(status.available as u64)),
);
details.insert(
"idle_connections".to_owned(),
serde_json::json!(status.available),
);
details.insert(
"slots".to_owned(),
serde_json::json!(self.shard.slots().len()),
);
if self.shard.replica_pool().is_some() {
details.insert(
"replica_ready".to_owned(),
serde_json::json!(self.shard.runtime().replica_ready()),
);
if let Some(detail) = self.shard.runtime().detail() {
details.insert("replica_detail".to_owned(), serde_json::json!(detail));
}
}
let primary_ok = match self.shard.primary_pool().get().await {
Ok(conn) => {
drop(conn);
true
}
Err(error) => {
details.insert(
"primary_detail".to_owned(),
serde_json::json!(format!("primary connection failed: {error}")),
);
false
}
};
details.insert("primary_ready".to_owned(), serde_json::json!(primary_ok));
let output = if primary_ok && self.shard.read_pool().is_some() {
crate::actuator::HealthCheckOutput::up()
} else {
crate::actuator::HealthCheckOutput::down()
};
output.with_details(details)
})
}
}
pub(crate) fn register_shard_health_indicators(
set: &ShardSet,
registry: &crate::actuator::HealthIndicatorRegistry,
) {
for shard in set.iter() {
let name = format!("db:shard:{}", shard.name());
if let Err(error) = registry.register(
name,
crate::actuator::IndicatorGroup::Readiness,
Arc::new(ShardHealthIndicator::new(shard.clone())),
) {
tracing::warn!("{error}");
}
}
}
#[derive(Debug, Clone)]
pub struct ShardKeyOverride(pub String);
pub struct Shards {
set: ShardSet,
ctx: crate::db::RequestDbContext,
}
impl<S> axum::extract::FromRequestParts<S> for Shards
where
S: crate::db::DbState + Send + Sync,
{
type Rejection = AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let set = state.shards().cloned().ok_or_else(no_shards_configured)?;
let ctx = crate::db::RequestDbContext::from_parts(parts, state);
Ok(Self { set, ctx })
}
}
fn no_shards_configured() -> AutumnError {
AutumnError::service_unavailable_msg(
"No shards configured: declare [[database.shards]] in autumn.toml \
(see docs/guide/sharding.md)",
)
}
fn cross_shard_seed(
set: &ShardSet,
ctx: &crate::db::RequestDbContext,
) -> Result<ShardRepositorySeed, AutumnError> {
let shard = set.iter().next().ok_or_else(no_shards_configured)?;
let mut seed =
ShardRepositorySeed::from_ctx(shard.primary_pool(), ctx, shard.name(), shard.read_route());
seed.route.clone_from(&ctx.route_key);
Ok(seed)
}
pub trait CrossShardRepository: Sized {
#[doc(hidden)]
fn __autumn_from_cross_shard(seed: ShardRepositorySeed, set: ShardSet) -> Self;
}
pub struct CrossShard<R>(pub R);
impl<R> std::ops::Deref for CrossShard<R> {
type Target = R;
fn deref(&self) -> &R {
&self.0
}
}
impl<R> std::ops::DerefMut for CrossShard<R> {
fn deref_mut(&mut self) -> &mut R {
&mut self.0
}
}
impl<S, R> axum::extract::FromRequestParts<S> for CrossShard<R>
where
S: crate::db::DbState + Send + Sync,
R: CrossShardRepository,
{
type Rejection = AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let shards =
<Shards as axum::extract::FromRequestParts<S>>::from_request_parts(parts, state)
.await?;
let seed = cross_shard_seed(&shards.set, &shards.ctx)?;
Ok(Self(R::__autumn_from_cross_shard(seed, shards.set)))
}
}
const FAN_OUT_CONCURRENCY: usize = 8;
impl Shards {
#[must_use]
pub const fn set(&self) -> &ShardSet {
&self.set
}
pub fn iter(&self) -> impl Iterator<Item = &Shard> {
self.set.iter()
}
pub async fn db_for<'k>(
&self,
key: impl Into<ShardKey<'k>>,
) -> Result<crate::db::Db, AutumnError> {
let shard = self.set.route(key).await?;
self.checkout_primary(shard).await
}
pub async fn read_for<'k>(
&self,
key: impl Into<ShardKey<'k>>,
) -> Result<crate::db::Db, AutumnError> {
let shard = self.set.route(key).await?;
let (pool, role) = shard.read_pool_with_role().ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"shard {:?} replica is not ready and replica_fallback = \"fail_readiness\"",
shard.name()
))
})?;
self.checkout(shard, pool, role).await
}
pub async fn read_replica_for<'k>(
&self,
key: impl Into<ShardKey<'k>>,
) -> Result<crate::db::Db, AutumnError> {
let shard = self.set.route(key).await?;
let pool = shard.replica_read_pool().ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"shard {:?} has no healthy replica; read_replica_for requires a \
configured, ready replica (no primary fallback)",
shard.name()
))
})?;
self.checkout(shard, pool, "replica").await
}
pub async fn db_on(&self, shard_name: &str) -> Result<crate::db::Db, AutumnError> {
let shard = self
.set
.by_name(shard_name)
.ok_or_else(|| AutumnError::bad_request_msg(format!("unknown shard {shard_name:?}")))?;
self.checkout_primary(shard).await
}
pub async fn each_shard<T, Fut, F>(&self, f: F) -> Vec<(ShardId, Result<T, AutumnError>)>
where
T: Send,
Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
{
use futures::StreamExt as _;
let mut results: Vec<Option<(ShardId, Result<T, AutumnError>)>> =
std::iter::repeat_with(|| None)
.take(self.set.len())
.collect();
let mut in_flight: futures::stream::FuturesUnordered<
futures::future::BoxFuture<'_, (ShardId, Result<T, AutumnError>)>,
> = futures::stream::FuturesUnordered::new();
for shard in self.set.iter() {
if in_flight.len() >= FAN_OUT_CONCURRENCY
&& let Some((id, result)) = in_flight.next().await
{
results[id.0] = Some((id, result));
}
in_flight.push(Box::pin(self.run_on_shard(shard, &f)));
}
while let Some((id, result)) = in_flight.next().await {
results[id.0] = Some((id, result));
}
results.into_iter().flatten().collect()
}
async fn run_on_shard<T, Fut, F>(
&self,
shard: &Shard,
f: &F,
) -> (ShardId, Result<T, AutumnError>)
where
T: Send,
Fut: std::future::Future<Output = Result<T, AutumnError>> + Send,
F: Fn(&Shard, crate::db::Db) -> Fut + Send + Sync,
{
let result = match self.checkout_primary(shard).await {
Ok(db) => f(shard, db).await,
Err(error) => Err(error),
};
(shard.id(), result)
}
async fn checkout_primary(&self, shard: &Shard) -> Result<crate::db::Db, AutumnError> {
self.checkout(shard, shard.primary_pool(), "primary").await
}
async fn checkout(
&self,
shard: &Shard,
pool: &Pool<RuntimeConnection>,
role: &str,
) -> Result<crate::db::Db, AutumnError> {
let ctx = self.ctx.clone();
crate::db::Db::checkout(crate::db::DbCheckoutParams {
pool,
pool_name: &format!("shard:{}:{role}", shard.name()),
shard: Some(shard.name()),
statement_timeout: ctx.statement_timeout,
route_key: ctx
.route_key
.map(|key| format!("{key} shard={}", shard.name())),
metrics: ctx.metrics,
slow_query_threshold: ctx.slow_query_threshold,
interceptors: ctx.interceptors,
})
.await
}
}
#[doc(hidden)]
#[derive(Clone)]
pub struct ShardRepositorySeed {
pub pool: Pool<RuntimeConnection>,
pub statement_timeout_ms: u64,
pub slow_query_threshold: std::time::Duration,
pub route: Option<String>,
pub read_route: crate::repository::ReadRoute,
}
impl ShardRepositorySeed {
pub(crate) fn from_ctx(
pool: &Pool<RuntimeConnection>,
ctx: &crate::db::RequestDbContext,
shard_name: &str,
read_route: crate::repository::ReadRoute,
) -> Self {
const PG_TIMEOUT_MAX_MS: u64 = i32::MAX as u64;
let statement_timeout_ms = ctx.statement_timeout.map_or(0, |d| {
u64::try_from(d.as_millis().min(u128::from(PG_TIMEOUT_MAX_MS)))
.unwrap_or(PG_TIMEOUT_MAX_MS)
});
Self {
pool: pool.clone(),
statement_timeout_ms,
slow_query_threshold: ctx.slow_query_threshold,
route: ctx
.route_key
.as_ref()
.map(|key| format!("{key} shard={shard_name}")),
read_route,
}
}
}
#[must_use]
pub fn reshard_route_label(parent: Option<&str>, shard_name: &str) -> Option<String> {
let parent = parent?;
let base = parent.rsplit_once(" shard=").map_or(parent, |(key, _)| key);
Some(format!("{base} shard={shard_name}"))
}
pub struct ShardedDb {
db: crate::db::Db,
shard_name: Arc<str>,
shard_id: ShardId,
repo_seed: ShardRepositorySeed,
shards: ShardSet,
}
impl ShardedDb {
#[must_use]
pub fn shard(&self) -> &str {
&self.shard_name
}
#[must_use]
pub const fn shard_id(&self) -> ShardId {
self.shard_id
}
#[must_use]
pub const fn span(&self) -> &tracing::Span {
self.db.span()
}
pub async fn tx<'a, T, E, F>(&'a mut self, f: F) -> Result<T, AutumnError>
where
T: Send + 'a,
E: From<diesel::result::Error> + Send + Sync + 'a,
AutumnError: From<E>,
F: for<'r> FnOnce(
&'r mut crate::db::PooledConnection,
) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
self.db.tx(f).await
}
pub async fn tx_with<'a, T, E, F>(
&'a mut self,
opts: crate::db::TxOptions,
f: F,
) -> Result<T, AutumnError>
where
T: Send + 'a,
E: From<diesel::result::Error> + Send + Sync + 'a,
AutumnError: From<E>,
F: for<'r> FnMut(
&'r mut crate::db::RuntimeConnection,
) -> scoped_futures::ScopedBoxFuture<'a, 'r, Result<T, E>>
+ Send
+ 'a,
{
self.db.tx_with(opts, f).await
}
pub const fn db_mut(&mut self) -> &mut crate::db::Db {
&mut self.db
}
#[doc(hidden)]
#[must_use]
pub const fn __autumn_repository_seed(&self) -> &ShardRepositorySeed {
&self.repo_seed
}
#[doc(hidden)]
#[must_use]
pub const fn __autumn_shard_set(&self) -> &ShardSet {
&self.shards
}
}
impl std::ops::Deref for ShardedDb {
type Target = RuntimeConnection;
fn deref(&self) -> &Self::Target {
&self.db
}
}
impl std::ops::DerefMut for ShardedDb {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.db
}
}
impl AsMut<crate::db::Db> for ShardedDb {
fn as_mut(&mut self) -> &mut crate::db::Db {
&mut self.db
}
}
#[doc(hidden)]
pub async fn __autumn_resolve_repo_seed(
parts: &mut axum::http::request::Parts,
state: &crate::AppState,
) -> Result<(ShardRepositorySeed, ShardSet), AutumnError> {
let shards = <Shards as axum::extract::FromRequestParts<crate::AppState>>::from_request_parts(
parts, state,
)
.await?;
let key = resolve_shard_key(parts, state).await?;
let shard = shards.set.route(&key).await?;
let shard_name = Arc::clone(&shard.name);
let seed = ShardRepositorySeed::from_ctx(
shard.primary_pool(),
&shards.ctx,
&shard_name,
shard.read_route(),
);
let set = shards.set.clone();
Ok((seed, set))
}
impl axum::extract::FromRequestParts<crate::AppState> for ShardedDb {
type Rejection = AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &crate::AppState,
) -> Result<Self, Self::Rejection> {
let shards = Shards::from_request_parts(parts, state).await?;
let key = resolve_shard_key(parts, state).await?;
let shard = shards.set.route(&key).await?;
let shard_name = Arc::clone(&shard.name);
let shard_id = shard.id();
let repo_seed = ShardRepositorySeed::from_ctx(
shard.primary_pool(),
&shards.ctx,
&shard_name,
shard.read_route(),
);
let shard_set = shards.set.clone();
let db = shards.checkout_primary(shard).await?;
crate::read_your_writes::mark_write();
Ok(Self {
db,
shard_name,
shard_id,
repo_seed,
shards: shard_set,
})
}
}
pub struct ShardedReadDb {
db: crate::db::Db,
shard_name: Arc<str>,
shard_id: ShardId,
}
impl ShardedReadDb {
#[must_use]
pub fn shard(&self) -> &str {
&self.shard_name
}
#[must_use]
pub const fn shard_id(&self) -> ShardId {
self.shard_id
}
#[must_use]
pub const fn span(&self) -> &tracing::Span {
self.db.span()
}
pub const fn db_mut(&mut self) -> &mut crate::db::Db {
&mut self.db
}
}
impl std::ops::Deref for ShardedReadDb {
type Target = RuntimeConnection;
fn deref(&self) -> &Self::Target {
&self.db
}
}
impl std::ops::DerefMut for ShardedReadDb {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.db
}
}
impl AsMut<crate::db::Db> for ShardedReadDb {
fn as_mut(&mut self) -> &mut crate::db::Db {
&mut self.db
}
}
impl axum::extract::FromRequestParts<crate::AppState> for ShardedReadDb {
type Rejection = AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &crate::AppState,
) -> Result<Self, Self::Rejection> {
let shards = Shards::from_request_parts(parts, state).await?;
let key = resolve_shard_key(parts, state).await?;
let shard = shards.set.route(&key).await?;
let shard_name = Arc::clone(&shard.name);
let shard_id = shard.id();
let pool = shard.replica_read_pool().ok_or_else(|| {
AutumnError::service_unavailable_msg(format!(
"shard {:?} has no healthy replica; ShardedReadDb requires a \
configured, ready replica (no primary fallback)",
shard.name()
))
})?;
let db = shards.checkout(shard, pool, "replica").await?;
Ok(Self {
db,
shard_name,
shard_id,
})
}
}
async fn resolve_shard_key(
parts: &mut axum::http::request::Parts,
state: &crate::AppState,
) -> Result<String, AutumnError> {
if let Some(overridden) = parts.extensions.get::<ShardKeyOverride>() {
return Ok(overridden.0.clone());
}
if let Ok(Some(tenant)) = crate::tenancy::CURRENT_TENANT.try_with(std::clone::Clone::clone) {
return Ok(tenant);
}
let config = state
.extension::<crate::config::AutumnConfig>()
.ok_or_else(|| AutumnError::service_unavailable_msg("Config is not available"))?;
crate::tenancy::extract_tenant_from_parts(parts, &config)
.await
.map_err(|error| {
AutumnError::bad_request_msg(format!(
"ShardedDb could not resolve a shard key: {error}. Enable [tenancy] so \
the tenant id can route the request, or insert a ShardKeyOverride \
request extension from middleware (see docs/guide/sharding.md)"
))
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{ShardConfig, SlotSpec};
#[test]
fn directory_invalidation_channel_and_interval_are_sane() {
assert_eq!(DIRECTORY_NOTIFY_CHANNEL, "autumn_shard_directory");
assert!(
DEFAULT_DIRECTORY_INVALIDATION_SWEEP_INTERVAL < DEFAULT_DIRECTORY_CACHE_TTL,
"sweep interval should beat the TTL"
);
}
fn shard_config(name: &str) -> ShardConfig {
ShardConfig {
name: name.to_owned(),
primary_url: format!("postgres://localhost/{name}"),
slots: None,
replica_url: None,
primary_pool_size: None,
replica_pool_size: None,
replica_fallback: None,
}
}
fn sharded_config(names: &[&str]) -> DatabaseConfig {
DatabaseConfig {
shards: names.iter().map(|name| shard_config(name)).collect(),
..Default::default()
}
}
fn shard_set(names: &[&str]) -> ShardSet {
create_shard_set(&sharded_config(names), Arc::new(HashShardRouter))
.expect("lazy pools should build")
.expect("shards configured")
}
#[test]
fn golden_vector_str_keys() {
let cases: &[(&str, u16)] = &[
("tenant-1", 12427),
("tenant-2", 12862),
("tenant-3", 13297),
("acme-corp", 11394),
("globex", 12846),
("initech", 11329),
("hooli", 3974),
("", 8997),
("a", 11404),
("00000000-0000-0000-0000-000000000001", 6206),
];
for (key, expected_slot) in cases {
assert_eq!(
slot_for_key(ShardKey::Str(key)),
SlotId(*expected_slot),
"key {key:?} must keep routing to slot {expected_slot} forever",
);
}
}
#[test]
fn golden_vector_int_keys() {
let cases: &[(i64, u16)] = &[
(0, 3503),
(1, 7361),
(2, 5838),
(42, 11925),
(1_000_000, 1511),
(-1, 11296),
(i64::MAX, 7847),
(i64::MIN, 13275),
];
for (key, expected_slot) in cases {
assert_eq!(
slot_for_key(ShardKey::Int(*key)),
SlotId(*expected_slot),
"key {key} must keep routing to slot {expected_slot} forever",
);
}
}
#[test]
fn golden_vector_bytes_match_equivalent_str() {
assert_eq!(
slot_for_key(ShardKey::Bytes(b"tenant-1")),
slot_for_key(ShardKey::Str("tenant-1")),
);
}
#[test]
fn slots_stay_in_range_and_spread_roughly_uniformly() {
let mut histogram = [0usize; 16];
for i in 0..10_000i64 {
let slot = slot_for_key(ShardKey::Int(i));
assert!(slot.0 < SLOT_COUNT);
histogram[usize::from(slot.0 / 1024)] += 1;
}
let expected = 10_000 / histogram.len();
for (bucket, count) in histogram.iter().enumerate() {
assert!(
*count > expected / 2 && *count < expected * 2,
"bucket {bucket} has {count} keys (expected ≈{expected})"
);
}
}
#[tokio::test]
async fn db_for_and_read_for_attempt_routed_checkouts() {
let shards = shards_handle(&["alpha"]);
let Err(error) = shards.db_for("tenant-1").await else {
panic!("checkout must fail without a server");
};
assert!(!error.to_string().contains("Unknown shard"));
let Err(error) = shards.read_for("tenant-1").await else {
panic!("checkout must fail without a server");
};
assert!(!error.to_string().contains("fail_readiness"));
}
#[test]
fn parity_recheck_is_throttled_per_window() {
let set = shard_set(&["a"]);
let runtime = set.get(ShardId(0)).expect("shard").runtime();
runtime.configure_migration_check(
"postgres://localhost/a".to_owned(),
"postgres://localhost/a_ro".to_owned(),
);
assert!(runtime.migration_check().is_some());
assert!(runtime.parity_check_due(), "first check claims the window");
assert!(
!runtime.parity_check_due(),
"checks within the window are suppressed"
);
}
#[tokio::test]
async fn route_is_deterministic_and_respects_slot_map() {
let mut config = sharded_config(&["a", "b"]);
config.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
config.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
let set = create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
for key in ["k1", "k2", "k3", "k4", "k5"] {
let slot = set.slot_for_key(key);
let expected = if slot.0 >= 8192 { "b" } else { "a" };
let routed = set.route(key).await.expect("route");
assert_eq!(routed.name(), expected, "key {key:?} slot {}", slot.0);
assert_eq!(set.route(key).await.expect("route").id(), routed.id());
}
}
#[tokio::test]
async fn arc_shard_router_delegates_to_inner() {
let config = sharded_config(&["a", "b"]);
let set = create_shard_set(&config, Arc::new(Arc::new(HashShardRouter)))
.expect("build")
.expect("configured");
let first = set.route("tenant-42").await.expect("route");
let again = set.route("tenant-42").await.expect("route");
assert_eq!(
first.id(),
again.id(),
"Arc<R> routes deterministically through its inner router"
);
}
#[tokio::test]
async fn moving_a_slot_in_config_moves_only_that_slot() {
let mut before = sharded_config(&["a", "b"]);
before.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
before.shards[1].slots = Some(vec![SlotSpec::Range("8192-16383".to_owned())]);
let mut after = sharded_config(&["a", "b", "c"]);
after.shards[0].slots = Some(vec![SlotSpec::Range("0-8191".to_owned())]);
after.shards[1].slots = Some(vec![SlotSpec::Range("8192-12287".to_owned())]);
after.shards[2].slots = Some(vec![SlotSpec::Range("12288-16383".to_owned())]);
let set_before = create_shard_set(&before, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
let set_after = create_shard_set(&after, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
let mut moved = 0;
for i in 0..200i64 {
let slot = set_before.slot_for_key(i);
assert_eq!(slot, set_after.slot_for_key(i), "key→slot never changes");
let before_shard = set_before.route(i).await.expect("route");
let after_shard = set_after.route(i).await.expect("route");
if slot.0 >= 12288 {
assert_eq!(before_shard.name(), "b");
assert_eq!(after_shard.name(), "c");
moved += 1;
} else {
assert_eq!(before_shard.name(), after_shard.name());
}
}
assert!(moved > 0, "some keys must exercise the moved slot range");
}
#[test]
fn by_name_and_get_resolve_shards() {
let set = shard_set(&["alpha", "beta"]);
assert_eq!(set.len(), 2);
assert_eq!(set.by_name("beta").expect("beta").id(), ShardId(1));
assert_eq!(set.get(ShardId(0)).expect("alpha").name(), "alpha");
assert!(set.by_name("gamma").is_none());
assert!(set.get(ShardId(9)).is_none());
let names: Vec<&str> = set.iter().map(Shard::name).collect();
assert_eq!(names, ["alpha", "beta"]);
}
#[test]
fn auto_split_assigns_contiguous_slots() {
let set = shard_set(&["a", "b"]);
assert_eq!(set.slot_count(), SLOT_COUNT);
assert_eq!(set.get(ShardId(0)).expect("a").slots().len(), 8192);
assert_eq!(
set.get(ShardId(1)).expect("b").slots(),
(8192..16384).collect::<Vec<u16>>()
);
}
#[test]
fn owns_key_agrees_with_route() {
let set = shard_set(&["a", "b"]);
assert!(
set.owns_key(ShardId(0), "hooli"),
"hooli (slot 3974) must be shard a"
);
assert!(
!set.owns_key(ShardId(1), "hooli"),
"hooli must not be shard b"
);
assert!(
set.owns_key(ShardId(1), "a"),
"key 'a' (slot 11404) must be shard b"
);
assert!(
!set.owns_key(ShardId(0), "a"),
"key 'a' must not be shard a"
);
}
#[test]
fn slots_for_shard_returns_correct_slice() {
let set = shard_set(&["a", "b"]);
let a_slots = set.slots_for_shard(ShardId(0)).expect("shard a exists");
let b_slots = set.slots_for_shard(ShardId(1)).expect("shard b exists");
assert_eq!(a_slots.len(), 8192);
assert_eq!(b_slots.len(), 8192);
assert!(a_slots.iter().all(|&s| s < 8192));
assert!(b_slots.iter().all(|&s| s >= 8192));
assert!(set.slots_for_shard(ShardId(9)).is_none());
}
#[test]
fn partition_by_shard_groups_golden_keys() {
let set = shard_set(&["a", "b"]);
let keys = ["hooli", "a", "tenant-1"]; let map = set.partition_by_shard(keys.iter().copied());
#[allow(clippy::similar_names)]
let keys_on_a = map.get(&ShardId(0)).map_or(&[][..], Vec::as_slice);
#[allow(clippy::similar_names)]
let keys_on_b = map.get(&ShardId(1)).map_or(&[][..], Vec::as_slice);
assert!(keys_on_a.contains(&"hooli"), "hooli must go to shard a");
assert!(keys_on_b.contains(&"a"), "key 'a' must go to shard b");
assert!(
keys_on_b.contains(&"tenant-1"),
"tenant-1 (slot 12427) must go to shard b"
);
assert_eq!(
keys_on_a.len() + keys_on_b.len(),
keys.len(),
"no key dropped"
);
}
#[test]
fn create_shard_set_returns_none_without_shards() {
let config = DatabaseConfig::default();
assert!(
create_shard_set(&config, Arc::new(HashShardRouter))
.expect("ok")
.is_none()
);
}
#[test]
fn build_shard_set_rejects_duplicate_names_without_config_validation() {
let config = sharded_config(&["twin", "twin"]);
let topologies = config
.shards
.iter()
.map(|shard| crate::db::create_shard_topology(shard, &config).expect("lazy pools"))
.collect();
let result = build_shard_set(&config, topologies, Arc::new(HashShardRouter));
let Err(ShardSetBuildError::Config(error)) = result else {
panic!("duplicate shard names must be rejected, got {result:?}");
};
assert!(error.to_string().contains("twin"));
}
#[test]
fn build_shard_set_rejects_topology_count_mismatch() {
let config = sharded_config(&["a", "b"]);
let result = build_shard_set(&config, Vec::new(), Arc::new(HashShardRouter));
assert!(matches!(
result,
Err(ShardSetBuildError::TopologyCountMismatch {
expected: 2,
actual: 0
})
));
}
fn shard_with_replica(fallback: ReplicaFallback) -> Shard {
let mut config = sharded_config(&["a"]);
config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
config.shards[0].replica_fallback = Some(fallback);
let set = create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
set.get(ShardId(0)).expect("shard").clone()
}
#[test]
fn read_pool_uses_primary_when_no_replica() {
let set = shard_set(&["a"]);
let shard = set.get(ShardId(0)).expect("shard");
assert!(shard.read_pool().is_some());
assert!(shard.replica_pool().is_none());
}
#[test]
fn read_pool_requires_readiness_check_before_replica_traffic() {
let shard = shard_with_replica(ReplicaFallback::Primary);
assert!(shard.read_pool().is_some());
assert!(shard.runtime().detail().is_some());
shard.runtime().mark_replica_connection_ready();
assert!(shard.runtime().replica_ready());
assert!(shard.read_pool().is_some());
assert!(shard.runtime().detail().is_none());
}
#[test]
fn read_pool_fails_closed_under_fail_readiness() {
let shard = shard_with_replica(ReplicaFallback::FailReadiness);
assert!(
shard.read_pool().is_none(),
"unchecked replica fails closed"
);
shard.runtime().mark_replica_connection_ready();
assert!(shard.read_pool().is_some());
shard
.runtime()
.mark_replica_migrations_unready("replica lags primary");
assert!(shard.read_pool().is_none());
assert!(shard.runtime().detail().expect("detail").contains("lags"));
}
const PRIMARY_SIZE: usize = 7;
const REPLICA_SIZE: usize = 3;
fn shard_with_sized_replica(fallback: ReplicaFallback) -> Shard {
let mut config = sharded_config(&["a"]);
config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
config.shards[0].replica_fallback = Some(fallback);
config.shards[0].primary_pool_size = Some(PRIMARY_SIZE);
config.shards[0].replica_pool_size = Some(REPLICA_SIZE);
let set = create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
set.get(ShardId(0)).expect("shard").clone()
}
fn read_pool_size(route: &crate::repository::ReadRoute) -> Option<usize> {
match route {
crate::repository::ReadRoute::ReadPool(pool) => Some(pool.status().max_size),
crate::repository::ReadRoute::Primary | crate::repository::ReadRoute::Unavailable => {
None
}
}
}
#[test]
fn read_route_is_primary_without_replica() {
let set = shard_set(&["a"]);
let shard = set.get(ShardId(0)).expect("shard");
assert!(
matches!(shard.read_route(), crate::repository::ReadRoute::Primary),
"a shard with no replica must keep reads on the primary"
);
}
#[test]
fn read_route_targets_replica_when_ready() {
let shard = shard_with_sized_replica(ReplicaFallback::Primary);
shard.runtime().mark_replica_connection_ready();
assert!(shard.runtime().replica_ready());
assert_eq!(
read_pool_size(&shard.read_route()),
Some(REPLICA_SIZE),
"a ready replica must route reads to the replica pool"
);
}
#[test]
fn read_route_falls_back_to_primary_when_unready_and_policy_allows() {
let shard = shard_with_sized_replica(ReplicaFallback::Primary);
assert_eq!(
read_pool_size(&shard.read_route()),
Some(PRIMARY_SIZE),
"primary fallback must route reads to the primary pool"
);
}
#[test]
fn read_route_is_unavailable_when_unready_and_fallback_forbidden() {
let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
assert!(
matches!(
shard.read_route(),
crate::repository::ReadRoute::Unavailable
),
"fail_readiness must not silently fall back to the primary"
);
}
#[test]
fn repository_seed_snapshots_the_shard_read_route() {
let shard = shard_with_sized_replica(ReplicaFallback::Primary);
shard.runtime().mark_replica_connection_ready();
let ctx = crate::db::RequestDbContext {
statement_timeout: None,
route_key: Some("GET /notes".to_owned()),
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
};
let seed = ShardRepositorySeed::from_ctx(
shard.primary_pool(),
&ctx,
shard.name(),
shard.read_route(),
);
assert_eq!(
read_pool_size(&seed.read_route),
Some(REPLICA_SIZE),
"the seed must carry the shard's read route for from_shard"
);
}
fn shards_handle(names: &[&str]) -> Shards {
Shards {
set: shard_set(names),
ctx: crate::db::RequestDbContext {
statement_timeout: None,
route_key: Some("GET /test".to_owned()),
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
},
}
}
#[tokio::test]
async fn db_on_rejects_unknown_shard_names() {
let shards = shards_handle(&["alpha"]);
let Err(error) = shards.db_on("beta").await else {
panic!("unknown shard name must be rejected");
};
assert!(error.to_string().contains("beta"));
}
#[tokio::test]
async fn read_for_fails_closed_without_checkout_under_fail_readiness() {
let mut config = sharded_config(&["a"]);
config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
config.shards[0].replica_fallback = Some(ReplicaFallback::FailReadiness);
let shards = Shards {
set: create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured"),
ctx: crate::db::RequestDbContext {
statement_timeout: None,
route_key: None,
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
},
};
let Err(error) = shards.read_for("tenant-1").await else {
panic!("unready replica under fail_readiness must be rejected");
};
assert!(error.to_string().contains("fail_readiness"));
}
#[test]
fn shards_exposes_set_and_iter() {
let shards = shards_handle(&["alpha", "beta"]);
assert_eq!(shards.set().len(), 2);
let names: Vec<&str> = shards.iter().map(Shard::name).collect();
assert_eq!(names, ["alpha", "beta"]);
}
#[tokio::test]
async fn route_rejects_out_of_range_router_results() {
struct BadRouter;
impl ShardRouter for BadRouter {
fn route<'a>(
&'a self,
_key: ShardKey<'a>,
_shards: &'a ShardSet,
) -> futures::future::BoxFuture<'a, Result<ShardId, AutumnError>> {
Box::pin(std::future::ready(Ok(ShardId(99))))
}
}
let set = create_shard_set(&sharded_config(&["a"]), Arc::new(BadRouter))
.expect("build")
.expect("configured");
let error = set.route("k").await.expect_err("out of range");
assert!(error.to_string().contains("out-of-range"));
}
#[test]
fn shard_key_from_impls_route_consistently() {
assert_eq!(
slot_for_key(ShardKey::from(42i32)),
slot_for_key(ShardKey::from(42i64)),
);
let owned = "tenant-1".to_owned();
let bytes: [u8; 16] = *b"0123456789abcdef";
assert_eq!(
slot_for_key(ShardKey::from(&owned)),
slot_for_key(ShardKey::from("tenant-1")),
);
assert_eq!(
slot_for_key(ShardKey::from(&bytes)),
slot_for_key(ShardKey::from(&b"0123456789abcdef"[..])),
);
}
#[test]
fn build_errors_and_debug_render_usefully() {
let error = ShardSetBuildError::TopologyCountMismatch {
expected: 2,
actual: 0,
};
assert!(error.to_string().contains("expected 2"));
let set = shard_set(&["alpha"]);
let debug = format!("{set:?}");
assert!(
debug.contains("alpha"),
"ShardSet Debug names shards: {debug}"
);
let shard_debug = format!("{:?}", set.get(ShardId(0)).expect("shard"));
assert!(shard_debug.contains("alpha"));
assert_eq!(
set.shard_for_slot(SlotId(0)).expect("owner").name(),
"alpha"
);
}
fn shard_with_unreachable_replica(fallback: ReplicaFallback) -> Shard {
let mut config = sharded_config(&["a"]);
config.connect_timeout_secs = 1;
config.shards[0].replica_url = Some("postgres://localhost:1/a_ro".to_owned());
config.shards[0].replica_fallback = Some(fallback);
let set = create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
set.get(ShardId(0)).expect("shard").clone()
}
#[tokio::test]
async fn shard_indicator_gates_readiness_for_fail_readiness_replica() {
use crate::actuator::HealthIndicator as _;
let shard = shard_with_unreachable_replica(ReplicaFallback::FailReadiness);
let indicator = ShardHealthIndicator::new(shard);
let output = indicator.check().await;
assert!(
!output.status.is_healthy(),
"unreachable replica under fail_readiness must report Down"
);
assert_eq!(output.details["replica_ready"], serde_json::json!(false));
assert!(output.details.contains_key("replica_detail"));
}
#[tokio::test]
async fn shard_indicator_reports_down_when_primary_unreachable() {
use crate::actuator::HealthIndicator as _;
let shard = shard_with_unreachable_replica(ReplicaFallback::Primary);
let indicator = ShardHealthIndicator::new(shard);
let output = indicator.check().await;
assert!(
!output.status.is_healthy(),
"unreachable primary must report Down even under primary fallback"
);
assert_eq!(output.details["primary_ready"], serde_json::json!(false));
assert!(output.details.contains_key("primary_detail"));
}
#[tokio::test]
async fn register_shard_health_indicators_names_components() {
let set = shard_set(&["alpha", "beta"]);
let registry = crate::actuator::HealthIndicatorRegistry::new();
register_shard_health_indicators(&set, ®istry);
register_shard_health_indicators(&set, ®istry);
let results = registry.run_all().await;
let mut names: Vec<&str> = results
.iter()
.map(|r| r.name.as_str())
.filter(|name| name.starts_with("db:shard:"))
.collect();
names.sort_unstable();
assert_eq!(names, ["db:shard:alpha", "db:shard:beta"]);
assert!(
results
.iter()
.filter(|r| r.name.starts_with("db:shard:"))
.all(|r| matches!(r.group, crate::actuator::IndicatorGroup::Readiness)),
"shard indicators gate readiness"
);
}
#[test]
fn total_max_connections_sums_every_pool() {
let mut config = sharded_config(&["a", "b"]);
config.pool_size = 7;
config.shards[1].replica_url = Some("postgres://localhost/b_ro".to_owned());
config.shards[1].replica_pool_size = Some(3);
let set = create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured");
assert_eq!(set.total_max_connections(), 17);
}
#[test]
fn repo_seed_from_ctx_preserves_statement_timeout() {
let set = shard_set(&["shard0"]);
let shard = set.get(ShardId(0)).expect("shard");
let ctx = crate::db::RequestDbContext {
statement_timeout: Some(std::time::Duration::from_secs(3)),
route_key: Some("GET /test".to_owned()),
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(200),
interceptors: Vec::new(),
};
let seed =
ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
assert_eq!(seed.statement_timeout_ms, 3_000, "timeout preserved as ms");
assert_eq!(
seed.slow_query_threshold,
std::time::Duration::from_millis(200),
"slow threshold preserved"
);
assert_eq!(
seed.route.as_deref(),
Some("GET /test shard=shard0"),
"route tagged with shard name"
);
}
#[test]
fn reshard_route_label_retags_with_target_shard() {
assert_eq!(
reshard_route_label(Some("GET /admin shard=shard0"), "shard2").as_deref(),
Some("GET /admin shard=shard2"),
);
assert_eq!(reshard_route_label(None, "shard2"), None);
assert_eq!(
reshard_route_label(Some("GET /admin"), "shard2").as_deref(),
Some("GET /admin shard=shard2"),
);
}
#[test]
fn cross_shard_wrapper_derefs_to_inner() {
let mut w = CrossShard(7i32);
assert_eq!(*w, 7); *w = 9; assert_eq!(w.0, 9);
}
#[test]
fn cross_shard_seed_is_tenant_free_and_untagged() {
let set = shard_set(&["shard0", "shard1"]);
let ctx = crate::db::RequestDbContext {
statement_timeout: Some(std::time::Duration::from_millis(1500)),
route_key: Some("GET /admin".to_owned()),
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(250),
interceptors: Vec::new(),
};
let seed = cross_shard_seed(&set, &ctx).expect("seed");
assert_eq!(seed.route.as_deref(), Some("GET /admin"));
assert!(!seed.route.as_deref().unwrap().contains("shard="));
assert_eq!(seed.statement_timeout_ms, 1500);
assert_eq!(
seed.slow_query_threshold,
std::time::Duration::from_millis(250)
);
}
#[test]
fn repo_seed_none_timeout_maps_to_zero() {
let set = shard_set(&["shard0"]);
let shard = set.get(ShardId(0)).expect("shard");
let ctx = crate::db::RequestDbContext {
statement_timeout: None,
route_key: None,
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
};
let seed =
ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
assert_eq!(seed.statement_timeout_ms, 0, "None timeout maps to 0");
assert!(seed.route.is_none(), "None route_key propagates as None");
}
#[test]
fn repo_seed_timeout_capped_at_i32_max() {
let set = shard_set(&["shard0"]);
let shard = set.get(ShardId(0)).expect("shard");
let ctx = crate::db::RequestDbContext {
statement_timeout: Some(std::time::Duration::from_secs(u64::MAX / 1_000)),
route_key: None,
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
};
let seed =
ShardRepositorySeed::from_ctx(shard.primary_pool(), &ctx, "shard0", shard.read_route());
assert_eq!(
seed.statement_timeout_ms,
i32::MAX as u64,
"timeout capped at i32::MAX ms"
);
}
#[test]
fn replica_read_pool_is_none_without_replica() {
let set = shard_set(&["a"]);
let shard = set.get(ShardId(0)).expect("shard");
assert!(
shard.replica_read_pool().is_none(),
"no replica configured → replica_read_pool must be None"
);
}
#[test]
fn replica_read_pool_is_none_when_unready_even_under_primary_fallback() {
let shard = shard_with_sized_replica(ReplicaFallback::Primary);
assert!(
shard.replica_read_pool().is_none(),
"unready replica under primary fallback must still return None for replica_read_pool"
);
}
#[test]
fn replica_read_pool_is_none_when_unready_under_fail_readiness() {
let shard = shard_with_sized_replica(ReplicaFallback::FailReadiness);
assert!(
shard.replica_read_pool().is_none(),
"unready replica under fail_readiness must return None"
);
}
#[test]
fn replica_read_pool_targets_replica_when_ready() {
let shard = shard_with_sized_replica(ReplicaFallback::Primary);
shard.runtime().mark_replica_connection_ready();
assert!(shard.runtime().replica_ready());
assert_eq!(
shard.replica_read_pool().map(|p| p.status().max_size),
Some(REPLICA_SIZE),
"a ready replica must be returned by replica_read_pool"
);
}
#[tokio::test]
async fn read_replica_for_fails_when_no_replica_configured() {
let shards = shards_handle(&["a"]);
let Err(error) = shards.read_replica_for("tenant-1").await else {
panic!("no replica configured must be rejected");
};
let msg = error.to_string();
assert!(
msg.contains("replica"),
"error must name the missing replica: {msg}"
);
assert!(
!msg.contains("fail_readiness"),
"error must not mention fallback policy: {msg}"
);
}
#[tokio::test]
async fn read_replica_for_fails_when_replica_unready_under_primary_fallback() {
let mut config = sharded_config(&["a"]);
config.shards[0].replica_url = Some("postgres://localhost/a_ro".to_owned());
config.shards[0].replica_fallback = Some(ReplicaFallback::Primary);
let shards = Shards {
set: create_shard_set(&config, Arc::new(HashShardRouter))
.expect("build")
.expect("configured"),
ctx: crate::db::RequestDbContext {
statement_timeout: None,
route_key: None,
metrics: None,
slow_query_threshold: std::time::Duration::from_millis(500),
interceptors: Vec::new(),
},
};
let Err(error) = shards.read_replica_for("tenant-1").await else {
panic!("unready replica must be rejected even under primary fallback");
};
assert!(
error.to_string().contains("replica"),
"must name the replica: {error}"
);
}
}