chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[smg][mesh] extract mesh to mesh crate to reduce compile time (#17907)

Browse Source
This commit is contained in:
Simo Lin
2026-01-29 10:58:42 +09:00
committed by GitHub
parent 0368ddf9ea
commit ac16d4450e
28 changed files with 31 additions and 10065 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-model-gateway/Cargo.toml
View File

@@ -84,6 +84,7 @@ wfaas = "1.0.0"
data-connector = "1.0.0"
smg-mcp = "1.0.0"
smg-wasm = "1.0.0"
smg-mesh = "1.0.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
rustls-pemfile = "2.2"
openssl = "0.10.73"

8
sgl-model-gateway/build.rs
View File

@@ -12,7 +12,6 @@ macro_rules! set_env {
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Rebuild triggers
println!("cargo:rerun-if-changed=src/mesh/proto/gossip.proto");
println!("cargo:rerun-if-changed=src/proto/sglang_scheduler.proto");
println!("cargo:rerun-if-changed=src/proto/vllm_engine.proto");
println!("cargo:rerun-if-changed=Cargo.toml");
@@ -31,13 +30,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
&["src/proto"],
)?;
// Compile gossip protobuf files
tonic_prost_build::configure()
// Generate both client and server code
.build_server(true)
.build_client(true)
.compile_protos(&["src/mesh/proto/gossip.proto"], &["src/mesh/proto"])?;
// Set version info environment variables
let version = read_cargo_version().unwrap_or_else(|_| DEFAULT_VERSION.to_string());
let target = std::env::var("TARGET").unwrap_or_else(|_| get_rustc_host().unwrap_or_default());

2
sgl-model-gateway/src/lib.rs
View File

@@ -5,7 +5,7 @@ pub mod core;
pub use data_connector;
pub mod grpc_client;
pub use smg_mcp as mcp;
pub mod mesh;
pub use smg_mesh as mesh;
pub mod middleware;
pub mod observability;
pub mod policies;

1123
sgl-model-gateway/src/mesh/README.md
View File

File diff suppressed because it is too large Load Diff

212
sgl-model-gateway/src/mesh/consistent_hash.rs
View File

@@ -1,212 +0,0 @@
//! Consistent hashing for rate-limit ownership
//!
//! Implements consistent hashing ring to determine K owners (K=1-3) for each rate-limit key.
//! Supports ownership transfer on node failures.
use std::{
collections::{hash_map::DefaultHasher, BTreeMap, HashSet},
hash::{Hash, Hasher},
};
/// Number of virtual nodes per physical node (for better distribution)
const VIRTUAL_NODES_PER_NODE: usize = 150;
/// Number of owners (K) for each key
const NUM_OWNERS: usize = 3;
/// Consistent hash ring
#[derive(Debug, Clone)]
pub struct ConsistentHashRing {
/// Ring: hash -> node_name
ring: BTreeMap<u64, String>,
/// Node -> set of virtual node hashes
node_hashes: BTreeMap<String, HashSet<u64>>,
}
impl ConsistentHashRing {
pub fn new() -> Self {
Self {
ring: BTreeMap::new(),
node_hashes: BTreeMap::new(),
}
}
/// Add a node to the ring
pub fn add_node(&mut self, node_name: &str) {
if self.node_hashes.contains_key(node_name) {
// Node already exists
return;
}
let mut hashes = HashSet::new();
for i in 0..VIRTUAL_NODES_PER_NODE {
let virtual_node = format!("{}:{}", node_name, i);
let hash = Self::hash(&virtual_node);
self.ring.insert(hash, node_name.to_string());
hashes.insert(hash);
}
self.node_hashes.insert(node_name.to_string(), hashes);
}
/// Remove a node from the ring
pub fn remove_node(&mut self, node_name: &str) {
if let Some(hashes) = self.node_hashes.remove(node_name) {
for hash in hashes {
self.ring.remove(&hash);
}
}
}
pub fn get_owners(&self, key: &str) -> Vec<&str> {
if self.ring.is_empty() {
return Vec::new();
}
let key_hash = Self::hash(key);
let mut owners = Vec::with_capacity(NUM_OWNERS);
let total_unique_nodes = self.node_hashes.len();
let mut iter = self.ring.range(key_hash..);
while owners.len() < NUM_OWNERS && owners.len() < total_unique_nodes {
if let Some((_, node)) = iter.next() {
if !owners.contains(&node.as_str()) {
owners.push(node.as_str());
}
} else {
iter = self.ring.range(..);
}
}
owners
}
/// Check if a node is an owner of a key
pub fn is_owner(&self, key: &str, node_name: &str) -> bool {
self.get_owners(key).contains(&node_name)
}
/// Get all nodes in the ring
pub fn get_nodes(&self) -> Vec<String> {
self.node_hashes.keys().cloned().collect()
}
/// Check if a node exists in the ring
pub fn has_node(&self, node_name: &str) -> bool {
self.node_hashes.contains_key(node_name)
}
/// Hash a string to u64
fn hash(s: &str) -> u64 {
let mut hasher = DefaultHasher::new();
s.hash(&mut hasher);
hasher.finish()
}
/// Update ring with current membership
pub fn update_membership(&mut self, nodes: &[String]) {
let current_nodes: HashSet<String> = self.node_hashes.keys().cloned().collect();
let new_nodes: HashSet<String> = nodes.iter().cloned().collect();
// Remove nodes that are no longer present
for node in current_nodes.difference(&new_nodes) {
self.remove_node(node);
}
// Add new nodes
for node in new_nodes.difference(&current_nodes) {
self.add_node(node);
}
}
}
impl Default for ConsistentHashRing {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_add_remove_node() {
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
assert!(ring.has_node("node1"));
assert_eq!(ring.get_nodes().len(), 1);
ring.add_node("node2");
assert_eq!(ring.get_nodes().len(), 2);
ring.remove_node("node1");
assert!(!ring.has_node("node1"));
assert_eq!(ring.get_nodes().len(), 1);
}
#[test]
fn test_get_owners() {
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
ring.add_node("node2");
ring.add_node("node3");
let owners = ring.get_owners("test_key");
assert_eq!(owners.len(), NUM_OWNERS);
assert!(owners.iter().all(|n| ring.has_node(n)));
}
#[test]
fn test_is_owner() {
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
ring.add_node("node2");
ring.add_node("node3");
let owners = ring.get_owners("test_key");
for owner in &owners {
assert!(ring.is_owner("test_key", owner));
}
}
#[test]
fn test_update_membership() {
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
ring.add_node("node2");
ring.update_membership(&["node2".to_string(), "node3".to_string()]);
assert!(!ring.has_node("node1"));
assert!(ring.has_node("node2"));
assert!(ring.has_node("node3"));
}
#[test]
fn test_get_owners_with_fewer_nodes_than_owners() {
// Test that the loop terminates correctly when there are fewer nodes than NUM_OWNERS
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
ring.add_node("node2");
// Only 2 nodes, but NUM_OWNERS is 3
let owners = ring.get_owners("test_key");
// Should return all available nodes (2) without infinite loop
assert_eq!(owners.len(), 2);
assert!(owners.contains(&"node1"));
assert!(owners.contains(&"node2"));
}
#[test]
fn test_get_owners_with_single_node() {
// Test with only one node
let mut ring = ConsistentHashRing::new();
ring.add_node("node1");
let owners = ring.get_owners("test_key");
// Should return the single node without infinite loop
assert_eq!(owners.len(), 1);
assert_eq!(owners[0], "node1");
}
}

266
sgl-model-gateway/src/mesh/controller.rs
View File

@@ -1,266 +0,0 @@
use std::{
collections::{BTreeMap, HashMap},
net::SocketAddr,
time::Duration,
};
use anyhow::Result;
use rand::seq::{IndexedRandom, SliceRandom};
use tracing as log;
use tracing::instrument;
use super::{
flow_control::RetryManager,
gossip::{gossip_message, NodeState, NodeStatus, Ping, PingReq, StateSync},
service::{broadcast_node_states, try_ping},
ClusterState,
};
pub struct MeshController {
state: ClusterState,
self_name: String,
self_addr: SocketAddr,
init_peer: Option<SocketAddr>,
}
impl MeshController {
pub fn new(
state: ClusterState,
self_addr: SocketAddr,
self_name: &str,
init_peer: Option<SocketAddr>,
) -> Self {
Self {
state,
self_name: self_name.to_string(),
self_addr,
init_peer,
}
}
#[instrument(fields(name = %self.self_name), skip(self, signal))]
pub async fn event_loop(self, mut signal: tokio::sync::watch::Receiver<()>) -> Result<()> {
let init_state = self.state.clone();
let read_state = self.state.clone();
let mut cnt: u64 = 0;
// Track retry managers for each peer
use std::collections::HashMap;
let mut retry_managers: HashMap<String, RetryManager> = HashMap::new();
loop {
log::info!("Round {} Status:{:?}", cnt, read_state.read());
// Get available peers from cluster state
let mut map = init_state.read().clone();
map.retain(|k, v| {
k.ne(&self.self_name.to_string())
&& v.status != NodeStatus::Down as i32
&& v.status != NodeStatus::Leaving as i32
});
let peer = if cnt == 0 && map.is_empty() {
// Only use init_peer if cluster state is empty (no service discovery)
self.init_peer.map(|init_peer| NodeState {
name: "init_peer".to_string(),
address: init_peer.to_string(),
status: NodeStatus::Suspected as i32,
version: 1,
metadata: HashMap::new(),
})
} else {
// Use nodes from cluster state (from service discovery or gossip)
let random_nodes = get_random_values_refs(&map, 1);
random_nodes.first().map(|&node| node.clone())
};
cnt += 1;
tokio::select! {
_ = signal.changed() => {
log::info!("Gossip app_server {} at {} is shutting down", self.self_name, self.self_addr);
break;
}
_ = tokio::time::sleep(Duration::from_secs(1)) => {
if let Some(peer) = peer {
let peer_name = peer.name.clone();
// Get or create retry manager for this peer
let retry_manager = retry_managers
.entry(peer_name.clone())
.or_default();
// Check if we should retry based on backoff
if retry_manager.should_retry() {
match self.connect_to_peer(peer.clone()).await {
Ok(_) => {
// Success - reset retry state
retry_manager.reset();
log::info!("Successfully connected to peer {}", peer_name);
}
Err(e) => {
// Failure - record attempt and calculate next delay
retry_manager.record_attempt();
let next_delay = retry_manager.next_delay();
let attempt = retry_manager.attempt_count();
log::warn!(
"Error connecting to peer {} (attempt {}): {}. Next retry in {:?}",
peer_name,
attempt,
e,
next_delay
);
}
}
} else {
// Still in backoff period, skip this attempt
let next_delay = retry_manager.next_delay();
log::debug!(
"Skipping connection to peer {} (backoff: {:?} remaining)",
peer_name,
next_delay
);
}
} else {
log::info!("No peer address available to connect");
}
}
}
}
Ok(())
}
async fn connect_to_peer(&self, peer: NodeState) -> Result<()> {
log::info!("Connecting to peer {} at {}", peer.name, peer.address);
let read_state = self.state.clone();
// TODO: Maybe we don't need to send the whole state.
let state_sync = StateSync {
nodes: read_state.read().values().cloned().collect(),
};
let peer_addr = peer.address.parse::<SocketAddr>()?;
let peer_name = peer.name.clone();
match try_ping(
&peer,
Some(gossip_message::Payload::Ping(Ping {
state_sync: Some(state_sync),
})),
)
.await
{
Ok(node_update) => {
log::info!("Received NodeUpdate from peer: {:?}", node_update);
// Update state for Alive or Leaving status
if node_update.status == NodeStatus::Alive as i32
|| node_update.status == NodeStatus::Leaving as i32
{
let mut s = read_state.write();
s.entry(node_update.name.clone())
.and_modify(|e| e.status = node_update.status)
.or_insert(NodeState {
name: node_update.name,
address: node_update.address,
status: node_update.status,
version: 1,
metadata: HashMap::new(),
});
}
}
Err(e) => {
log::info!("Failed to connect to peer: {}, now try ping-req", e);
let mut map = read_state.read().clone();
map.retain(|k, v| {
k.ne(&self.self_name)
&& k.ne(&peer_name)
&& v.status == NodeStatus::Alive as i32
&& v.status != NodeStatus::Leaving as i32
});
let random_nodes = get_random_values_refs(&map, 3);
let mut reachable = false;
for node in random_nodes {
log::info!(
"Trying to ping-req node {}, req target: {}",
node.address,
peer_addr
);
if try_ping(
node,
Some(gossip_message::Payload::PingReq(PingReq {
node: Some(peer.clone()),
})),
)
.await
.is_ok()
{
reachable = true;
break;
}
}
if !reachable {
let mut target = read_state.read().clone();
// Broadcast only the unreachable node's status is enough.
if let Some(mut unreachable_node) = target.remove(&peer_name) {
if unreachable_node.status == NodeStatus::Suspected as i32 {
unreachable_node.status = NodeStatus::Down as i32
} else {
unreachable_node.status = NodeStatus::Suspected as i32
}
unreachable_node.version += 1;
// Broadcast target nodes should include self.
let target_nodes: Vec<NodeState> = target
.values()
.filter(|v| {
v.name.ne(&peer_name)
&& v.status == NodeStatus::Alive as i32
&& v.status != NodeStatus::Leaving as i32
})
.cloned()
.collect();
log::info!(
"Broadcasting node status to {} alive nodes, new_state: {:?}",
target_nodes.len(),
unreachable_node
);
let (success_count, total_count) = broadcast_node_states(
vec![unreachable_node],
target_nodes,
None, // Use default timeout
)
.await;
log::info!(
"Broadcast node status: {}/{} successful",
success_count,
total_count
);
}
}
}
}
log::info!("Successfully connected to peer {}", peer_addr);
Ok(())
}
}
// TODO: Support weighted random selection. e.g. nodes in INIT state should be more likely to be selected.
fn get_random_values_refs<K, V>(map: &BTreeMap<K, V>, k: usize) -> Vec<&V> {
let values: Vec<&V> = map.values().collect();
if k >= values.len() {
let mut all_values = values;
all_values.shuffle(&mut rand::rng());
return all_values;
}
let mut rng = rand::rng();
values.choose_multiple(&mut rng, k).cloned().collect()
}

962
sgl-model-gateway/src/mesh/crdt.rs
View File

@@ -1,962 +0,0 @@
//! CRDT (Conflict-free Replicated Data Types) wrapper for HA state synchronization
//!
//! This module provides CRDT data structures for eventual consistency:
//! - Map<SKey, LWWReg> for Last-Write-Wins Register maps
//! - PNCounter for rate-limit and load balance aggregates
use std::{
collections::BTreeMap,
sync::Arc,
time::{SystemTime, UNIX_EPOCH},
};
use crdts::{CmRDT, CvRDT, PNCounter};
use num_bigint::BigInt;
use num_traits::ToPrimitive;
use parking_lot::RwLock;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
/// State key for CRDT maps
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct SKey(pub String);
impl SKey {
pub fn new(key: String) -> Self {
Self(key)
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<String> for SKey {
fn from(s: String) -> Self {
Self(s)
}
}
impl From<&str> for SKey {
fn from(s: &str) -> Self {
Self(s.to_string())
}
}
/// Last-Write-Wins Register wrapper
/// Simplified implementation using timestamp and version
#[derive(Debug, Clone, serde::Serialize)]
#[serde(bound(serialize = "T: Serialize"))]
#[derive(serde::Deserialize)]
#[serde(bound(deserialize = "T: DeserializeOwned"))]
pub struct LWWRegister<T: Clone + Serialize + DeserializeOwned> {
value: T,
timestamp: u64,
version: u64,
actor: String,
}
impl<T: Clone + Serialize + DeserializeOwned> LWWRegister<T> {
pub fn new(value: T, actor: String) -> Self {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
Self {
value,
timestamp,
version: 1,
actor,
}
}
pub fn read(&self) -> &T {
&self.value
}
pub fn write(&mut self, value: T, actor: String) {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
self.value = value;
self.timestamp = timestamp;
self.version += 1;
self.actor = actor;
}
pub fn merge(&mut self, other: &Self) {
// Last-Write-Wins: choose the one with higher timestamp, or higher version if equal
if other.timestamp > self.timestamp
|| (other.timestamp == self.timestamp && other.version > self.version)
{
self.value = other.value.clone();
self.timestamp = other.timestamp;
self.version = other.version;
self.actor = other.actor.clone();
}
}
}
/// CRDT Map wrapper using LWWRegister for values
/// Simplified implementation using BTreeMap with LWWRegister values
#[derive(Debug, Clone, serde::Serialize)]
#[serde(bound(serialize = "T: Serialize + DeserializeOwned"))]
#[derive(serde::Deserialize)]
#[serde(bound(deserialize = "T: DeserializeOwned"))]
pub struct CRDTMap<T: Clone + Serialize + DeserializeOwned> {
inner: BTreeMap<SKey, LWWRegister<T>>,
}
impl<T: Clone + Serialize + DeserializeOwned> Default for CRDTMap<T> {
fn default() -> Self {
Self {
inner: BTreeMap::new(),
}
}
}
impl<T: Clone + Serialize + DeserializeOwned> CRDTMap<T> {
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, key: &SKey) -> Option<&T> {
self.inner.get(key).map(|reg| reg.read())
}
pub fn insert(&mut self, key: SKey, value: T, actor: String) {
// Check if key already exists to preserve version
if let Some(existing_reg) = self.inner.get_mut(&key) {
// Update existing register, which will increment version
existing_reg.write(value, actor);
} else {
// New entry, start with version 1
let reg = LWWRegister::new(value, actor);
self.inner.insert(key, reg);
}
}
/// Get the version and actor for a key
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner
.get(key)
.map(|reg| (reg.version, reg.actor.clone()))
}
pub fn remove(&mut self, key: &SKey) {
self.inner.remove(key);
}
pub fn contains_key(&self, key: &SKey) -> bool {
self.inner.contains_key(key)
}
pub fn iter(&self) -> impl Iterator<Item = (&SKey, &T)> {
self.inner.iter().map(|(k, v)| (k, v.read()))
}
pub fn keys(&self) -> impl Iterator<Item = &SKey> {
self.inner.keys()
}
pub fn values(&self) -> impl Iterator<Item = &T> {
self.inner.values().map(|v| v.read())
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn merge(&mut self, other: &Self) {
for (key, other_reg) in &other.inner {
match self.inner.get_mut(key) {
Some(self_reg) => {
self_reg.merge(other_reg);
}
None => {
self.inner.insert(key.clone(), other_reg.clone());
}
}
}
}
pub fn to_map(&self) -> BTreeMap<SKey, T> {
self.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
}
}
/// Positive-Negative Counter for rate-limit and load balance aggregates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CRDTPNCounter {
inner: PNCounter<String>,
}
impl Default for CRDTPNCounter {
fn default() -> Self {
Self {
inner: PNCounter::new(),
}
}
}
impl CRDTPNCounter {
pub fn new() -> Self {
Self::default()
}
pub fn inc(&mut self, actor: String, delta: i64) {
// PNCounter API: inc(actor) and dec(actor) return operations that need to be applied
// In crdts 7.3, we need to call apply() to actually modify the counter
if delta > 0 {
for i in 0..delta as u64 {
// Use a unique actor for each increment to ensure they're all counted
let unique_actor = format!("{}:{}", actor, i);
let op = self.inner.inc(unique_actor);
self.inner.apply(op);
}
} else if delta < 0 {
for i in 0..(-delta) as u64 {
// Use a unique actor for each decrement
let unique_actor = format!("{}:{}", actor, i);
let op = self.inner.dec(unique_actor);
self.inner.apply(op);
}
}
}
pub fn value(&self) -> i64 {
// PNCounter read() returns BigInt in crdts 7.3
let val: BigInt = self.inner.read();
// Convert BigInt to i64, clamping to i64::MAX/i64::MIN if value is out of range
val.to_i64().unwrap_or_else(|| {
// If value is too large, clamp to i64::MAX
if val > BigInt::from(i64::MAX) {
i64::MAX
} else if val < BigInt::from(i64::MIN) {
i64::MIN
} else {
0
}
})
}
pub fn merge(&mut self, other: &Self) {
// Merge PNCounter using CvRDT trait
// CvRDT::merge takes &mut self and other by value, but we need to clone
let other_clone = other.inner.clone();
<PNCounter<String> as CvRDT>::merge(&mut self.inner, other_clone);
}
}
/// Thread-safe wrapper for CRDT Map
#[derive(Debug, Clone)]
pub struct SyncCRDTMap<T: Clone + Serialize + DeserializeOwned> {
inner: Arc<RwLock<CRDTMap<T>>>,
}
impl<T: Clone + Serialize + DeserializeOwned> Default for SyncCRDTMap<T> {
fn default() -> Self {
Self {
inner: Arc::new(RwLock::new(CRDTMap::new())),
}
}
}
impl<T: Clone + Serialize + DeserializeOwned> SyncCRDTMap<T> {
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, key: &SKey) -> Option<T> {
self.inner.read().get(key).cloned()
}
pub fn insert(&self, key: SKey, value: T, actor: String) {
self.inner.write().insert(key, value, actor);
}
/// Get the version and actor for a key
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner.read().get_metadata(key)
}
pub fn remove(&self, key: &SKey) {
self.inner.write().remove(key);
}
pub fn contains_key(&self, key: &SKey) -> bool {
self.inner.read().contains_key(key)
}
pub fn merge(&self, other: &CRDTMap<T>) {
self.inner.write().merge(other);
}
pub fn snapshot(&self) -> CRDTMap<T> {
self.inner.read().clone()
}
pub fn len(&self) -> usize {
self.inner.read().len()
}
pub fn is_empty(&self) -> bool {
self.inner.read().is_empty()
}
}
/// Thread-safe wrapper for PNCounter
#[derive(Debug, Clone)]
pub struct SyncPNCounter {
inner: Arc<RwLock<CRDTPNCounter>>,
}
impl Default for SyncPNCounter {
fn default() -> Self {
Self {
inner: Arc::new(RwLock::new(CRDTPNCounter::new())),
}
}
}
impl SyncPNCounter {
pub fn new() -> Self {
Self::default()
}
pub fn inc(&self, actor: String, delta: i64) {
self.inner.write().inc(actor, delta);
}
pub fn value(&self) -> i64 {
self.inner.read().value()
}
pub fn merge(&self, other: &CRDTPNCounter) {
let mut inner = self.inner.write();
inner.merge(other);
}
pub fn snapshot(&self) -> CRDTPNCounter {
self.inner.read().clone()
}
}
#[cfg(test)]
mod tests {
use std::{thread, time::Duration};
use super::*;
#[test]
fn test_crdt_pncounter_inc_and_value() {
let mut counter = CRDTPNCounter::new();
assert_eq!(counter.value(), 0);
// Test direct PNCounter usage
use crdts::{CmRDT, PNCounter};
let mut pn = PNCounter::new();
let op = pn.inc("actor1".to_string());
pn.apply(op);
let pn_val: BigInt = pn.read();
println!("Direct PNCounter value after inc(1): {:?}", pn_val);
counter.inc("actor1".to_string(), 5);
let val = counter.value();
println!("Counter value after inc(5): {}", val);
println!("Counter inner read(): {:?}", counter.inner.read());
assert!(val > 0, "Counter should be incremented, got: {}", val);
counter.inc("actor2".to_string(), 3);
let val2 = counter.value();
println!("Counter value after inc(3): {}", val2);
assert!(val2 > val, "Counter should be incremented further");
}
// SKey tests
#[test]
fn test_skey_new() {
let key = SKey::new("test_key".to_string());
assert_eq!(key.as_str(), "test_key");
}
#[test]
fn test_skey_from_string() {
let key: SKey = "test_key".to_string().into();
assert_eq!(key.as_str(), "test_key");
}
#[test]
fn test_skey_from_str() {
let key: SKey = "test_key".into();
assert_eq!(key.as_str(), "test_key");
}
#[test]
fn test_skey_ordering() {
let key1 = SKey::new("a".to_string());
let key2 = SKey::new("b".to_string());
assert!(key1 < key2);
}
// LWWRegister tests with i32
#[test]
fn test_lww_register_new() {
let reg = LWWRegister::new(42, "actor1".to_string());
assert_eq!(*reg.read(), 42);
assert_eq!(reg.actor, "actor1");
assert_eq!(reg.version, 1);
}
#[test]
fn test_lww_register_write() {
let mut reg = LWWRegister::new(42, "actor1".to_string());
let old_version = reg.version;
reg.write(100, "actor2".to_string());
assert_eq!(*reg.read(), 100);
assert_eq!(reg.actor, "actor2");
assert_eq!(reg.version, old_version + 1);
}
#[test]
fn test_lww_register_merge_newer_wins() {
let mut reg1 = LWWRegister::new(42, "actor1".to_string());
thread::sleep(Duration::from_millis(1));
let reg2 = LWWRegister::new(100, "actor2".to_string());
reg1.merge(&reg2);
assert_eq!(*reg1.read(), 100);
assert_eq!(reg1.actor, "actor2");
}
// LWWRegister tests with String
#[test]
fn test_lww_register_create_and_read() {
let reg = LWWRegister::new("value1".to_string(), "actor1".to_string());
assert_eq!(reg.read(), "value1");
assert_eq!(reg.version, 1);
assert_eq!(reg.actor, "actor1");
}
#[test]
fn test_lww_register_version_increment() {
let mut reg = LWWRegister::new("value1".to_string(), "actor1".to_string());
let initial_version = reg.version;
reg.write("value2".to_string(), "actor2".to_string());
assert_eq!(reg.version, initial_version + 1);
assert_eq!(reg.read(), "value2");
assert_eq!(reg.actor, "actor2");
}
#[test]
fn test_lww_register_merge_timestamp_priority() {
let mut reg1 = LWWRegister::new("value1".to_string(), "actor1".to_string());
thread::sleep(Duration::from_millis(10)); // Ensure different timestamp
let reg2 = LWWRegister::new("value2".to_string(), "actor2".to_string());
// reg2 has newer timestamp, should win
reg1.merge(&reg2);
assert_eq!(reg1.read(), "value2");
assert_eq!(reg1.actor, "actor2");
}
#[test]
fn test_lww_register_merge_older_loses() {
let reg1 = LWWRegister::new(42, "actor1".to_string());
thread::sleep(Duration::from_millis(1));
let reg2 = LWWRegister::new(100, "actor2".to_string());
let mut reg2_clone = reg2.clone();
reg2_clone.merge(&reg1);
assert_eq!(*reg2_clone.read(), 100);
assert_eq!(reg2_clone.actor, "actor2");
}
#[test]
fn test_lww_register_merge_version_priority() {
let mut reg1 = LWWRegister::new("value1".to_string(), "actor1".to_string());
let mut reg2 = LWWRegister::new("value2".to_string(), "actor2".to_string());
// Set same timestamp but different versions
reg2.timestamp = reg1.timestamp;
reg2.version = reg1.version + 1;
reg1.merge(&reg2);
assert_eq!(reg1.read(), "value2");
assert_eq!(reg1.version, reg2.version);
}
#[test]
fn test_lww_register_concurrent_merge() {
let mut reg1 = LWWRegister::new("value1".to_string(), "actor1".to_string());
thread::sleep(Duration::from_millis(10));
let reg2 = LWWRegister::new("value2".to_string(), "actor2".to_string());
thread::sleep(Duration::from_millis(10));
let reg3 = LWWRegister::new("value3".to_string(), "actor3".to_string());
// Merge in different orders should give same result (latest wins)
reg1.merge(&reg2);
reg1.merge(&reg3);
assert_eq!(reg1.read(), "value3");
let mut reg4 = LWWRegister::new("value1".to_string(), "actor1".to_string());
thread::sleep(Duration::from_millis(10));
let reg5 = LWWRegister::new("value2".to_string(), "actor2".to_string());
thread::sleep(Duration::from_millis(10));
let reg6 = LWWRegister::new("value3".to_string(), "actor3".to_string());
reg4.merge(&reg6);
reg4.merge(&reg5);
// reg6 should win (latest timestamp)
assert_eq!(reg4.read(), "value3");
}
// CRDTMap tests with i32
#[test]
fn test_crdt_map_new() {
let map: CRDTMap<i32> = CRDTMap::new();
assert!(map.is_empty());
assert_eq!(map.len(), 0);
}
#[test]
fn test_crdt_map_insert_get() {
let mut map = CRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), 42, "actor1".to_string());
assert_eq!(map.get(&key), Some(&42));
assert_eq!(map.len(), 1);
assert!(!map.is_empty());
}
#[test]
fn test_crdt_map_remove() {
let mut map = CRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), 42, "actor1".to_string());
assert_eq!(map.len(), 1);
map.remove(&key);
assert_eq!(map.get(&key), None);
assert_eq!(map.len(), 0);
assert!(map.is_empty());
}
#[test]
fn test_crdt_map_contains_key() {
let mut map = CRDTMap::new();
let key = SKey::new("key1".to_string());
assert!(!map.contains_key(&key));
map.insert(key.clone(), 42, "actor1".to_string());
assert!(map.contains_key(&key));
}
#[test]
fn test_crdt_map_iter() {
let mut map = CRDTMap::new();
map.insert(SKey::new("key1".to_string()), 1, "actor1".to_string());
map.insert(SKey::new("key2".to_string()), 2, "actor1".to_string());
map.insert(SKey::new("key3".to_string()), 3, "actor1".to_string());
let mut values: Vec<i32> = map.values().cloned().collect();
values.sort();
assert_eq!(values, vec![1, 2, 3]);
}
// CRDTMap tests with String
#[test]
fn test_crdt_map_insert_get_remove_string() {
let mut map = CRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), "value1".to_string(), "actor1".to_string());
assert_eq!(map.get(&key), Some(&"value1".to_string()));
assert_eq!(map.len(), 1);
map.remove(&key);
assert_eq!(map.get(&key), None);
assert_eq!(map.len(), 0);
}
#[test]
fn test_crdt_map_version_management() {
let mut map = CRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), "value1".to_string(), "actor1".to_string());
let (version1, actor1) = map.get_metadata(&key).unwrap();
assert_eq!(version1, 1);
assert_eq!(actor1, "actor1");
map.insert(key.clone(), "value2".to_string(), "actor2".to_string());
let (version2, actor2) = map.get_metadata(&key).unwrap();
assert_eq!(version2, 2);
assert_eq!(actor2, "actor2");
}
#[test]
fn test_crdt_map_merge() {
let mut map1 = CRDTMap::new();
map1.insert(SKey::new("key1".to_string()), 1, "actor1".to_string());
map1.insert(SKey::new("key2".to_string()), 2, "actor1".to_string());
let mut map2 = CRDTMap::new();
map2.insert(SKey::new("key2".to_string()), 20, "actor2".to_string());
map2.insert(SKey::new("key3".to_string()), 3, "actor2".to_string());
// Wait a bit to ensure map2 has newer timestamps
thread::sleep(Duration::from_millis(1));
map1.merge(&map2);
assert_eq!(map1.get(&SKey::new("key1".to_string())), Some(&1));
assert_eq!(map1.get(&SKey::new("key2".to_string())), Some(&20)); // Newer value wins
assert_eq!(map1.get(&SKey::new("key3".to_string())), Some(&3));
assert_eq!(map1.len(), 3);
}
#[test]
fn test_crdt_map_merge_string() {
let mut map1 = CRDTMap::new();
let mut map2 = CRDTMap::new();
let key1 = SKey::new("key1".to_string());
let key2 = SKey::new("key2".to_string());
map1.insert(key1.clone(), "value1".to_string(), "actor1".to_string());
map2.insert(key2.clone(), "value2".to_string(), "actor2".to_string());
map1.merge(&map2);
assert_eq!(map1.get(&key1), Some(&"value1".to_string()));
assert_eq!(map1.get(&key2), Some(&"value2".to_string()));
assert_eq!(map1.len(), 2);
}
#[test]
fn test_crdt_map_merge_conflict_resolution() {
let mut map1 = CRDTMap::new();
let mut map2 = CRDTMap::new();
let key = SKey::new("key1".to_string());
map1.insert(key.clone(), "value1".to_string(), "actor1".to_string());
thread::sleep(Duration::from_millis(10));
map2.insert(key.clone(), "value2".to_string(), "actor2".to_string());
// map2 has newer timestamp, should win
map1.merge(&map2);
assert_eq!(map1.get(&key), Some(&"value2".to_string()));
}
#[test]
fn test_crdt_map_to_map() {
let mut map = CRDTMap::new();
map.insert(SKey::new("key1".to_string()), 1, "actor1".to_string());
map.insert(SKey::new("key2".to_string()), 2, "actor1".to_string());
let btree_map = map.to_map();
assert_eq!(btree_map.len(), 2);
assert_eq!(btree_map.get(&SKey::new("key1".to_string())), Some(&1));
assert_eq!(btree_map.get(&SKey::new("key2".to_string())), Some(&2));
}
// CRDTPNCounter tests
#[test]
fn test_pn_counter_new() {
let counter = CRDTPNCounter::new();
assert_eq!(counter.value(), 0);
}
#[test]
fn test_pn_counter_inc_positive() {
let mut counter = CRDTPNCounter::new();
counter.inc("actor1".to_string(), 5);
// Note: PNCounter read() may require ReadCtx or have different behavior
// This test verifies the inc() method works, value() conversion may need adjustment
let val = counter.value();
// For now, just verify inc() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic and return non-negative
}
#[test]
fn test_pn_counter_inc_negative() {
let mut counter = CRDTPNCounter::new();
counter.inc("actor1".to_string(), 10);
counter.inc("actor1".to_string(), -3);
// Note: PNCounter read() may require ReadCtx or have different behavior
// This test verifies the inc() method works with negative deltas
let val = counter.value();
// For now, just verify inc() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_pn_counter_inc_dec() {
let mut counter = CRDTPNCounter::new();
assert_eq!(counter.value(), 0);
counter.inc("actor1".to_string(), 5);
assert_eq!(counter.value(), 5);
counter.inc("actor2".to_string(), 3);
assert_eq!(counter.value(), 8);
counter.inc("actor1".to_string(), -2);
assert_eq!(counter.value(), 6);
}
#[test]
fn test_pn_counter_merge() {
let mut counter1 = CRDTPNCounter::new();
counter1.inc("actor1".to_string(), 5);
let mut counter2 = CRDTPNCounter::new();
counter2.inc("actor2".to_string(), 3);
counter1.merge(&counter2);
// Note: PNCounter read() may require ReadCtx or have different behavior
// This test verifies the merge() method works
let val = counter1.value();
// For now, just verify merge() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_pn_counter_merge_exact() {
let mut counter1 = CRDTPNCounter::new();
let mut counter2 = CRDTPNCounter::new();
counter1.inc("actor1".to_string(), 10);
counter2.inc("actor2".to_string(), 5);
counter1.merge(&counter2);
assert_eq!(counter1.value(), 15);
}
#[test]
fn test_pn_counter_merge_idempotent() {
let mut counter1 = CRDTPNCounter::new();
let mut counter2 = CRDTPNCounter::new();
counter1.inc("actor1".to_string(), 10);
counter2.inc("actor1".to_string(), 10);
counter1.merge(&counter2);
// Merging same operations should not double count
assert_eq!(counter1.value(), 10);
}
#[test]
fn test_pn_counter_multiple_actors() {
let mut counter = CRDTPNCounter::new();
counter.inc("actor1".to_string(), 5);
counter.inc("actor2".to_string(), 3);
counter.inc("actor1".to_string(), -2);
// Note: PNCounter read() may require ReadCtx or have different behavior
// This test verifies multiple actors work
let val = counter.value();
// For now, just verify inc() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
// SyncCRDTMap tests
#[test]
fn test_sync_crdt_map_new() {
let map: SyncCRDTMap<i32> = SyncCRDTMap::new();
assert!(map.is_empty());
assert_eq!(map.len(), 0);
}
#[test]
fn test_sync_crdt_map_insert_get() {
let map = SyncCRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), 42, "actor1".to_string());
assert_eq!(map.get(&key), Some(42));
assert_eq!(map.len(), 1);
}
#[test]
fn test_sync_crdt_map() {
let map = SyncCRDTMap::new();
let key = SKey::new("key1".to_string());
map.insert(key.clone(), "value1".to_string(), "actor1".to_string());
assert_eq!(map.get(&key), Some("value1".to_string()));
let (version, actor) = map.get_metadata(&key).unwrap();
assert_eq!(version, 1);
assert_eq!(actor, "actor1");
}
#[test]
fn test_sync_crdt_map_concurrent_access() {
let map = Arc::new(SyncCRDTMap::new());
let mut handles = vec![];
for i in 0..10 {
let map_clone = map.clone();
let handle = thread::spawn(move || {
let key = SKey::new(format!("key{}", i));
map_clone.insert(key.clone(), i, format!("actor{}", i));
assert_eq!(map_clone.get(&key), Some(i));
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
assert_eq!(map.len(), 10);
}
#[test]
fn test_sync_crdt_map_snapshot() {
let map = SyncCRDTMap::new();
map.insert(SKey::new("key1".to_string()), 1, "actor1".to_string());
map.insert(SKey::new("key2".to_string()), 2, "actor1".to_string());
let snapshot = map.snapshot();
assert_eq!(snapshot.len(), 2);
assert_eq!(snapshot.get(&SKey::new("key1".to_string())), Some(&1));
}
#[test]
fn test_sync_crdt_map_merge() {
let map = SyncCRDTMap::new();
map.insert(SKey::new("key1".to_string()), 1, "actor1".to_string());
let mut other = CRDTMap::new();
thread::sleep(Duration::from_millis(1));
other.insert(SKey::new("key2".to_string()), 2, "actor2".to_string());
map.merge(&other);
assert_eq!(map.len(), 2);
assert_eq!(map.get(&SKey::new("key1".to_string())), Some(1));
assert_eq!(map.get(&SKey::new("key2".to_string())), Some(2));
}
// SyncPNCounter tests
#[test]
fn test_sync_pn_counter_new() {
let counter = SyncPNCounter::new();
assert_eq!(counter.value(), 0);
}
#[test]
fn test_sync_pn_counter_inc() {
let counter = SyncPNCounter::new();
counter.inc("actor1".to_string(), 5);
// Note: PNCounter read() may require ReadCtx or have different behavior
let val = counter.value();
// For now, just verify inc() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_sync_pn_counter() {
let counter = SyncPNCounter::new();
assert_eq!(counter.value(), 0);
counter.inc("actor1".to_string(), 10);
assert_eq!(counter.value(), 10);
let snapshot = counter.snapshot();
let counter2 = SyncPNCounter::new();
counter2.merge(&snapshot);
assert_eq!(counter2.value(), 10);
}
#[test]
fn test_sync_pn_counter_concurrent_access() {
let counter = Arc::new(SyncPNCounter::new());
let mut handles = vec![];
for i in 0..10 {
let counter_clone = counter.clone();
let handle = thread::spawn(move || {
counter_clone.inc(format!("actor{}", i), 1);
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
// Note: PNCounter read() may require ReadCtx or have different behavior
let val = counter.value();
// For now, just verify concurrent access doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_sync_pn_counter_concurrent() {
let counter = Arc::new(SyncPNCounter::new());
let mut handles = vec![];
for i in 0..10 {
let counter_clone = counter.clone();
let handle = thread::spawn(move || {
counter_clone.inc(format!("actor{}", i), 1);
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let val = counter.value();
assert!(val >= 0);
}
#[test]
fn test_sync_pn_counter_merge() {
let counter = SyncPNCounter::new();
counter.inc("actor1".to_string(), 5);
let mut other = CRDTPNCounter::new();
other.inc("actor2".to_string(), 3);
counter.merge(&other);
// Note: PNCounter read() may require ReadCtx or have different behavior
let val = counter.value();
// For now, just verify merge() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_sync_pn_counter_snapshot() {
let counter = SyncPNCounter::new();
counter.inc("actor1".to_string(), 5);
let snapshot = counter.snapshot();
// Note: PNCounter read() may require ReadCtx or have different behavior
let val = snapshot.value();
// For now, just verify snapshot() doesn't panic
// TODO: Fix value() method to properly read PNCounter value
assert!(val >= 0); // At minimum, should not panic
}
#[test]
fn test_sync_pn_counter_value() {
let counter = SyncPNCounter::new();
counter.inc("actor1".to_string(), 10);
assert_eq!(counter.value(), 10);
}
}

195
sgl-model-gateway/src/mesh/flow_control.rs
View File

@@ -1,195 +0,0 @@
//! Flow control for mesh cluster communication
//!
//! Provides:
//! - Backpressure control (channel capacity monitoring)
//! - Message size limits and validation
//! - Exponential backoff for reconnection
use std::{
sync::Arc,
time::{Duration, Instant},
};
use parking_lot::RwLock;
/// Maximum message size in bytes (default: 10MB)
pub const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024;
/// Channel capacity threshold for backpressure (default: 20% remaining)
pub const BACKPRESSURE_THRESHOLD: usize = 25; // 25 out of 128 = ~20%
/// Backpressure controller for managing channel capacity
#[derive(Debug, Clone)]
pub struct BackpressureController {
channel_capacity: usize,
threshold: usize,
}
impl BackpressureController {
pub fn new(channel_capacity: usize, threshold: usize) -> Self {
Self {
channel_capacity,
threshold,
}
}
/// Check if channel has capacity for sending
pub fn can_send(&self, current_len: usize) -> bool {
let remaining = self.channel_capacity.saturating_sub(current_len);
remaining > self.threshold
}
/// Get remaining capacity
pub fn remaining_capacity(&self, current_len: usize) -> usize {
self.channel_capacity.saturating_sub(current_len)
}
}
impl Default for BackpressureController {
fn default() -> Self {
Self::new(128, BACKPRESSURE_THRESHOLD)
}
}
/// Message size validator
#[derive(Debug, Clone)]
pub struct MessageSizeValidator {
max_size: usize,
}
impl MessageSizeValidator {
pub fn new(max_size: usize) -> Self {
Self { max_size }
}
/// Validate message size
pub fn validate(&self, size: usize) -> Result<(), MessageSizeError> {
if size > self.max_size {
Err(MessageSizeError::TooLarge {
size,
max: self.max_size,
})
} else {
Ok(())
}
}
/// Get maximum allowed size
pub fn max_size(&self) -> usize {
self.max_size
}
}
impl Default for MessageSizeValidator {
fn default() -> Self {
Self::new(MAX_MESSAGE_SIZE)
}
}
/// Message size validation error
#[derive(Debug, Clone)]
pub enum MessageSizeError {
TooLarge { size: usize, max: usize },
}
impl std::fmt::Display for MessageSizeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MessageSizeError::TooLarge { size, max } => {
write!(f, "Message size {} exceeds maximum {}", size, max)
}
}
}
}
impl std::error::Error for MessageSizeError {}
/// Exponential backoff calculator for reconnection
#[derive(Debug, Clone)]
pub struct ExponentialBackoff {
initial_delay: Duration,
max_delay: Duration,
multiplier: f64,
}
impl ExponentialBackoff {
pub fn new(initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
Self {
initial_delay,
max_delay,
multiplier,
}
}
/// Calculate delay for attempt number (0-indexed)
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
let delay_secs = self.initial_delay.as_secs_f64() * self.multiplier.powi(attempt as i32);
let delay = Duration::from_secs_f64(delay_secs);
delay.min(self.max_delay)
}
}
impl Default for ExponentialBackoff {
fn default() -> Self {
Self::new(Duration::from_secs(1), Duration::from_secs(60), 2.0)
}
}
/// Connection retry manager with exponential backoff
#[derive(Debug)]
pub struct RetryManager {
backoff: ExponentialBackoff,
last_attempt: Arc<RwLock<Option<Instant>>>,
attempt_count: Arc<RwLock<u32>>,
}
impl RetryManager {
pub fn new(backoff: ExponentialBackoff) -> Self {
Self {
backoff,
last_attempt: Arc::new(RwLock::new(None)),
attempt_count: Arc::new(RwLock::new(0)),
}
}
/// Check if we should retry now (based on backoff delay)
pub fn should_retry(&self) -> bool {
let last = self.last_attempt.read();
if let Some(last_attempt) = *last {
let attempt = *self.attempt_count.read();
let delay = self.backoff.delay_for_attempt(attempt);
last_attempt.elapsed() >= delay
} else {
true // First attempt
}
}
/// Record a retry attempt
pub fn record_attempt(&self) {
*self.last_attempt.write() = Some(Instant::now());
*self.attempt_count.write() += 1;
}
/// Reset retry state (on successful connection)
pub fn reset(&self) {
*self.last_attempt.write() = None;
*self.attempt_count.write() = 0;
}
/// Get current attempt count
pub fn attempt_count(&self) -> u32 {
*self.attempt_count.read()
}
/// Get next retry delay
pub fn next_delay(&self) -> Duration {
let attempt = *self.attempt_count.read();
self.backoff.delay_for_attempt(attempt)
}
}
impl Default for RetryManager {
fn default() -> Self {
Self::new(ExponentialBackoff::default())
}
}

544
sgl-model-gateway/src/mesh/incremental.rs
View File

@@ -1,544 +0,0 @@
//! Incremental update collection and batching
//!
//! Collects local state changes and batches them for efficient transmission
use std::{
collections::HashMap,
sync::Arc,
time::{SystemTime, UNIX_EPOCH},
};
use parking_lot::RwLock;
use tracing::{debug, trace};
use super::{
gossip::StateUpdate,
stores::{MembershipState, PolicyState, StateStores, StoreType, WorkerState},
SKey,
};
/// Tracks the last sent version for each key in each store
#[derive(Debug, Clone, Default)]
struct LastSentVersions {
worker: HashMap<String, u64>,
policy: HashMap<String, u64>,
app: HashMap<String, u64>,
membership: HashMap<String, u64>,
rate_limit: HashMap<String, u64>, // Track last sent timestamp for rate limit counters
}
/// Incremental update collector
pub struct IncrementalUpdateCollector {
stores: Arc<StateStores>,
self_name: String,
last_sent: Arc<RwLock<LastSentVersions>>,
}
impl IncrementalUpdateCollector {
pub fn new(stores: Arc<StateStores>, self_name: String) -> Self {
Self {
stores,
self_name,
last_sent: Arc::new(RwLock::new(LastSentVersions::default())),
}
}
/// Get current timestamp in nanoseconds
fn current_timestamp() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos() as u64
}
/// Helper function to collect updates for stores with serializable state
fn collect_serializable_updates<S>(
&self,
all_items: std::collections::BTreeMap<SKey, S>,
get_version: impl Fn(&SKey) -> u64,
last_sent_map: &mut HashMap<String, u64>,
store_name: &str,
get_id: impl Fn(&S) -> String,
) -> Vec<StateUpdate>
where
S: serde::Serialize,
{
let mut updates = Vec::new();
let timestamp = Self::current_timestamp();
for (key, state) in all_items {
let key_str = key.as_str().to_string();
let current_version = get_version(&key);
let last_sent_version = last_sent_map.get(&key_str).copied().unwrap_or(0);
if current_version > last_sent_version {
if let Ok(serialized) = serde_json::to_vec(&state) {
updates.push(StateUpdate {
key: key_str.clone(),
value: serialized,
version: current_version,
actor: self.self_name.clone(),
timestamp,
});
last_sent_map.insert(key_str, current_version);
trace!(
"Collected {} update: {} (version: {})",
store_name,
get_id(&state),
current_version
);
}
}
}
updates
}
/// Collect incremental updates for a specific store type
pub fn collect_updates_for_store(&self, store_type: StoreType) -> Vec<StateUpdate> {
let mut updates = Vec::new();
let mut last_sent = self.last_sent.write();
match store_type {
StoreType::Worker => {
let all_workers = self.stores.worker.all();
let get_version = |key: &SKey| {
self.stores
.worker
.get_metadata(key)
.map(|(v, _)| v)
.unwrap_or(0)
};
updates = self.collect_serializable_updates(
all_workers,
get_version,
&mut last_sent.worker,
"worker",
|state: &WorkerState| state.worker_id.clone(),
);
}
StoreType::Policy => {
let all_policies = self.stores.policy.all();
let get_version = |key: &SKey| {
self.stores
.policy
.get_metadata(key)
.map(|(v, _)| v)
.unwrap_or(0)
};
updates = self.collect_serializable_updates(
all_policies,
get_version,
&mut last_sent.policy,
"policy",
|state: &PolicyState| state.model_id.clone(),
);
}
StoreType::App => {
let all_apps = self.stores.app.all();
let timestamp = Self::current_timestamp();
for (key, state) in all_apps {
let key_str = key.as_str().to_string();
let current_version = self
.stores
.app
.get_metadata(&key)
.map(|(v, _)| v)
.unwrap_or(0);
let last_sent_version = last_sent.app.get(&key_str).copied().unwrap_or(0);
if current_version > last_sent_version {
updates.push(StateUpdate {
key: key_str.clone(),
value: state.value.clone(),
version: current_version,
actor: self.self_name.clone(),
timestamp,
});
last_sent.app.insert(key_str, current_version);
trace!(
"Collected app update: {} (version: {})",
state.key,
current_version
);
}
}
}
StoreType::Membership => {
let all_members = self.stores.membership.all();
let get_version = |key: &SKey| {
self.stores
.membership
.get_metadata(key)
.map(|(v, _)| v)
.unwrap_or(0)
};
updates = self.collect_serializable_updates(
all_members,
get_version,
&mut last_sent.membership,
"membership",
|state: &MembershipState| state.name.clone(),
);
}
StoreType::RateLimit => {
let rate_limit_keys = self.stores.rate_limit.keys();
let current_timestamp = Self::current_timestamp();
for key in rate_limit_keys {
if self.stores.rate_limit.is_owner(&key) {
if let Some(counter) = self.stores.rate_limit.get_counter(&key) {
let last_sent_timestamp =
last_sent.rate_limit.get(&key).copied().unwrap_or(0);
// Only send if at least 1 second has passed since last send
if current_timestamp > last_sent_timestamp + 1_000_000_000 {
if let Ok(serialized) = serde_json::to_vec(&counter.snapshot()) {
let key_str = key.clone();
updates.push(StateUpdate {
key: key_str.clone(),
value: serialized,
version: current_timestamp,
actor: self.self_name.clone(),
timestamp: current_timestamp,
});
last_sent.rate_limit.insert(key_str, current_timestamp);
trace!("Collected rate limit counter update: {}", key);
}
}
}
}
}
}
}
debug!(
"Collected {} incremental updates for store {:?}",
updates.len(),
store_type
);
updates
}
/// Collect all incremental updates across all stores
pub fn collect_all_updates(&self) -> Vec<(StoreType, Vec<StateUpdate>)> {
let mut all_updates = Vec::new();
for store_type in [
StoreType::Worker,
StoreType::Policy,
StoreType::App,
StoreType::Membership,
StoreType::RateLimit,
] {
let updates = self.collect_updates_for_store(store_type);
if !updates.is_empty() {
all_updates.push((store_type, updates));
}
}
all_updates
}
/// Mark updates as sent (called after successful transmission)
pub fn mark_sent(&self, store_type: StoreType, updates: &[StateUpdate]) {
let mut last_sent = self.last_sent.write();
let target_map = match store_type {
StoreType::Worker => &mut last_sent.worker,
StoreType::Policy => &mut last_sent.policy,
StoreType::App => &mut last_sent.app,
StoreType::Membership => &mut last_sent.membership,
StoreType::RateLimit => &mut last_sent.rate_limit,
};
for update in updates {
target_map.insert(update.key.clone(), update.version);
}
}
}
#[cfg(test)]
mod tests {
use std::{thread, time::Duration};
use super::*;
use crate::mesh::stores::{AppState, MembershipState, PolicyState, StateStores, WorkerState};
fn create_test_collector(self_name: String) -> IncrementalUpdateCollector {
let stores = Arc::new(StateStores::with_self_name(self_name.clone()));
IncrementalUpdateCollector::new(stores, self_name)
}
#[test]
fn test_collect_worker_updates() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
// Insert a worker state
let key = SKey::new("worker1".to_string());
let worker_state = WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1,
};
stores.worker.insert(key, worker_state, "node1".to_string());
// Collect updates
let updates = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates.len(), 1);
assert_eq!(updates[0].key, "worker1");
assert_eq!(updates[0].version, 1);
assert_eq!(updates[0].actor, "node1");
// Collect again - should be empty (already sent)
let updates2 = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates2.len(), 0);
// Update worker state
let key2 = SKey::new("worker1".to_string());
let worker_state2 = WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: false,
load: 0.8,
version: 2,
};
stores
.worker
.insert(key2, worker_state2, "node1".to_string());
// Should collect new version
let updates3 = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates3.len(), 1);
assert_eq!(updates3[0].version, 2);
}
#[test]
fn test_collect_policy_updates() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
let key = SKey::new("policy:model1".to_string());
let policy_state = PolicyState {
model_id: "model1".to_string(),
policy_type: "cache_aware".to_string(),
config: b"config_data".to_vec(),
version: 1,
};
stores.policy.insert(key, policy_state, "node1".to_string());
let updates = collector.collect_updates_for_store(StoreType::Policy);
assert_eq!(updates.len(), 1);
assert_eq!(updates[0].key, "policy:model1");
}
#[test]
fn test_collect_app_updates() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
let key = SKey::new("app_key1".to_string());
let app_state = AppState {
key: "app_key1".to_string(),
value: b"app_value".to_vec(),
version: 1,
};
stores.app.insert(key, app_state, "node1".to_string());
let updates = collector.collect_updates_for_store(StoreType::App);
assert_eq!(updates.len(), 1);
assert_eq!(updates[0].key, "app_key1");
}
#[test]
fn test_collect_membership_updates() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
let key = SKey::new("node2".to_string());
let membership_state = MembershipState {
name: "node2".to_string(),
address: "127.0.0.1:8001".to_string(),
status: 1, // Alive
version: 1,
metadata: std::collections::BTreeMap::new(),
};
stores
.membership
.insert(key, membership_state, "node1".to_string());
let updates = collector.collect_updates_for_store(StoreType::Membership);
assert_eq!(updates.len(), 1);
assert_eq!(updates[0].key, "node2");
}
#[test]
fn test_collect_all_updates() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
// Insert into multiple stores
let worker_key = SKey::new("worker1".to_string());
stores.worker.insert(
worker_key,
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1,
},
"node1".to_string(),
);
let policy_key = SKey::new("policy:model1".to_string());
stores.policy.insert(
policy_key,
PolicyState {
model_id: "model1".to_string(),
policy_type: "cache_aware".to_string(),
config: vec![],
version: 1,
},
"node1".to_string(),
);
let all_updates = collector.collect_all_updates();
assert_eq!(all_updates.len(), 2); // Worker and Policy
}
#[test]
fn test_mark_sent() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
// Insert and collect
let key = SKey::new("worker1".to_string());
stores.worker.insert(
key,
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1,
},
"node1".to_string(),
);
let updates = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates.len(), 1);
// Mark as sent
collector.mark_sent(StoreType::Worker, &updates);
// Should not collect again
let updates2 = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates2.len(), 0);
}
#[test]
fn test_rate_limit_timestamp_filtering() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
// Update membership to make node1 an owner
stores.rate_limit.update_membership(&["node1".to_string()]);
// Insert a counter (node1 should be owner)
let test_key = "test_rate_limit_key".to_string();
if stores.rate_limit.is_owner(&test_key) {
stores
.rate_limit
.inc(test_key.clone(), "node1".to_string(), 1);
}
// Collect immediately - should be filtered by timestamp
let _updates = collector.collect_updates_for_store(StoreType::RateLimit);
// May be empty if timestamp check fails, or may have one update
// The exact behavior depends on timing
// Wait a bit and try again
thread::sleep(Duration::from_secs(2));
// Now should collect (enough time has passed)
let updates2 = collector.collect_updates_for_store(StoreType::RateLimit);
// Should have at least one update if node1 is owner
if stores.rate_limit.is_owner(&test_key) {
// Updates may be 0 or 1 depending on timing
let _ = updates2;
}
}
#[test]
fn test_version_tracking() {
let collector = create_test_collector("node1".to_string());
let stores = collector.stores.clone();
let key = SKey::new("worker1".to_string());
// Insert first version (will be version 1 in store)
stores.worker.insert(
key.clone(),
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1, // Note: CRDT will use this but increment internally
},
"node1".to_string(),
);
let updates1 = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates1.len(), 1);
let version1 = updates1[0].version;
assert!(version1 >= 1);
// Insert second version (will increment from version1)
stores.worker.insert(
key.clone(),
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: false,
load: 0.8,
version: 2, // Note: CRDT will increment internally
},
"node1".to_string(),
);
let updates2 = collector.collect_updates_for_store(StoreType::Worker);
assert_eq!(updates2.len(), 1);
let version2 = updates2[0].version;
assert!(version2 > version1);
// Insert again - should increment version and be collected
stores.worker.insert(
key,
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.3,
version: 1, // Note: CRDT ignores this and increments internally
},
"node1".to_string(),
);
let updates3 = collector.collect_updates_for_store(StoreType::Worker);
// Should collect because version was incremented (version2 + 1 > version2)
assert_eq!(updates3.len(), 1);
let version3 = updates3[0].version;
assert!(version3 > version2);
}
}

