diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 422d2f8c8..cc35ad5f8 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -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" diff --git a/sgl-model-gateway/build.rs b/sgl-model-gateway/build.rs index 7dffd8c16..57a02656b 100644 --- a/sgl-model-gateway/build.rs +++ b/sgl-model-gateway/build.rs @@ -12,7 +12,6 @@ macro_rules! set_env { fn main() -> Result<(), Box> { // 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> { &["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()); diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index 83a5a1704..cd1af3ecd 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -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; diff --git a/sgl-model-gateway/src/mesh/README.md b/sgl-model-gateway/src/mesh/README.md deleted file mode 100644 index 04c88aefb..000000000 --- a/sgl-model-gateway/src/mesh/README.md +++ /dev/null @@ -1,1123 +0,0 @@ -# Mesh Module - -The Mesh module provides a distributed, eventually consistent state synchronization system for high-availability (HA) clusters. It enables multiple router nodes to share state information, coordinate rate limiting, and synchronize cache-aware routing policies across the cluster. - -## Table of Contents - -- [Overview](#overview) -- [Architecture](#architecture) -- [HTTP API Reference](#http-api-reference) -- [Rate Limiting](#rate-limiting) -- [Cache-Aware Routing](#cache-aware-routing) -- [Service Discovery Integration](#service-discovery-integration) -- [Usage Examples](#usage-examples) - -## Overview - -The Mesh module implements a gossip-based protocol for state synchronization across cluster nodes. It uses Conflict-free Replicated Data Types (CRDTs) to ensure eventual consistency without requiring strong coordination or consensus protocols. - -### Key Features - -- **Eventual Consistency**: State converges across all nodes using CRDTs -- **Gossip Protocol**: Efficient peer-to-peer state propagation -- **Rate Limiting**: Distributed rate limiting with consistent hashing -- **Cache-Aware Routing**: Synchronized cache state for optimal routing -- **Service Discovery**: Integration with service discovery for dynamic membership -- **Topology Management**: Supports both full mesh and sparse mesh topologies - -## Architecture - -``` -┌─────────────────────────────────────────────────────────────┐ -│ Mesh Module │ -├─────────────────────────────────────────────────────────────┤ -│ │ -│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ -│ │ Gossip │ │ CRDT │ │ Stores │ │ -│ │ Protocol │ │ Layer │ │ Layer │ │ -│ └──────────────┘ └──────────────┘ └──────────────┘ │ -│ │ │ │ │ -│ └─────────────────┴─────────────────┘ │ -│ │ │ -│ ┌──────────────────────────────────────────────┐ │ -│ │ MeshSyncManager │ │ -│ │ - Worker State Sync │ │ -│ │ - Policy State Sync │ │ -│ │ - Rate Limit Coordination │ │ -│ │ - Tree Operation Sync │ │ -│ └──────────────────────────────────────────────┘ │ -│ │ │ -│ ┌──────────────────────────────────────────────┐ │ -│ │ TopologyManager │ │ -│ │ - Full Mesh (≤ threshold) │ │ -│ │ - Sparse Mesh (> threshold, by region/AZ) │ │ -│ └──────────────────────────────────────────────┘ │ -│ │ -└─────────────────────────────────────────────────────────────┘ -``` - -## HTTP API Reference - -All mesh endpoints are prefixed with `/mesh` and require authentication. - -### Cluster Status - -#### GET `/mesh/status` - -Get the current cluster status including node information and store counts. - -**Response:** -```json -{ - "node_name": "node1", - "node_count": 3, - "nodes": [ - { - "name": "node1", - "address": "127.0.0.1:8000", - "status": "Alive", - "version": 1 - }, - { - "name": "node2", - "address": "127.0.0.1:8001", - "status": "Alive", - "version": 1 - } - ], - "stores": { - "membership_count": 3, - "worker_count": 0, - "policy_count": 0, - "app_count": 0 - } -} -``` - -**Status Codes:** -- `200 OK`: Success -- `503 Service Unavailable`: Mesh not enabled - -### Health Check - -#### GET `/mesh/health` - -Get the health status of the mesh cluster. - -**Response:** -```json -{ - "status": "healthy", - "node_name": "node1", - "cluster_size": 3, - "stores_healthy": true -} -``` - -**Status Codes:** -- `200 OK`: Success -- `503 Service Unavailable`: Mesh not enabled - -### Worker States - -#### GET `/mesh/workers` - -Get all worker states from the mesh store. - -**Response:** -```json -[ - { - "worker_id": "worker-1", - "model_id": "model-1", - "url": "http://worker1:8000", - "health": true, - "load": 0.75, - "version": 1 - }, - { - "worker_id": "worker-2", - "model_id": "model-1", - "url": "http://worker2:8000", - "health": true, - "load": 0.50, - "version": 1 - } -] -``` - -**Status Codes:** -- `200 OK`: Success -- `503 Service Unavailable`: Mesh sync manager not available - -#### GET `/mesh/workers/{worker_id}` - -Get a specific worker state by worker ID. - -**Path Parameters:** -- `worker_id` (string): The worker identifier - -**Response:** -```json -{ - "worker_id": "worker-1", - "model_id": "model-1", - "url": "http://worker1:8000", - "health": true, - "load": 0.75, - "version": 1 -} -``` - -**Status Codes:** -- `200 OK`: Success -- `404 Not Found`: Worker not found -- `503 Service Unavailable`: Mesh sync manager not available - -### Policy States - -#### GET `/mesh/policies` - -Get all policy states from the mesh store. - -**Response:** -```json -[ - { - "model_id": "model-1", - "policy_type": "cache_aware", - "config": "...", - "version": 1 - } -] -``` - -**Status Codes:** -- `200 OK`: Success -- `503 Service Unavailable`: Mesh sync manager not available - -#### GET `/mesh/policies/{model_id}` - -Get a specific policy state by model ID. - -**Path Parameters:** -- `model_id` (string): The model identifier - -**Response:** -```json -{ - "model_id": "model-1", - "policy_type": "cache_aware", - "config": "...", - "version": 1 -} -``` - -**Status Codes:** -- `200 OK`: Success -- `404 Not Found`: Policy not found -- `503 Service Unavailable`: Mesh sync manager not available - -### Application Configuration - -#### GET `/mesh/config/{key}` - -Get application configuration by key. - -**Path Parameters:** -- `key` (string): The configuration key - -**Response:** -```json -{ - "key": "config_key", - "value": "68656c6c6f", // Hex-encoded value - "format": "hex" -} -``` - -**Status Codes:** -- `200 OK`: Success -- `404 Not Found`: Config not found -- `503 Service Unavailable`: Mesh not enabled - -#### POST `/mesh/config` - -Update application configuration. - -**Request Body:** -```json -{ - "key": "config_key", - "value": "68656c6c6f" // Hex-encoded string (even length) -} -``` - -**Response:** -```json -{ - "status": "updated", - "key": "config_key" -} -``` - -**Status Codes:** -- `200 OK`: Success -- `400 Bad Request`: Invalid hex encoding or odd-length string -- `503 Service Unavailable`: Mesh not enabled - -### Rate Limiting - -#### POST `/mesh/rate-limit` - -Set the global rate limit configuration. - -**Request Body:** -```json -{ - "limit_per_second": 1000 -} -``` - -**Response:** -```json -{ - "status": "updated", - "limit_per_second": 1000 -} -``` - -**Status Codes:** -- `200 OK`: Success -- `500 Internal Server Error`: Failed to serialize config -- `503 Service Unavailable`: Mesh not enabled - -**Note:** Setting `limit_per_second` to `0` disables rate limiting. - -#### GET `/mesh/rate-limit` - -Get the global rate limit configuration. - -**Response:** -```json -{ - "limit_per_second": 1000 -} -``` - -**Status Codes:** -- `200 OK`: Success -- `404 Not Found`: Global rate limit not configured -- `500 Internal Server Error`: Failed to deserialize config -- `503 Service Unavailable`: Mesh not enabled - -#### GET `/mesh/rate-limit/stats` - -Get global rate limit statistics including current count and remaining capacity. - -**Response:** -```json -{ - "limit_per_second": 1000, - "current_count": 750, - "remaining": 250 -} -``` - -**Response Fields:** -- `limit_per_second`: The configured rate limit (0 means unlimited) -- `current_count`: Current aggregated count across all nodes -- `remaining`: Remaining capacity (`-1` if unlimited) - -**Status Codes:** -- `200 OK`: Success -- `503 Service Unavailable`: Mesh sync manager not available - -**How It Works:** -- Rate limit counters are distributed across cluster nodes using consistent hashing -- Each key is assigned to specific owner nodes -- Only owner nodes can increment counters -- Counter values are aggregated using CRDT (PNCounter) across all owners -- Counters are automatically reset periodically (default: every 1 second) - -### Graceful Shutdown - -#### POST `/mesh/shutdown` - -Trigger a graceful shutdown of the mesh node. - -**Response:** -```json -{ - "status": "shutdown initiated" -} -``` - -**Status Codes:** -- `202 Accepted`: Shutdown initiated -- `503 Service Unavailable`: Mesh not enabled - -**Note:** This endpoint initiates the shutdown process asynchronously. The node will gracefully leave the cluster. - -## Rate Limiting - -The Mesh module provides distributed rate limiting using consistent hashing and CRDT counters. - -### How It Works - -1. **Consistent Hashing**: Each rate limit key is assigned to specific nodes (owners) based on hash -2. **Owner-Only Updates**: Only owner nodes can increment counters for their assigned keys -3. **CRDT Aggregation**: Counter values are merged across all owners using PNCounter (Positive-Negative Counter) -4. **Time Window Reset**: Counters are periodically reset (default: every 1 second) - -### Usage Examples - -#### Setting Global Rate Limit - -Set a global rate limit of 1000 requests per second: - -```bash -curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 1000}' -``` - -**Response:** -```json -{ - "status": "updated", - "limit_per_second": 1000 -} -``` - -#### Checking Rate Limit Configuration - -Get the current rate limit configuration: - -```bash -curl http://localhost:8000/mesh/rate-limit \ - -H "Authorization: Bearer " -``` - -**Response:** -```json -{ - "limit_per_second": 1000 -} -``` - -#### Monitoring Rate Limit Statistics - -Check current usage and remaining capacity: - -```bash -curl http://localhost:8000/mesh/rate-limit/stats \ - -H "Authorization: Bearer " -``` - -**Response:** -```json -{ - "limit_per_second": 1000, - "current_count": 750, - "remaining": 250 -} -``` - -#### Disabling Rate Limiting - -To disable rate limiting, set `limit_per_second` to `0`: - -```bash -curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 0}' -``` - -**Response:** -```json -{ - "limit_per_second": 0, - "current_count": 0, - "remaining": -1 -} -``` - -### Complete Example: Rate Limiting Workflow - -1. **Initialize Rate Limit:** - ```bash - # Set rate limit to 500 requests per second - curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 500}' - ``` - -2. **Monitor Usage:** - ```bash - # Check current statistics - curl http://localhost:8000/mesh/rate-limit/stats \ - -H "Authorization: Bearer " - ``` - -3. **Adjust Rate Limit:** - ```bash - # Increase to 2000 requests per second - curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 2000}' - ``` - -### Key Concepts - -- **Distributed Counters**: Counters are sharded across nodes, not replicated -- **Eventual Consistency**: Counter values converge across all nodes -- **Automatic Reset**: Counters reset periodically to implement time-window rate limiting -- **Membership Updates**: When nodes join/leave, ownership is automatically recalculated - -## Cache-Aware Routing - -The Mesh module synchronizes cache-aware routing tree operations across cluster nodes. This enables cache-aware routing policies to share cache state information using a **global radix tree**. - -### Global Radix Tree: How Global State is Achieved - -The cache-aware routing policy uses a **radix tree** (prefix tree) to track which workers have cached which request prefixes. The key innovation is making this tree **global** across all mesh nodes through state synchronization. - -#### What is a Radix Tree? - -A radix tree is a data structure that stores strings as a tree of character-based nodes. Each node represents a prefix segment, and the tree efficiently tracks which workers (tenants) have processed which request prefixes. - -**Simple Example:** -``` -Tree stores: "Hello" → worker1, "Help" → worker2 - -Structure: -Root -└── "H" → "ello" [worker1] - └── "p" [worker2] -``` - -When routing a new request "Hello world", the tree finds the longest matching prefix ("Hello") and returns worker1, indicating worker1 likely has this cached. - -#### How Global State is Achieved - -The radix tree becomes **global** through mesh synchronization. Each node maintains a local copy of the tree, and all tree operations are synchronized across the cluster. - -**Architecture:** - -``` -┌─────────────┐ ┌─────────────┐ ┌─────────────┐ -│ Node 1 │ │ Node 2 │ │ Node 3 │ -│ │ │ │ │ │ -│ Local Tree │◄───────►│ Local Tree │◄───────►│ Local Tree │ -│ │ Mesh │ │ Mesh │ │ -│ │ Sync │ │ Sync │ │ -└─────────────┘ └─────────────┘ └─────────────┘ - │ │ │ - └───────────────────────┴───────────────────────┘ - Global Tree State - (Eventual Consistency via CRDT) -``` - -#### Synchronization Mechanism - -**1. Operation-Based Synchronization:** - -Instead of synchronizing the entire tree structure, only **tree operations** are synchronized: - -- **Insert Operation**: `TreeOperation::Insert { text: "Hello", tenant: "worker1" }` -- **Remove Operation**: `TreeOperation::Remove { tenant: "worker1" }` - -**2. Synchronization Flow:** - -``` -Node 1: Request arrives → Route to worker1 - ↓ - Insert into local tree: "Hello" → worker1 - ↓ - Generate operation: Insert("Hello", "worker1") - ↓ - Sync to mesh via MeshSyncManager - ↓ - Operation stored in PolicyStore with key: "tree:model-1" - ↓ - Gossip protocol propagates to other nodes - ↓ -Node 2: Receives operation via gossip - ↓ - Applies operation to local tree - ↓ - Local tree now matches Node 1's tree -``` - -**3. State Storage:** - -Tree state is stored in the `PolicyStore` (a CRDT map) with keys in the format `tree:{model_id}`: - -```rust -// Tree state structure -TreeState { - model_id: "llama-3", - operations: [ - Insert("Hello", "worker1"), - Insert("Help", "worker2"), - Insert("World", "worker1"), - ], - version: 5 -} -``` - -**4. Incremental Updates:** - -The mesh uses incremental update collection to only send new operations: - -- Each operation has a version number -- Only operations with versions > last_sent_version are transmitted -- Reduces network overhead and ensures efficient synchronization - -**5. Eventual Consistency:** - -- All nodes apply the same sequence of operations -- Operations are idempotent (can be applied multiple times safely) -- CRDT properties ensure convergence even with network partitions -- All nodes eventually have the same tree state - -#### Complete Synchronization Example - -**Initial State (All Nodes):** -``` -All nodes have empty trees -``` - -**Step 1: Node 1 processes request** - -``` -Node 1: -1. Request: "Hello world" → Route to worker1 -2. Insert locally: tree.insert("Hello world", "worker1") -3. Generate operation: Insert("Hello world", "worker1") -4. Sync to mesh: mesh_sync.sync_tree_operation("model-1", operation) -5. Operation stored in PolicyStore with version 1 -``` - -**Step 2: Gossip propagation** - -``` -Gossip Protocol: -1. Node 1 sends state sync message to Node 2 -2. Node 2 receives TreeState with operations: [Insert("Hello world", "worker1")] -3. Node 2 applies operation to local tree -4. Node 2's tree now matches Node 1's tree -5. Node 2 forwards to Node 3 (gossip continues) -``` - -**Step 3: Node 3 processes similar request** - -``` -Node 3: -1. Request: "Hello" arrives -2. Prefix match finds: "Hello" (partial match of "Hello world") -3. Match rate: 5/5 = 1.0 > cache_threshold -4. Route to worker1 (knows worker1 has "Hello" cached) -5. No new operation needed (already in tree) -``` - -**Step 4: Worker failure** - -``` -All Nodes: -1. Worker1 fails -2. Node 1 detects failure -3. Remove locally: tree.remove_tenant("worker1") -4. Generate operation: Remove("worker1") -5. Sync to mesh: mesh_sync.sync_tree_operation("model-1", operation) -6. All nodes receive and apply removal -7. All trees updated consistently -``` - -#### State Restoration on Startup - -When a node restarts or joins the cluster: - -``` -1. Node starts with empty tree -2. CacheAwarePolicy.restore_tree_state_from_mesh() called -3. Retrieves TreeState from PolicyStore via mesh -4. Applies all operations sequentially: - for operation in tree_state.operations { - match operation { - Insert(text, tenant) => tree.insert(text, tenant), - Remove(tenant) => tree.remove_tenant(tenant), - } - } -5. Tree rebuilt to match cluster state -6. Node ready to route with full cache knowledge -``` - -#### Benefits of Global Synchronization - -1. **Shared Cache Knowledge**: All nodes know which workers have cached which prefixes -2. **Optimal Routing**: Any node can route to the best worker based on cache state -3. **Fault Tolerance**: Tree state persists across node failures via mesh storage -4. **Automatic Recovery**: New nodes automatically get full cache state -5. **Eventual Consistency**: All nodes converge to the same view without coordination - -### How It Works - -1. **Tree Operations**: When cache-aware routing makes routing decisions, tree operations (insert/remove) are generated -2. **Mesh Synchronization**: Tree operations are automatically synchronized to the mesh via the sync manager -3. **State Restoration**: On startup, cache-aware policies restore tree state from the mesh -4. **Eventual Consistency**: Tree states converge across all nodes - -### Integration - -Cache-aware routing is automatically integrated with the mesh when: -- Mesh is enabled in the router configuration -- Cache-aware policy is configured for a model -- Tree operations are performed during routing - -### Usage Examples - -#### Checking Policy State - -View the cache-aware policy state for a specific model: - -```bash -curl http://localhost:8000/mesh/policies/model-1 \ - -H "Authorization: Bearer " -``` - -**Response:** -```json -{ - "model_id": "model-1", - "policy_type": "cache_aware", - "config": "...", - "version": 5 -} -``` - -#### Viewing All Policy States - -List all policy states across the cluster: - -```bash -curl http://localhost:8000/mesh/policies \ - -H "Authorization: Bearer " -``` - -**Response:** -```json -[ - { - "model_id": "model-1", - "policy_type": "cache_aware", - "config": "...", - "version": 5 - }, - { - "model_id": "model-2", - "policy_type": "cache_aware", - "config": "...", - "version": 3 - } -] -``` - -### Complete Example: Global Tree Synchronization Workflow - -This example demonstrates how the radix tree becomes global through mesh synchronization. - -**Initial Setup:** -- 3-node cluster: node1, node2, node3 -- Model: llama-3 -- All nodes start with empty trees - -**Step 1: Node 1 processes first request** - -``` -Node 1: -Request: "The quick brown fox" -Model: llama-3 - -1. Local tree is empty → no prefix match -2. Route to worker1 (smallest tree) -3. Insert locally: tree.insert("The quick brown fox", "worker1") -4. Generate operation: Insert("The quick brown fox", "worker1") -5. Sync to mesh: - mesh_sync.sync_tree_operation("llama-3", operation) - ↓ - Operation stored in PolicyStore: tree:llama-3 - Version: 1 -``` - -**Step 2: Mesh propagates to other nodes** - -``` -Gossip Protocol: -Node 1 → Node 2: Sends TreeState { operations: [Insert(...)], version: 1 } -Node 2 → Node 3: Forwards TreeState - -Node 2: -1. Receives TreeState via gossip -2. Applies operation: tree.insert("The quick brown fox", "worker1") -3. Local tree now matches Node 1 - -Node 3: -1. Receives TreeState via gossip -2. Applies operation: tree.insert("The quick brown fox", "worker1") -3. Local tree now matches Node 1 and Node 2 -``` - -**Step 3: Node 2 processes similar request** - -``` -Node 2: -Request: "The quick brown" -Model: llama-3 - -1. Prefix match finds: "The quick brown" (partial match) -2. Match rate: 17/19 = 0.89 > cache_threshold (0.7) -3. Route to worker1 (knows worker1 has this cached) -4. No new operation (tree already has this prefix) -5. All nodes already have this in their trees -``` - -**Step 4: Node 3 processes different request** - -``` -Node 3: -Request: "Hello world" -Model: llama-3 - -1. Prefix match: "" (no match) -2. Route to worker2 (smallest tree) -3. Insert locally: tree.insert("Hello world", "worker2") -4. Generate operation: Insert("Hello world", "worker2") -5. Sync to mesh: - mesh_sync.sync_tree_operation("llama-3", operation) - ↓ - Operation stored in PolicyStore: tree:llama-3 - Version: 2 -6. Gossip propagates to Node 1 and Node 2 -7. All nodes now have both prefixes in their trees -``` - -**Step 5: Node 1 restarts** - -``` -Node 1 (after restart): -1. CacheAwarePolicy initializes -2. Calls restore_tree_state_from_mesh() -3. Retrieves TreeState from PolicyStore: - { - model_id: "llama-3", - operations: [ - Insert("The quick brown fox", "worker1"), - Insert("Hello world", "worker2") - ], - version: 2 - } -4. Applies all operations sequentially: - tree.insert("The quick brown fox", "worker1") - tree.insert("Hello world", "worker2") -5. Tree rebuilt to match cluster state -6. Node 1 has full cache knowledge again -``` - -**Result:** -- All nodes have identical tree state -- Any node can route optimally based on cache -- State persists across restarts -- New nodes automatically get full state - -### Tree State Storage - -Tree states are stored in the PolicyStore with keys in the format: `tree:{model_id}` - -The tree state contains: -- `model_id`: The model identifier -- `operations`: Sequence of tree operations (insert/remove) -- `version`: Version number for conflict resolution - -### Benefits - -- **Shared Cache Knowledge**: All nodes know which workers have cached which request prefixes -- **Optimal Routing**: Routes requests to workers with relevant cache data -- **Automatic Synchronization**: No manual intervention required -- **Fault Tolerance**: Tree state is preserved across node failures - -### Example: Global Synchronization in Action - -This example shows how tree operations are synchronized across nodes to create a global view. - -**Scenario:** 3-node cluster, model "llama-3" - -**Timeline:** - -``` -T0: All nodes have empty trees - Node1: [] - Node2: [] - Node3: [] - -T1: Node1 processes "Hello world" → worker1 - Node1: [Insert("Hello world", "worker1")] → syncs to mesh - Node2: [] (not yet received) - Node3: [] (not yet received) - -T2: Gossip propagates (Node1 → Node2 → Node3) - Node1: [Insert("Hello world", "worker1")] - Node2: [Insert("Hello world", "worker1")] ← applied from mesh - Node3: [Insert("Hello world", "worker1")] ← applied from mesh - -T3: Node2 processes "Hello" → worker1 (cache hit, no sync needed) - Node3 processes "Help" → worker2 - Node3: [Insert("Hello world", "worker1"), Insert("Help", "worker2")] → syncs to mesh - -T4: Gossip propagates - Node1: [Insert("Hello world", "worker1"), Insert("Help", "worker2")] ← applied - Node2: [Insert("Hello world", "worker1"), Insert("Help", "worker2")] ← applied - Node3: [Insert("Hello world", "worker1"), Insert("Help", "worker2")] - -T5: All nodes have identical tree state - All nodes can route "Hello" → worker1 (cache hit) - All nodes can route "Help" → worker2 (cache hit) -``` - -**Key Points:** - -1. **Operations are the source of truth**: Tree structure is derived from operations -2. **Gossip ensures propagation**: Operations spread to all nodes automatically -3. **Eventual consistency**: All nodes converge to the same state -4. **No coordination needed**: Each node applies operations independently -5. **State persists**: Operations stored in PolicyStore survive node restarts - -## Service Discovery Integration - -The Mesh module integrates with service discovery systems to maintain cluster membership dynamically. - -### Membership Updates - -When service discovery detects membership changes (nodes joining or leaving), the mesh: - -1. **Updates Hash Rings**: Rate limit hash rings are recalculated -2. **Updates Membership Store**: Node information is synchronized -3. **Ownership Transfer**: Rate limit ownership is transferred for failed nodes -4. **Topology Recalculation**: Peer connections are recalculated - -### Topology Management - -The mesh supports two topology modes: - -**Full Mesh** (≤ threshold nodes, default: 10): -- All nodes connect to all other nodes -- Best for small clusters -- Maximum redundancy - -**Sparse Mesh** (> threshold nodes): -- Nodes connect based on region/availability zone -- Reduces connection overhead -- Suitable for large clusters - -### Node States - -Nodes can be in different states: - -- **Alive**: Node is healthy and reachable -- **Suspected**: Node may be unreachable (gossip-based detection) -- **Down**: Node is confirmed unreachable -- **Leaving**: Node is gracefully shutting down - -### Service Discovery Flow - -``` -Service Discovery → Membership Update → Hash Ring Update → Ownership Transfer - ↓ - Topology Recalculation - ↓ - Peer Connection Update -``` - -### Configuration - -Topology configuration can be set via: -- Region identifier (for sparse mesh) -- Availability zone identifier (for sparse mesh) -- Full mesh threshold (default: 10 nodes) - -## Error Handling - -All endpoints return appropriate HTTP status codes: - -- `200 OK`: Successful operation -- `202 Accepted`: Operation accepted (async) -- `400 Bad Request`: Invalid request format -- `404 Not Found`: Resource not found -- `500 Internal Server Error`: Server error -- `503 Service Unavailable`: Mesh not enabled or service unavailable - -Error responses follow this format: -```json -{ - "error": "Error message description" -} -``` - -## Usage Examples - -### Complete Workflow: Setting Up Mesh with Rate Limiting and Cache-Aware Routing - -This example demonstrates how to set up and use mesh features in a production environment. - -#### 1. Enable Mesh in Configuration - -Configure mesh in your router configuration file: - -```yaml -mesh: - enabled: true - self_name: "router-node-1" - self_addr: "0.0.0.0:8000" - init_peer: "router-node-2:8000" # Optional: initial peer for bootstrap -``` - -#### 2. Set Up Global Rate Limiting - -```bash -# Set initial rate limit -curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 1000}' - -# Verify configuration -curl http://localhost:8000/mesh/rate-limit \ - -H "Authorization: Bearer " -``` - -#### 3. Configure Cache-Aware Policy - -When configuring a model with cache-aware routing, the mesh automatically handles synchronization: - -```yaml -models: - - model_id: "llama-3" - policy: - type: "cache_aware" - config: - cache_threshold: 0.7 - balance_abs_threshold: 10 - balance_rel_threshold: 1.5 - eviction_interval_secs: 300 - max_tree_size: 10000 -``` - -#### 4. Monitor Cluster Status - -```bash -# Check cluster health -curl http://localhost:8000/mesh/health \ - -H "Authorization: Bearer " - -# View cluster status -curl http://localhost:8000/mesh/status \ - -H "Authorization: Bearer " -``` - -#### 5. Monitor Rate Limiting - -```bash -# Check current rate limit statistics -curl http://localhost:8000/mesh/rate-limit/stats \ - -H "Authorization: Bearer " - -# Adjust rate limit based on load -curl -X POST http://localhost:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 2000}' -``` - -#### 6. Monitor Cache-Aware Policy State - -```bash -# View policy state for a model -curl http://localhost:8000/mesh/policies/llama-3 \ - -H "Authorization: Bearer " - -# View all policy states -curl http://localhost:8000/mesh/policies \ - -H "Authorization: Bearer " -``` - -#### 7. View Worker States - -```bash -# List all worker states -curl http://localhost:8000/mesh/workers \ - -H "Authorization: Bearer " - -# Get specific worker state -curl http://localhost:8000/mesh/workers/worker-1 \ - -H "Authorization: Bearer " -``` - -### Example: Multi-Node Setup - -In a 3-node cluster setup: - -**Node 1 (router-node-1):** -```bash -# Set rate limit -curl -X POST http://router-node-1:8000/mesh/rate-limit \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{"limit_per_second": 1000}' -``` - -**Node 2 and Node 3:** -- Automatically receive rate limit configuration via mesh synchronization -- Share cache-aware tree state across all nodes -- Maintain consistent state without manual configuration - -**Verify Consistency:** -```bash -# Check rate limit on all nodes (should be consistent) -curl http://router-node-1:8000/mesh/rate-limit/stats -curl http://router-node-2:8000/mesh/rate-limit/stats -curl http://router-node-3:8000/mesh/rate-limit/stats -``` - -### Example: Dynamic Configuration Updates - -Update application configuration dynamically: - -```bash -# Store custom configuration -curl -X POST http://localhost:8000/mesh/config \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{ - "key": "custom_feature_flag", - "value": "74727565" # hex for "true" - }' - -# Retrieve configuration -curl http://localhost:8000/mesh/config/custom_feature_flag \ - -H "Authorization: Bearer " -``` - -## Authentication - -All mesh endpoints require authentication. Configure authentication via the router's authentication middleware. - -## See Also - -- [CRDT Documentation](https://github.com/rust-crdt/rust-crdt) -- [Gossip Protocol](https://en.wikipedia.org/wiki/Gossip_protocol) -- [Consistent Hashing](https://en.wikipedia.org/wiki/Consistent_hashing) diff --git a/sgl-model-gateway/src/mesh/consistent_hash.rs b/sgl-model-gateway/src/mesh/consistent_hash.rs deleted file mode 100644 index 4846fcff1..000000000 --- a/sgl-model-gateway/src/mesh/consistent_hash.rs +++ /dev/null @@ -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, - /// Node -> set of virtual node hashes - node_hashes: BTreeMap>, -} - -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 { - 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 = self.node_hashes.keys().cloned().collect(); - let new_nodes: HashSet = 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(¤t_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"); - } -} diff --git a/sgl-model-gateway/src/mesh/controller.rs b/sgl-model-gateway/src/mesh/controller.rs deleted file mode 100644 index 3ad9c544f..000000000 --- a/sgl-model-gateway/src/mesh/controller.rs +++ /dev/null @@ -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, -} - -impl MeshController { - pub fn new( - state: ClusterState, - self_addr: SocketAddr, - self_name: &str, - init_peer: Option, - ) -> 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 = 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::()?; - 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 = 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(map: &BTreeMap, 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() -} diff --git a/sgl-model-gateway/src/mesh/crdt.rs b/sgl-model-gateway/src/mesh/crdt.rs deleted file mode 100644 index 81dc9cdab..000000000 --- a/sgl-model-gateway/src/mesh/crdt.rs +++ /dev/null @@ -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 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 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 { - value: T, - timestamp: u64, - version: u64, - actor: String, -} - -impl LWWRegister { - 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 { - inner: BTreeMap>, -} - -impl Default for CRDTMap { - fn default() -> Self { - Self { - inner: BTreeMap::new(), - } - } -} - -impl CRDTMap { - 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 { - self.inner.iter().map(|(k, v)| (k, v.read())) - } - - pub fn keys(&self) -> impl Iterator { - self.inner.keys() - } - - pub fn values(&self) -> impl Iterator { - 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 { - 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, -} - -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(); - as CvRDT>::merge(&mut self.inner, other_clone); - } -} - -/// Thread-safe wrapper for CRDT Map -#[derive(Debug, Clone)] -pub struct SyncCRDTMap { - inner: Arc>>, -} - -impl Default for SyncCRDTMap { - fn default() -> Self { - Self { - inner: Arc::new(RwLock::new(CRDTMap::new())), - } - } -} - -impl SyncCRDTMap { - pub fn new() -> Self { - Self::default() - } - - pub fn get(&self, key: &SKey) -> Option { - 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) { - self.inner.write().merge(other); - } - - pub fn snapshot(&self) -> CRDTMap { - 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>, -} - -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(®2); - 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(®2); - 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(®1); - 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(®2); - 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(®2); - reg1.merge(®3); - 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(®6); - reg4.merge(®5); - // reg6 should win (latest timestamp) - assert_eq!(reg4.read(), "value3"); - } - - // CRDTMap tests with i32 - #[test] - fn test_crdt_map_new() { - let map: CRDTMap = 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 = 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 = 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); - } -} diff --git a/sgl-model-gateway/src/mesh/flow_control.rs b/sgl-model-gateway/src/mesh/flow_control.rs deleted file mode 100644 index a8d802a19..000000000 --- a/sgl-model-gateway/src/mesh/flow_control.rs +++ /dev/null @@ -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>>, - attempt_count: Arc>, -} - -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()) - } -} diff --git a/sgl-model-gateway/src/mesh/incremental.rs b/sgl-model-gateway/src/mesh/incremental.rs deleted file mode 100644 index c696556fd..000000000 --- a/sgl-model-gateway/src/mesh/incremental.rs +++ /dev/null @@ -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, - policy: HashMap, - app: HashMap, - membership: HashMap, - rate_limit: HashMap, // Track last sent timestamp for rate limit counters -} - -/// Incremental update collector -pub struct IncrementalUpdateCollector { - stores: Arc, - self_name: String, - last_sent: Arc>, -} - -impl IncrementalUpdateCollector { - pub fn new(stores: Arc, 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( - &self, - all_items: std::collections::BTreeMap, - get_version: impl Fn(&SKey) -> u64, - last_sent_map: &mut HashMap, - store_name: &str, - get_id: impl Fn(&S) -> String, - ) -> Vec - 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 { - 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)> { - 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); - } -} diff --git a/sgl-model-gateway/src/mesh/metrics.rs b/sgl-model-gateway/src/mesh/metrics.rs deleted file mode 100644 index 900f9b9a1..000000000 --- a/sgl-model-gateway/src/mesh/metrics.rs +++ /dev/null @@ -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() - } -} diff --git a/sgl-model-gateway/src/mesh/mod.rs b/sgl-model-gateway/src/mesh/mod.rs deleted file mode 100644 index 258246d88..000000000 --- a/sgl-model-gateway/src/mesh/mod.rs +++ /dev/null @@ -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}; diff --git a/sgl-model-gateway/src/mesh/mtls.rs b/sgl-model-gateway/src/mesh/mtls.rs deleted file mode 100644 index 5fe8f7ca0..000000000 --- a/sgl-model-gateway/src/mesh/mtls.rs +++ /dev/null @@ -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>>>, - client_config: Arc>>>, -} - -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> { - 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> { - 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>> { - let cert_data = fs::read(path).await?; - let certs = certs(&mut cert_data.as_slice()).collect::, _>>()?; - Ok(certs) - } - - /// Load private key from file - async fn load_private_key(&self, path: &Path) -> Result> { - let key_data = fs::read(path).await?; - let mut keys = - pkcs8_private_keys(&mut key_data.as_slice()).collect::, _>>()?; - - 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>>>, - _client_config: &Arc>>>, - ) -> 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> { - self.server_config.read().await.clone() - } - - /// Get current client config (for use with tonic) - pub async fn get_client_config(&self) -> Option> { - self.client_config.read().await.clone() - } -} - -/// Optional mTLS manager -pub type OptionalMTLSManager = Option>; diff --git a/sgl-model-gateway/src/mesh/node_state_machine.rs b/sgl-model-gateway/src/mesh/node_state_machine.rs deleted file mode 100644 index 716ea41bb..000000000 --- a/sgl-model-gateway/src/mesh/node_state_machine.rs +++ /dev/null @@ -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, - stable_update_count: usize, - last_state_hash: Option, -} - -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>, - config: ConvergenceConfig, - convergence_tracker: Arc>, - snapshot_start_time: Arc>>, - stores: Arc, -} - -impl NodeStateMachine { - pub fn new(stores: Arc, 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 { - 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()); - } -} diff --git a/sgl-model-gateway/src/mesh/partition.rs b/sgl-model-gateway/src/mesh/partition.rs deleted file mode 100644 index f856b22a6..000000000 --- a/sgl-model-gateway/src/mesh/partition.rs +++ /dev/null @@ -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>>, - current_state: Arc>, -} - -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) -> 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, - ) -> Vec { - 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 - ); - } -} diff --git a/sgl-model-gateway/src/mesh/ping_server.rs b/sgl-model-gateway/src/mesh/ping_server.rs deleted file mode 100644 index b52c620c5..000000000 --- a/sgl-model-gateway/src/mesh/ping_server.rs +++ /dev/null @@ -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>, // Optional state stores for CRDT-based sync - sync_manager: Option>, // Optional sync manager for applying remote updates - state_machine: Option>, - partition_detector: Option>, -} - -impl GossipService { - /// Create snapshot chunks for a store - async fn create_snapshot_chunks( - &self, - store_type: LocalStoreType, - chunk_size: usize, - ) -> Vec { - 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)> = 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 = 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) -> 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) -> Self { - self.sync_manager = Some(sync_manager); - self - } - - pub fn with_partition_detector(mut self, partition_detector: Arc) -> Self { - self.partition_detector = Some(partition_detector); - self - } - - pub async fn serve_ping_with_shutdown>( - 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) -> 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> + Send + 'static>>; - - #[instrument(fields(name = %self.self_name), skip(self, request))] - async fn ping_server( - &self, - request: tonic::Request, - ) -> std::result::Result, 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>, - ) -> Result, 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::>(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> = 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::(); - 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::(&entry.value) { - stores.membership.insert(key, membership_state, entry.actor.clone()); - } - } - LocalStoreType::App => { - if let Ok(app_state) = serde_json::from_slice::(&entry.value) { - stores.app.insert(key, app_state, entry.actor.clone()); - } - } - LocalStoreType::Worker => { - if let Ok(worker_state) = serde_json::from_slice::(&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::(&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::(&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 - )) - } -} diff --git a/sgl-model-gateway/src/mesh/proto/gossip.proto b/sgl-model-gateway/src/mesh/proto/gossip.proto deleted file mode 100644 index aeef98c76..000000000 --- a/sgl-model-gateway/src/mesh/proto/gossip.proto +++ /dev/null @@ -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 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; -} diff --git a/sgl-model-gateway/src/mesh/rate_limit_window.rs b/sgl-model-gateway/src/mesh/rate_limit_window.rs deleted file mode 100644 index 8a74799d3..000000000 --- a/sgl-model-gateway/src/mesh/rate_limit_window.rs +++ /dev/null @@ -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, - window_seconds: u64, -} - -impl RateLimitWindow { - pub fn new(sync_manager: Arc, 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); - } - } -} diff --git a/sgl-model-gateway/src/mesh/service.rs b/sgl-model-gateway/src/mesh/service.rs deleted file mode 100644 index c21c5292c..000000000 --- a/sgl-model-gateway/src/mesh/service.rs +++ /dev/null @@ -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>>; - -pub struct MeshServerConfig { - pub self_name: String, - pub self_addr: SocketAddr, - pub init_peer: Option, -} - -/// 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>, - state_machine: Option>, -} - -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>, - ) -> 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> { - self.partition_detector.as_ref() - } - - /// Get state machine - pub fn state_machine(&self) -> Option<&Arc> { - 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::>(); - (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) { - 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> { - 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, -} - -impl MeshServerBuilder { - pub fn new(self_name: String, self_addr: SocketAddr, init_peer: Option) -> 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>, - ) -> (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, - signal_rx: tokio::sync::watch::Receiver<()>, -} - -impl MeshServer { - pub fn new( - state: ClusterState, - self_name: &str, - self_addr: SocketAddr, - init_peer: Option, - 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>, - sync_manager: Option>, - partition_detector: Option>, - ) -> 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, - target_nodes: Vec, - timeout: Option, -) -> (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, -) -> Result { - 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) - ); - } -} diff --git a/sgl-model-gateway/src/mesh/stores.rs b/sgl-model-gateway/src/mesh/stores.rs deleted file mode 100644 index 6429e1969..000000000 --- a/sgl-model-gateway/src/mesh/stores.rs +++ /dev/null @@ -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>, -} - -/// App state entry (application configuration) -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)] -pub struct AppState { - pub key: String, - pub value: Vec, // 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(&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, // 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, -} - -impl MembershipStore { - pub fn new() -> Self { - Self { - inner: SyncCRDTMap::new(), - } - } - - pub fn get(&self, key: &SKey) -> Option { - 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) { - self.inner.merge(other); - } - - pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap { - 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 { - 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, -} - -impl AppStore { - pub fn new() -> Self { - Self { - inner: SyncCRDTMap::new(), - } - } - - pub fn get(&self, key: &SKey) -> Option { - 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) { - self.inner.merge(other); - } - - pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap { - 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 { - 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, -} - -impl WorkerStore { - pub fn new() -> Self { - Self { - inner: SyncCRDTMap::new(), - } - } - - pub fn get(&self, key: &SKey) -> Option { - 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) { - self.inner.merge(other); - } - - pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap { - 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 { - 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, -} - -impl PolicyStore { - pub fn new() -> Self { - Self { - inner: SyncCRDTMap::new(), - } - } - - pub fn get(&self, key: &SKey) -> Option { - 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) { - self.inner.merge(other); - } - - pub fn snapshot(&self) -> crate::mesh::crdt::CRDTMap { - 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 { - 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>>, // key -> counter - hash_ring: Arc>, - 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 { - 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 { - 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 { - 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 { - 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 { - 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 { - 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)); - } -} diff --git a/sgl-model-gateway/src/mesh/sync.rs b/sgl-model-gateway/src/mesh/sync.rs deleted file mode 100644 index 884b7bdab..000000000 --- a/sgl-model-gateway/src/mesh/sync.rs +++ /dev/null @@ -1,1177 +0,0 @@ -//! Mesh state synchronization module -//! -//! Handles synchronization of worker and policy states across mesh cluster nodes - -use std::sync::Arc; - -use tracing::debug; - -use super::{ - crdt::SKey, - gossip::NodeStatus, - stores::{ - tree_state_key, PolicyState, RateLimitConfig, StateStores, WorkerState, - GLOBAL_RATE_LIMIT_COUNTER_KEY, GLOBAL_RATE_LIMIT_KEY, - }, - tree_ops::{TreeOperation, TreeState}, -}; - -/// Mesh sync manager for coordinating state synchronization -#[derive(Clone, Debug)] -pub struct MeshSyncManager { - pub(crate) stores: Arc, - self_name: String, -} - -impl MeshSyncManager { - pub fn new(stores: Arc, self_name: String) -> Self { - Self { stores, self_name } - } - - /// Get the node name (actor) for this sync manager - pub fn self_name(&self) -> &str { - &self.self_name - } - - /// Sync worker state to mesh stores - pub fn sync_worker_state( - &self, - worker_id: String, - model_id: String, - url: String, - health: bool, - load: f64, - ) { - let key = SKey::new(worker_id.clone()); - - // Get current version if exists, otherwise start at 1 - let current_version = self - .stores - .worker - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - let new_version = current_version + 1; - - let state = WorkerState { - worker_id: worker_id.clone(), - model_id, - url, - health, - load, - version: new_version, - }; - - // Use self node name as actor - let actor = self.self_name.clone(); - self.stores.worker.insert(key, state, actor); - debug!( - "Synced worker state to mesh {} (version: {})", - worker_id, new_version - ); - } - - /// Remove worker state from mesh stores - pub fn remove_worker_state(&self, worker_id: &str) { - let key = SKey::new(worker_id.to_string()); - self.stores.worker.remove(&key); - debug!("Removed worker state from mesh {}", worker_id); - } - - /// Sync policy state to mesh stores - pub fn sync_policy_state(&self, model_id: String, policy_type: String, config: Vec) { - let key = SKey::new(format!("policy:{}", model_id)); - - // Get current version if exists, otherwise start at 1 - let current_version = self - .stores - .policy - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - let new_version = current_version + 1; - - let state = PolicyState { - model_id: model_id.clone(), - policy_type, - config, - version: new_version, - }; - - // Use self node name as actor - let actor = self.self_name.clone(); - self.stores.policy.insert(key, state, actor); - debug!( - "Synced policy state to mesh model={} (version: {})", - model_id, new_version - ); - } - - /// Remove policy state from mesh stores - pub fn remove_policy_state(&self, model_id: &str) { - let key = SKey::new(format!("policy:{}", model_id)); - self.stores.policy.remove(&key); - debug!("Removed policy state from mesh model={}", model_id); - } - - /// Get worker state from mesh stores - pub fn get_worker_state(&self, worker_id: &str) -> Option { - let key = SKey::new(worker_id.to_string()); - self.stores.worker.get(&key) - } - - /// Get all worker states from mesh stores - pub fn get_all_worker_states(&self) -> Vec { - self.stores.worker.all().into_values().collect() - } - - /// Get policy state from mesh stores - pub fn get_policy_state(&self, model_id: &str) -> Option { - let key = SKey::new(format!("policy:{}", model_id)); - self.stores.policy.get(&key) - } - - /// Get all policy states from mesh stores - pub fn get_all_policy_states(&self) -> Vec { - self.stores.policy.all().into_values().collect() - } - - /// Apply worker state update from remote node - /// The actor should be extracted from the state update context (e.g., from StateUpdate message) - pub fn apply_remote_worker_state(&self, state: WorkerState, actor: Option) { - let key = SKey::new(state.worker_id.clone()); - // Use provided actor, or fallback to a default if not available - // In practice, actor should come from the StateUpdate message - let actor = actor.unwrap_or_else(|| "remote".to_string()); - - // Check if we should update based on version - let current_version = self - .stores - .worker - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - - if state.version > current_version { - self.stores.worker.insert(key, state.clone(), actor.clone()); - debug!( - "Applied remote worker state update: {} (version: {} -> {})", - state.worker_id, current_version, state.version - ); - } else { - debug!( - "Skipped remote worker state update: {} (version {} <= current {})", - state.worker_id, state.version, current_version - ); - } - } - - /// Apply policy state update from remote node - /// The actor should be extracted from the state update context (e.g., from StateUpdate message) - pub fn apply_remote_policy_state(&self, state: PolicyState, actor: Option) { - let key = SKey::new(format!("policy:{}", state.model_id)); - // Use provided actor, or fallback to a default if not available - let actor = actor.unwrap_or_else(|| "remote".to_string()); - - // Check if we should update based on version - let current_version = self - .stores - .policy - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - - if state.version > current_version { - self.stores.policy.insert(key, state.clone(), actor.clone()); - debug!( - "Applied remote policy state update: {} (version: {} -> {})", - state.model_id, current_version, state.version - ); - } else { - debug!( - "Skipped remote policy state update: {} (version {} <= current {})", - state.model_id, state.version, current_version - ); - } - } - - /// Update rate-limit hash ring with current membership - pub fn update_rate_limit_membership(&self) { - // Get all alive nodes from membership store - let all_members = self.stores.membership.all(); - let alive_nodes: Vec = all_members - .values() - .filter(|m| m.status == NodeStatus::Alive as i32) - .map(|m| m.name.clone()) - .collect(); - - self.stores.rate_limit.update_membership(&alive_nodes); - debug!( - "Updated rate-limit hash ring with {} alive nodes", - alive_nodes.len() - ); - } - - /// Handle node failure and transfer rate-limit ownership - pub fn handle_node_failure(&self, failed_nodes: &[String]) { - if failed_nodes.is_empty() { - return; - } - - debug!("Handling node failure for rate-limit: {:?}", failed_nodes); - - // Check which keys need ownership transfer - let affected_keys = self - .stores - .rate_limit - .check_ownership_transfer(failed_nodes); - - if !affected_keys.is_empty() { - debug!( - "Ownership transfer needed for {} rate-limit keys", - affected_keys.len() - ); - - // Update membership to reflect node failures - self.update_rate_limit_membership(); - - // For each affected key, we may need to initialize counters if we're now an owner - for key in &affected_keys { - if self.stores.rate_limit.is_owner(key) { - debug!("This node is now owner of rate-limit key: {}", key); - // Counter will be created on first inc() call - } - } - } - } - - /// Sync rate-limit counter increment (only if this node is an owner) - pub fn sync_rate_limit_inc(&self, key: String, delta: i64) { - if !self.stores.rate_limit.is_owner(&key) { - // Not an owner, skip - return; - } - - self.stores - .rate_limit - .inc(key.clone(), self.self_name.clone(), delta); - debug!("Synced rate-limit increment: key={}, delta={}", key, delta); - } - - /// Apply remote rate-limit counter update (merge CRDT) - pub fn apply_remote_rate_limit_counter( - &self, - key: String, - counter: &super::crdt::SyncPNCounter, - ) { - // Merge counter regardless of ownership (for CRDT consistency) - self.stores.rate_limit.merge_counter(key.clone(), counter); - debug!("Applied remote rate-limit counter update: key={}", key); - } - - /// Get rate-limit value (aggregate from all owners) - pub fn get_rate_limit_value(&self, key: &str) -> Option { - self.stores.rate_limit.value(key) - } - - /// Get global rate limit configuration from AppStore - pub fn get_global_rate_limit_config(&self) -> Option { - let key = SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string()); - self.stores - .app - .get(&key) - .and_then(|app_state| serde_json::from_slice::(&app_state.value).ok()) - } - - /// Check if global rate limit is exceeded - /// Returns (is_exceeded, current_count, limit) - pub fn check_global_rate_limit(&self) -> (bool, i64, u64) { - let config = self.get_global_rate_limit_config().unwrap_or_default(); - - if config.limit_per_second == 0 { - // Rate limit disabled - return (false, 0, 0); - } - - // Increment counter if this node is an owner - self.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 1); - - // Get aggregated counter value from all owners - let current_count = self - .get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY) - .unwrap_or(0); - - let is_exceeded = current_count > config.limit_per_second as i64; - (is_exceeded, current_count, config.limit_per_second) - } - - /// Reset global rate limit counter (called periodically for time window reset) - pub fn reset_global_rate_limit_counter(&self) { - // Reset by decrementing the current value - // Since we use PNCounter, we can't directly reset, but we can track the window - // For simplicity, we'll use a time-based approach where counters are reset periodically - // The actual reset logic will be handled by the window manager - let current_count = self - .get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY) - .unwrap_or(0); - - if current_count > 0 { - // Decrement by current count to effectively reset - // Note: This is a workaround since PNCounter doesn't support direct reset - // In production, you might want to use a different approach like timestamped counters - self.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), -current_count); - } - } - - /// Sync tree operation to mesh stores - /// This adds a tree operation (insert or remove) to the tree state for a specific model - pub fn sync_tree_operation( - &self, - model_id: String, - operation: TreeOperation, - ) -> Result<(), String> { - let key = SKey::new(tree_state_key(&model_id)); - - // Get current tree state or create new one - let mut tree_state = if let Some(policy_state) = self.stores.policy.get(&key) { - // Deserialize existing tree state - serde_json::from_slice::(&policy_state.config) - .unwrap_or_else(|_| TreeState::new(model_id.clone())) - } else { - TreeState::new(model_id.clone()) - }; - - // Add the new operation - tree_state.add_operation(operation); - - // Serialize and store back - let serialized = serde_json::to_vec(&tree_state) - .map_err(|e| format!("Failed to serialize tree state: {}", e))?; - - // Get current version if exists - let current_version = self - .stores - .policy - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - let new_version = current_version + 1; - - let state = PolicyState { - model_id: model_id.clone(), - policy_type: "tree_state".to_string(), - config: serialized, - version: new_version, - }; - - let actor = self.self_name.clone(); - self.stores.policy.insert(key, state, actor); - debug!( - "Synced tree operation to mesh: model={} (version: {})", - model_id, new_version - ); - - Ok(()) - } - - /// Get tree state for a model from mesh stores - pub fn get_tree_state(&self, model_id: &str) -> Option { - let key = SKey::new(tree_state_key(model_id)); - self.stores - .policy - .get(&key) - .and_then(|policy_state| serde_json::from_slice::(&policy_state.config).ok()) - } - - /// Apply remote tree operation to local policy - /// This is called when receiving tree state updates from other nodes - pub fn apply_remote_tree_operation( - &self, - model_id: String, - tree_state: TreeState, - actor: Option, - ) { - let key = SKey::new(tree_state_key(&model_id)); - let actor = actor.unwrap_or_else(|| "remote".to_string()); - - // Check if we should update based on version - let current_version = self - .stores - .policy - .get_metadata(&key) - .map(|(v, _)| v) - .unwrap_or(0); - - if tree_state.version > current_version { - // Serialize tree state - if let Ok(serialized) = serde_json::to_vec(&tree_state) { - let state = PolicyState { - model_id: model_id.clone(), - policy_type: "tree_state".to_string(), - config: serialized, - version: tree_state.version, - }; - - self.stores.policy.insert(key, state, actor.clone()); - debug!( - "Applied remote tree state update: model={} (version: {} -> {})", - model_id, current_version, tree_state.version - ); - } else { - debug!( - "Failed to serialize remote tree state for model={}", - model_id - ); - } - } else { - debug!( - "Skipped remote tree state update: model={} (version {} <= current {})", - model_id, tree_state.version, current_version - ); - } - } -} - -/// Optional mesh sync manager (can be None if mesh is not enabled) -pub type OptionalMeshSyncManager = Option>; - -#[cfg(test)] -mod tests { - use std::collections::BTreeMap; - - use super::*; - use crate::mesh::stores::{ - AppState, MembershipState, RateLimitConfig, StateStores, GLOBAL_RATE_LIMIT_COUNTER_KEY, - GLOBAL_RATE_LIMIT_KEY, - }; - - fn create_test_sync_manager() -> MeshSyncManager { - let stores = Arc::new(StateStores::new()); - MeshSyncManager::new(stores, "test_node".to_string()) - } - - fn create_test_manager(self_name: String) -> MeshSyncManager { - let stores = Arc::new(StateStores::with_self_name(self_name.clone())); - MeshSyncManager::new(stores, self_name) - } - - #[test] - fn test_sync_manager_new() { - let manager = create_test_sync_manager(); - // Should create without panicking - assert_eq!(manager.get_all_worker_states().len(), 0); - assert_eq!(manager.get_all_policy_states().len(), 0); - } - - #[test] - fn test_sync_worker_state() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - let state = manager.get_worker_state("worker1").unwrap(); - assert_eq!(state.worker_id, "worker1"); - assert_eq!(state.model_id, "model1"); - assert_eq!(state.url, "http://localhost:8000"); - assert!(state.health); - assert_eq!(state.load, 0.5); - assert_eq!(state.version, 1); - } - - #[test] - fn test_sync_multiple_worker_states() { - let manager = create_test_sync_manager(); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - manager.sync_worker_state( - "worker2".to_string(), - "model1".to_string(), - "http://localhost:8001".to_string(), - false, - 0.8, - ); - - manager.sync_worker_state( - "worker3".to_string(), - "model2".to_string(), - "http://localhost:8002".to_string(), - true, - 0.3, - ); - - let all_states = manager.get_all_worker_states(); - assert_eq!(all_states.len(), 3); - - let worker1 = manager.get_worker_state("worker1").unwrap(); - assert_eq!(worker1.worker_id, "worker1"); - assert!(worker1.health); - - let worker2 = manager.get_worker_state("worker2").unwrap(); - assert_eq!(worker2.worker_id, "worker2"); - assert!(!worker2.health); - - let worker3 = manager.get_worker_state("worker3").unwrap(); - assert_eq!(worker3.worker_id, "worker3"); - assert_eq!(worker3.model_id, "model2"); - } - - #[test] - fn test_sync_worker_state_version_increment() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - let state1 = manager.get_worker_state("worker1").unwrap(); - assert_eq!(state1.version, 1); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - false, - 0.8, - ); - - let state2 = manager.get_worker_state("worker1").unwrap(); - assert_eq!(state2.version, 2); - assert!(!state2.health); - assert_eq!(state2.load, 0.8); - } - - #[test] - fn test_remove_worker_state() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - assert!(manager.get_worker_state("worker1").is_some()); - - manager.remove_worker_state("worker1"); - - assert!(manager.get_worker_state("worker1").is_none()); - assert_eq!(manager.get_all_worker_states().len(), 0); - } - - #[test] - fn test_remove_nonexistent_worker_state() { - let manager = create_test_sync_manager(); - - // Should not panic - manager.remove_worker_state("nonexistent"); - assert!(manager.get_worker_state("nonexistent").is_none()); - } - - #[test] - fn test_sync_policy_state() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_policy_state( - "model1".to_string(), - "cache_aware".to_string(), - b"config_data".to_vec(), - ); - - let state = manager.get_policy_state("model1").unwrap(); - assert_eq!(state.model_id, "model1"); - assert_eq!(state.policy_type, "cache_aware"); - assert_eq!(state.config, b"config_data"); - assert_eq!(state.version, 1); - } - - #[test] - fn test_sync_multiple_policy_states() { - let manager = create_test_sync_manager(); - - manager.sync_policy_state( - "model1".to_string(), - "round_robin".to_string(), - b"config1".to_vec(), - ); - - manager.sync_policy_state( - "model2".to_string(), - "random".to_string(), - b"config2".to_vec(), - ); - - manager.sync_policy_state( - "model3".to_string(), - "consistent_hash".to_string(), - b"config3".to_vec(), - ); - - let all_states = manager.get_all_policy_states(); - assert_eq!(all_states.len(), 3); - - let policy1 = manager.get_policy_state("model1").unwrap(); - assert_eq!(policy1.model_id, "model1"); - assert_eq!(policy1.policy_type, "round_robin"); - - let policy2 = manager.get_policy_state("model2").unwrap(); - assert_eq!(policy2.model_id, "model2"); - assert_eq!(policy2.policy_type, "random"); - } - - #[test] - fn test_remove_policy_state() { - let manager = create_test_sync_manager(); - - manager.sync_policy_state( - "model1".to_string(), - "round_robin".to_string(), - b"config".to_vec(), - ); - - assert!(manager.get_policy_state("model1").is_some()); - - manager.remove_policy_state("model1"); - - assert!(manager.get_policy_state("model1").is_none()); - assert_eq!(manager.get_all_policy_states().len(), 0); - } - - #[test] - fn test_remove_nonexistent_policy_state() { - let manager = create_test_sync_manager(); - - // Should not panic - manager.remove_policy_state("nonexistent"); - assert!(manager.get_policy_state("nonexistent").is_none()); - } - - #[test] - fn test_apply_remote_worker_state() { - let manager = create_test_manager("node1".to_string()); - - // Apply remote state with higher version - let remote_state = WorkerState { - worker_id: "worker1".to_string(), - model_id: "model1".to_string(), - url: "http://localhost:8000".to_string(), - health: true, - load: 0.5, - version: 5, - }; - - manager.apply_remote_worker_state(remote_state.clone(), Some("node2".to_string())); - - let state = manager.get_worker_state("worker1").unwrap(); - assert_eq!(state.version, 5); - } - - #[test] - fn test_apply_remote_worker_state_basic() { - let manager = create_test_sync_manager(); - - let remote_state = WorkerState { - worker_id: "remote_worker1".to_string(), - model_id: "model1".to_string(), - url: "http://localhost:8000".to_string(), - health: true, - load: 0.6, - version: 1, - }; - - manager.apply_remote_worker_state(remote_state.clone(), None); - - let state = manager.get_worker_state("remote_worker1"); - assert!(state.is_some()); - let state = state.unwrap(); - assert_eq!(state.worker_id, "remote_worker1"); - assert_eq!(state.model_id, "model1"); - assert_eq!(state.url, "http://localhost:8000"); - assert!(state.health); - assert_eq!(state.load, 0.6); - } - - #[test] - fn test_apply_remote_worker_state_version_check() { - let manager = create_test_manager("node1".to_string()); - - // First insert local state - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - // Try to apply older version - should be skipped - 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 - }; - - manager.apply_remote_worker_state(old_state, Some("node2".to_string())); - - // Should still have version 1 - let state = manager.get_worker_state("worker1").unwrap(); - assert_eq!(state.version, 1); - assert!(state.health); // Not updated - } - - #[test] - fn test_apply_remote_policy_state() { - let manager = create_test_sync_manager(); - - let remote_state = PolicyState { - model_id: "model1".to_string(), - policy_type: "remote_policy".to_string(), - config: b"remote_config".to_vec(), - version: 1, - }; - - manager.apply_remote_policy_state(remote_state.clone(), None); - - let state = manager.get_policy_state("model1"); - assert!(state.is_some()); - let state = state.unwrap(); - assert_eq!(state.model_id, "model1"); - assert_eq!(state.policy_type, "remote_policy"); - assert_eq!(state.config, b"remote_config"); - } - - #[test] - fn test_mixed_local_and_remote_states() { - let manager = create_test_sync_manager(); - - // Add local worker - manager.sync_worker_state( - "local_worker".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - // Add remote worker - let remote_state = WorkerState { - worker_id: "remote_worker".to_string(), - model_id: "model1".to_string(), - url: "http://localhost:8001".to_string(), - health: true, - load: 0.7, - version: 1, - }; - manager.apply_remote_worker_state(remote_state, None); - - let all_states = manager.get_all_worker_states(); - assert_eq!(all_states.len(), 2); - - assert!(manager.get_worker_state("local_worker").is_some()); - assert!(manager.get_worker_state("remote_worker").is_some()); - } - - #[test] - fn test_update_worker_state() { - let manager = create_test_sync_manager(); - - // Initial state - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - // Update state - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - false, - 0.9, - ); - - let state = manager.get_worker_state("worker1").unwrap(); - assert!(!state.health); - assert_eq!(state.load, 0.9); - assert_eq!(manager.get_all_worker_states().len(), 1); - } - - #[test] - fn test_update_policy_state() { - let manager = create_test_sync_manager(); - - // Initial state - manager.sync_policy_state( - "model1".to_string(), - "round_robin".to_string(), - b"config1".to_vec(), - ); - - // Update state - manager.sync_policy_state( - "model1".to_string(), - "random".to_string(), - b"config2".to_vec(), - ); - - let state = manager.get_policy_state("model1").unwrap(); - assert_eq!(state.policy_type, "random"); - assert_eq!(state.config, b"config2"); - assert_eq!(manager.get_all_policy_states().len(), 1); - } - - #[test] - fn test_get_all_worker_states_empty() { - let manager = create_test_sync_manager(); - let states = manager.get_all_worker_states(); - assert!(states.is_empty()); - } - - #[test] - fn test_get_all_policy_states_empty() { - let manager = create_test_sync_manager(); - let states = manager.get_all_policy_states(); - assert!(states.is_empty()); - } - - #[test] - fn test_update_rate_limit_membership() { - let manager = create_test_manager("node1".to_string()); - - // Add membership nodes - let key1 = SKey::new("node1".to_string()); - manager.stores.membership.insert( - key1, - MembershipState { - name: "node1".to_string(), - address: "127.0.0.1:8000".to_string(), - status: NodeStatus::Alive as i32, - version: 1, - metadata: BTreeMap::new(), - }, - "node1".to_string(), - ); - - let key2 = SKey::new("node2".to_string()); - manager.stores.membership.insert( - key2, - MembershipState { - name: "node2".to_string(), - address: "127.0.0.1:8001".to_string(), - status: NodeStatus::Alive as i32, - version: 1, - metadata: BTreeMap::new(), - }, - "node1".to_string(), - ); - - manager.update_rate_limit_membership(); - - // Check that hash ring was updated - let owners = manager.stores.rate_limit.get_owners("test_key"); - assert!(!owners.is_empty()); - } - - #[test] - fn test_handle_node_failure() { - let manager = create_test_manager("node1".to_string()); - - // Setup membership - let key1 = SKey::new("node1".to_string()); - manager.stores.membership.insert( - key1, - MembershipState { - name: "node1".to_string(), - address: "127.0.0.1:8000".to_string(), - status: NodeStatus::Alive as i32, - version: 1, - metadata: BTreeMap::new(), - }, - "node1".to_string(), - ); - - let key2 = SKey::new("node2".to_string()); - manager.stores.membership.insert( - key2, - MembershipState { - name: "node2".to_string(), - address: "127.0.0.1:8001".to_string(), - status: NodeStatus::Alive as i32, - version: 1, - metadata: BTreeMap::new(), - }, - "node1".to_string(), - ); - - manager.update_rate_limit_membership(); - - // Handle node failure - manager.handle_node_failure(&["node2".to_string()]); - - // Membership should be updated - manager.update_rate_limit_membership(); - } - - #[test] - fn test_sync_rate_limit_inc() { - let manager = create_test_manager("node1".to_string()); - - // Setup membership to make node1 an owner - manager - .stores - .rate_limit - .update_membership(&["node1".to_string()]); - - let test_key = "test_key".to_string(); - if manager.stores.rate_limit.is_owner(&test_key) { - manager.sync_rate_limit_inc(test_key.clone(), 5); - - let value = manager.get_rate_limit_value(&test_key); - assert_eq!(value, Some(5)); - } - } - - #[test] - fn test_sync_rate_limit_inc_non_owner() { - let manager = create_test_manager("node1".to_string()); - - // Setup membership without node1 - manager - .stores - .rate_limit - .update_membership(&["node2".to_string(), "node3".to_string()]); - - let test_key = "test_key".to_string(); - if !manager.stores.rate_limit.is_owner(&test_key) { - manager.sync_rate_limit_inc(test_key.clone(), 5); - - // Should not increment if not owner - let value = manager.get_rate_limit_value(&test_key); - assert_eq!(value, None); - } - } - - #[test] - fn test_get_global_rate_limit_config() { - let manager = create_test_manager("node1".to_string()); - - // Initially should be None - assert!(manager.get_global_rate_limit_config().is_none()); - - // Set config - let key = SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string()); - let config = RateLimitConfig { - limit_per_second: 100, - }; - let serialized = serde_json::to_vec(&config).unwrap(); - manager.stores.app.insert( - key, - AppState { - key: GLOBAL_RATE_LIMIT_KEY.to_string(), - value: serialized, - version: 1, - }, - "node1".to_string(), - ); - - let retrieved = manager.get_global_rate_limit_config().unwrap(); - assert_eq!(retrieved.limit_per_second, 100); - } - - #[test] - fn test_check_global_rate_limit() { - let manager = create_test_manager("node1".to_string()); - - // Setup config - let key = SKey::new(GLOBAL_RATE_LIMIT_KEY.to_string()); - let config = RateLimitConfig { - limit_per_second: 10, - }; - let serialized = serde_json::to_vec(&config).unwrap(); - manager.stores.app.insert( - key, - AppState { - key: GLOBAL_RATE_LIMIT_KEY.to_string(), - value: serialized, - version: 1, - }, - "node1".to_string(), - ); - - // Setup membership - manager - .stores - .rate_limit - .update_membership(&["node1".to_string()]); - - // Check rate limit - let (is_exceeded, _current_count, limit) = manager.check_global_rate_limit(); - assert!(!is_exceeded); // First check should not exceed - assert_eq!(limit, 10); - - // Increment multiple times - for _ in 0..15 { - manager.check_global_rate_limit(); - } - - let (is_exceeded2, current_count2, _) = manager.check_global_rate_limit(); - // Should exceed after many increments - assert!(is_exceeded2 || current_count2 > 10); - } - - #[test] - fn test_reset_global_rate_limit_counter() { - let manager = create_test_manager("node1".to_string()); - - // Setup membership - manager - .stores - .rate_limit - .update_membership(&["node1".to_string()]); - - // Increment counter - if manager - .stores - .rate_limit - .is_owner(GLOBAL_RATE_LIMIT_COUNTER_KEY) - { - manager.sync_rate_limit_inc(GLOBAL_RATE_LIMIT_COUNTER_KEY.to_string(), 10); - let value = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY); - assert!(value.is_some() && value.unwrap() > 0); - - // Reset - manager.reset_global_rate_limit_counter(); - let value_after = manager.get_rate_limit_value(GLOBAL_RATE_LIMIT_COUNTER_KEY); - // Should be reset (0 or negative) - assert!(value_after.is_none() || value_after.unwrap() <= 0); - } - } - - #[test] - fn test_sync_tree_operation() { - let manager = create_test_manager("node1".to_string()); - - use crate::mesh::tree_ops::{TreeInsertOp, TreeOperation}; - - let op = TreeOperation::Insert(TreeInsertOp { - text: "test_text".to_string(), - tenant: "http://localhost:8000".to_string(), - }); - - let result = manager.sync_tree_operation("model1".to_string(), op); - assert!(result.is_ok()); - - // Verify tree state was stored - let tree_state = manager.get_tree_state("model1"); - assert!(tree_state.is_some()); - let tree = tree_state.unwrap(); - assert_eq!(tree.model_id, "model1"); - assert_eq!(tree.operations.len(), 1); - } - - #[test] - fn test_get_tree_state() { - let manager = create_test_manager("node1".to_string()); - - // Initially should be None - assert!(manager.get_tree_state("model1").is_none()); - - // Sync an operation - use crate::mesh::tree_ops::{TreeInsertOp, TreeOperation}; - let op = TreeOperation::Insert(TreeInsertOp { - text: "test_text".to_string(), - tenant: "http://localhost:8000".to_string(), - }); - manager - .sync_tree_operation("model1".to_string(), op) - .unwrap(); - - let tree_state = manager.get_tree_state("model1"); - assert!(tree_state.is_some()); - } - - #[test] - fn test_apply_remote_tree_operation() { - let manager = create_test_manager("node1".to_string()); - - use crate::mesh::tree_ops::{TreeInsertOp, TreeOperation, TreeState}; - - let mut tree_state = TreeState::new("model1".to_string()); - tree_state.version = 5; - tree_state.add_operation(TreeOperation::Insert(TreeInsertOp { - text: "remote_text".to_string(), - tenant: "http://localhost:8001".to_string(), - })); - // add_operation increments version, so version is now 6 - - manager.apply_remote_tree_operation( - "model1".to_string(), - tree_state, - Some("node2".to_string()), - ); - - let retrieved = manager.get_tree_state("model1").unwrap(); - assert_eq!(retrieved.version, 6); // add_operation increments version from 5 to 6 - assert_eq!(retrieved.operations.len(), 1); - } - - #[test] - fn test_get_all_worker_states() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_worker_state( - "worker1".to_string(), - "model1".to_string(), - "http://localhost:8000".to_string(), - true, - 0.5, - ); - - manager.sync_worker_state( - "worker2".to_string(), - "model2".to_string(), - "http://localhost:8001".to_string(), - false, - 0.8, - ); - - let all_states = manager.get_all_worker_states(); - assert_eq!(all_states.len(), 2); - } - - #[test] - fn test_get_all_policy_states() { - let manager = create_test_manager("node1".to_string()); - - manager.sync_policy_state("model1".to_string(), "cache_aware".to_string(), vec![]); - - manager.sync_policy_state("model2".to_string(), "round_robin".to_string(), vec![]); - - let all_states = manager.get_all_policy_states(); - assert_eq!(all_states.len(), 2); - } -} diff --git a/sgl-model-gateway/src/mesh/test_utils.rs b/sgl-model-gateway/src/mesh/test_utils.rs deleted file mode 100644 index d800cc6a6..000000000 --- a/sgl-model-gateway/src/mesh/test_utils.rs +++ /dev/null @@ -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 { - Arc::new(StateStores::with_self_name(self_name)) -} - -/// Create test MeshSyncManager -pub fn create_test_sync_manager(self_name: String) -> Arc { - 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")); - } -} diff --git a/sgl-model-gateway/src/mesh/topology.rs b/sgl-model-gateway/src/mesh/topology.rs deleted file mode 100644 index 038ba9f8e..000000000 --- a/sgl-model-gateway/src/mesh/topology.rs +++ /dev/null @@ -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, - /// Availability zone identifier (for sparse mesh) - pub availability_zone: Option, -} - -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>>, -} - -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 { - 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, - count: usize, - ) -> Vec { - 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> = 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 - } - - /// Get peers for sparse mesh topology (by region/AZ) - fn get_sparse_mesh_peers( - &self, - state: &BTreeMap, - count: usize, - ) -> Vec { - 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> = 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> = 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()); - } -} diff --git a/sgl-model-gateway/src/mesh/tree_ops.rs b/sgl-model-gateway/src/mesh/tree_ops.rs deleted file mode 100644 index 17e91785c..000000000 --- a/sgl-model-gateway/src/mesh/tree_ops.rs +++ /dev/null @@ -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, - 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); - } -} diff --git a/sgl-model-gateway/src/mesh/endpoints.rs b/sgl-model-gateway/src/routers/mesh/handlers.rs similarity index 98% rename from sgl-model-gateway/src/mesh/endpoints.rs rename to sgl-model-gateway/src/routers/mesh/handlers.rs index 5534613be..8be810dad 100644 --- a/sgl-model-gateway/src/mesh/endpoints.rs +++ b/sgl-model-gateway/src/routers/mesh/handlers.rs @@ -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>) // 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>) - ) .into_response() } - -use std::sync::Arc; - -use crate::{ - mesh::stores::{RateLimitConfig, GLOBAL_RATE_LIMIT_KEY}, - server::AppState, -}; diff --git a/sgl-model-gateway/src/routers/mesh/mod.rs b/sgl-model-gateway/src/routers/mesh/mod.rs new file mode 100644 index 000000000..d88008372 --- /dev/null +++ b/sgl-model-gateway/src/routers/mesh/mod.rs @@ -0,0 +1,7 @@ +//! Mesh cluster management HTTP handlers +//! +//! This module provides HTTP API endpoints for mesh cluster management. + +mod handlers; + +pub use handlers::*; diff --git a/sgl-model-gateway/src/routers/mod.rs b/sgl-model-gateway/src/routers/mod.rs index 7301ce2d6..11d34a68b 100644 --- a/sgl-model-gateway/src/routers/mod.rs +++ b/sgl-model-gateway/src/routers/mod.rs @@ -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; diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index cfbfc9755..6a4d5c531 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -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}, diff --git a/sgl-model-gateway/tests/mesh_integration_test.rs b/sgl-model-gateway/tests/mesh_integration_test.rs deleted file mode 100644 index 2aa4d5d0c..000000000 --- a/sgl-model-gateway/tests/mesh_integration_test.rs +++ /dev/null @@ -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 { - Arc::new(StateStores::with_self_name(node_name)) -} - -/// Create test sync manager for a node -fn create_test_sync_manager(node_name: String) -> Arc { - 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); - } -}