diff --git a/crates/services/src/state_manager/push.rs b/crates/services/src/state_manager/push.rs index ecc585ab..cdd4644d 100644 --- a/crates/services/src/state_manager/push.rs +++ b/crates/services/src/state_manager/push.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use super::{ @@ -178,9 +180,14 @@ pub fn spawn_push_manager(inner: Arc) -> mpsc::Sender { match event_or_timeout { Ok(Some(event)) => match event { Event::Update { account_id } => { - if server.core.jmap.push_total_shards > 1 - && account_id % server.core.jmap.push_total_shards - != server.registry().cluster_push_shard() + // INBUXA: the startup load honours the node's role, but + // subscription updates are broadcast to every node, so + // without this check each node registered the + // subscription and the client was pushed once per node. + if !server.core.network.roles.push_notifications + || (server.core.jmap.push_total_shards > 1 + && account_id % server.core.jmap.push_total_shards + != server.registry().cluster_push_shard()) { continue; } diff --git a/crates/store/src/backend/mod.rs b/crates/store/src/backend/mod.rs index 5af6fb21..195d2eec 100644 --- a/crates/store/src/backend/mod.rs +++ b/crates/store/src/backend/mod.rs @@ -34,7 +34,6 @@ pub mod scaleout; #[cfg(any(feature = "postgres", feature = "mysql"))] pub mod query_timeout; - pub const MAX_TOKEN_LENGTH: usize = (u8::MAX >> 1) as usize; pub const MAX_TOKEN_MASK: usize = MAX_TOKEN_LENGTH - 1; diff --git a/crates/store/src/backend/mysql/main.rs b/crates/store/src/backend/mysql/main.rs index deb8aa20..1cd3f17c 100644 --- a/crates/store/src/backend/mysql/main.rs +++ b/crates/store/src/backend/mysql/main.rs @@ -23,7 +23,11 @@ use mysql_async::{ impl MysqlStore { pub async fn open(config: structs::MySqlStore) -> Result { // inbuxa: ST-15: where the primary is, to tell a replica from it - let primary_location = (config.host.clone(), config.port as u16, config.database.clone()); + let primary_location = ( + config.host.clone(), + config.port as u16, + config.database.clone(), + ); let mut opts = OptsBuilder::default() .ip_or_hostname(config.host) .user(config.auth_username) diff --git a/crates/store/src/backend/postgres/main.rs b/crates/store/src/backend/postgres/main.rs index a5f71cbf..743d2d5c 100644 --- a/crates/store/src/backend/postgres/main.rs +++ b/crates/store/src/backend/postgres/main.rs @@ -53,7 +53,11 @@ pub(crate) fn pool_timeouts(connect_timeout: Option) -> Timeouts { impl PostgresStore { pub async fn open(config: structs::PostgreSqlStore) -> Result { // inbuxa: ST-15: where the primary is, to tell a replica from it - let primary_location = (config.host.clone(), config.port as u16, config.database.clone()); + let primary_location = ( + config.host.clone(), + config.port as u16, + config.database.clone(), + ); let mut cfg = Config::new(); cfg.dbname = config.database.into(); cfg.host = config.host.into(); diff --git a/crates/store/src/build/blob.rs b/crates/store/src/build/blob.rs index 079beca3..3170ca97 100644 --- a/crates/store/src/build/blob.rs +++ b/crates/store/src/build/blob.rs @@ -66,5 +66,4 @@ impl BlobStore { } } } - } diff --git a/crates/store/src/build/cluster.rs b/crates/store/src/build/cluster.rs index 447c523f..65d366ce 100644 --- a/crates/store/src/build/cluster.rs +++ b/crates/store/src/build/cluster.rs @@ -19,9 +19,9 @@ use registry::{ schema::{enums::ClusterNodeStatus, structs::ClusterNode}, types::datetime::UTCDateTime, }; -use std::time::Duration; -use trc::AddContext; -use utils::snowflake::MAX_NODE_ID; +use std::{sync::atomic::Ordering, time::Duration}; +use trc::{AddContext, ClusterEvent}; +use utils::snowflake::{MAX_NODE_ID, SnowflakeIdGenerator}; const STALE_NODE_TIMEOUT: u64 = 60 * 60; // 1 hour const DEAD_NODE_TIMEOUT: u64 = 60 * 60 * 24; // 24 hours @@ -50,46 +50,14 @@ struct NodeClaim { impl RegistryStoreInner { pub(super) async fn acquire_node_id(&mut self) -> Result<(), String> { - let mut retry_count = 0; - let slots = loop { - let now = now(); - let slots = NodeSlot::list(&self.store, now) - .await - .map_err(|err| format!("Failed to iterate store: {err}"))?; - let claim = NodeSlot::claim(&slots, &self.env_hostname)?; - let mut batch = BatchBuilder::new(); - - batch - .assert_value(ValueClass::NodeId(claim.node_id), claim.assert) - .set( - ValueClass::NodeId(claim.node_id), - KeySerializer::new(self.env_hostname.len() + U64_LEN) - .write(now) - .write(&self.env_hostname) - .finalize(), - ); - - match self.store.write(batch.build_all()).await { - Ok(_) => { - self.node_id = claim.node_id; - break slots; - } - Err(err) => { - if err.is_assertion_failure() && retry_count < MAX_LEASE_RETRIES { - retry_count += 1; - continue; - } else { - return Err(format!("Failed to write node id to store: {err}")); - } - } - } - }; + let (node_id, slots) = NodeSlot::acquire(&self.store, &self.env_hostname).await?; + self.node_id.store(node_id, Ordering::Relaxed); if let Err(err) = NodeSlot::release( &self.store, slots .iter() - .filter(|slot| slot.node_id != self.node_id && slot.is_dead()), + .filter(|slot| slot.node_id != node_id && slot.is_dead()), ) .await { @@ -102,7 +70,7 @@ impl RegistryStoreInner { impl RegistryStore { pub fn node_id(&self) -> u16 { - self.0.node_id + self.0.node_id.load(Ordering::Relaxed) } pub fn refresh_node_id_interval(&self) -> Duration { @@ -116,7 +84,7 @@ impl RegistryStore { } pub async fn refresh_node_id_lease(&self) -> trc::Result<()> { - let node_id = self.0.node_id; + let node_id = self.node_id(); let assert = match NodeSlot::list(&self.0.store, now()) .await .caused_by(trc::location!())? @@ -125,11 +93,30 @@ impl RegistryStore { { Some(slot) if slot.is_owned_by(&self.0.env_hostname) => AssertValue::Hash(slot.hash), Some(slot) => { - return Err(trc::StoreEvent::AssertValueFailed - .into_err() - .details("Node id lease is held by another host") - .ctx(trc::Key::Id, node_id) - .ctx(trc::Key::Hostname, slot.hostname)); + // INBUXA: another host took this id after the hour of silence, + // so this node was cut off from the store for longer than + // STALE_NODE_TIMEOUT. Keeping the id would mint snowflake ids + // that collide with that host's. Step down and claim a fresh + // id instead of only logging it. + let (new_node_id, _) = NodeSlot::acquire(&self.0.store, &self.0.env_hostname) + .await + .map_err(|err| { + trc::StoreEvent::UnexpectedError + .into_err() + .details(err) + .ctx(trc::Key::Id, node_id) + })?; + self.0.node_id.store(new_node_id, Ordering::Relaxed); + SnowflakeIdGenerator::set_node_id(new_node_id as u64); + + trc::event!( + Cluster(ClusterEvent::NodeIdReassigned), + Id = new_node_id, + Hostname = slot.hostname, + Details = format!("Node id {node_id} lease is held by another host"), + ); + + return Ok(()); } None => AssertValue::None, }; @@ -152,7 +139,7 @@ impl RegistryStore { } pub async fn purge_dead_nodes(&self) -> trc::Result<()> { - let node_id = self.0.node_id; + let node_id = self.node_id(); let slots = NodeSlot::list(&self.0.store, now()) .await .caused_by(trc::location!())?; @@ -174,6 +161,43 @@ impl RegistryStore { } impl NodeSlot { + // Claims a node id for this hostname: its own slot if it still has one, + // otherwise a stale one or the lowest free id. Returns the id and the + // slots seen while claiming it. + async fn acquire(store: &Store, hostname: &str) -> Result<(u16, Vec), String> { + let mut retry_count = 0; + loop { + let now = now(); + let slots = NodeSlot::list(store, now) + .await + .map_err(|err| format!("Failed to iterate store: {err}"))?; + let claim = NodeSlot::claim(&slots, hostname)?; + let mut batch = BatchBuilder::new(); + + batch + .assert_value(ValueClass::NodeId(claim.node_id), claim.assert) + .set( + ValueClass::NodeId(claim.node_id), + KeySerializer::new(hostname.len() + U64_LEN) + .write(now) + .write(hostname) + .finalize(), + ); + + match store.write(batch.build_all()).await { + Ok(_) => return Ok((claim.node_id, slots)), + Err(err) => { + if err.is_assertion_failure() && retry_count < MAX_LEASE_RETRIES { + retry_count += 1; + continue; + } else { + return Err(format!("Failed to write node id to store: {err}")); + } + } + } + } + } + async fn list(store: &Store, now: u64) -> trc::Result> { let mut slots = Vec::new(); @@ -346,10 +370,19 @@ mod tests { #[test] fn status_follows_the_heartbeat() { assert_eq!(slot(0).status(), ClusterNodeStatus::Active); - assert_eq!(slot(UNRESPONSIVE_NODE_TIMEOUT).status(), ClusterNodeStatus::Active); - assert_eq!(slot(UNRESPONSIVE_NODE_TIMEOUT + 1).status(), ClusterNodeStatus::Stale); + assert_eq!( + slot(UNRESPONSIVE_NODE_TIMEOUT).status(), + ClusterNodeStatus::Active + ); + assert_eq!( + slot(UNRESPONSIVE_NODE_TIMEOUT + 1).status(), + ClusterNodeStatus::Stale + ); assert_eq!(slot(DEAD_NODE_TIMEOUT).status(), ClusterNodeStatus::Stale); - assert_eq!(slot(DEAD_NODE_TIMEOUT + 1).status(), ClusterNodeStatus::Inactive); + assert_eq!( + slot(DEAD_NODE_TIMEOUT + 1).status(), + ClusterNodeStatus::Inactive + ); } #[test] @@ -361,6 +394,39 @@ mod tests { assert!(slot(STALE_NODE_TIMEOUT + 1).is_stale()); } + fn owned(node_id: u16, hostname: &str, elapsed: u64) -> NodeSlot { + NodeSlot { + node_id, + hostname: hostname.into(), + last_renewal: 0, + elapsed, + hash: node_id as u64, + } + } + + #[test] + fn a_node_whose_id_was_taken_claims_a_fresh_one() { + // mx2 held id 1, was silent for over an hour, and mx3 took it over. + // Back online, mx2 must not keep minting ids as node 1. + let slots = vec![ + owned(0, "mail.example.org", 5), + owned(1, "mx3.example.org", 5), + ]; + let claim = NodeSlot::claim(&slots, "mx2.example.org").unwrap(); + assert_eq!(claim.node_id, 2); + assert!(matches!(claim.assert, AssertValue::None)); + + // Its own slot is preferred whenever it still has one. + let slots = vec![ + owned(0, "mail.example.org", 5), + owned(3, "mx2.example.org", 5), + ]; + assert_eq!( + NodeSlot::claim(&slots, "mx2.example.org").unwrap().node_id, + 3 + ); + } + #[test] fn several_renewals_fit_before_a_node_looks_unresponsive() { assert!(UNRESPONSIVE_NODE_TIMEOUT >= 3 * HEARTBEAT_INTERVAL); diff --git a/crates/store/src/build/memory.rs b/crates/store/src/build/memory.rs index cec5f12e..ac9dfe85 100644 --- a/crates/store/src/build/memory.rs +++ b/crates/store/src/build/memory.rs @@ -54,5 +54,4 @@ impl InMemoryStore { } } } - } diff --git a/crates/store/src/build/registry.rs b/crates/store/src/build/registry.rs index 5751a745..15ea1f6e 100644 --- a/crates/store/src/build/registry.rs +++ b/crates/store/src/build/registry.rs @@ -28,8 +28,12 @@ impl RegistryStore { if !inner.env_recovery_mode && inner.env_recovery_admin.take().is_some() { eprintln!(); eprintln!("⚠️ INBUXA_RECOVERY_ADMIN is set, but the"); - eprintln!(" server is configured and not in recovery mode, so it is ignored."); - eprintln!(" Remove it from the environment. To use it for recovery, also set"); + eprintln!( + " server is configured and not in recovery mode, so it is ignored." + ); + eprintln!( + " Remove it from the environment. To use it for recovery, also set" + ); eprintln!(" INBUXA_RECOVERY_MODE=1."); eprintln!(); } @@ -164,7 +168,7 @@ impl RegistryStore { RegistryStoreInner { local_path: PathBuf::from(path), store, - node_id: 0, + node_id: Default::default(), env_recovery_mode: false, env_recovery_admin: Some(("admin".to_string(), "popolna_zapora".to_string())), env_cluster_role: cluster_role, diff --git a/crates/store/src/dispatch/search.rs b/crates/store/src/dispatch/search.rs index c65e1244..b8afd7e3 100644 --- a/crates/store/src/dispatch/search.rs +++ b/crates/store/src/dispatch/search.rs @@ -219,7 +219,9 @@ impl SearchStore { crate::sql_backend!(&store.primary, db => db.query(index, filters, sort).await) } }, - None => crate::sql_backend!(&store.primary, db => db.query(index, filters, sort).await), + None => { + crate::sql_backend!(&store.primary, db => db.query(index, filters, sort).await) + } }; } match self { diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index 589347bb..d65b2a98 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -295,7 +295,9 @@ impl Store { #[cfg(feature = "rocks")] Self::RocksDb(store) => store.purge_store().await, Self::Ephemeral(store) => store.purge_store().await, - Self::Replicated(store) => crate::sql_backend!(&store.primary, db => db.purge_store().await), + Self::Replicated(store) => { + crate::sql_backend!(&store.primary, db => db.purge_store().await) + } Self::None => Err(trc::StoreEvent::NotConfigured.into()), } .caused_by(trc::location!()) diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index a332cfc2..4a88fc5f 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -29,7 +29,11 @@ pub use xxhash_rust; use crate::backend::{elastic::ElasticSearchStore, meili::MeiliSearchStore}; use ahash::AHashMap; use backend::{ephemeral::EphemeralStore, fs::FsStore, http::HttpStore, memory::StaticMemoryStore}; -use std::{borrow::Cow, path::PathBuf, sync::Arc}; +use std::{ + borrow::Cow, + path::PathBuf, + sync::{Arc, atomic::AtomicU16}, +}; use write::ValueClass; pub trait Deserialize: Sized + Sync + Send { @@ -204,7 +208,8 @@ pub struct RegistryStore(pub(crate) Arc); pub struct RegistryStoreInner { pub(crate) local_path: PathBuf, pub(crate) store: Store, - pub(crate) node_id: u16, + // inbuxa: atomic so a node can step down to a fresh id after losing its lease + pub(crate) node_id: Arc, pub(crate) env_recovery_mode: bool, pub(crate) env_recovery_admin: Option<(String, String)>, pub(crate) env_cluster_role: Option, @@ -708,7 +713,6 @@ impl Store { pub fn is_ephemeral(&self) -> bool { matches!(self, Self::Ephemeral(_)) } - } impl std::fmt::Debug for Store { diff --git a/crates/store/src/registry/bootstrap.rs b/crates/store/src/registry/bootstrap.rs index 87494bdc..585a74e9 100644 --- a/crates/store/src/registry/bootstrap.rs +++ b/crates/store/src/registry/bootstrap.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::{RegistryStore, Store, registry::RegistryObject}; @@ -189,7 +191,7 @@ impl Bootstrap { } pub fn node_id(&self) -> u16 { - self.registry.0.node_id + self.registry.node_id() } pub fn log_errors(&self) { diff --git a/crates/store/src/registry/local.rs b/crates/store/src/registry/local.rs index 6f6e0653..03b0d2d3 100644 --- a/crates/store/src/registry/local.rs +++ b/crates/store/src/registry/local.rs @@ -36,7 +36,7 @@ impl RegistryStoreInner { local_path, store: Store::None, id_generator: SnowflakeIdGenerator::new(), - node_id: 0, + node_id: Default::default(), env_recovery_mode: types::branding::env_var("RECOVERY_MODE") .ok() .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) diff --git a/crates/store/src/registry/write.rs b/crates/store/src/registry/write.rs index 99c5d1f0..301699ec 100644 --- a/crates/store/src/registry/write.rs +++ b/crates/store/src/registry/write.rs @@ -87,12 +87,19 @@ impl RegistryStore { return self.write_unhooked(write).await; }; let (object_type, id, before, after) = match &write { - RegistryWrite::Insert { object, id } => (object.object_type(), *id, None, Some(*object)), + RegistryWrite::Insert { object, id } => { + (object.object_type(), *id, None, Some(*object)) + } RegistryWrite::Update { object, id, old_object, - } => (object.object_type(), Some(*id), Some(*old_object), Some(*object)), + } => ( + object.object_type(), + Some(*id), + Some(*old_object), + Some(*object), + ), RegistryWrite::Delete { object_id, object, .. } => (object_id.object(), Some(object_id.id()), *object, None), diff --git a/crates/trc/src/event/enums.rs b/crates/trc/src/event/enums.rs index 2a7132a8..91bce264 100644 --- a/crates/trc/src/event/enums.rs +++ b/crates/trc/src/event/enums.rs @@ -152,6 +152,7 @@ pub enum ClusterEvent { MessageSkipped = 47, MessageInvalid = 49, NodeIdRenewed = 275, + NodeIdReassigned = 649, // inbuxa: the coordinator's connection CoordinatorConnected = 644, CoordinatorDisconnected = 645, diff --git a/crates/trc/src/event/enums_impl.rs b/crates/trc/src/event/enums_impl.rs index 3e7c1a55..86cfe72f 100644 --- a/crates/trc/src/event/enums_impl.rs +++ b/crates/trc/src/event/enums_impl.rs @@ -81,6 +81,7 @@ impl EventType { b"cluster.message-skipped" => EventType::Cluster(ClusterEvent::MessageSkipped), b"cluster.message-invalid" => EventType::Cluster(ClusterEvent::MessageInvalid), b"cluster.node-id-renewed" => EventType::Cluster(ClusterEvent::NodeIdRenewed), + b"cluster.node-id-reassigned" => EventType::Cluster(ClusterEvent::NodeIdReassigned), // inbuxa: coordinator connection b"cluster.coordinator-connected" => EventType::Cluster(ClusterEvent::CoordinatorConnected), b"cluster.coordinator-disconnected" => EventType::Cluster(ClusterEvent::CoordinatorDisconnected), @@ -749,6 +750,7 @@ impl EventType { EventType::Cluster(ClusterEvent::MessageSkipped) => "cluster.message-skipped", EventType::Cluster(ClusterEvent::MessageInvalid) => "cluster.message-invalid", EventType::Cluster(ClusterEvent::NodeIdRenewed) => "cluster.node-id-renewed", + EventType::Cluster(ClusterEvent::NodeIdReassigned) => "cluster.node-id-reassigned", // inbuxa: coordinator connection EventType::Cluster(ClusterEvent::CoordinatorConnected) => { "cluster.coordinator-connected" @@ -1542,6 +1544,7 @@ impl EventType { EventType::Cluster(ClusterEvent::MessageSkipped) => 47, EventType::Cluster(ClusterEvent::MessageInvalid) => 49, EventType::Cluster(ClusterEvent::NodeIdRenewed) => 275, + EventType::Cluster(ClusterEvent::NodeIdReassigned) => 649, // inbuxa: coordinator connection EventType::Cluster(ClusterEvent::CoordinatorConnected) => 644, EventType::Cluster(ClusterEvent::CoordinatorDisconnected) => 645, @@ -2201,6 +2204,7 @@ impl EventType { 47 => Some(EventType::Cluster(ClusterEvent::MessageSkipped)), 49 => Some(EventType::Cluster(ClusterEvent::MessageInvalid)), 275 => Some(EventType::Cluster(ClusterEvent::NodeIdRenewed)), + 649 => Some(EventType::Cluster(ClusterEvent::NodeIdReassigned)), // inbuxa: coordinator connection 644 => Some(EventType::Cluster(ClusterEvent::CoordinatorConnected)), 645 => Some(EventType::Cluster(ClusterEvent::CoordinatorDisconnected)), @@ -3283,6 +3287,7 @@ impl EventType { EventType::Cluster(ClusterEvent::MessageSkipped) => "PubSub message skipped", EventType::Cluster(ClusterEvent::MessageInvalid) => "Invalid PubSub message", EventType::Cluster(ClusterEvent::NodeIdRenewed) => "Node ID renewed", + EventType::Cluster(ClusterEvent::NodeIdReassigned) => "Node ID reassigned after its lease was taken", // inbuxa: coordinator connection EventType::Cluster(ClusterEvent::CoordinatorConnected) => "Coordinator connected", EventType::Cluster(ClusterEvent::CoordinatorDisconnected) => "Coordinator unavailable", @@ -4375,6 +4380,7 @@ impl EventType { EventType::Cluster(ClusterEvent::MessageSkipped), EventType::Cluster(ClusterEvent::MessageInvalid), EventType::Cluster(ClusterEvent::NodeIdRenewed), + EventType::Cluster(ClusterEvent::NodeIdReassigned), // inbuxa: coordinator connection EventType::Cluster(ClusterEvent::CoordinatorConnected), EventType::Cluster(ClusterEvent::CoordinatorDisconnected),