230
sgl-model-gateway/src/mesh/metrics.rs
View File

@@ -1,230 +0,0 @@
//! Mesh cluster metrics for Prometheus
//!
//! Implements all metrics required by issue #10839:
//! - Convergence latency
//! - Traffic metrics (batches, bytes)
//! - Snapshot metrics
//! - Peer health metrics
//! - State integrity metrics
//! - Rate-limit/LB drift metrics
use std::time::{Duration, Instant};
use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram};
/// Initialize mesh metrics descriptions
pub fn init_mesh_metrics() {
// Convergence latency
describe_histogram!(
"router_mesh_convergence_ms",
"Time for state to converge across mesh in milliseconds"
);
// Traffic metrics
describe_counter!(
"router_mesh_batches_total",
"Total number of state update batches sent/received"
);
describe_counter!("router_mesh_bytes_total", "Total bytes transmitted in mesh");
// Snapshot metrics
describe_counter!(
"router_mesh_snapshot_trigger_total",
"Total number of snapshot triggers"
);
describe_histogram!(
"router_mesh_snapshot_duration_seconds",
"Time to generate and send snapshot"
);
describe_counter!(
"router_mesh_snapshot_bytes_total",
"Total bytes in snapshots"
);
// Peer health metrics
describe_gauge!(
"router_mesh_peer_connections",
"Number of active peer connections"
);
describe_counter!(
"router_mesh_peer_reconnects_total",
"Total number of peer reconnections"
);
describe_counter!("router_mesh_peer_ack_total", "Total number of ACK messages");
describe_counter!(
"router_mesh_peer_nack_total",
"Total number of NACK messages"
);
// State integrity metrics
describe_gauge!(
"router_mesh_store_cardinality",
"Number of entries in each store"
);
describe_gauge!(
"router_mesh_store_hash",
"Hash of store state for integrity checking"
);
// Rate-limit and LB drift metrics
describe_gauge!(
"router_rl_drift_ratio",
"Rate-limit drift ratio (actual vs expected)"
);
describe_gauge!(
"router_lb_drift_ratio",
"Load balance drift ratio (actual vs expected)"
);
}
/// Record convergence latency
pub fn record_convergence_latency(duration: Duration) {
histogram!("router_mesh_convergence_ms",
"quantile" => "p50"
)
.record(duration.as_millis() as f64);
}
/// Record batch transmission
pub fn record_batch_sent(peer: &str, batch_size: usize) {
counter!("router_mesh_batches_total",
"direction" => "sent",
"peer" => peer.to_string()
)
.increment(1);
counter!("router_mesh_bytes_total",
"direction" => "sent",
"peer" => peer.to_string()
)
.increment(batch_size as u64);
}
/// Record batch reception
pub fn record_batch_received(peer: &str, batch_size: usize) {
counter!("router_mesh_batches_total",
"direction" => "received",
"peer" => peer.to_string()
)
.increment(1);
counter!("router_mesh_bytes_total",
"direction" => "received",
"peer" => peer.to_string()
)
.increment(batch_size as u64);
}
/// Record snapshot trigger
pub fn record_snapshot_trigger(store: &str, reason: &str) {
counter!("router_mesh_snapshot_trigger_total",
"store" => store.to_string(),
"reason" => reason.to_string()
)
.increment(1);
}
/// Record snapshot generation duration
pub fn record_snapshot_duration(store: &str, duration: Duration) {
histogram!("router_mesh_snapshot_duration_seconds",
"store" => store.to_string()
)
.record(duration.as_secs_f64());
}
/// Record snapshot bytes
pub fn record_snapshot_bytes(store: &str, direction: &str, bytes: usize) {
counter!("router_mesh_snapshot_bytes_total",
"store" => store.to_string(),
"direction" => direction.to_string()
)
.increment(bytes as u64);
}
/// Update peer connection status
pub fn update_peer_connections(peer: &str, connected: bool) {
gauge!("router_mesh_peer_connections",
"peer" => peer.to_string()
)
.set(if connected { 1.0 } else { 0.0 });
}
/// Record peer reconnection
pub fn record_peer_reconnect(peer: &str) {
counter!("router_mesh_peer_reconnects_total",
"peer" => peer.to_string()
)
.increment(1);
}
/// Record ACK
pub fn record_ack(peer: &str, success: bool) {
let status = if success { "success" } else { "failure" };
counter!("router_mesh_peer_ack_total",
"peer" => peer.to_string(),
"status" => status.to_string()
)
.increment(1);
}
/// Record NACK
pub fn record_nack(peer: &str) {
counter!("router_mesh_peer_nack_total",
"peer" => peer.to_string()
)
.increment(1);
}
/// Update store cardinality
pub fn update_store_cardinality(store: &str, count: usize) {
gauge!("router_mesh_store_cardinality",
"store" => store.to_string()
)
.set(count as f64);
}
/// Update store hash (for integrity checking)
pub fn update_store_hash(store: &str, hash: u64) {
gauge!("router_mesh_store_hash",
"store" => store.to_string()
)
.set(hash as f64);
}
/// Update rate-limit drift ratio
pub fn update_rl_drift_ratio(key: &str, ratio: f64) {
gauge!("router_rl_drift_ratio",
"key" => key.to_string()
)
.set(ratio);
}
/// Update load balance drift ratio
pub fn update_lb_drift_ratio(model: &str, ratio: f64) {
gauge!("router_lb_drift_ratio",
"model" => model.to_string()
)
.set(ratio);
}
/// Helper struct for tracking convergence time
pub struct ConvergenceTracker {
start_time: Instant,
}
impl ConvergenceTracker {
pub fn new() -> Self {
Self {
start_time: Instant::now(),
}
}
pub fn record_convergence(&self) {
let duration = self.start_time.elapsed();
record_convergence_latency(duration);
}
}
impl Default for ConvergenceTracker {
fn default() -> Self {
Self::new()
}
}

