Import upstream v0.16.22, stripped
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
This commit is contained in:
@@ -0,0 +1,316 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use crate::{
|
||||
IterateParams, RegistryStore, RegistryStoreInner, Store, U16_LEN, U32_LEN, U64_LEN, ValueKey,
|
||||
write::{
|
||||
BatchBuilder, ValueClass,
|
||||
assert::AssertValue,
|
||||
key::{DeserializeBigEndian, KeySerializer},
|
||||
now,
|
||||
},
|
||||
};
|
||||
use registry::{
|
||||
schema::{enums::ClusterNodeStatus, structs::ClusterNode},
|
||||
types::datetime::UTCDateTime,
|
||||
};
|
||||
use std::time::Duration;
|
||||
use trc::AddContext;
|
||||
use utils::snowflake::MAX_NODE_ID;
|
||||
|
||||
const STALE_NODE_TIMEOUT: u64 = 60 * 60; // 1 hour
|
||||
const DEAD_NODE_TIMEOUT: u64 = 60 * 60 * 24; // 24 hours
|
||||
const MAX_LEASE_RETRIES: u32 = 5;
|
||||
|
||||
struct NodeSlot {
|
||||
node_id: u16,
|
||||
hostname: String,
|
||||
last_renewal: u64,
|
||||
elapsed: u64,
|
||||
hash: u64,
|
||||
}
|
||||
|
||||
struct NodeClaim {
|
||||
node_id: u16,
|
||||
assert: AssertValue,
|
||||
}
|
||||
|
||||
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}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(err) = NodeSlot::release(
|
||||
&self.store,
|
||||
slots
|
||||
.iter()
|
||||
.filter(|slot| slot.node_id != self.node_id && slot.is_dead()),
|
||||
)
|
||||
.await
|
||||
{
|
||||
trc::error!(err.details("Failed to release expired node id leases"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl RegistryStore {
|
||||
pub fn node_id(&self) -> u16 {
|
||||
self.0.node_id
|
||||
}
|
||||
|
||||
pub fn refresh_node_id_interval(&self) -> Duration {
|
||||
Duration::from_secs(STALE_NODE_TIMEOUT / 2)
|
||||
}
|
||||
|
||||
pub async fn cluster_node_list(&self) -> trc::Result<Vec<ClusterNode>> {
|
||||
NodeSlot::list(&self.0.store, now())
|
||||
.await
|
||||
.map(|slots| slots.into_iter().map(ClusterNode::from).collect())
|
||||
}
|
||||
|
||||
pub async fn refresh_node_id_lease(&self) -> trc::Result<()> {
|
||||
let node_id = self.0.node_id;
|
||||
let assert = match NodeSlot::list(&self.0.store, now())
|
||||
.await
|
||||
.caused_by(trc::location!())?
|
||||
.into_iter()
|
||||
.find(|slot| slot.node_id == node_id)
|
||||
{
|
||||
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));
|
||||
}
|
||||
None => AssertValue::None,
|
||||
};
|
||||
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.assert_value(ValueClass::NodeId(node_id), assert).set(
|
||||
ValueClass::NodeId(node_id),
|
||||
KeySerializer::new(self.0.env_hostname.len() + U64_LEN)
|
||||
.write(now())
|
||||
.write(&self.0.env_hostname)
|
||||
.finalize(),
|
||||
);
|
||||
|
||||
self.0
|
||||
.store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub async fn purge_dead_nodes(&self) -> trc::Result<()> {
|
||||
let node_id = self.0.node_id;
|
||||
let slots = NodeSlot::list(&self.0.store, now())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
if !slots.iter().any(|slot| {
|
||||
slot.node_id == node_id && slot.is_owned_by(&self.0.env_hostname) && !slot.is_stale()
|
||||
}) {
|
||||
Ok(())
|
||||
} else {
|
||||
NodeSlot::release(
|
||||
&self.0.store,
|
||||
slots
|
||||
.iter()
|
||||
.filter(|slot| slot.node_id != node_id && slot.is_dead()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NodeSlot {
|
||||
async fn list(store: &Store, now: u64) -> trc::Result<Vec<NodeSlot>> {
|
||||
let mut slots = Vec::new();
|
||||
|
||||
store
|
||||
.iterate(
|
||||
IterateParams::new(
|
||||
ValueKey::from(ValueClass::NodeId(0)),
|
||||
ValueKey::from(ValueClass::NodeId(u16::MAX)),
|
||||
)
|
||||
.ascending(),
|
||||
|key, value| {
|
||||
if key.len() == U16_LEN * 3 {
|
||||
let node_id = key.deserialize_be_u16(U32_LEN)?;
|
||||
|
||||
match (
|
||||
value.deserialize_be_u64(0),
|
||||
value
|
||||
.get(U64_LEN..)
|
||||
.and_then(|bytes| std::str::from_utf8(bytes).ok())
|
||||
.filter(|text| !text.is_empty()),
|
||||
) {
|
||||
(Ok(last_renewal), Some(hostname)) => {
|
||||
slots.push(NodeSlot {
|
||||
node_id,
|
||||
hostname: hostname.to_string(),
|
||||
last_renewal,
|
||||
elapsed: now.saturating_sub(last_renewal),
|
||||
hash: xxhash_rust::xxh3::xxh3_64(value),
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
trc::error!(
|
||||
trc::StoreEvent::DataCorruption
|
||||
.into_err()
|
||||
.details("Invalid node id lease")
|
||||
.ctx(trc::Key::Id, node_id)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map(|_| slots)
|
||||
}
|
||||
|
||||
fn claim(slots: &[NodeSlot], hostname: &str) -> Result<NodeClaim, String> {
|
||||
if let Some(slot) = slots
|
||||
.iter()
|
||||
.find(|slot| slot.is_owned_by(hostname) && slot.is_assignable())
|
||||
.or_else(|| {
|
||||
slots
|
||||
.iter()
|
||||
.find(|slot| slot.is_stale() && slot.is_assignable())
|
||||
})
|
||||
{
|
||||
return Ok(NodeClaim {
|
||||
node_id: slot.node_id,
|
||||
assert: AssertValue::Hash(slot.hash),
|
||||
});
|
||||
}
|
||||
|
||||
let mut leased = slots
|
||||
.iter()
|
||||
.filter(|slot| !slot.is_stale())
|
||||
.map(|slot| slot.node_id)
|
||||
.collect::<Vec<_>>();
|
||||
leased.sort_unstable();
|
||||
|
||||
let mut node_id = 0;
|
||||
for leased_id in leased {
|
||||
if leased_id > node_id {
|
||||
break;
|
||||
}
|
||||
node_id = leased_id.saturating_add(1);
|
||||
if node_id > MAX_NODE_ID {
|
||||
return Err(format!(
|
||||
"Failed to obtain a node id: all {} ids are leased by active nodes",
|
||||
MAX_NODE_ID as u32 + 1
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(NodeClaim {
|
||||
node_id,
|
||||
assert: AssertValue::None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn release<'x>(
|
||||
store: &Store,
|
||||
slots: impl Iterator<Item = &'x NodeSlot>,
|
||||
) -> trc::Result<()> {
|
||||
for slot in slots {
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch
|
||||
.assert_value(
|
||||
ValueClass::NodeId(slot.node_id),
|
||||
AssertValue::Hash(slot.hash),
|
||||
)
|
||||
.clear(ValueClass::NodeId(slot.node_id));
|
||||
|
||||
if let Err(err) = store.write(batch.build_all()).await
|
||||
&& !err.is_assertion_failure()
|
||||
{
|
||||
return Err(err.caused_by(trc::location!()));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_owned_by(&self, hostname: &str) -> bool {
|
||||
self.hostname == hostname
|
||||
}
|
||||
|
||||
fn is_stale(&self) -> bool {
|
||||
self.elapsed > STALE_NODE_TIMEOUT
|
||||
}
|
||||
|
||||
fn is_dead(&self) -> bool {
|
||||
self.elapsed > DEAD_NODE_TIMEOUT
|
||||
}
|
||||
|
||||
fn is_assignable(&self) -> bool {
|
||||
self.node_id <= MAX_NODE_ID
|
||||
}
|
||||
|
||||
fn status(&self) -> ClusterNodeStatus {
|
||||
if self.is_dead() {
|
||||
ClusterNodeStatus::Inactive
|
||||
} else if self.is_stale() {
|
||||
ClusterNodeStatus::Stale
|
||||
} else {
|
||||
ClusterNodeStatus::Active
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<NodeSlot> for ClusterNode {
|
||||
fn from(slot: NodeSlot) -> Self {
|
||||
ClusterNode {
|
||||
status: slot.status(),
|
||||
last_renewal: UTCDateTime::from_timestamp(slot.last_renewal.cast_signed()),
|
||||
node_id: slot.node_id as u64,
|
||||
hostname: slot.hostname,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user