33
sgl-model-gateway/src/mesh/mod.rs
View File

@@ -1,33 +0,0 @@
pub mod consistent_hash;
pub mod controller;
pub mod crdt;
pub mod endpoints;
pub mod flow_control;
pub mod incremental;
pub mod metrics;
pub mod mtls;
pub mod node_state_machine;
pub mod partition;
mod ping_server;
pub mod rate_limit_window;
pub mod service;
pub mod stores;
pub mod sync;
pub mod topology;
pub mod tree_ops;
#[cfg(test)]
mod test_utils;
pub use crdt::{CRDTMap, CRDTPNCounter, SKey, SyncCRDTMap, SyncPNCounter};
pub use endpoints::{
get_app_config, get_cluster_status, get_mesh_health, get_policy_state, get_policy_states,
get_worker_state, get_worker_states, trigger_graceful_shutdown, update_app_config,
};
pub use service::{broadcast_node_states, gossip, try_ping, ClusterState};
pub use stores::{
tree_state_key, AppState, AppStore, MembershipState, MembershipStore, PolicyState, PolicyStore,
RateLimitStore, StateStores, StoreType, WorkerState, WorkerStore,
};
pub use sync::{MeshSyncManager, OptionalMeshSyncManager};
pub use tree_ops::{TreeInsertOp, TreeOperation, TreeRemoveOp, TreeState};

180
sgl-model-gateway/src/mesh/mtls.rs
View File

@@ -1,180 +0,0 @@
//! mTLS (mutual TLS) support for mesh cluster communication
//!
//! Provides optional mTLS encryption for gRPC mesh connections using rustls.
//! Supports certificate rotation without restart.
use std::{
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};
use anyhow::Result;
use rustls::{
pki_types::{CertificateDer, PrivateKeyDer},
ClientConfig, RootCertStore, ServerConfig,
};
use rustls_pemfile::{certs, pkcs8_private_keys};
use tokio::{fs, sync::RwLock};
use tracing::{info, warn};
/// mTLS configuration
#[derive(Debug, Clone)]
pub struct MTLSConfig {
/// Path to CA certificate file
pub ca_cert_path: PathBuf,
/// Path to server certificate file
pub server_cert_path: PathBuf,
/// Path to server private key file
pub server_key_path: PathBuf,
/// Whether to require client certificates
pub require_client_cert: bool,
/// Certificate rotation check interval
pub rotation_check_interval: Duration,
}
impl Default for MTLSConfig {
fn default() -> Self {
Self {
ca_cert_path: PathBuf::from("/etc/ssl/certs/ca-certificates.crt"),
server_cert_path: PathBuf::from("/etc/ssl/certs/server.crt"),
server_key_path: PathBuf::from("/etc/ssl/private/server.key"),
require_client_cert: true,
rotation_check_interval: Duration::from_secs(300), // 5 minutes
}
}
}
/// mTLS certificate manager
pub struct MTLSManager {
config: MTLSConfig,
server_config: Arc<RwLock<Option<Arc<ServerConfig>>>>,
client_config: Arc<RwLock<Option<Arc<ClientConfig>>>>,
}
impl MTLSManager {
/// Create a new mTLS manager
pub fn new(config: MTLSConfig) -> Self {
Self {
config,
server_config: Arc::new(RwLock::new(None)),
client_config: Arc::new(RwLock::new(None)),
}
}
/// Load server TLS configuration
pub async fn load_server_config(&self) -> Result<Arc<ServerConfig>> {
let certs = self.load_certs(&self.config.server_cert_path).await?;
let key = self.load_private_key(&self.config.server_key_path).await?;
let mut server_config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(certs, key)?;
// Enable ALPN for HTTP/2
server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
let config = Arc::new(server_config);
*self.server_config.write().await = Some(config.clone());
Ok(config)
}
/// Load client TLS configuration
pub async fn load_client_config(&self) -> Result<Arc<ClientConfig>> {
let mut root_store = RootCertStore::empty();
// Load CA certificate
let ca_certs = self.load_certs(&self.config.ca_cert_path).await?;
for cert in ca_certs {
root_store.add(cert)?;
}
let mut client_config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
// Enable ALPN for HTTP/2
client_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
let config = Arc::new(client_config);
*self.client_config.write().await = Some(config.clone());
Ok(config)
}
/// Load certificates from file
async fn load_certs(&self, path: &Path) -> Result<Vec<CertificateDer<'static>>> {
let cert_data = fs::read(path).await?;
let certs = certs(&mut cert_data.as_slice()).collect::<Result<Vec<_>, _>>()?;
Ok(certs)
}
/// Load private key from file
async fn load_private_key(&self, path: &Path) -> Result<PrivateKeyDer<'static>> {
let key_data = fs::read(path).await?;
let mut keys =
pkcs8_private_keys(&mut key_data.as_slice()).collect::<Result<Vec<_>, _>>()?;
if keys.is_empty() {
return Err(anyhow::anyhow!("No private key found in file"));
}
Ok(PrivateKeyDer::Pkcs8(keys.remove(0)))
}
/// Start certificate rotation monitoring
pub async fn start_rotation_monitor(&self) {
let config = self.config.clone();
let server_config = self.server_config.clone();
let client_config = self.client_config.clone();
tokio::spawn(async move {
let mut interval = tokio::time::interval(config.rotation_check_interval);
loop {
interval.tick().await;
// Check if certificates have changed
if let Err(e) =
Self::check_and_reload_certs(&config, &server_config, &client_config).await
{
warn!("Error checking certificate rotation: {}", e);
}
}
});
}
/// Check and reload certificates if they have changed
async fn check_and_reload_certs(
config: &MTLSConfig,
_server_config: &Arc<RwLock<Option<Arc<ServerConfig>>>>,
_client_config: &Arc<RwLock<Option<Arc<ClientConfig>>>>,
) -> Result<()> {
// Get file modification times
let server_cert_mtime = fs::metadata(&config.server_cert_path).await?.modified()?;
let server_key_mtime = fs::metadata(&config.server_key_path).await?.modified()?;
let ca_cert_mtime = fs::metadata(&config.ca_cert_path).await?.modified()?;
// TODO: Compare with cached modification times
// For now, we'll just log that rotation monitoring is active
info!(
"Certificate rotation check: server_cert={:?}, server_key={:?}, ca_cert={:?}",
server_cert_mtime, server_key_mtime, ca_cert_mtime
);
// Reload if certificates have changed
// This is a simplified version - in production, you'd compare mtimes
Ok(())
}
/// Get current server config (for use with tonic)
pub async fn get_server_config(&self) -> Option<Arc<ServerConfig>> {
self.server_config.read().await.clone()
}
/// Get current client config (for use with tonic)
pub async fn get_client_config(&self) -> Option<Arc<ClientConfig>> {
self.client_config.read().await.clone()
}
}
/// Optional mTLS manager
pub type OptionalMTLSManager = Option<Arc<MTLSManager>>;

549
sgl-model-gateway/src/mesh/node_state_machine.rs
View File

@@ -1,549 +0,0 @@
//! Node state machine for cold start
//!
//! Manages node lifecycle: NotReady -> Joining -> SnapshotPull -> Converging -> Ready
use std::{
sync::Arc,
time::{Duration, Instant},
};
use parking_lot::RwLock;
use tracing::info;
use super::stores::StateStores;
/// Node readiness state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeReadiness {
/// Node is not ready (initial state)
NotReady,
/// Node is joining the cluster
Joining,
/// Node is pulling snapshot from peers
SnapshotPull,
/// Node is converging (applying state updates)
Converging,
/// Node is ready to serve traffic
Ready,
}
impl NodeReadiness {
pub fn as_str(&self) -> &'static str {
match self {
NodeReadiness::NotReady => "not_ready",
NodeReadiness::Joining => "joining",
NodeReadiness::SnapshotPull => "snapshot_pull",
NodeReadiness::Converging => "converging",
NodeReadiness::Ready => "ready",
}
}
}
/// Convergence detection configuration
#[derive(Debug, Clone)]
pub struct ConvergenceConfig {
/// Time window for convergence detection (seconds)
pub convergence_window: Duration,
/// Minimum number of state updates without changes to consider converged
pub min_stable_updates: usize,
/// Timeout for snapshot pull (seconds)
pub snapshot_timeout: Duration,
}
impl Default for ConvergenceConfig {
fn default() -> Self {
Self {
convergence_window: Duration::from_secs(10),
min_stable_updates: 5,
snapshot_timeout: Duration::from_secs(60),
}
}
}
/// Convergence tracker
#[derive(Debug)]
struct ConvergenceTracker {
last_update_time: Option<Instant>,
stable_update_count: usize,
last_state_hash: Option<u64>,
}
impl ConvergenceTracker {
fn new() -> Self {
Self {
last_update_time: None,
stable_update_count: 0,
last_state_hash: None,
}
}
fn record_update(&mut self, state_hash: u64, config: &ConvergenceConfig) -> bool {
let now = Instant::now();
if let Some(last_hash) = self.last_state_hash {
if last_hash == state_hash {
// State unchanged
self.stable_update_count += 1;
} else {
// State changed, reset counter
self.stable_update_count = 0;
}
} else {
// First update
self.stable_update_count = 0;
}
self.last_state_hash = Some(state_hash);
self.last_update_time = Some(now);
// Check if we've been stable long enough
if let Some(last_time) = self.last_update_time {
let elapsed = now.duration_since(last_time);
if elapsed >= config.convergence_window
&& self.stable_update_count >= config.min_stable_updates
{
return true;
}
}
false
}
fn reset(&mut self) {
self.last_update_time = None;
self.stable_update_count = 0;
self.last_state_hash = None;
}
}
/// Node state machine for managing cold start
#[derive(Debug)]
pub struct NodeStateMachine {
readiness: Arc<RwLock<NodeReadiness>>,
config: ConvergenceConfig,
convergence_tracker: Arc<RwLock<ConvergenceTracker>>,
snapshot_start_time: Arc<RwLock<Option<Instant>>>,
stores: Arc<StateStores>,
}
impl NodeStateMachine {
pub fn new(stores: Arc<StateStores>, config: ConvergenceConfig) -> Self {
Self {
readiness: Arc::new(RwLock::new(NodeReadiness::NotReady)),
config,
convergence_tracker: Arc::new(RwLock::new(ConvergenceTracker::new())),
snapshot_start_time: Arc::new(RwLock::new(None)),
stores,
}
}
/// Get current readiness state
pub fn readiness(&self) -> NodeReadiness {
*self.readiness.read()
}
/// Transition to joining state
pub fn start_joining(&self) {
let mut readiness = self.readiness.write();
if *readiness == NodeReadiness::NotReady {
*readiness = NodeReadiness::Joining;
info!("Node state: NotReady -> Joining");
}
}
/// Transition to snapshot pull state
pub fn start_snapshot_pull(&self) {
let mut readiness = self.readiness.write();
if *readiness == NodeReadiness::Joining {
*readiness = NodeReadiness::SnapshotPull;
*self.snapshot_start_time.write() = Some(Instant::now());
info!("Node state: Joining -> SnapshotPull");
}
}
/// Check if snapshot pull has timed out
pub fn is_snapshot_timeout(&self) -> bool {
if let Some(start_time) = *self.snapshot_start_time.read() {
start_time.elapsed() > self.config.snapshot_timeout
} else {
false
}
}
/// Transition to converging state
pub fn start_converging(&self) {
let mut readiness = self.readiness.write();
if *readiness == NodeReadiness::SnapshotPull {
*readiness = NodeReadiness::Converging;
*self.snapshot_start_time.write() = None;
self.convergence_tracker.write().reset();
info!("Node state: SnapshotPull -> Converging");
}
}
/// Record a state update and check for convergence
pub fn record_state_update(&self) -> bool {
if self.readiness() != NodeReadiness::Converging {
return false;
}
// Calculate a simple hash of store states
let state_hash = self.calculate_state_hash();
let mut tracker = self.convergence_tracker.write();
let converged = tracker.record_update(state_hash, &self.config);
if converged {
self.transition_to_ready();
return true;
}
false
}
/// Transition to ready state
pub fn transition_to_ready(&self) {
let mut readiness = self.readiness.write();
if *readiness == NodeReadiness::Converging {
*readiness = NodeReadiness::Ready;
info!("Node state: Converging -> Ready");
}
}
/// Check if node is ready
pub fn is_ready(&self) -> bool {
self.readiness() == NodeReadiness::Ready
}
/// Check if stores are empty (need snapshot)
pub fn needs_snapshot(&self) -> bool {
self.stores.membership.is_empty()
|| self.stores.worker.is_empty()
|| self.stores.policy.is_empty()
}
/// Calculate a simple hash of current state (for convergence detection)
fn calculate_state_hash(&self) -> u64 {
use std::{
collections::hash_map::DefaultHasher,
hash::{Hash, Hasher},
};
let mut hasher = DefaultHasher::new();
self.stores.membership.len().hash(&mut hasher);
self.stores.worker.len().hash(&mut hasher);
self.stores.policy.len().hash(&mut hasher);
self.stores.app.len().hash(&mut hasher);
hasher.finish()
}
/// Reset state machine (for testing or recovery)
pub fn reset(&self) {
*self.readiness.write() = NodeReadiness::NotReady;
self.convergence_tracker.write().reset();
*self.snapshot_start_time.write() = None;
}
}
impl Default for NodeStateMachine {
fn default() -> Self {
Self::new(
Arc::new(StateStores::default()),
ConvergenceConfig::default(),
)
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
fn create_test_stores() -> Arc<StateStores> {
Arc::new(StateStores::default())
}
fn create_test_config() -> ConvergenceConfig {
ConvergenceConfig {
convergence_window: Duration::from_millis(100),
min_stable_updates: 3,
snapshot_timeout: Duration::from_secs(1),
}
}
#[test]
fn test_node_readiness_as_str() {
assert_eq!(NodeReadiness::NotReady.as_str(), "not_ready");
assert_eq!(NodeReadiness::Joining.as_str(), "joining");
assert_eq!(NodeReadiness::SnapshotPull.as_str(), "snapshot_pull");
assert_eq!(NodeReadiness::Converging.as_str(), "converging");
assert_eq!(NodeReadiness::Ready.as_str(), "ready");
}
#[test]
fn test_convergence_config_default() {
let config = ConvergenceConfig::default();
assert_eq!(config.convergence_window, Duration::from_secs(10));
assert_eq!(config.min_stable_updates, 5);
assert_eq!(config.snapshot_timeout, Duration::from_secs(60));
}
#[test]
fn test_node_state_machine_initial_state() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores, config);
assert_eq!(sm.readiness(), NodeReadiness::NotReady);
assert!(!sm.is_ready());
}
#[test]
fn test_state_transition_flow() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores, config);
// Start joining
sm.start_joining();
assert_eq!(sm.readiness(), NodeReadiness::Joining);
// Start snapshot pull
sm.start_snapshot_pull();
assert_eq!(sm.readiness(), NodeReadiness::SnapshotPull);
assert!(!sm.is_snapshot_timeout());
// Start converging
sm.start_converging();
assert_eq!(sm.readiness(), NodeReadiness::Converging);
// Transition to ready
sm.transition_to_ready();
assert_eq!(sm.readiness(), NodeReadiness::Ready);
assert!(sm.is_ready());
}
#[test]
fn test_state_transition_guards() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores, config);
// Cannot start snapshot pull without joining first
sm.start_snapshot_pull();
assert_eq!(sm.readiness(), NodeReadiness::NotReady);
// Cannot start converging without snapshot pull
sm.start_joining();
sm.start_converging();
assert_eq!(sm.readiness(), NodeReadiness::Joining);
// Cannot transition to ready without converging
sm.transition_to_ready();
assert_eq!(sm.readiness(), NodeReadiness::Joining);
}
#[test]
fn test_snapshot_timeout() {
let stores = create_test_stores();
let mut config = create_test_config();
config.snapshot_timeout = Duration::from_millis(50);
let sm = NodeStateMachine::new(stores, config);
sm.start_joining();
sm.start_snapshot_pull();
assert!(!sm.is_snapshot_timeout());
// Wait for timeout
std::thread::sleep(Duration::from_millis(100));
assert!(sm.is_snapshot_timeout());
}
#[test]
fn test_needs_snapshot() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores.clone(), config);
// Empty stores need snapshot
assert!(sm.needs_snapshot());
// Add some data to stores
use super::super::{
crdt::SKey,
stores::{MembershipState, PolicyState, WorkerState},
};
stores.membership.insert(
SKey::from("node1"),
MembershipState {
name: "node1".to_string(),
address: "127.0.0.1:8080".to_string(),
status: 1,
version: 1,
metadata: Default::default(),
},
"test".to_string(),
);
stores.worker.insert(
SKey::from("worker1"),
WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1,
},
"test".to_string(),
);
stores.policy.insert(
SKey::from("policy1"),
PolicyState {
model_id: "model1".to_string(),
policy_type: "round_robin".to_string(),
config: vec![],
version: 1,
},
"test".to_string(),
);
// Now should not need snapshot
assert!(!sm.needs_snapshot());
}
#[test]
fn test_record_state_update_not_converging() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores, config);
// Should return false when not in converging state
assert!(!sm.record_state_update());
assert_eq!(sm.readiness(), NodeReadiness::NotReady);
}
#[test]
fn test_convergence_detection() {
let stores = create_test_stores();
let mut config = create_test_config();
config.convergence_window = Duration::from_millis(50);
config.min_stable_updates = 2;
let sm = NodeStateMachine::new(stores, config);
// Transition to converging state
sm.start_joining();
sm.start_snapshot_pull();
sm.start_converging();
assert_eq!(sm.readiness(), NodeReadiness::Converging);
// Record multiple updates with same state
let converged1 = sm.record_state_update();
assert!(!converged1);
// Wait a bit and record more updates
std::thread::sleep(Duration::from_millis(60));
let converged2 = sm.record_state_update();
assert!(!converged2); // Still not enough stable updates
// Record more stable updates
std::thread::sleep(Duration::from_millis(10));
let converged3 = sm.record_state_update();
// Should converge after enough stable updates within window
if converged3 {
assert_eq!(sm.readiness(), NodeReadiness::Ready);
}
}
#[test]
fn test_convergence_reset_on_state_change() {
let stores = create_test_stores();
let mut config = create_test_config();
config.convergence_window = Duration::from_millis(100);
config.min_stable_updates = 2;
let sm = NodeStateMachine::new(stores.clone(), config);
sm.start_joining();
sm.start_snapshot_pull();
sm.start_converging();
// Record update
sm.record_state_update();
// Change state by adding data
use super::super::{crdt::SKey, stores::AppState};
stores.app.insert(
SKey::from("app1"),
AppState {
key: "app1".to_string(),
value: vec![1, 2, 3],
version: 1,
},
"test".to_string(),
);
// Record update with changed state
sm.record_state_update();
// The stable count should be reset
std::thread::sleep(Duration::from_millis(110));
let converged = sm.record_state_update();
// Should not converge immediately after state change
assert!(!converged || sm.readiness() == NodeReadiness::Converging);
}
#[test]
fn test_reset() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores, config);
// Go through states
sm.start_joining();
sm.start_snapshot_pull();
sm.start_converging();
sm.transition_to_ready();
assert_eq!(sm.readiness(), NodeReadiness::Ready);
// Reset
sm.reset();
assert_eq!(sm.readiness(), NodeReadiness::NotReady);
assert!(!sm.is_ready());
assert!(!sm.is_snapshot_timeout());
}
#[test]
fn test_calculate_state_hash() {
let stores = create_test_stores();
let config = create_test_config();
let sm = NodeStateMachine::new(stores.clone(), config);
let hash1 = sm.calculate_state_hash();
// Add some data
use super::super::{crdt::SKey, stores::AppState};
stores.app.insert(
SKey::from("app1"),
AppState {
key: "app1".to_string(),
value: vec![],
version: 1,
},
"test".to_string(),
);
// Hash should change
let hash2 = sm.calculate_state_hash();
assert_ne!(hash1, hash2);
}
#[test]
fn test_default_implementation() {
let sm = NodeStateMachine::default();
assert_eq!(sm.readiness(), NodeReadiness::NotReady);
assert!(!sm.is_ready());
}
}

507
sgl-model-gateway/src/mesh/partition.rs
View File

@@ -1,507 +0,0 @@
//! Partition detection and handling
//!
//! Detects network partitions and handles state isolation and recovery
use std::{
collections::{BTreeMap, HashSet},
sync::Arc,
time::{Duration, Instant},
};
use parking_lot::RwLock;
use tracing::warn;
use super::gossip::{NodeState, NodeStatus};
/// Partition detection configuration
#[derive(Debug, Clone)]
pub struct PartitionConfig {
/// Timeout for considering a node unreachable (seconds)
pub unreachable_timeout: Duration,
/// Minimum cluster size to consider a partition
pub min_cluster_size: usize,
/// Quorum threshold (minimum nodes needed for quorum)
pub quorum_threshold: usize,
}
impl Default for PartitionConfig {
fn default() -> Self {
Self {
unreachable_timeout: Duration::from_secs(30),
min_cluster_size: 3,
quorum_threshold: 2,
}
}
}
/// Partition state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PartitionState {
/// Normal operation, no partition detected
Normal,
/// Partition detected, but we have quorum
PartitionedWithQuorum,
/// Partition detected, we don't have quorum
PartitionedWithoutQuorum,
}
/// Partition detector
#[derive(Debug)]
pub struct PartitionDetector {
config: PartitionConfig,
last_seen: Arc<RwLock<BTreeMap<String, Instant>>>,
current_state: Arc<RwLock<PartitionState>>,
}
impl PartitionDetector {
pub fn new(config: PartitionConfig) -> Self {
Self {
config,
last_seen: Arc::new(RwLock::new(BTreeMap::new())),
current_state: Arc::new(RwLock::new(PartitionState::Normal)),
}
}
/// Update last seen time for a node
pub fn update_last_seen(&self, node_name: &str) {
let mut last_seen = self.last_seen.write();
last_seen.insert(node_name.to_string(), Instant::now());
}
/// Detect partition based on current cluster state
pub fn detect_partition(&self, cluster_state: &BTreeMap<String, NodeState>) -> PartitionState {
let now = Instant::now();
let last_seen = self.last_seen.read();
// Count alive nodes and unreachable nodes
let mut alive_count = 0;
let mut unreachable_count = 0;
let mut reachable_nodes = HashSet::new();
for (name, node) in cluster_state.iter() {
if node.status == NodeStatus::Alive as i32 {
alive_count += 1;
// Check if we've seen this node recently
if let Some(last_seen_time) = last_seen.get(name) {
if now.duration_since(*last_seen_time) < self.config.unreachable_timeout {
reachable_nodes.insert(name.clone());
} else {
unreachable_count += 1;
warn!(
"Node {} unreachable for {:?}",
name,
now.duration_since(*last_seen_time)
);
}
} else {
// New node, consider it reachable for now
reachable_nodes.insert(name.clone());
}
}
}
let reachable_count = reachable_nodes.len();
// Determine partition state
let state = if unreachable_count == 0 {
PartitionState::Normal
} else if reachable_count >= self.config.quorum_threshold {
PartitionState::PartitionedWithQuorum
} else {
PartitionState::PartitionedWithoutQuorum
};
// Update current state
*self.current_state.write() = state.clone();
if state != PartitionState::Normal {
warn!(
"Partition detected: state={:?}, reachable={}, unreachable={}, total_alive={}",
state, reachable_count, unreachable_count, alive_count
);
}
state
}
/// Get current partition state
pub fn current_state(&self) -> PartitionState {
self.current_state.read().clone()
}
/// Check if we have quorum
pub fn has_quorum(&self, reachable_count: usize) -> bool {
reachable_count >= self.config.quorum_threshold
}
/// Get unreachable nodes
pub fn get_unreachable_nodes(
&self,
cluster_state: &BTreeMap<String, NodeState>,
) -> Vec<String> {
let now = Instant::now();
let last_seen = self.last_seen.read();
let mut unreachable = Vec::new();
for (name, node) in cluster_state.iter() {
if node.status == NodeStatus::Alive as i32 {
if let Some(last_seen_time) = last_seen.get(name) {
if now.duration_since(*last_seen_time) >= self.config.unreachable_timeout {
unreachable.push(name.clone());
}
}
}
}
unreachable
}
/// Check if we should continue serving (have quorum)
pub fn should_serve(&self) -> bool {
let state = self.current_state.read();
matches!(
*state,
PartitionState::Normal | PartitionState::PartitionedWithQuorum
)
}
}
impl Default for PartitionDetector {
fn default() -> Self {
Self::new(PartitionConfig::default())
}
}
#[cfg(test)]
mod tests {
use std::{collections::BTreeMap, time::Duration};
use super::*;
// Import NodeState and NodeStatus from gossip module
use crate::mesh::service::gossip::{NodeState, NodeStatus};
fn create_test_config() -> PartitionConfig {
PartitionConfig {
unreachable_timeout: Duration::from_millis(100),
min_cluster_size: 3,
quorum_threshold: 2,
}
}
fn create_node_state(name: &str, address: &str, status: NodeStatus) -> NodeState {
NodeState {
name: name.to_string(),
address: address.to_string(),
status: status as i32,
version: 1,
metadata: std::collections::HashMap::new(),
}
}
#[test]
fn test_partition_config_default() {
let config = PartitionConfig::default();
assert_eq!(config.unreachable_timeout, Duration::from_secs(30));
assert_eq!(config.min_cluster_size, 3);
assert_eq!(config.quorum_threshold, 2);
}
#[test]
fn test_partition_detector_initial_state() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
assert_eq!(detector.current_state(), PartitionState::Normal);
assert!(detector.should_serve());
}
#[test]
fn test_update_last_seen() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
detector.update_last_seen("node1");
detector.update_last_seen("node2");
// Verify nodes are tracked
let cluster_state = BTreeMap::new();
let state = detector.detect_partition(&cluster_state);
assert_eq!(state, PartitionState::Normal);
}
#[test]
fn test_detect_partition_normal() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Alive),
);
cluster_state.insert(
"node3".to_string(),
create_node_state("node3", "127.0.0.1:8082", NodeStatus::Alive),
);
// Update last seen for all nodes
detector.update_last_seen("node1");
detector.update_last_seen("node2");
detector.update_last_seen("node3");
let state = detector.detect_partition(&cluster_state);
assert_eq!(state, PartitionState::Normal);
assert!(detector.should_serve());
}
#[test]
fn test_detect_partition_with_quorum() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Alive),
);
cluster_state.insert(
"node3".to_string(),
create_node_state("node3", "127.0.0.1:8082", NodeStatus::Alive),
);
// Update last seen for node1 and node2 (quorum)
detector.update_last_seen("node1");
detector.update_last_seen("node2");
// Don't update node3, but wait for it to be considered unreachable
// Since node3 is new, it's initially considered reachable
// We need to update it first, then wait for timeout
detector.update_last_seen("node3");
std::thread::sleep(Duration::from_millis(150));
// Update node1 and node2 again to keep them reachable
detector.update_last_seen("node1");
detector.update_last_seen("node2");
let state = detector.detect_partition(&cluster_state);
// node1 and node2 are still reachable (quorum of 2), node3 is unreachable
assert_eq!(state, PartitionState::PartitionedWithQuorum);
assert!(detector.should_serve());
}
#[test]
fn test_detect_partition_without_quorum() {
let mut config = create_test_config();
config.quorum_threshold = 2;
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Alive),
);
cluster_state.insert(
"node3".to_string(),
create_node_state("node3", "127.0.0.1:8082", NodeStatus::Alive),
);
// Update last seen for all nodes first
detector.update_last_seen("node1");
detector.update_last_seen("node2");
detector.update_last_seen("node3");
// Wait for node2 and node3 to become unreachable
std::thread::sleep(Duration::from_millis(150));
// Only update node1 again to keep it reachable
detector.update_last_seen("node1");
let state = detector.detect_partition(&cluster_state);
// Only node1 is reachable (below quorum of 2)
assert_eq!(state, PartitionState::PartitionedWithoutQuorum);
assert!(!detector.should_serve());
}
#[test]
fn test_has_quorum() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
assert!(detector.has_quorum(2));
assert!(detector.has_quorum(3));
assert!(!detector.has_quorum(1));
assert!(!detector.has_quorum(0));
}
#[test]
fn test_get_unreachable_nodes() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Alive),
);
cluster_state.insert(
"node3".to_string(),
create_node_state("node3", "127.0.0.1:8082", NodeStatus::Alive),
);
// Update last seen for all nodes
detector.update_last_seen("node1");
detector.update_last_seen("node2");
detector.update_last_seen("node3");
// Initially no unreachable nodes
let unreachable = detector.get_unreachable_nodes(&cluster_state);
assert!(unreachable.is_empty());
// Wait for timeout
std::thread::sleep(Duration::from_millis(150));
// All nodes should be unreachable now
let unreachable = detector.get_unreachable_nodes(&cluster_state);
assert_eq!(unreachable.len(), 3);
assert!(unreachable.contains(&"node1".to_string()));
assert!(unreachable.contains(&"node2".to_string()));
assert!(unreachable.contains(&"node3".to_string()));
}
#[test]
fn test_get_unreachable_nodes_with_recent_updates() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Alive),
);
// Update node1 first (old)
detector.update_last_seen("node1");
std::thread::sleep(Duration::from_millis(50));
// Update node2 later (more recent)
detector.update_last_seen("node2");
// Wait for node1 to timeout but node2 should still be reachable
std::thread::sleep(Duration::from_millis(60));
let unreachable = detector.get_unreachable_nodes(&cluster_state);
// node1 should be unreachable (updated 110ms ago), node2 should still be reachable (updated 60ms ago)
assert!(unreachable.contains(&"node1".to_string()));
assert!(!unreachable.contains(&"node2".to_string()));
}
#[test]
fn test_detect_partition_ignores_non_alive_nodes() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"node2".to_string(),
create_node_state("node2", "127.0.0.1:8081", NodeStatus::Down),
);
cluster_state.insert(
"node3".to_string(),
create_node_state("node3", "127.0.0.1:8082", NodeStatus::Suspected),
);
detector.update_last_seen("node1");
let state = detector.detect_partition(&cluster_state);
// Only node1 is alive and reachable
// Since node2 and node3 are not alive, they don't count as unreachable
// If all alive nodes are reachable (unreachable_count == 0), state is Normal
assert_eq!(state, PartitionState::Normal);
}
#[test]
fn test_new_node_considered_reachable() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
let mut cluster_state = BTreeMap::new();
cluster_state.insert(
"node1".to_string(),
create_node_state("node1", "127.0.0.1:8080", NodeStatus::Alive),
);
cluster_state.insert(
"new_node".to_string(),
create_node_state("new_node", "127.0.0.1:8083", NodeStatus::Alive),
);
// Don't update last_seen for new_node, it should be considered reachable
detector.update_last_seen("node1");
let state = detector.detect_partition(&cluster_state);
// Both nodes should be considered reachable (node1 explicitly, new_node by default)
assert_eq!(state, PartitionState::Normal);
}
#[test]
fn test_should_serve() {
let config = create_test_config();
let detector = PartitionDetector::new(config);
// Normal state should serve
*detector.current_state.write() = PartitionState::Normal;
assert!(detector.should_serve());
// Partitioned with quorum should serve
*detector.current_state.write() = PartitionState::PartitionedWithQuorum;
assert!(detector.should_serve());
// Partitioned without quorum should not serve
*detector.current_state.write() = PartitionState::PartitionedWithoutQuorum;
assert!(!detector.should_serve());
}
#[test]
fn test_default_implementation() {
let detector = PartitionDetector::default();
assert_eq!(detector.current_state(), PartitionState::Normal);
assert!(detector.should_serve());
}
#[test]
fn test_partition_state_equality() {
assert_eq!(PartitionState::Normal, PartitionState::Normal);
assert_ne!(
PartitionState::Normal,
PartitionState::PartitionedWithQuorum
);
assert_ne!(
PartitionState::PartitionedWithQuorum,
PartitionState::PartitionedWithoutQuorum
);
}
}

964
sgl-model-gateway/src/mesh/ping_server.rs
View File

@@ -1,964 +0,0 @@
use std::{
net::SocketAddr,
pin::Pin,
sync::Arc,
time::{Duration, Instant},
};
use anyhow::Result;
use futures::Stream;
use tokio_stream::StreamExt;
use tonic::{transport::Server, Response, Status};
use tracing as log;
use tracing::instrument;
use super::{
crdt::SKey,
flow_control::MessageSizeValidator,
gossip::{
self,
gossip_server::{Gossip, GossipServer},
GossipMessage, IncrementalUpdate, NodeState, NodeStatus, NodeUpdate, PingReq,
SnapshotChunk, SnapshotRequest, StateUpdate, StreamAck, StreamMessage, StreamMessageType,
},
incremental::IncrementalUpdateCollector,
metrics::{
record_ack, record_batch_sent, record_nack, record_peer_reconnect, record_snapshot_bytes,
record_snapshot_duration, record_snapshot_trigger, update_peer_connections,
ConvergenceTracker,
},
node_state_machine::NodeStateMachine,
partition::PartitionDetector,
stores::{StateStores, StoreType as LocalStoreType},
sync::MeshSyncManager,
try_ping, ClusterState,
};
#[derive(Debug)]
pub struct GossipService {
state: ClusterState,
self_addr: SocketAddr,
self_name: String,
stores: Option<Arc<StateStores>>, // Optional state stores for CRDT-based sync
sync_manager: Option<Arc<MeshSyncManager>>, // Optional sync manager for applying remote updates
state_machine: Option<Arc<NodeStateMachine>>,
partition_detector: Option<Arc<PartitionDetector>>,
}
impl GossipService {
/// Create snapshot chunks for a store
async fn create_snapshot_chunks(
&self,
store_type: LocalStoreType,
chunk_size: usize,
) -> Vec<SnapshotChunk> {
let stores = match self.stores.as_ref() {
Some(s) => s,
None => {
log::warn!("State stores not available for snapshot generation");
return vec![];
}
};
let proto_store_type = match store_type {
LocalStoreType::Membership => gossip::StoreType::Membership as i32,
LocalStoreType::App => gossip::StoreType::App as i32,
LocalStoreType::Worker => gossip::StoreType::Worker as i32,
LocalStoreType::Policy => gossip::StoreType::Policy as i32,
LocalStoreType::RateLimit => gossip::StoreType::RateLimit as i32,
};
// Get all entries from the store
let entries: Vec<(SKey, Vec<u8>)> = match store_type {
LocalStoreType::Membership => stores
.membership
.all()
.into_iter()
.map(|(k, v)| {
let serialized = serde_json::to_vec(&v).unwrap_or_else(|e| {
log::error!("Failed to serialize membership state: {}", e);
vec![]
});
(k, serialized)
})
.collect(),
LocalStoreType::App => stores
.app
.all()
.into_iter()
.map(|(k, v)| {
let serialized = serde_json::to_vec(&v).unwrap_or_else(|e| {
log::error!("Failed to serialize app state: {}", e);
vec![]
});
(k, serialized)
})
.collect(),
LocalStoreType::Worker => stores
.worker
.all()
.into_iter()
.map(|(k, v)| {
let serialized = serde_json::to_vec(&v).unwrap_or_else(|e| {
log::error!("Failed to serialize worker state: {}", e);
vec![]
});
(k, serialized)
})
.collect(),
LocalStoreType::Policy => stores
.policy
.all()
.into_iter()
.map(|(k, v)| {
let serialized = serde_json::to_vec(&v).unwrap_or_else(|e| {
log::error!("Failed to serialize policy state: {}", e);
vec![]
});
(k, serialized)
})
.collect(),
LocalStoreType::RateLimit => {
// For rate limit, serialize all counters from owners
stores
.rate_limit
.keys()
.into_iter()
.filter_map(|key| {
if stores.rate_limit.is_owner(&key) {
stores.rate_limit.get_counter(&key).map(|counter| {
let serialized = serde_json::to_vec(&counter.snapshot())
.unwrap_or_else(|e| {
log::error!(
"Failed to serialize rate limit counter: {}",
e
);
vec![]
});
(SKey::new(key.clone()), serialized)
})
} else {
None
}
})
.collect()
}
};
if entries.is_empty() {
return vec![];
}
// Split entries into chunks
let mut chunks = Vec::new();
let total_chunks = entries.len().div_ceil(chunk_size);
for (chunk_idx, chunk_entries) in entries.chunks(chunk_size).enumerate() {
let state_updates: Vec<StateUpdate> = chunk_entries
.iter()
.map(|(key, value)| {
// Get actual version from CRDT metadata
let version = match store_type {
LocalStoreType::Membership => stores
.membership
.get_metadata(key)
.map(|(v, _)| v)
.unwrap_or(1),
LocalStoreType::App => {
stores.app.get_metadata(key).map(|(v, _)| v).unwrap_or(1)
}
LocalStoreType::Worker => {
stores.worker.get_metadata(key).map(|(v, _)| v).unwrap_or(1)
}
LocalStoreType::Policy => {
stores.policy.get_metadata(key).map(|(v, _)| v).unwrap_or(1)
}
LocalStoreType::RateLimit => {
// For rate limit, use timestamp as version
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64
}
};
StateUpdate {
key: key.as_str().to_string(),
value: value.clone(),
version,
actor: self.self_name.clone(),
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64,
}
})
.collect();
// Calculate checksum for integrity verification
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
for update in &state_updates {
update.key.hash(&mut hasher);
update.value.hash(&mut hasher);
}
let checksum = hasher.finish().to_le_bytes().to_vec();
chunks.push(SnapshotChunk {
store: proto_store_type,
chunk_index: chunk_idx as u64,
total_chunks: total_chunks as u64,
entries: state_updates,
checksum,
});
}
log::info!(
"Generated {} snapshot chunks for store {:?}",
chunks.len(),
store_type
);
chunks
}
}
impl GossipService {
pub fn new(state: ClusterState, self_addr: SocketAddr, self_name: &str) -> Self {
Self {
state,
self_addr,
self_name: self_name.to_string(),
stores: None,
sync_manager: None,
state_machine: None,
partition_detector: None,
}
}
pub fn with_stores(mut self, stores: Arc<StateStores>) -> Self {
self.stores = Some(stores.clone());
// Create state machine if stores are provided
if self.state_machine.is_none() {
use super::node_state_machine::ConvergenceConfig;
self.state_machine = Some(Arc::new(NodeStateMachine::new(
stores,
ConvergenceConfig::default(),
)));
}
self
}
pub fn with_sync_manager(mut self, sync_manager: Arc<MeshSyncManager>) -> Self {
self.sync_manager = Some(sync_manager);
self
}
pub fn with_partition_detector(mut self, partition_detector: Arc<PartitionDetector>) -> Self {
self.partition_detector = Some(partition_detector);
self
}
pub async fn serve_ping_with_shutdown<F: std::future::Future<Output = ()>>(
self,
signal: F,
) -> Result<()> {
let listen_addr = self.self_addr;
let service = GossipServer::new(self);
Server::builder()
.add_service(service)
.serve_with_shutdown(listen_addr, signal)
.await?;
Ok(())
}
async fn merge_state(&self, incoming_nodes: Vec<NodeState>) -> bool {
let mut state = self.state.write();
let mut updated = false;
for node in incoming_nodes {
state
.entry(node.name.clone())
.and_modify(|entry| {
if node.version > entry.version {
*entry = node.clone();
updated = true;
}
})
.or_insert_with(|| {
updated = true;
node
});
}
if updated {
log::info!("Cluster state updated. Current nodes: {}", state.len());
}
updated
}
}
#[tonic::async_trait]
impl Gossip for GossipService {
type SyncStreamStream =
Pin<Box<dyn Stream<Item = Result<StreamMessage, Status>> + Send + 'static>>;
#[instrument(fields(name = %self.self_name), skip(self, request))]
async fn ping_server(
&self,
request: tonic::Request<GossipMessage>,
) -> std::result::Result<Response<NodeUpdate>, Status> {
let message = request.into_inner();
match message.payload {
Some(gossip::gossip_message::Payload::Ping(ping)) => {
log::info!("Received {:?}", ping);
if let Some(stat_sync) = ping.state_sync {
log::info!("Merging state from Ping: {} nodes", stat_sync.nodes.len());
self.merge_state(stat_sync.nodes).await;
}
// Return current status of self node (could be Alive or Leaving)
let current_status = {
let state = self.state.read();
state
.get(&self.self_name)
.map(|n| n.status)
.unwrap_or(NodeStatus::Alive as i32)
};
Ok(Response::new(NodeUpdate {
name: self.self_name.clone(),
address: self.self_addr.to_string(),
status: current_status,
}))
}
Some(gossip::gossip_message::Payload::PingReq(PingReq { node: Some(node) })) => {
log::info!("PingReq to node {} addr:{}", node.name, node.address);
let res = try_ping(&node, None).await?;
Ok(Response::new(res))
}
_ => Err(Status::invalid_argument("Invalid message payload")),
}
}
#[instrument(fields(name = %self.self_name), skip(self, request))]
async fn sync_stream(
&self,
request: tonic::Request<tonic::Streaming<StreamMessage>>,
) -> Result<Response<Self::SyncStreamStream>, Status> {
let mut incoming = request.into_inner();
let self_name = self.self_name.clone();
let state = self.state.clone();
let stores = self.stores.clone();
let sync_manager = self.sync_manager.clone();
// Create output stream with flow control
const CHANNEL_CAPACITY: usize = 128;
let (tx, rx) =
tokio::sync::mpsc::channel::<Result<StreamMessage, Status>>(CHANNEL_CAPACITY);
let size_validator = MessageSizeValidator::default();
// Create incremental update collector if stores are available
let collector = stores.as_ref().map(|stores| {
Arc::new(IncrementalUpdateCollector::new(
stores.clone(),
self_name.clone(),
))
});
// Spawn task to periodically send incremental updates
if let Some(collector) = collector {
let tx_incremental = tx.clone();
let self_name_incremental = self_name.clone();
let size_validator_clone = size_validator.clone();
tokio::spawn(async move {
// Use 1 second interval for rate limit counter sync (faster than other stores)
let mut interval = tokio::time::interval(Duration::from_secs(1)); // Send every 1 second
let mut sequence_counter: u64 = 0;
loop {
interval.tick().await;
// Collect all incremental updates
let all_updates = collector.collect_all_updates();
if !all_updates.is_empty() {
for (store_type, updates) in all_updates {
let proto_store_type = match store_type {
LocalStoreType::Membership => gossip::StoreType::Membership as i32,
LocalStoreType::App => gossip::StoreType::App as i32,
LocalStoreType::Worker => gossip::StoreType::Worker as i32,
LocalStoreType::Policy => gossip::StoreType::Policy as i32,
LocalStoreType::RateLimit => gossip::StoreType::RateLimit as i32,
};
sequence_counter += 1;
let batch_size: usize = updates.iter().map(|u| u.value.len()).sum();
// Validate message size
if let Err(e) = size_validator_clone.validate(batch_size) {
log::warn!(
"Incremental update too large, skipping: {} (max: {} bytes)",
e,
size_validator_clone.max_size()
);
continue;
}
let incremental_update = StreamMessage {
message_type: StreamMessageType::IncrementalUpdate as i32,
payload: Some(gossip::stream_message::Payload::Incremental(
IncrementalUpdate {
store: proto_store_type,
updates: updates.clone(),
version: 0, // Version is tracked per key in StateUpdate
},
)),
sequence: sequence_counter,
peer_id: self_name_incremental.clone(),
};
// Check backpressure using try_send (mpsc::Sender doesn't have len())
match tx_incremental.try_send(Ok(incremental_update)) {
Ok(_) => {
// Successfully queued
// Record metrics
record_batch_sent(&self_name_incremental, batch_size);
// Mark as sent after successful transmission
collector.mark_sent(store_type, &updates);
}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
log::debug!(
"Backpressure: channel full, skipping send (will retry next interval)"
);
// Don't mark as sent, will retry next interval
continue;
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
log::warn!(
"Channel closed, stopping incremental update sender"
);
break;
}
}
log::debug!(
"Sent incremental update: store={:?}, {} updates",
store_type,
updates.len()
);
}
}
}
});
}
// Spawn task to handle incoming messages
let mut sequence: u64 = 0;
let _convergence_tracker = ConvergenceTracker::new();
// Track snapshot reception state: (store_type, total_chunks) -> received_chunks
use std::collections::HashMap;
let mut snapshot_state: HashMap<(LocalStoreType, u64), Vec<SnapshotChunk>> = HashMap::new();
tokio::spawn(async move {
let mut peer_id = String::new();
update_peer_connections(&peer_id, true);
// Check if we need to request snapshots on connection
// This happens when:
// 1. We're a new node joining (stores are empty or very small)
// 2. We detect a version gap
if let Some(ref stores) = stores {
for store_type in [
LocalStoreType::Membership,
LocalStoreType::App,
LocalStoreType::Worker,
LocalStoreType::Policy,
LocalStoreType::RateLimit,
] {
let store_len = match store_type {
LocalStoreType::Membership => stores.membership.len(),
LocalStoreType::App => stores.app.len(),
LocalStoreType::Worker => stores.worker.len(),
LocalStoreType::Policy => stores.policy.len(),
LocalStoreType::RateLimit => stores.rate_limit.keys().len(),
};
// If store is empty or very small, request snapshot
if store_len == 0 {
log::info!(
"Store {:?} is empty, requesting snapshot from {}",
store_type,
peer_id
);
let proto_store_type = match store_type {
LocalStoreType::Membership => gossip::StoreType::Membership as i32,
LocalStoreType::App => gossip::StoreType::App as i32,
LocalStoreType::Worker => gossip::StoreType::Worker as i32,
LocalStoreType::Policy => gossip::StoreType::Policy as i32,
LocalStoreType::RateLimit => gossip::StoreType::RateLimit as i32,
};
let snapshot_request = StreamMessage {
message_type: StreamMessageType::SnapshotRequest as i32,
payload: Some(gossip::stream_message::Payload::SnapshotRequest(
SnapshotRequest {
store: proto_store_type,
from_version: 0, // Request from beginning
},
)),
sequence: 0,
peer_id: self_name.clone(),
};
if tx.send(Ok(snapshot_request)).await.is_err() {
log::warn!("Failed to send snapshot request");
}
}
}
}
while let Some(msg_result) = incoming.next().await {
match msg_result {
Ok(msg) => {
sequence += 1;
peer_id = msg.peer_id.clone();
match msg.message_type() {
StreamMessageType::IncrementalUpdate => {
if let Some(gossip::stream_message::Payload::Incremental(update)) =
&msg.payload
{
// Validate message size
let msg_size: usize =
update.updates.iter().map(|u| u.value.len()).sum();
if let Err(e) = size_validator.validate(msg_size) {
log::warn!(
"Received oversized incremental update from {}: {} (max: {} bytes), rejecting",
peer_id, e, size_validator.max_size()
);
let nack = StreamMessage {
message_type: StreamMessageType::Nack as i32,
payload: Some(gossip::stream_message::Payload::Ack(
StreamAck {
sequence: msg.sequence,
success: false,
error_message: format!(
"Message too large: {}",
e
),
},
)),
sequence,
peer_id: self_name.clone(),
};
if tx.send(Ok(nack)).await.is_err() {
break;
}
record_nack(&peer_id);
continue;
}
let store_type = LocalStoreType::from_proto(update.store);
log::info!("Received incremental update from {}: store={:?}, {} updates",
peer_id, store_type, update.updates.len());
// Apply incremental updates to state stores
// This will be handled by the sync manager if available
// For now, we acknowledge and the sync manager will handle it
if let Some(ref sync_manager) = sync_manager {
for state_update in &update.updates {
match store_type {
LocalStoreType::Worker => {
// Deserialize and apply worker state
if let Ok(worker_state) = serde_json::from_slice::<
super::stores::WorkerState,
>(
&state_update.value
) {
// Extract actor from StateUpdate
let actor =
Some(state_update.actor.clone());
sync_manager.apply_remote_worker_state(
worker_state,
actor,
);
}
}
LocalStoreType::Policy => {
// Deserialize and apply policy state
if let Ok(policy_state) = serde_json::from_slice::<
super::stores::PolicyState,
>(
&state_update.value
) {
// Extract actor from StateUpdate
let actor =
Some(state_update.actor.clone());
// Check if this is a tree state update
if policy_state.policy_type == "tree_state"
{
// Deserialize tree state
if let Ok(tree_state) =
serde_json::from_slice::<
super::tree_ops::TreeState,
>(
&policy_state.config
)
{
sync_manager
.apply_remote_tree_operation(
policy_state
.model_id
.clone(),
tree_state,
actor,
);
}
} else {
// Regular policy state update
sync_manager.apply_remote_policy_state(
policy_state,
actor,
);
}
}
}
LocalStoreType::RateLimit => {
// Deserialize and apply rate limit counter
if let Ok(counter) = serde_json::from_slice::<
super::crdt::CRDTPNCounter,
>(
&state_update.value
) {
// Convert CRDTPNCounter to SyncPNCounter for merging
let sync_counter =
super::crdt::SyncPNCounter::new();
sync_counter.merge(&counter);
sync_manager
.apply_remote_rate_limit_counter(
state_update.key.clone(),
&sync_counter,
);
}
}
_ => {
// Other store types handled elsewhere
}
}
}
}
let ack = StreamMessage {
message_type: StreamMessageType::Ack as i32,
payload: Some(gossip::stream_message::Payload::Ack(
StreamAck {
sequence: msg.sequence,
success: true,
error_message: String::new(),
},
)),
sequence,
peer_id: self_name.clone(),
};
if tx.send(Ok(ack)).await.is_err() {
break;
}
}
}
StreamMessageType::SnapshotRequest => {
if let Some(gossip::stream_message::Payload::SnapshotRequest(req)) =
&msg.payload
{
let store_type = LocalStoreType::from_proto(req.store);
let store_name = store_type.as_str();
log::info!("Received snapshot request from {}: store={:?}, from_version={}",
peer_id, store_type, req.from_version);
record_snapshot_trigger(store_name, "request");
let snapshot_start = Instant::now();
// Generate and send snapshot chunks
let service = GossipService {
state: state.clone(),
self_addr: SocketAddr::from(([0, 0, 0, 0], 0)), // Not used in snapshot generation
self_name: self_name.clone(),
stores: stores.clone(),
sync_manager: sync_manager.clone(),
state_machine: None,
partition_detector: None,
};
let chunks =
service.create_snapshot_chunks(store_type, 100).await; // chunk_size = 100 entries
let total_chunks = chunks.len() as u64;
let mut total_bytes = 0;
for (idx, chunk) in chunks.into_iter().enumerate() {
let chunk_bytes = chunk
.entries
.iter()
.map(|e| e.value.len())
.sum::<usize>();
total_bytes += chunk_bytes;
let mut chunk_msg = StreamMessage {
message_type: StreamMessageType::SnapshotChunk as i32,
payload: Some(
gossip::stream_message::Payload::SnapshotChunk(
chunk,
),
),
sequence: sequence + idx as u64 + 1,
peer_id: self_name.clone(),
};
// Update chunk metadata
if let Some(
gossip::stream_message::Payload::SnapshotChunk(
ref mut c,
),
) = chunk_msg.payload
{
c.chunk_index = idx as u64;
c.total_chunks = total_chunks;
}
// Check backpressure using try_send
match tx.try_send(Ok(chunk_msg)) {
Ok(_) => {
// Successfully queued
}
Err(tokio::sync::mpsc::error::TrySendError::Full(
msg,
)) => {
log::debug!(
"Backpressure: channel full, waiting for drain"
);
// Wait a bit for channel to drain, then use blocking send
tokio::time::sleep(Duration::from_millis(100))
.await;
if tx.send(msg).await.is_err() {
log::warn!("Backpressure: channel closed, stopping snapshot");
break;
}
}
Err(
tokio::sync::mpsc::error::TrySendError::Closed(_),
) => {
log::warn!("Channel closed, stopping snapshot");
break;
}
}
}
record_snapshot_duration(store_name, snapshot_start.elapsed());
record_snapshot_bytes(store_name, "sent", total_bytes);
// Send snapshot complete message
let complete = StreamMessage {
message_type: StreamMessageType::SnapshotComplete as i32,
payload: None,
sequence: sequence + total_chunks + 1,
peer_id: self_name.clone(),
};
if tx.send(Ok(complete)).await.is_err() {
break;
}
// Send ACK
let ack = StreamMessage {
message_type: StreamMessageType::Ack as i32,
payload: Some(gossip::stream_message::Payload::Ack(
StreamAck {
sequence: msg.sequence,
success: true,
error_message: String::new(),
},
)),
sequence,
peer_id: self_name.clone(),
};
record_ack(&peer_id, true);
if tx.send(Ok(ack)).await.is_err() {
break;
}
}
}
StreamMessageType::SnapshotChunk => {
if let Some(gossip::stream_message::Payload::SnapshotChunk(chunk)) =
&msg.payload
{
let store_type = LocalStoreType::from_proto(chunk.store);
let store_name = store_type.as_str();
log::info!(
"Received snapshot chunk from {}: store={:?}, chunk={}/{}",
peer_id,
store_type,
chunk.chunk_index,
chunk.total_chunks
);
// Record metrics
let chunk_bytes: usize =
chunk.entries.iter().map(|e| e.value.len()).sum();
record_snapshot_bytes(store_name, "received", chunk_bytes);
// Store chunk for later application
let chunk_key = (store_type, chunk.total_chunks);
snapshot_state
.entry(chunk_key)
.or_default()
.push(chunk.clone());
// Check if we've received all chunks
if let Some(received_chunks) = snapshot_state.get(&chunk_key) {
if received_chunks.len() as u64 == chunk.total_chunks {
// All chunks received, apply snapshot
log::info!("All {} chunks received for store {:?}, applying snapshot",
chunk.total_chunks, store_type);
if let Some(ref stores) = stores {
// Sort chunks by index
let mut sorted_chunks = received_chunks.clone();
sorted_chunks.sort_by_key(|c| c.chunk_index);
// Apply all entries from chunks
for chunk in &sorted_chunks {
for entry in &chunk.entries {
let key = SKey::new(entry.key.clone());
match store_type {
LocalStoreType::Membership => {
if let Ok(membership_state) = serde_json::from_slice::<super::stores::MembershipState>(&entry.value) {
stores.membership.insert(key, membership_state, entry.actor.clone());
}
}
LocalStoreType::App => {
if let Ok(app_state) = serde_json::from_slice::<super::stores::AppState>(&entry.value) {
stores.app.insert(key, app_state, entry.actor.clone());
}
}
LocalStoreType::Worker => {
if let Ok(worker_state) = serde_json::from_slice::<super::stores::WorkerState>(&entry.value) {
stores.worker.insert(key, worker_state.clone(), entry.actor.clone());
// Also update sync manager if available
if let Some(ref sync_manager) = sync_manager {
sync_manager.apply_remote_worker_state(worker_state, Some(entry.actor.clone()));
}
}
}
LocalStoreType::Policy => {
if let Ok(policy_state) = serde_json::from_slice::<super::stores::PolicyState>(&entry.value) {
stores.policy.insert(key, policy_state.clone(), entry.actor.clone());
// Also update sync manager if available
if let Some(ref sync_manager) = sync_manager {
// Check if this is a tree state update
if policy_state.policy_type == "tree_state" {
// Deserialize tree state
if let Ok(tree_state) = serde_json::from_slice::<
super::tree_ops::TreeState,
>(
&policy_state.config
) {
sync_manager.apply_remote_tree_operation(
policy_state.model_id.clone(),
tree_state,
Some(entry.actor.clone()),
);
}
} else {
sync_manager.apply_remote_policy_state(policy_state, Some(entry.actor.clone()));
}
}
}
}
LocalStoreType::RateLimit => {
// For rate limit counters, deserialize and merge
if let Ok(counter) = serde_json::from_slice::<super::crdt::CRDTPNCounter>(&entry.value) {
if let Some(ref sync_manager) = sync_manager {
let sync_counter = super::crdt::SyncPNCounter::new();
sync_counter.merge(&counter);
sync_manager.apply_remote_rate_limit_counter(entry.key.clone(), &sync_counter);
}
}
}
}
}
}
// Clear snapshot state
snapshot_state.remove(&chunk_key);
log::info!(
"Snapshot applied successfully for store {:?}",
store_type
);
}
}
}
let ack = StreamMessage {
message_type: StreamMessageType::Ack as i32,
payload: Some(gossip::stream_message::Payload::Ack(
StreamAck {
sequence: msg.sequence,
success: true,
error_message: String::new(),
},
)),
sequence,
peer_id: self_name.clone(),
};
record_ack(&peer_id, true);
if tx.send(Ok(ack)).await.is_err() {
break;
}
}
}
StreamMessageType::Ack => {
log::debug!(
"Received ACK from {}: sequence={}",
peer_id,
msg.sequence
);
if let Some(gossip::stream_message::Payload::Ack(ack)) =
&msg.payload
{
record_ack(&peer_id, ack.success);
}
}
StreamMessageType::Heartbeat => {
// Send heartbeat back
let heartbeat = StreamMessage {
message_type: StreamMessageType::Heartbeat as i32,
payload: None,
sequence,
peer_id: self_name.clone(),
};
if tx.send(Ok(heartbeat)).await.is_err() {
break;
}
}
_ => {
log::warn!(
"Unknown message type from {}: {:?}",
peer_id,
msg.message_type
);
}
}
}
Err(e) => {
log::error!("Error receiving stream message: {}", e);
record_nack(&peer_id);
update_peer_connections(&peer_id, false);
record_peer_reconnect(&peer_id);
break;
}
}
}
log::info!("Stream from {} closed", peer_id);
update_peer_connections(&peer_id, false);
});
// Convert receiver to stream
let output_stream = tokio_stream::wrappers::ReceiverStream::new(rx);
Ok(Response::new(
Box::pin(output_stream) as Self::SyncStreamStream
))
}
}

122
sgl-model-gateway/src/mesh/proto/gossip.proto
View File

@@ -1,122 +0,0 @@
syntax = "proto3";
package sglang.mesh.gossip;
service Gossip {
rpc PingServer(GossipMessage) returns (NodeUpdate);
// Bidirectional streaming for state synchronization
rpc SyncStream(stream StreamMessage) returns (stream StreamMessage);
}
enum NodeStatus {
INIT = 0;
ALIVE = 1;
SUSPECTED = 2;
DOWN = 3;
LEAVING = 4;
}
message NodeState {
string name = 1;
string address = 2;
NodeStatus status = 3;
uint64 version = 4;
map<string, bytes> metadata = 5;
}
message StateSync {
repeated NodeState nodes = 1;
}
message Ping {
StateSync state_sync = 1;
}
message PingReq {
NodeState node = 1;
}
message Ack {
uint64 timestamp = 1;
}
message NodeUpdate {
string name = 1;
string address = 2;
NodeStatus status = 3;
}
message GossipMessage {
oneof payload {
Ping ping = 1;
PingReq ping_req = 2;
}
}
// Stream message types for bidirectional streaming
enum StreamMessageType {
INCREMENTAL_UPDATE = 0;
SNAPSHOT_REQUEST = 1;
SNAPSHOT_CHUNK = 2;
SNAPSHOT_COMPLETE = 3;
ACK = 4;
NACK = 5;
HEARTBEAT = 6;
}
message StreamMessage {
StreamMessageType message_type = 1;
oneof payload {
IncrementalUpdate incremental = 2;
SnapshotRequest snapshot_request = 3;
SnapshotChunk snapshot_chunk = 4;
StreamAck ack = 5;
}
uint64 sequence = 6; // Sequence number for ordering
string peer_id = 7; // Sender peer ID
}
// Incremental state update (steady-state)
message IncrementalUpdate {
StoreType store = 1;
repeated StateUpdate updates = 2;
uint64 version = 3;
}
message StateUpdate {
string key = 1;
bytes value = 2; // Serialized state value
uint64 version = 3;
string actor = 4;
uint64 timestamp = 5;
}
// Snapshot request (on gap or new join)
message SnapshotRequest {
StoreType store = 1;
uint64 from_version = 2; // Request snapshot from this version
}
// Snapshot chunk (per-store chunking)
message SnapshotChunk {
StoreType store = 1;
uint64 chunk_index = 2;
uint64 total_chunks = 3;
repeated StateUpdate entries = 4;
bytes checksum = 5; // For integrity verification
}
message StreamAck {
uint64 sequence = 1;
bool success = 2;
string error_message = 3;
}
enum StoreType {
MEMBERSHIP = 0;
APP = 1;
WORKER = 2;
POLICY = 3;
RATE_LIMIT = 4;
}

257
sgl-model-gateway/src/mesh/rate_limit_window.rs
View File

@@ -1,257 +0,0 @@
//! Rate limit time window management
//!
//! Manages time windows for global rate limiting with periodic counter resets
use std::{sync::Arc, time::Duration};
use tokio::time::interval;
use tracing::{debug, info};
use super::sync::MeshSyncManager;
/// Rate limit window manager
/// Handles periodic reset of rate limit counters for time window management
pub struct RateLimitWindow {
sync_manager: Arc<MeshSyncManager>,
window_seconds: u64,
}
impl RateLimitWindow {
pub fn new(sync_manager: Arc<MeshSyncManager>, window_seconds: u64) -> Self {
Self {
sync_manager,
window_seconds,
}
}
/// Start the window reset task
/// This task periodically resets the global rate limit counter
pub async fn start_reset_task(self) {
let mut interval_timer = interval(Duration::from_secs(self.window_seconds));
info!(
"Starting rate limit window reset task with {}s interval",
self.window_seconds
);
loop {
interval_timer.tick().await;
debug!("Resetting global rate limit counter");
self.sync_manager.reset_global_rate_limit_counter();
}
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use tokio::time::sleep;
use super::*;
use crate::mesh::stores::{
RateLimitConfig, StateStores, GLOBAL_RATE_LIMIT_COUNTER_KEY, GLOBAL_RATE_LIMIT_KEY,
};
#[test]
fn test_rate_limit_window_new() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
let window = RateLimitWindow::new(sync_manager, 60);
// Should create without panicking
assert_eq!(window.window_seconds, 60);
}
#[test]
fn test_rate_limit_window_different_intervals() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
let window1 = RateLimitWindow::new(sync_manager.clone(), 30);
assert_eq!(window1.window_seconds, 30);
let window2 = RateLimitWindow::new(sync_manager, 120);
assert_eq!(window2.window_seconds, 120);
}
#[tokio::test]
async fn test_rate_limit_window_reset_task_interval() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
// Set a very short window for testing (1 second)
let window = RateLimitWindow::new(sync_manager, 1);
// Spawn the reset task
let task_handle = tokio::spawn(async move {
window.start_reset_task().await;
});
// Wait a bit to allow the task to run
sleep(Duration::from_millis(1500)).await;
// Cancel the task
task_handle.abort();
// The task should have started (we can't easily verify it ran without
// more complex mocking, but we can verify it doesn't panic)
// In a real scenario, you'd use a mock to track reset calls
}
#[tokio::test]
async fn test_rate_limit_window_reset_task() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores.clone(), "node1".to_string()));
// Setup membership
stores.rate_limit.update_membership(&["node1".to_string()]);
// Setup config
let key = crate::mesh::crdt::SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string());
let config = RateLimitConfig {
limit_per_second: 100,
};
let serialized = serde_json::to_vec(&config).unwrap();
stores.app.insert(
key,
crate::mesh::stores::AppState {
key: GLOBAL_RATE_LIMIT_KEY.to_string(),
value: serialized,
version: 1,
},
"node1".to_string(),
);
// Increment counter
if stores.rate_limit.is_owner(GLOBAL_RATE_LIMIT_COUNTER_KEY) {
sync_manager.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 10);
let value_before = sync_manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
assert!(value_before.is_some() && value_before.unwrap() > 0);
// Create window manager with short interval for testing
let window = RateLimitWindow::new(sync_manager.clone(), 1); // 1 second
// Start reset task in background
let reset_handle = tokio::spawn(async move {
window.start_reset_task().await;
});
// Wait a bit for reset to happen
sleep(Duration::from_millis(1500)).await;
// Check that counter was reset (or at least decremented)
let _value_after = sync_manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
// Counter should be reset or significantly reduced
// Note: The exact value depends on timing, but it should be less than initial
// Cancel the task
reset_handle.abort();
}
}
#[tokio::test]
async fn test_rate_limit_window_reset_with_counter() {
use crate::mesh::{crdt::SKey, stores::MembershipState};
// Use with_self_name to ensure RateLimitStore uses the same self_name
let stores = Arc::new(StateStores::with_self_name("test_node".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(
stores.clone(),
"test_node".to_string(),
));
// First, add this node to membership so it can be an owner
let membership_key = SKey::new("test_node".to_string());
let membership_state = MembershipState {
name: "test_node".to_string(),
address: "127.0.0.1:8080".to_string(),
status: 1, // NodeStatus::Alive
version: 1,
metadata: Default::default(),
};
stores
.membership
.insert(membership_key, membership_state, "test_node".to_string());
// Update rate limit membership so this node becomes an owner
sync_manager.update_rate_limit_membership();
// Check if node is owner before incrementing
let key = GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string();
let is_owner = stores.rate_limit.is_owner(&key);
assert!(is_owner, "Node should be owner of the rate limit key");
// Set up a rate limit counter via sync_manager
// This should increment the counter if the node is an owner
sync_manager.sync_rate_limit_inc(key.clone(), 10);
// Verify counter exists (was created)
// Note: The actual value might be 0 due to PNCounter implementation details,
// but the counter should exist after inc is called
let counter_opt = stores.rate_limit.get_counter(&key);
assert!(counter_opt.is_some(), "Counter should exist after inc call");
// Verify counter was created after inc call
// Note: The actual value depends on PNCounter implementation,
// but the counter should exist after inc is called
// Reset the counter
sync_manager.reset_global_rate_limit_counter();
// Verify reset was called (counter should still exist)
// The reset implementation decrements by current count,
// so the value should be 0 or negative after reset
let reset_value = stores.rate_limit.value(&key).unwrap_or(0);
// After reset, value should be <= 0 (since we decrement by current count)
assert!(
reset_value <= 0,
"Counter should be reset to 0 or less, got: {}",
reset_value
);
}
#[test]
fn test_rate_limit_window_zero_seconds() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
// Should handle zero seconds (though not recommended in practice)
let window = RateLimitWindow::new(sync_manager, 0);
assert_eq!(window.window_seconds, 0);
}
#[test]
fn test_rate_limit_window_large_interval() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
// Test with a large interval
let window = RateLimitWindow::new(sync_manager, 86400); // 24 hours
assert_eq!(window.window_seconds, 86400);
}
#[tokio::test]
async fn test_reset_global_rate_limit_counter_logic() {
let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
let sync_manager = Arc::new(MeshSyncManager::new(stores.clone(), "node1".to_string()));
// Setup membership
stores.rate_limit.update_membership(&["node1".to_string()]);
if stores.rate_limit.is_owner(GLOBAL_RATE_LIMIT_COUNTER_KEY) {
// Increment counter
sync_manager.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 20);
let value_before = sync_manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
assert!(value_before.is_some() && value_before.unwrap() > 0);
// Reset
sync_manager.reset_global_rate_limit_counter();
// Check that counter was reset
let value_after = sync_manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
// Should be 0 or negative after reset
assert!(value_after.is_none() || value_after.unwrap() <= 0);
}
}
}

600
sgl-model-gateway/src/mesh/service.rs
View File

@@ -1,600 +0,0 @@
use std::{
collections::{BTreeMap, HashMap},
net::SocketAddr,
str::FromStr,
sync::Arc,
time::Duration,
};
use anyhow::Result;
use parking_lot::RwLock;
use tonic::Request;
use tracing as log;
pub mod gossip {
#![allow(unused_qualifications)]
tonic::include_proto!("sglang.mesh.gossip");
}
use gossip::{
gossip_client, gossip_message, GossipMessage, NodeState, NodeStatus, NodeUpdate, Ping,
StateSync,
};
use crate::mesh::{
controller::MeshController,
node_state_machine::{ConvergenceConfig, NodeStateMachine},
partition::PartitionDetector,
ping_server::GossipService,
};
pub type ClusterState = Arc<RwLock<BTreeMap<String, NodeState>>>;
pub struct MeshServerConfig {
pub self_name: String,
pub self_addr: SocketAddr,
pub init_peer: Option<SocketAddr>,
}
/// MeshServerHandler
/// It is the handler for the mesh server, which is responsible for the node management.
/// Includes some basic node management logic, like shutdown,
/// node discovery(TODO), node status update(TODO), etc.
pub struct MeshServerHandler {
pub state: ClusterState,
pub self_name: String,
_self_addr: SocketAddr,
signal_tx: tokio::sync::watch::Sender<()>,
partition_detector: Option<Arc<PartitionDetector>>,
state_machine: Option<Arc<NodeStateMachine>>,
}
impl MeshServerHandler {
pub fn new(
state: ClusterState,
self_name: &str,
self_addr: SocketAddr,
signal_tx: tokio::sync::watch::Sender<()>,
) -> Self {
Self {
state,
self_name: self_name.to_string(),
_self_addr: self_addr,
signal_tx,
partition_detector: None,
state_machine: None,
}
}
/// Create with partition detector and state machine
pub fn with_partition_and_state_machine(
state: ClusterState,
self_name: &str,
self_addr: SocketAddr,
signal_tx: tokio::sync::watch::Sender<()>,
stores: Option<Arc<super::stores::StateStores>>,
) -> Self {
let partition_detector = Some(Arc::new(PartitionDetector::default()));
let state_machine =
stores.map(|s| Arc::new(NodeStateMachine::new(s, ConvergenceConfig::default())));
Self {
state,
self_name: self_name.to_string(),
_self_addr: self_addr,
signal_tx,
partition_detector,
state_machine,
}
}
/// Get partition detector
pub fn partition_detector(&self) -> Option<&Arc<PartitionDetector>> {
self.partition_detector.as_ref()
}
/// Get state machine
pub fn state_machine(&self) -> Option<&Arc<NodeStateMachine>> {
self.state_machine.as_ref()
}
/// Check if node is ready
pub fn is_ready(&self) -> bool {
self.state_machine
.as_ref()
.map(|sm| sm.is_ready())
.unwrap_or(true) // If no state machine, consider ready
}
/// Check if we should serve (have quorum)
pub fn should_serve(&self) -> bool {
self.partition_detector
.as_ref()
.map(|pd| pd.should_serve())
.unwrap_or(true) // If no partition detector, consider should serve
}
/// Shutdown immediately without graceful shutdown
pub fn shutdown(&self) {
self.signal_tx.send(()).ok();
}
/// Graceful shutdown: broadcast LEAVING status to all alive nodes,
/// wait for propagation, then shutdown
pub async fn graceful_shutdown(&self) -> Result<()> {
log::info!("Starting graceful shutdown for node {}", self.self_name);
let (leaving_node, alive_nodes) = {
let state = self.state.read();
let mut self_node = if let Some(self_node) = state.get(&self.self_name) {
self_node.clone()
} else {
self.signal_tx.send(()).ok();
return Ok(());
};
if self_node.status != NodeStatus::Leaving as i32 {
self_node.status = NodeStatus::Leaving as i32;
self_node.version += 1;
let alive_nodes = state
.values()
.filter(|node| {
node.status == NodeStatus::Alive as i32 // include self
})
.cloned()
.collect::<Vec<NodeState>>();
(self_node.clone(), alive_nodes)
} else {
self.signal_tx.send(()).ok();
return Ok(());
}
};
log::info!(
"Broadcasting LEAVING status to {} alive nodes",
alive_nodes.len()
);
// Broadcast LEAVING status to all alive nodes
let (success_count, total_count) = broadcast_node_states(
vec![leaving_node],
alive_nodes,
Some(Duration::from_secs(3)),
)
.await;
log::info!(
"Broadcast LEAVING status: {}/{} successful",
success_count,
total_count
);
// Wait a bit more for state propagation
let propagation_delay = Duration::from_secs(1);
log::info!(
"Waiting {} seconds for LEAVING status propagation",
propagation_delay.as_secs()
);
tokio::time::sleep(propagation_delay).await;
log::info!("Sending shutdown signal");
self.signal_tx.send(()).ok();
Ok(())
}
pub fn write_data(&self, key: String, value: Vec<u8>) {
let mut state = self.state.write();
state.entry(self.self_name.clone()).and_modify(|e| {
e.version += 1;
e.metadata.insert(key, value);
});
}
pub fn read_data(&self, key: String) -> Option<Vec<u8>> {
let state = self.state.read();
state
.get(&self.self_name)
.and_then(|e| e.metadata.get(&key).cloned())
}
}
pub struct MeshServerBuilder {
state: ClusterState,
self_name: String,
self_addr: SocketAddr,
init_peer: Option<SocketAddr>,
}
impl MeshServerBuilder {
pub fn new(self_name: String, self_addr: SocketAddr, init_peer: Option<SocketAddr>) -> Self {
let state = Arc::new(RwLock::new(BTreeMap::from([(
self_name.clone(),
NodeState {
name: self_name.clone(),
address: self_addr.to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: HashMap::new(),
},
)])));
Self {
state,
self_name,
self_addr,
init_peer,
}
}
pub fn build(&self) -> (MeshServer, MeshServerHandler) {
self.build_with_stores(None)
}
pub fn build_with_stores(
&self,
stores: Option<Arc<super::stores::StateStores>>,
) -> (MeshServer, MeshServerHandler) {
let (signal_tx, signal_rx) = tokio::sync::watch::channel(());
(
MeshServer::new(
self.state.clone(),
&self.self_name,
self.self_addr,
self.init_peer,
signal_rx,
),
MeshServerHandler::with_partition_and_state_machine(
self.state.clone(),
&self.self_name,
self.self_addr,
signal_tx,
stores,
),
)
}
}
pub struct MeshServer {
state: ClusterState,
self_name: String,
self_addr: SocketAddr,
init_peer: Option<SocketAddr>,
signal_rx: tokio::sync::watch::Receiver<()>,
}
impl MeshServer {
pub fn new(
state: ClusterState,
self_name: &str,
self_addr: SocketAddr,
init_peer: Option<SocketAddr>,
signal_rx: tokio::sync::watch::Receiver<()>,
) -> Self {
MeshServer {
state,
self_name: self_name.to_string(),
self_addr,
init_peer,
signal_rx,
}
}
pub fn build_ping_server(&self) -> GossipService {
GossipService::new(self.state.clone(), self.self_addr, &self.self_name)
}
pub fn build_controller(&self) -> MeshController {
MeshController::new(
self.state.clone(),
self.self_addr,
&self.self_name,
self.init_peer,
)
}
pub async fn start_serve(self) -> Result<()> {
log::info!("Mesh server listening on {}", self.self_addr);
let self_name = self.self_name.clone();
let self_address = self.self_addr;
let service = self.build_ping_server();
let controller = self.build_controller();
let mut service_shutdown = self.signal_rx.clone();
let listener = tokio::spawn(service.serve_ping_with_shutdown(async move {
_ = service_shutdown.changed().await;
}));
tokio::time::sleep(Duration::from_secs(1)).await;
let app_handle = tokio::spawn(controller.event_loop(self.signal_rx.clone()));
tokio::select! {
res = listener => res??,
res = app_handle => res??,
}
log::info!(
"Mesh server {} at {} is shutting down",
self_name,
self_address
);
Ok(())
}
pub async fn start_serve_with_stores(
self,
stores: Option<Arc<super::stores::StateStores>>,
sync_manager: Option<Arc<super::sync::MeshSyncManager>>,
partition_detector: Option<Arc<PartitionDetector>>,
) -> Result<()> {
log::info!("Mesh server listening on {}", self.self_addr);
let self_name = self.self_name.clone();
let self_address = self.self_addr;
let mut service = self.build_ping_server();
if let Some(stores) = stores {
service = service.with_stores(stores);
}
if let Some(sync_manager) = sync_manager {
service = service.with_sync_manager(sync_manager);
}
if let Some(partition_detector) = partition_detector {
service = service.with_partition_detector(partition_detector);
}
let controller = self.build_controller();
let mut service_shutdown = self.signal_rx.clone();
let listener = tokio::spawn(service.serve_ping_with_shutdown(async move {
_ = service_shutdown.changed().await;
}));
tokio::time::sleep(Duration::from_secs(1)).await;
let app_handle = tokio::spawn(controller.event_loop(self.signal_rx.clone()));
tokio::select! {
res = listener => res??,
res = app_handle => res??,
}
log::info!(
"Mesh server {} at {} is shutting down",
self_name,
self_address
);
Ok(())
}
}
/// Broadcast node state updates to target nodes
/// Returns (success_count, total_count)
pub async fn broadcast_node_states(
nodes_to_broadcast: Vec<NodeState>,
target_nodes: Vec<NodeState>,
timeout: Option<Duration>,
) -> (usize, usize) {
if nodes_to_broadcast.is_empty() || target_nodes.is_empty() {
log::debug!(
"Nothing to broadcast: nodes_to_broadcast={}, target_nodes={}",
nodes_to_broadcast.len(),
target_nodes.len()
);
return (0, target_nodes.len());
}
let mut broadcast_tasks = Vec::new();
for target_node in &target_nodes {
let target_node_clone = target_node.clone();
let nodes_for_task = nodes_to_broadcast.clone();
let task = tokio::spawn(async move {
let state_sync = StateSync {
nodes: nodes_for_task,
};
let ping_payload = gossip_message::Payload::Ping(Ping {
state_sync: Some(state_sync),
});
match try_ping(&target_node_clone, Some(ping_payload)).await {
Ok(_) => {
log::debug!("Successfully broadcasted to {}", target_node_clone.name);
Ok(())
}
Err(e) => {
log::warn!("Failed to broadcast to {}: {}", target_node_clone.name, e);
Err(e)
}
}
});
broadcast_tasks.push(task);
}
let timeout_duration = timeout.unwrap_or(Duration::from_secs(3));
let broadcast_result = tokio::time::timeout(timeout_duration, async {
futures::future::join_all(broadcast_tasks).await
})
.await;
match broadcast_result {
Ok(results) => {
let success_count = results.iter().filter(|r| r.is_ok()).count();
let total_count = target_nodes.len();
log::info!(
"Broadcast completed: {}/{} successful",
success_count,
total_count
);
(success_count, total_count)
}
Err(_) => {
log::warn!(
"Broadcast timeout after {} seconds",
timeout_duration.as_secs()
);
(0, target_nodes.len())
}
}
}
pub async fn try_ping(
peer_node: &NodeState,
payload: Option<gossip_message::Payload>,
) -> Result<NodeUpdate, tonic::Status> {
let peer_name = peer_node.name.clone();
let peer_addr = SocketAddr::from_str(&peer_node.address).map_err(|e| {
tonic::Status::invalid_argument(format!(
"Invalid address for node {}: {}, {}",
peer_name, peer_node.address, e
))
})?;
let mut client = gossip_client::GossipClient::connect(format!("http://{}", peer_addr))
.await
.map_err(|e| {
log::warn!(
"Failed to connect to peer {} {}: {}.",
peer_name,
peer_addr,
e
);
tonic::Status::unavailable("Failed to connect to peer")
})?;
let ping_message = GossipMessage { payload };
let response = client.ping_server(Request::new(ping_message)).await?;
Ok(response.into_inner())
}
#[macro_export]
macro_rules! mesh_run {
($addr:expr, $init_peer:expr) => {{
mesh_run!($addr.to_string(), $addr, $init_peer)
}};
($name:expr, $addr:expr, $init_peer:expr) => {{
tracing::info!("Starting mesh server : {}", $addr);
let (server, handler) =
$crate::mesh::service::MeshServerBuilder::new($name.to_string(), $addr, $init_peer)
.build();
tokio::spawn(async move {
if let Err(e) = server.start_serve().await {
tracing::error!("Mesh server failed: {}", e);
}
});
handler
}};
}
#[cfg(test)]
mod tests {
use std::sync::Once;
use tokio::net::TcpListener;
use tracing as log;
use tracing_subscriber::{
filter::LevelFilter, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
};
use super::*;
static INIT: Once = Once::new();
fn init() {
INIT.call_once(|| {
let _ = tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer())
.with(
EnvFilter::builder()
.with_default_directive(LevelFilter::INFO.into())
.from_env_lossy(),
)
.try_init();
});
}
async fn find_free_port() -> (TcpListener, u16) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
log::info!("Found free port: {}", port);
(listener, port)
}
async fn get_node() -> SocketAddr {
let (_listener, port) = find_free_port().await;
format!("127.0.0.1:{}", port).parse().unwrap()
}
fn print_state(handler: &MeshServerHandler) -> String {
let state = handler.state.read();
let mut res = vec![];
for (k, v) in state.iter() {
res.push(format!(
"{}: {:?} - {:?}",
k,
NodeStatus::try_from(v.status).unwrap(),
v.metadata
));
}
res.join(", ")
}
#[tokio::test]
async fn test_state_synchronization() {
init();
log::info!("Starting test_state_synchronization");
// 1. setup node A and B for initial cluster
let addr_a = get_node().await;
let handler_a = mesh_run!("A", addr_a, None);
let addr_b = get_node().await;
let handler_b = mesh_run!("B", addr_b, Some(addr_a));
// 2. wait for node A and B to sync and write some data
tokio::time::sleep(Duration::from_secs(2)).await;
handler_a.write_data("hello".into(), "world".into());
log::info!("================================================");
// 3. add node C and D and wait for them to sync
let addr_c = get_node().await;
let handler_c = mesh_run!("C", addr_c, Some(addr_a));
let addr_d = get_node().await;
let handler_d = mesh_run!("D", addr_d, Some(addr_c));
tokio::time::sleep(Duration::from_secs(2)).await;
log::info!("================================================");
// 4. add node E and wait for it to sync and kill it
{
let addr_e = get_node().await;
let handler_e = mesh_run!("E", addr_e, Some(addr_d));
tokio::time::sleep(Duration::from_secs(3)).await;
log::info!("State E: {:?}", print_state(&handler_e));
// killing_button.send(()).unwrap();
handler_e.shutdown();
}
handler_d.graceful_shutdown().await.unwrap();
tokio::time::sleep(Duration::from_secs(2)).await;
log::info!("================================================");
// 5. wait for node status to sync
tokio::time::sleep(Duration::from_secs(8)).await;
log::info!("================================================");
// 6. verify node status, status of all nodes should be same, and node E should be down
let final_state = String::from("A: Alive - {\"hello\": [119, 111, 114, 108, 100]}, B: Alive - {}, C: Alive - {}, D: Leaving - {}, E: Down - {}");
assert_eq!(
print_state(&handler_a),
final_state,
"State A: {:?}",
print_state(&handler_a)
);
assert_eq!(
print_state(&handler_b),
final_state,
"State B: {:?}",
print_state(&handler_b)
);
assert_eq!(
print_state(&handler_c),
final_state,
"State C: {:?}",
print_state(&handler_c)
);
}
}

742
sgl-model-gateway/src/mesh/stores.rs
View File

@@ -1,742 +0,0 @@
//! State stores for mesh cluster synchronization
//!
//! Four types of state stores:
//! - MembershipStore: Router node membership
//! - AppStore: Application configuration, rate-limiting rules, LB algorithms
//! - WorkerStore: Worker status, load, health
//! - PolicyStore: Routing policy internal state
use std::{collections::BTreeMap, sync::Arc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use tracing::debug;
use super::{
consistent_hash::ConsistentHashRing,
crdt::{SKey, SyncCRDTMap, SyncPNCounter},
};
/// Store type identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum StoreType {
Membership,
App,
Worker,
Policy,
RateLimit,
}
impl StoreType {
pub fn as_str(&self) -> &'static str {
match self {
StoreType::Membership => "membership",
StoreType::App => "app",
StoreType::Worker => "worker",
StoreType::Policy => "policy",
StoreType::RateLimit => "rate_limit",
}
}
/// Convert from proto StoreType (i32) to local StoreType
pub fn from_proto(proto_value: i32) -> Self {
match proto_value {
0 => StoreType::Membership,
1 => StoreType::App,
2 => StoreType::Worker,
3 => StoreType::Policy,
4 => StoreType::RateLimit,
_ => StoreType::Membership, // Default fallback
}
}
}
/// Membership state entry
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
pub struct MembershipState {
pub name: String,
pub address: String,
pub status: i32, // NodeStatus enum value
pub version: u64,
pub metadata: BTreeMap<String, Vec<u8>>,
}
/// App state entry (application configuration)
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
pub struct AppState {
pub key: String,
pub value: Vec<u8>, // Serialized config
pub version: u64,
}
/// Global rate limit configuration
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct RateLimitConfig {
pub limit_per_second: u64,
}
/// Key for global rate limit configuration in AppStore
pub const GLOBAL_RATE_LIMIT_KEY: &str = "global_rate_limit";
/// Key for global rate limit counter in RateLimitStore
pub const GLOBAL_RATE_LIMIT_COUNTER_KEY: &str = "global";
/// Worker state entry
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct WorkerState {
pub worker_id: String,
pub model_id: String,
pub url: String,
pub health: bool,
pub load: f64,
pub version: u64,
}
// Implement Hash manually for WorkerState (excluding f64)
impl std::hash::Hash for WorkerState {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.worker_id.hash(state);
self.model_id.hash(state);
self.url.hash(state);
self.health.hash(state);
// f64 cannot be hashed directly, use a workaround
(self.load as i64).hash(state);
self.version.hash(state);
}
}
// Implement Eq manually (f64 comparison with epsilon)
impl Eq for WorkerState {}
/// Policy state entry
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
pub struct PolicyState {
pub model_id: String,
pub policy_type: String,
pub config: Vec<u8>, // Serialized policy config
pub version: u64,
}
/// Helper function to get tree state key for a model
pub fn tree_state_key(model_id: &str) -> String {
format!("tree:{}", model_id)
}
/// Membership store
#[derive(Debug, Clone)]
pub struct MembershipStore {
inner: SyncCRDTMap<MembershipState>,
}
impl MembershipStore {
pub fn new() -> Self {
Self {
inner: SyncCRDTMap::new(),
}
}
pub fn get(&self, key: &SKey) -> Option<MembershipState> {
self.inner.get(key)
}
pub fn insert(&self, key: SKey, value: MembershipState, actor: String) {
self.inner.insert(key, value, actor);
}
pub fn remove(&self, key: &SKey) {
self.inner.remove(key);
}
pub fn merge(&self, other: &crate::mesh::crdt::CRDTMap<MembershipState>) {
self.inner.merge(other);
}
pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap<MembershipState> {
self.inner.snapshot()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.len() == 0
}
pub fn all(&self) -> BTreeMap<SKey, MembershipState> {
self.inner.snapshot().to_map()
}
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner.get_metadata(key)
}
}
impl Default for MembershipStore {
fn default() -> Self {
Self::new()
}
}
/// App store
#[derive(Debug, Clone)]
pub struct AppStore {
inner: SyncCRDTMap<AppState>,
}
impl AppStore {
pub fn new() -> Self {
Self {
inner: SyncCRDTMap::new(),
}
}
pub fn get(&self, key: &SKey) -> Option<AppState> {
self.inner.get(key)
}
pub fn insert(&self, key: SKey, value: AppState, actor: String) {
self.inner.insert(key, value, actor);
}
pub fn remove(&self, key: &SKey) {
self.inner.remove(key);
}
pub fn merge(&self, other: &crate::mesh::crdt::CRDTMap<AppState>) {
self.inner.merge(other);
}
pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap<AppState> {
self.inner.snapshot()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.len() == 0
}
pub fn all(&self) -> BTreeMap<SKey, AppState> {
self.inner.snapshot().to_map()
}
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner.get_metadata(key)
}
}
impl Default for AppStore {
fn default() -> Self {
Self::new()
}
}
/// Worker store
#[derive(Debug, Clone)]
pub struct WorkerStore {
inner: SyncCRDTMap<WorkerState>,
}
impl WorkerStore {
pub fn new() -> Self {
Self {
inner: SyncCRDTMap::new(),
}
}
pub fn get(&self, key: &SKey) -> Option<WorkerState> {
self.inner.get(key)
}
pub fn insert(&self, key: SKey, value: WorkerState, actor: String) {
self.inner.insert(key, value, actor);
}
pub fn remove(&self, key: &SKey) {
self.inner.remove(key);
}
pub fn merge(&self, other: &crate::mesh::crdt::CRDTMap<WorkerState>) {
self.inner.merge(other);
}
pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap<WorkerState> {
self.inner.snapshot()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.len() == 0
}
pub fn all(&self) -> BTreeMap<SKey, WorkerState> {
self.inner.snapshot().to_map()
}
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner.get_metadata(key)
}
}
impl Default for WorkerStore {
fn default() -> Self {
Self::new()
}
}
/// Policy store
#[derive(Debug, Clone)]
pub struct PolicyStore {
inner: SyncCRDTMap<PolicyState>,
}
impl PolicyStore {
pub fn new() -> Self {
Self {
inner: SyncCRDTMap::new(),
}
}
pub fn get(&self, key: &SKey) -> Option<PolicyState> {
self.inner.get(key)
}
pub fn insert(&self, key: SKey, value: PolicyState, actor: String) {
self.inner.insert(key, value, actor);
}
pub fn remove(&self, key: &SKey) {
self.inner.remove(key);
}
pub fn merge(&self, other: &crate::mesh::crdt::CRDTMap<PolicyState>) {
self.inner.merge(other);
}
pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap<PolicyState> {
self.inner.snapshot()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.len() == 0
}
pub fn all(&self) -> BTreeMap<SKey, PolicyState> {
self.inner.snapshot().to_map()
}
pub fn get_metadata(&self, key: &SKey) -> Option<(u64, String)> {
self.inner.get_metadata(key)
}
}
impl Default for PolicyStore {
fn default() -> Self {
Self::new()
}
}
/// Rate-limit counter store (using PNCounter with consistent hashing)
#[derive(Debug, Clone)]
pub struct RateLimitStore {
counters: Arc<RwLock<BTreeMap<String, SyncPNCounter>>>, // key -> counter
hash_ring: Arc<RwLock<ConsistentHashRing>>,
self_name: String,
}
impl RateLimitStore {
pub fn new(self_name: String) -> Self {
Self {
counters: Arc::new(RwLock::new(BTreeMap::new())),
hash_ring: Arc::new(RwLock::new(ConsistentHashRing::new())),
self_name,
}
}
/// Update the hash ring with current membership
pub fn update_membership(&self, nodes: &[String]) {
let mut ring = self.hash_ring.write();
ring.update_membership(nodes);
debug!("Updated rate-limit hash ring with {} nodes", nodes.len());
}
/// Check if this node is an owner of a key
pub fn is_owner(&self, key: &str) -> bool {
let ring = self.hash_ring.read();
ring.is_owner(key, &self.self_name)
}
/// Get owners for a key
pub fn get_owners(&self, key: &str) -> Vec<String> {
let ring = self.hash_ring.read();
ring.get_owners(key)
.into_iter()
.map(|s| s.to_string())
.collect()
}
/// Get or create counter (only if this node is an owner)
#[allow(dead_code)]
fn get_or_create_counter_internal(&self, key: String) -> Option<SyncPNCounter> {
if !self.is_owner(&key) {
return None;
}
let mut counters = self.counters.write();
Some(counters.entry(key.clone()).or_default().clone())
}
pub fn get_counter(&self, key: &str) -> Option<SyncPNCounter> {
if !self.is_owner(key) {
return None;
}
let counters = self.counters.read();
counters.get(key).cloned()
}
/// Increment counter (only if this node is an owner)
pub fn inc(&self, key: String, actor: String, delta: i64) {
if !self.is_owner(&key) {
// Not an owner, skip
return;
}
let mut counters = self.counters.write();
let counter = counters.entry(key.clone()).or_default();
counter.inc(actor, delta);
}
/// Get counter value (aggregate from all owners via CRDT merge)
pub fn value(&self, key: &str) -> Option<i64> {
let counters = self.counters.read();
counters.get(key).map(|c| c.value())
}
/// Merge counter from another node (for CRDT synchronization)
pub fn merge_counter(&self, key: String, other: &SyncPNCounter) {
let mut counters = self.counters.write();
let counter = counters.entry(key).or_default();
// Get the inner CRDTPNCounter from other SyncPNCounter
let other_inner = other.snapshot();
counter.merge(&other_inner);
}
/// Get all counter keys
pub fn keys(&self) -> Vec<String> {
let counters = self.counters.read();
counters.keys().cloned().collect()
}
/// Check if we need to transfer ownership due to node failure
pub fn check_ownership_transfer(&self, failed_nodes: &[String]) -> Vec<String> {
let mut affected_keys = Vec::new();
let ring = self.hash_ring.read();
let counters = self.counters.read();
for key in counters.keys() {
let owners = ring.get_owners(key);
if owners
.iter()
.any(|&owner| failed_nodes.iter().any(|f| f == owner))
{
// Check if we are now an owner
if ring.is_owner(key, &self.self_name) {
affected_keys.push(key.clone());
}
}
}
affected_keys
}
}
impl Default for RateLimitStore {
fn default() -> Self {
Self::new("default".to_string())
}
}
/// All state stores container
#[derive(Debug, Clone)]
pub struct StateStores {
pub membership: MembershipStore,
pub app: AppStore,
pub worker: WorkerStore,
pub policy: PolicyStore,
pub rate_limit: RateLimitStore,
}
impl StateStores {
pub fn new() -> Self {
Self {
membership: MembershipStore::new(),
app: AppStore::new(),
worker: WorkerStore::new(),
policy: PolicyStore::new(),
rate_limit: RateLimitStore::new("default".to_string()),
}
}
pub fn with_self_name(self_name: String) -> Self {
Self {
membership: MembershipStore::new(),
app: AppStore::new(),
worker: WorkerStore::new(),
policy: PolicyStore::new(),
rate_limit: RateLimitStore::new(self_name),
}
}
}
impl Default for StateStores {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::*;
use crate::mesh::service::gossip::NodeStatus;
#[test]
fn test_membership_store() {
let store = MembershipStore::new();
let key = SKey::new("node1".to_string());
let state = MembershipState {
name: "node1".to_string(),
address: "127.0.0.1:8000".to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: BTreeMap::new(),
};
store.insert(key.clone(), state.clone(), "node1".to_string());
assert_eq!(store.get(&key).unwrap().name, "node1");
assert_eq!(store.len(), 1);
store.remove(&key);
assert!(store.get(&key).is_none());
}
#[test]
fn test_app_store() {
let store = AppStore::new();
let key = SKey::new("app_key1".to_string());
let state = AppState {
key: "app_key1".to_string(),
value: b"app_value".to_vec(),
version: 1,
};
store.insert(key.clone(), state.clone(), "node1".to_string());
assert_eq!(store.get(&key).unwrap().key, "app_key1");
assert_eq!(store.len(), 1);
}
#[test]
fn test_worker_store() {
let store = WorkerStore::new();
let key = SKey::new("worker1".to_string());
let state = WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: true,
load: 0.5,
version: 1,
};
store.insert(key.clone(), state.clone(), "node1".to_string());
assert_eq!(store.get(&key).unwrap().worker_id, "worker1");
assert_eq!(store.len(), 1);
}
#[test]
fn test_policy_store() {
let store = PolicyStore::new();
let key = SKey::new("policy:model1".to_string());
let state = PolicyState {
model_id: "model1".to_string(),
policy_type: "cache_aware".to_string(),
config: b"config_data".to_vec(),
version: 1,
};
store.insert(key.clone(), state.clone(), "node1".to_string());
assert_eq!(store.get(&key).unwrap().model_id, "model1");
assert_eq!(store.len(), 1);
}
#[test]
fn test_rate_limit_store_update_membership() {
let store = RateLimitStore::new("node1".to_string());
store.update_membership(&[
"node1".to_string(),
"node2".to_string(),
"node3".to_string(),
]);
let owners = store.get_owners("test_key");
assert_eq!(owners.len(), 3);
assert!(
owners.contains(&"node1".to_string())
|| owners.contains(&"node2".to_string())
|| owners.contains(&"node3".to_string())
);
}
#[test]
fn test_rate_limit_store_is_owner() {
let store = RateLimitStore::new("node1".to_string());
store.update_membership(&["node1".to_string()]);
let test_key = "test_key".to_string();
let is_owner = store.is_owner(&test_key);
// node1 should be owner since it's the only node
assert!(is_owner);
}
#[test]
fn test_rate_limit_store_inc_only_owner() {
let store = RateLimitStore::new("node1".to_string());
store.update_membership(&["node1".to_string()]);
let test_key = "test_key".to_string();
if store.is_owner(&test_key) {
store.inc(test_key.clone(), "node1".to_string(), 5);
let value = store.value(&test_key);
assert_eq!(value, Some(5));
}
}
#[test]
fn test_rate_limit_store_inc_non_owner() {
let store = RateLimitStore::new("node1".to_string());
// Setup membership without node1 as owner
store.update_membership(&["node2".to_string(), "node3".to_string()]);
let test_key = "test_key".to_string();
if !store.is_owner(&test_key) {
store.inc(test_key.clone(), "node1".to_string(), 5);
// Should not increment if not owner
let value = store.value(&test_key);
assert_eq!(value, None);
}
}
#[test]
fn test_rate_limit_store_merge_counter() {
let store1 = RateLimitStore::new("node1".to_string());
let store2 = RateLimitStore::new("node2".to_string());
store1.update_membership(&["node1".to_string()]);
store2.update_membership(&["node2".to_string()]);
let test_key = "test_key".to_string();
// Both nodes increment their counters
if store1.is_owner(&test_key) {
store1.inc(test_key.clone(), "node1".to_string(), 10);
}
if store2.is_owner(&test_key) {
store2.inc(test_key.clone(), "node2".to_string(), 5);
}
// Merge counter from store2 into store1
if let Some(counter2) = store2.get_counter(&test_key) {
store1.merge_counter(test_key.clone(), &counter2);
}
// Get aggregated value (if node1 is owner)
if store1.is_owner(&test_key) {
let value = store1.value(&test_key);
// Should include merged value
assert!(value.is_some());
}
}
#[test]
fn test_rate_limit_store_check_ownership_transfer() {
let store = RateLimitStore::new("node1".to_string());
store.update_membership(&[
"node1".to_string(),
"node2".to_string(),
"node3".to_string(),
]);
let test_key = "test_key".to_string();
// Setup a counter (if node1 is owner)
if store.is_owner(&test_key) {
store.inc(test_key.clone(), "node1".to_string(), 10);
}
// Check ownership transfer when node2 fails
let affected = store.check_ownership_transfer(&["node2".to_string()]);
// Should detect if node2 was an owner
let _ = affected;
}
#[test]
fn test_rate_limit_store_keys() {
let store = RateLimitStore::new("node1".to_string());
store.update_membership(&["node1".to_string()]);
let key1 = "key1".to_string();
let key2 = "key2".to_string();
if store.is_owner(&key1) {
store.inc(key1.clone(), "node1".to_string(), 1);
}
if store.is_owner(&key2) {
store.inc(key2.clone(), "node1".to_string(), 1);
}
let keys = store.keys();
// Should contain keys where node1 is owner
let _ = keys;
}
#[test]
fn test_state_stores_new() {
let stores = StateStores::new();
assert_eq!(stores.membership.len(), 0);
assert_eq!(stores.app.len(), 0);
assert_eq!(stores.worker.len(), 0);
assert_eq!(stores.policy.len(), 0);
}
#[test]
fn test_state_stores_with_self_name() {
let stores = StateStores::with_self_name("test_node".to_string());
// Rate limit store should have the self_name
let test_key = "test_key".to_string();
stores
.rate_limit
.update_membership(&["test_node".to_string()]);
assert!(stores.rate_limit.is_owner(&test_key));
}
}

1177
sgl-model-gateway/src/mesh/sync.rs
View File

File diff suppressed because it is too large Load Diff

86
sgl-model-gateway/src/mesh/test_utils.rs
View File

@@ -1,86 +0,0 @@
//! Test utilities for mesh module
use std::{
collections::{BTreeMap, HashMap},
sync::Arc,
};
use parking_lot::RwLock;
use super::{
service::{gossip::NodeState, ClusterState},
stores::{MembershipState, StateStores},
sync::MeshSyncManager,
};
/// Create test StateStores with a given node name
pub fn create_test_stores(self_name: String) -> Arc<StateStores> {
Arc::new(StateStores::with_self_name(self_name))
}
/// Create test MeshSyncManager
pub fn create_test_sync_manager(self_name: String) -> Arc<MeshSyncManager> {
let stores = create_test_stores(self_name.clone());
Arc::new(MeshSyncManager::new(stores, self_name))
}
/// Create test cluster state with given nodes
pub fn create_test_cluster_state(
nodes: Vec<(String, String, i32)>, // (name, address, status)
) -> ClusterState {
let mut state = BTreeMap::new();
for (name, address, status) in nodes {
state.insert(
name.clone(),
NodeState {
name: name.clone(),
address,
status,
version: 1,
metadata: HashMap::new(),
},
);
}
Arc::new(RwLock::new(state))
}
/// Create test membership state
#[allow(dead_code)]
pub fn create_test_membership_state(name: String, address: String, status: i32) -> MembershipState {
MembershipState {
name,
address,
status,
version: 1,
metadata: BTreeMap::new(),
}
}
#[cfg(test)]
mod test_utils_tests {
use super::*;
#[test]
fn test_create_test_stores() {
let stores = create_test_stores("test_node".to_string());
assert!(!stores.rate_limit.is_owner("key1"));
}
#[test]
fn test_create_test_sync_manager() {
let manager = create_test_sync_manager("test_node".to_string());
assert_eq!(manager.self_name(), "test_node");
}
#[test]
fn test_create_test_cluster_state() {
let state = create_test_cluster_state(vec![
("node1".to_string(), "127.0.0.1:8000".to_string(), 1),
("node2".to_string(), "127.0.0.1:8001".to_string(), 1),
]);
let read_state = state.read();
assert_eq!(read_state.len(), 2);
assert!(read_state.contains_key("node1"));
assert!(read_state.contains_key("node2"));
}
}

629
sgl-model-gateway/src/mesh/topology.rs
View File

@@ -1,629 +0,0 @@
//! Topology management for mesh cluster
//!
//! Supports:
//! - Full mesh for small/medium clusters
//! - Sparse mesh for large clusters (by region/AZ)
use std::{
collections::{BTreeMap, HashSet},
sync::Arc,
};
use parking_lot::RwLock;
use tracing::debug;
use super::{service::ClusterState, stores::MembershipState};
/// Topology configuration
#[derive(Debug, Clone)]
pub struct TopologyConfig {
/// Maximum nodes for full mesh (beyond this, use sparse)
pub full_mesh_threshold: usize,
/// Region identifier (for sparse mesh)
pub region: Option<String>,
/// Availability zone identifier (for sparse mesh)
pub availability_zone: Option<String>,
}
impl Default for TopologyConfig {
fn default() -> Self {
Self {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
}
}
}
/// Topology manager
pub struct TopologyManager {
config: TopologyConfig,
state: ClusterState,
self_name: String,
/// Active peer connections (for sparse mesh)
active_peers: Arc<RwLock<HashSet<String>>>,
}
impl TopologyManager {
pub fn new(config: TopologyConfig, state: ClusterState, self_name: String) -> Self {
Self {
config,
state,
self_name,
active_peers: Arc::new(RwLock::new(HashSet::new())),
}
}
/// Get peers to connect to based on topology
pub fn get_peers(&self, count: usize) -> Vec<MembershipState> {
let state = self.state.read();
let total_nodes = state.len();
if total_nodes <= self.config.full_mesh_threshold {
// Full mesh: connect to all nodes
self.get_full_mesh_peers(&state, count)
} else {
// Sparse mesh: connect based on region/AZ
self.get_sparse_mesh_peers(&state, count)
}
}
/// Get peers for full mesh topology
fn get_full_mesh_peers(
&self,
state: &BTreeMap<String, super::gossip::NodeState>,
count: usize,
) -> Vec<MembershipState> {
let mut peers = Vec::new();
let active = self.active_peers.read();
for (name, node) in state.iter() {
if name != &self.self_name
&& node.status == super::gossip::NodeStatus::Alive as i32
&& !active.contains(name)
{
let metadata: BTreeMap<String, Vec<u8>> = node
.metadata
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect::<BTreeMap<_, _>>();
peers.push(MembershipState {
name: node.name.clone(),
address: node.address.clone(),
status: node.status,
version: node.version,
metadata,
});
if peers.len() >= count {
break;
}
}
}
peers
}
/// Get peers for sparse mesh topology (by region/AZ)
fn get_sparse_mesh_peers(
&self,
state: &BTreeMap<String, super::gossip::NodeState>,
count: usize,
) -> Vec<MembershipState> {
let mut peers = Vec::new();
let active = self.active_peers.read();
// First, try to connect to nodes in same region/AZ
if let (Some(ref region), Some(ref az)) =
(&self.config.region, &self.config.availability_zone)
{
for (name, node) in state.iter() {
if name != &self.self_name
&& node.status == super::gossip::NodeStatus::Alive as i32
&& !active.contains(name)
{
// Check if node is in same region/AZ (from metadata)
let node_region = node
.metadata
.get("region")
.and_then(|v| String::from_utf8(v.clone()).ok());
let node_az = node
.metadata
.get("availability_zone")
.and_then(|v| String::from_utf8(v.clone()).ok());
if node_region.as_ref() == Some(region) && node_az.as_ref() == Some(az) {
let metadata: BTreeMap<String, Vec<u8>> = node
.metadata
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
peers.push(MembershipState {
name: node.name.clone(),
address: node.address.clone(),
status: node.status,
version: node.version,
metadata,
});
if peers.len() >= count {
break;
}
}
}
}
}
// If not enough peers, add from other regions
if peers.len() < count {
for (name, node) in state.iter() {
if name != &self.self_name
&& node.status == super::gossip::NodeStatus::Alive as i32
&& !active.contains(name)
&& !peers.iter().any(|p| p.name == node.name)
{
let metadata: BTreeMap<String, Vec<u8>> = node
.metadata
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
peers.push(MembershipState {
name: node.name.clone(),
address: node.address.clone(),
status: node.status,
version: node.version,
metadata,
});
if peers.len() >= count {
break;
}
}
}
}
peers
}
/// Mark peer as active
pub fn mark_peer_active(&self, peer_name: &str) {
self.active_peers.write().insert(peer_name.to_string());
debug!("Marked peer {} as active", peer_name);
}
/// Mark peer as inactive
pub fn mark_peer_inactive(&self, peer_name: &str) {
self.active_peers.write().remove(peer_name);
debug!("Marked peer {} as inactive", peer_name);
}
/// Get number of active peers
pub fn active_peer_count(&self) -> usize {
self.active_peers.read().len()
}
/// Check if we should use full mesh
pub fn is_full_mesh(&self) -> bool {
let state = self.state.read();
state.len() <= self.config.full_mesh_threshold
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::*;
use crate::mesh::service::gossip::{NodeState, NodeStatus};
fn create_test_cluster_state(nodes: Vec<(String, String, i32)>) -> ClusterState {
let mut state = BTreeMap::new();
for (name, address, status) in nodes {
state.insert(
name.clone(),
NodeState {
name: name.clone(),
address,
status,
version: 1,
metadata: std::collections::HashMap::new(),
},
);
}
Arc::new(RwLock::new(state))
}
#[test]
fn test_full_mesh_topology() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
(
"node3".to_string(),
"127.0.0.1:8002".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
let peers = manager.get_peers(5);
// Should return all available peers (node2 and node3)
assert_eq!(peers.len(), 2);
assert!(peers.iter().any(|p| p.name == "node2"));
assert!(peers.iter().any(|p| p.name == "node3"));
}
#[test]
fn test_full_mesh_topology_excludes_self() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
let peers = manager.get_peers(5);
// Should not include self (node1)
assert_eq!(peers.len(), 1);
assert_eq!(peers[0].name, "node2");
}
#[test]
fn test_full_mesh_topology_filters_down_nodes() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Down as i32,
),
(
"node3".to_string(),
"127.0.0.1:8002".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
let peers = manager.get_peers(5);
// Should only return alive nodes (node3)
assert_eq!(peers.len(), 1);
assert_eq!(peers[0].name, "node3");
}
#[test]
fn test_sparse_mesh_topology() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
(
"node3".to_string(),
"127.0.0.1:8002".to_string(),
NodeStatus::Alive as i32,
),
(
"node4".to_string(),
"127.0.0.1:8003".to_string(),
NodeStatus::Alive as i32,
),
(
"node5".to_string(),
"127.0.0.1:8004".to_string(),
NodeStatus::Alive as i32,
),
(
"node6".to_string(),
"127.0.0.1:8005".to_string(),
NodeStatus::Alive as i32,
),
(
"node7".to_string(),
"127.0.0.1:8006".to_string(),
NodeStatus::Alive as i32,
),
(
"node8".to_string(),
"127.0.0.1:8007".to_string(),
NodeStatus::Alive as i32,
),
(
"node9".to_string(),
"127.0.0.1:8008".to_string(),
NodeStatus::Alive as i32,
),
(
"node10".to_string(),
"127.0.0.1:8009".to_string(),
NodeStatus::Alive as i32,
),
(
"node11".to_string(),
"127.0.0.1:8010".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10, // 11 nodes > 10, should use sparse
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
let peers = manager.get_peers(5);
// Should return peers (sparse mesh mode)
assert!(!peers.is_empty());
assert!(peers.len() <= 5);
}
#[test]
fn test_sparse_mesh_with_region_az() {
let mut state_map = BTreeMap::new();
// Create nodes with region/AZ metadata
let mut node1_metadata = std::collections::HashMap::new();
node1_metadata.insert("region".to_string(), b"us-west".to_vec());
node1_metadata.insert("availability_zone".to_string(), b"us-west-1a".to_vec());
state_map.insert(
"node1".to_string(),
NodeState {
name: "node1".to_string(),
address: "127.0.0.1:8000".to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: node1_metadata.clone(),
},
);
let mut node2_metadata = std::collections::HashMap::new();
node2_metadata.insert("region".to_string(), b"us-west".to_vec());
node2_metadata.insert("availability_zone".to_string(), b"us-west-1a".to_vec());
state_map.insert(
"node2".to_string(),
NodeState {
name: "node2".to_string(),
address: "127.0.0.1:8001".to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: node2_metadata,
},
);
let mut node3_metadata = std::collections::HashMap::new();
node3_metadata.insert("region".to_string(), b"us-east".to_vec());
node3_metadata.insert("availability_zone".to_string(), b"us-east-1a".to_vec());
state_map.insert(
"node3".to_string(),
NodeState {
name: "node3".to_string(),
address: "127.0.0.1:8002".to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: node3_metadata,
},
);
let state = Arc::new(RwLock::new(state_map));
let config = TopologyConfig {
full_mesh_threshold: 2,
region: Some("us-west".to_string()),
availability_zone: Some("us-west-1a".to_string()),
};
let manager = TopologyManager::new(config, state, "node1".to_string());
let peers = manager.get_peers(5);
// Should prefer nodes in same region/AZ (node2)
assert!(!peers.is_empty());
// node2 should be in the list (same region/AZ)
assert!(peers.iter().any(|p| p.name == "node2"));
}
#[test]
fn test_mark_peer_active_inactive() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
assert_eq!(manager.active_peer_count(), 0);
manager.mark_peer_active("node2");
assert_eq!(manager.active_peer_count(), 1);
manager.mark_peer_inactive("node2");
assert_eq!(manager.active_peer_count(), 0);
}
#[test]
fn test_get_peers_excludes_active_peers() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
(
"node3".to_string(),
"127.0.0.1:8002".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
manager.mark_peer_active("node2");
let peers = manager.get_peers(5);
// Should exclude node2 (already active)
assert!(!peers.iter().any(|p| p.name == "node2"));
// Should include node3
assert!(peers.iter().any(|p| p.name == "node3"));
}
#[test]
fn test_is_full_mesh() {
let state = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
]);
let config = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager = TopologyManager::new(config, state, "node1".to_string());
assert!(manager.is_full_mesh());
let state2 = create_test_cluster_state(vec![
(
"node1".to_string(),
"127.0.0.1:8000".to_string(),
NodeStatus::Alive as i32,
),
(
"node2".to_string(),
"127.0.0.1:8001".to_string(),
NodeStatus::Alive as i32,
),
(
"node3".to_string(),
"127.0.0.1:8002".to_string(),
NodeStatus::Alive as i32,
),
(
"node4".to_string(),
"127.0.0.1:8003".to_string(),
NodeStatus::Alive as i32,
),
(
"node5".to_string(),
"127.0.0.1:8004".to_string(),
NodeStatus::Alive as i32,
),
(
"node6".to_string(),
"127.0.0.1:8005".to_string(),
NodeStatus::Alive as i32,
),
(
"node7".to_string(),
"127.0.0.1:8006".to_string(),
NodeStatus::Alive as i32,
),
(
"node8".to_string(),
"127.0.0.1:8007".to_string(),
NodeStatus::Alive as i32,
),
(
"node9".to_string(),
"127.0.0.1:8008".to_string(),
NodeStatus::Alive as i32,
),
(
"node10".to_string(),
"127.0.0.1:8009".to_string(),
NodeStatus::Alive as i32,
),
(
"node11".to_string(),
"127.0.0.1:8010".to_string(),
NodeStatus::Alive as i32,
),
]);
let config2 = TopologyConfig {
full_mesh_threshold: 10,
region: None,
availability_zone: None,
};
let manager2 = TopologyManager::new(config2, state2, "node1".to_string());
assert!(!manager2.is_full_mesh());
}
}

274
sgl-model-gateway/src/mesh/tree_ops.rs
View File

@@ -1,274 +0,0 @@
//! Tree operation definitions for mesh synchronization
//!
//! Defines serializable tree operations that can be synchronized across mesh cluster nodes
use serde::{Deserialize, Serialize};
/// Tree insert operation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct TreeInsertOp {
pub text: String,
pub tenant: String, // worker URL
}
/// Tree remove operation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct TreeRemoveOp {
pub tenant: String, // worker URL
}
/// Tree operation type
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum TreeOperation {
Insert(TreeInsertOp),
Remove(TreeRemoveOp),
}
/// Tree state for a specific model
/// Contains a sequence of operations that can be applied to reconstruct the tree
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
pub struct TreeState {
pub model_id: String,
pub operations: Vec<TreeOperation>,
pub version: u64,
}
impl TreeState {
pub fn new(model_id: String) -> Self {
Self {
model_id,
operations: Vec::new(),
version: 0,
}
}
pub fn add_operation(&mut self, operation: TreeOperation) {
self.operations.push(operation);
self.version += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tree_insert_op_creation() {
let op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
assert_eq!(op.text, "test_text");
assert_eq!(op.tenant, "http://worker1:8000");
}
#[test]
fn test_tree_remove_op_creation() {
let op = TreeRemoveOp {
tenant: "http://worker1:8000".to_string(),
};
assert_eq!(op.tenant, "http://worker1:8000");
}
#[test]
fn test_tree_operation_insert() {
let insert_op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
let operation = TreeOperation::Insert(insert_op.clone());
match &operation {
TreeOperation::Insert(op) => {
assert_eq!(op.text, "test_text");
assert_eq!(op.tenant, "http://worker1:8000");
}
TreeOperation::Remove(_) => panic!("Expected Insert operation"),
}
}
#[test]
fn test_tree_operation_remove() {
let remove_op = TreeRemoveOp {
tenant: "http://worker1:8000".to_string(),
};
let operation = TreeOperation::Remove(remove_op.clone());
match &operation {
TreeOperation::Insert(_) => panic!("Expected Remove operation"),
TreeOperation::Remove(op) => {
assert_eq!(op.tenant, "http://worker1:8000");
}
}
}
#[test]
fn test_tree_operation_serialization() {
let insert_op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
let operation = TreeOperation::Insert(insert_op);
let serialized = serde_json::to_string(&operation).unwrap();
let deserialized: TreeOperation = serde_json::from_str(&serialized).unwrap();
match (&operation, &deserialized) {
(TreeOperation::Insert(a), TreeOperation::Insert(b)) => {
assert_eq!(a.text, b.text);
assert_eq!(a.tenant, b.tenant);
}
_ => panic!("Operations should match"),
}
}
#[test]
fn test_tree_operation_remove_serialization() {
let remove_op = TreeRemoveOp {
tenant: "http://worker1:8000".to_string(),
};
let operation = TreeOperation::Remove(remove_op);
let serialized = serde_json::to_string(&operation).unwrap();
let deserialized: TreeOperation = serde_json::from_str(&serialized).unwrap();
match (&operation, &deserialized) {
(TreeOperation::Remove(a), TreeOperation::Remove(b)) => {
assert_eq!(a.tenant, b.tenant);
}
_ => panic!("Operations should match"),
}
}
#[test]
fn test_tree_state_new() {
let state = TreeState::new("model1".to_string());
assert_eq!(state.model_id, "model1");
assert_eq!(state.operations.len(), 0);
assert_eq!(state.version, 0);
}
#[test]
fn test_tree_state_default() {
let state = TreeState::default();
assert_eq!(state.model_id, "");
assert_eq!(state.operations.len(), 0);
assert_eq!(state.version, 0);
}
#[test]
fn test_tree_state_add_operation() {
let mut state = TreeState::new("model1".to_string());
let insert_op = TreeInsertOp {
text: "text1".to_string(),
tenant: "http://worker1:8000".to_string(),
};
state.add_operation(TreeOperation::Insert(insert_op));
assert_eq!(state.operations.len(), 1);
assert_eq!(state.version, 1);
let remove_op = TreeRemoveOp {
tenant: "http://worker1:8000".to_string(),
};
state.add_operation(TreeOperation::Remove(remove_op));
assert_eq!(state.operations.len(), 2);
assert_eq!(state.version, 2);
}
#[test]
fn test_tree_state_add_multiple_operations() {
let mut state = TreeState::new("model1".to_string());
for i in 0..5 {
let insert_op = TreeInsertOp {
text: format!("text_{}", i),
tenant: format!("http://worker{}:8000", i),
};
state.add_operation(TreeOperation::Insert(insert_op));
}
assert_eq!(state.operations.len(), 5);
assert_eq!(state.version, 5);
}
#[test]
fn test_tree_state_serialization() {
let mut state = TreeState::new("model1".to_string());
let insert_op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
state.add_operation(TreeOperation::Insert(insert_op));
let remove_op = TreeRemoveOp {
tenant: "http://worker1:8000".to_string(),
};
state.add_operation(TreeOperation::Remove(remove_op));
let serialized = serde_json::to_string(&state).unwrap();
let deserialized: TreeState = serde_json::from_str(&serialized).unwrap();
assert_eq!(state.model_id, deserialized.model_id);
assert_eq!(state.operations.len(), deserialized.operations.len());
assert_eq!(state.version, deserialized.version);
}
#[test]
fn test_tree_state_clone() {
let mut state = TreeState::new("model1".to_string());
let insert_op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
state.add_operation(TreeOperation::Insert(insert_op));
let cloned = state.clone();
assert_eq!(state.model_id, cloned.model_id);
assert_eq!(state.operations.len(), cloned.operations.len());
assert_eq!(state.version, cloned.version);
}
#[test]
fn test_tree_state_equality() {
let mut state1 = TreeState::new("model1".to_string());
let mut state2 = TreeState::new("model1".to_string());
let insert_op = TreeInsertOp {
text: "test_text".to_string(),
tenant: "http://worker1:8000".to_string(),
};
state1.add_operation(TreeOperation::Insert(insert_op.clone()));
state2.add_operation(TreeOperation::Insert(insert_op));
assert_eq!(state1, state2);
}
#[test]
fn test_tree_operation_hash() {
use std::collections::HashSet;
let insert_op1 = TreeInsertOp {
text: "text1".to_string(),
tenant: "http://worker1:8000".to_string(),
};
let insert_op2 = TreeInsertOp {
text: "text1".to_string(),
tenant: "http://worker1:8000".to_string(),
};
let op1 = TreeOperation::Insert(insert_op1);
let op2 = TreeOperation::Insert(insert_op2);
let mut set = HashSet::new();
set.insert(op1.clone());
set.insert(op2.clone());
// Same operations should be considered equal
assert_eq!(set.len(), 1);
}
}

16
sgl-model-gateway/src/mesh/endpoints.rs → sgl-model-gateway/src/routers/mesh/handlers.rs
View File

@@ -1,10 +1,12 @@
//! Mesh management endpoints
//! Mesh management HTTP handlers
//!
//! Provides REST API for mesh cluster management:
//! - Configuration CRUD operations
//! - Health checks
//! - Cluster status
use std::sync::Arc;
use axum::{
extract::{Path, State},
http::StatusCode,
@@ -13,8 +15,11 @@ use axum::{
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use smg_mesh::{RateLimitConfig, GLOBAL_RATE_LIMIT_COUNTER_KEY, GLOBAL_RATE_LIMIT_KEY};
use tracing::{info, warn};
use crate::server::AppState;
/// Mesh cluster status response
#[derive(Debug, Serialize, Deserialize)]
pub struct ClusterStatusResponse {
@@ -392,7 +397,7 @@ pub async fn get_global_rate_limit_stats(State(app_state): State<Arc<AppState>>)
// Get current counter value
let current_count = sync_manager
.get_rate_limit_value(crate::mesh::stores::GLOBAL_RATE_LIMIT_COUNTER_KEY)
.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY)
.unwrap_or(0);
(
@@ -434,10 +439,3 @@ pub async fn trigger_graceful_shutdown(State(app_state): State<Arc<AppState>>) -
)
.into_response()
}
use std::sync::Arc;
use crate::{
mesh::stores::{RateLimitConfig, GLOBAL_RATE_LIMIT_KEY},
server::AppState,
};

7
sgl-model-gateway/src/routers/mesh/mod.rs Normal file
View File

@@ -0,0 +1,7 @@
//! Mesh cluster management HTTP handlers
//!
//! This module provides HTTP API endpoints for mesh cluster management.
mod handlers;
pub use handlers::*;

1
sgl-model-gateway/src/routers/mod.rs
View File

@@ -27,6 +27,7 @@ pub mod grpc;
pub mod header_utils;
pub mod http;
pub mod mcp_utils;
pub mod mesh;
pub mod openai;
pub mod parse;
pub mod persistence_utils;

25
sgl-model-gateway/src/server.rs
View File

@@ -16,6 +16,9 @@ use axum::{
use rustls::crypto::ring;
use serde::Deserialize;
use serde_json::{json, Value};
use smg_mesh::{
rate_limit_window::RateLimitWindow, MeshServerConfig, MeshServerHandler, MeshSyncManager,
};
use tokio::{signal, spawn};
use tracing::{debug, error, info, warn, Level};
@@ -29,16 +32,6 @@ use crate::{
worker_manager::WorkerManager,
Job,
},
mesh::{
endpoints::{
get_app_config, get_cluster_status, get_global_rate_limit, get_global_rate_limit_stats,
get_mesh_health, get_policy_state, get_policy_states, get_worker_state,
get_worker_states, set_global_rate_limit, trigger_graceful_shutdown, update_app_config,
},
rate_limit_window::RateLimitWindow,
service::{MeshServerConfig, MeshServerHandler},
sync::MeshSyncManager,
},
middleware::{self, AuthConfig, QueuedRequest},
observability::{
logging::{self, LoggingConfig},
@@ -58,7 +51,17 @@ use crate::{
validated::ValidatedJson,
worker_spec::{WorkerConfigRequest, WorkerUpdateRequest},
},
routers::{conversations, parse, router_manager::RouterManager, tokenize, RouterTrait},
routers::{
conversations,
mesh::{
get_app_config, get_cluster_status, get_global_rate_limit, get_global_rate_limit_stats,
get_mesh_health, get_policy_state, get_policy_states, get_worker_state,
get_worker_states, set_global_rate_limit, trigger_graceful_shutdown, update_app_config,
},
parse,
router_manager::RouterManager,
tokenize, RouterTrait,
},
service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
tokenizer::TokenizerRegistry,
wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},

384
sgl-model-gateway/tests/mesh_integration_test.rs
View File

@@ -1,384 +0,0 @@
//! Integration tests for mesh functionality
//!
//! Tests multi-node scenarios including state synchronization,
//! rate limiting, and cache-aware routing across cluster nodes.
use std::sync::Arc;
use smg::mesh::{
crdt::SKey,
gossip::NodeStatus,
stores::{
AppState, MembershipState, RateLimitConfig, StateStores, WorkerState,
GLOBAL_RATE_LIMIT_COUNTER_KEY, GLOBAL_RATE_LIMIT_KEY,
},
sync::MeshSyncManager,
tree_ops::{TreeInsertOp, TreeOperation},
};
/// Create test stores for a node
fn create_test_stores(node_name: String) -> Arc<StateStores> {
Arc::new(StateStores::with_self_name(node_name))
}
/// Create test sync manager for a node
fn create_test_sync_manager(node_name: String) -> Arc<MeshSyncManager> {
let stores = create_test_stores(node_name.clone());
Arc::new(MeshSyncManager::new(stores, node_name))
}
#[tokio::test]
async fn test_multi_node_state_synchronization() {
// Create three nodes
let manager1 = create_test_sync_manager("node1".to_string());
let manager2 = create_test_sync_manager("node2".to_string());
let manager3 = create_test_sync_manager("node3".to_string());
// Node1 syncs a worker state
manager1.sync_worker_state(
"worker1".to_string(),
"model1".to_string(),
"http://localhost:8000".to_string(),
true,
0.5,
);
// Simulate synchronization: Node2 and Node3 receive the update
let worker_state = manager1.get_worker_state("worker1").unwrap();
manager2.apply_remote_worker_state(worker_state.clone(), Some("node1".to_string()));
manager3.apply_remote_worker_state(worker_state, Some("node1".to_string()));
// Verify all nodes have the same state
let state1 = manager1.get_worker_state("worker1").unwrap();
let state2 = manager2.get_worker_state("worker1").unwrap();
let state3 = manager3.get_worker_state("worker1").unwrap();
assert_eq!(state1.worker_id, state2.worker_id);
assert_eq!(state2.worker_id, state3.worker_id);
assert_eq!(state1.version, state2.version);
assert_eq!(state2.version, state3.version);
}
#[tokio::test]
async fn test_node_join_and_leave() {
let manager1 = create_test_sync_manager("node1".to_string());
let manager2 = create_test_sync_manager("node2".to_string());
// Node1 has some state
manager1.sync_worker_state(
"worker1".to_string(),
"model1".to_string(),
"http://localhost:8000".to_string(),
true,
0.5,
);
manager1.sync_policy_state(
"model1".to_string(),
"cache_aware".to_string(),
b"config".to_vec(),
);
// Node2 joins and receives state
let worker_state = manager1.get_worker_state("worker1").unwrap();
manager2.apply_remote_worker_state(worker_state, Some("node1".to_string()));
let policy_state = manager1.get_policy_state("model1").unwrap();
manager2.apply_remote_policy_state(policy_state, Some("node1".to_string()));
// Verify Node2 has the state
assert!(manager2.get_worker_state("worker1").is_some());
assert!(manager2.get_policy_state("model1").is_some());
// Node1 removes worker
manager1.remove_worker_state("worker1");
// In a real scenario, this would be propagated via gossip
// For test, we verify the removal happened
assert!(manager1.get_worker_state("worker1").is_none());
}
#[tokio::test]
async fn test_rate_limit_cluster_consistency() {
// Create stores and managers
let stores1 = create_test_stores("node1".to_string());
let stores2 = create_test_stores("node2".to_string());
let stores3 = create_test_stores("node3".to_string());
// Add all nodes to membership store (required for rate limit hash ring)
let node_names = ["node1", "node2", "node3"];
let node_addresses = ["127.0.0.1:8001", "127.0.0.1:8002", "127.0.0.1:8003"];
for stores in [&stores1, &stores2, &stores3] {
for (i, &name) in node_names.iter().enumerate() {
let key = SKey::new(name.to_string());
stores.membership.insert(
key,
MembershipState {
name: name.to_string(),
address: node_addresses[i].to_string(),
status: NodeStatus::Alive as i32,
version: 1,
metadata: std::collections::BTreeMap::new(),
},
name.to_string(),
);
}
}
// Setup global rate limit config
let config = RateLimitConfig {
limit_per_second: 100,
};
let serialized = serde_json::to_vec(&config).unwrap();
let key = SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string());
for stores in [&stores1, &stores2, &stores3] {
stores.app.insert(
key.clone(),
AppState {
key: GLOBAL_RATE_LIMIT_KEY.to_string(),
value: serialized.clone(),
version: 1,
},
"node1".to_string(),
);
}
// Create managers with updated stores
let manager1 = Arc::new(MeshSyncManager::new(stores1.clone(), "node1".to_string()));
let manager2 = Arc::new(MeshSyncManager::new(stores2.clone(), "node2".to_string()));
let manager3 = Arc::new(MeshSyncManager::new(stores3.clone(), "node3".to_string()));
// Update rate limit membership (reads from membership store)
manager1.update_rate_limit_membership();
manager2.update_rate_limit_membership();
manager3.update_rate_limit_membership();
// Each node increments the counter (if it's an owner)
let test_key = GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string();
manager1.sync_rate_limit_inc(test_key.clone(), 10);
manager2.sync_rate_limit_inc(test_key.clone(), 5);
manager3.sync_rate_limit_inc(test_key.clone(), 3);
// Simulate counter merging (in real scenario, this happens via gossip)
// Get counters from each node and merge them into all nodes
if let Some(counter2) = stores2.rate_limit.get_counter(&test_key) {
manager1.apply_remote_rate_limit_counter(test_key.clone(), &counter2);
manager3.apply_remote_rate_limit_counter(test_key.clone(), &counter2);
}
if let Some(counter3) = stores3.rate_limit.get_counter(&test_key) {
manager1.apply_remote_rate_limit_counter(test_key.clone(), &counter3);
manager2.apply_remote_rate_limit_counter(test_key.clone(), &counter3);
}
if let Some(counter1) = stores1.rate_limit.get_counter(&test_key) {
manager2.apply_remote_rate_limit_counter(test_key.clone(), &counter1);
manager3.apply_remote_rate_limit_counter(test_key.clone(), &counter1);
}
// Check aggregated value
let value = manager1.get_rate_limit_value(&test_key);
// Should have aggregated value from all owners
assert!(value.is_some());
// The value should be the sum of all increments (10 + 5 + 3 = 18)
// But note: only owners actually increment, so the sum depends on ownership
let value = value.unwrap();
assert!(value > 0, "Counter value should be greater than 0");
}
#[tokio::test]
async fn test_rate_limit_node_failure() {
let manager1 = create_test_sync_manager("node1".to_string());
let _manager2 = create_test_sync_manager("node2".to_string());
let _manager3 = create_test_sync_manager("node3".to_string());
// Setup membership through sync manager
// In a real scenario, membership would be updated through gossip protocol
manager1.update_rate_limit_membership();
// Simulate node2 failure
manager1.handle_node_failure(&["node2".to_string()]);
// Update membership to reflect failure
manager1.update_rate_limit_membership();
// Verify system continues to work
let test_key = "test_key".to_string();
manager1.sync_rate_limit_inc(test_key.clone(), 1);
let _value = manager1.get_rate_limit_value(&test_key);
// Value may be None if not owner, which is acceptable
// In a real scenario, ownership would be redistributed after node failure
}
#[tokio::test]
async fn test_cache_aware_tree_synchronization() {
let manager1 = create_test_sync_manager("node1".to_string());
let manager2 = create_test_sync_manager("node2".to_string());
// Node1 syncs tree operations
let op1 = TreeOperation::Insert(TreeInsertOp {
text: "request1".to_string(),
tenant: "http://worker1:8000".to_string(),
});
manager1
.sync_tree_operation("model1".to_string(), op1)
.unwrap();
let op2 = TreeOperation::Insert(TreeInsertOp {
text: "request2".to_string(),
tenant: "http://worker2:8000".to_string(),
});
manager1
.sync_tree_operation("model1".to_string(), op2)
.unwrap();
// Node2 receives tree state (simulated)
let tree_state = manager1.get_tree_state("model1").unwrap();
manager2.apply_remote_tree_operation(
"model1".to_string(),
tree_state,
Some("node1".to_string()),
);
// Verify Node2 has the tree state
let tree_state2 = manager2.get_tree_state("model1");
assert!(tree_state2.is_some());
let tree = tree_state2.unwrap();
assert_eq!(tree.operations.len(), 2);
}
#[tokio::test]
async fn test_version_conflict_resolution() {
let manager1 = create_test_sync_manager("node1".to_string());
let manager2 = create_test_sync_manager("node2".to_string());
// Both nodes update the same worker with different versions
manager1.sync_worker_state(
"worker1".to_string(),
"model1".to_string(),
"http://localhost:8000".to_string(),
true,
0.5,
);
// Node2 tries to apply an older version
let old_state = WorkerState {
worker_id: "worker1".to_string(),
model_id: "model1".to_string(),
url: "http://localhost:8000".to_string(),
health: false,
load: 0.8,
version: 0, // Older version
};
manager2.apply_remote_worker_state(old_state, Some("node2".to_string()));
// Node2 should not have the state (version too old)
// But if it does, it should have version 0
let state2 = manager2.get_worker_state("worker1");
if let Some(s) = state2 {
// If state exists, it should be from node1 (version 1)
assert!(s.version >= 1);
}
// Node1 applies newer version to Node2
let new_state = manager1.get_worker_state("worker1").unwrap();
manager2.apply_remote_worker_state(new_state, Some("node1".to_string()));
// Now Node2 should have the correct state
let final_state = manager2.get_worker_state("worker1").unwrap();
assert_eq!(final_state.version, 1);
assert!(final_state.health);
}
#[tokio::test]
async fn test_concurrent_updates() {
let manager1 = create_test_sync_manager("node1".to_string());
let manager2 = create_test_sync_manager("node2".to_string());
let manager3 = create_test_sync_manager("node3".to_string());
// All nodes update different workers concurrently
manager1.sync_worker_state(
"worker1".to_string(),
"model1".to_string(),
"http://localhost:8000".to_string(),
true,
0.5,
);
manager2.sync_worker_state(
"worker2".to_string(),
"model1".to_string(),
"http://localhost:8001".to_string(),
true,
0.6,
);
manager3.sync_worker_state(
"worker3".to_string(),
"model1".to_string(),
"http://localhost:8002".to_string(),
true,
0.7,
);
// Simulate synchronization: all nodes receive all updates
let worker1_state = manager1.get_worker_state("worker1").unwrap();
let worker2_state = manager2.get_worker_state("worker2").unwrap();
let worker3_state = manager3.get_worker_state("worker3").unwrap();
manager2.apply_remote_worker_state(worker1_state.clone(), Some("node1".to_string()));
manager3.apply_remote_worker_state(worker1_state, Some("node1".to_string()));
manager1.apply_remote_worker_state(worker2_state.clone(), Some("node2".to_string()));
manager3.apply_remote_worker_state(worker2_state, Some("node2".to_string()));
manager1.apply_remote_worker_state(worker3_state.clone(), Some("node3".to_string()));
manager2.apply_remote_worker_state(worker3_state, Some("node3".to_string()));
// All nodes should have all workers
assert_eq!(manager1.get_all_worker_states().len(), 3);
assert_eq!(manager2.get_all_worker_states().len(), 3);
assert_eq!(manager3.get_all_worker_states().len(), 3);
}
#[tokio::test]
async fn test_rate_limit_window_reset() {
let manager = create_test_sync_manager("node1".to_string());
// Setup membership
manager.update_rate_limit_membership();
// Setup config through stores (for testing)
let stores = create_test_stores("node1".to_string());
let config = RateLimitConfig {
limit_per_second: 100,
};
let serialized = serde_json::to_vec(&config).unwrap();
let key = SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string());
stores.app.insert(
key,
AppState {
key: GLOBAL_RATE_LIMIT_KEY.to_string(),
value: serialized,
version: 1,
},
"node1".to_string(),
);
// Recreate manager with updated stores
let manager = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
// Increment counter (if owner)
manager.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 50);
let value_before = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
// Value may be None if not owner, or Some if owner
if let Some(val) = value_before {
assert!(val > 0);
// Reset counter
manager.reset_global_rate_limit_counter();
let value_after = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY);
// Should be reset
assert!(value_after.is_none() || value_after.unwrap_or(0) <= 0);
}
}