From 7fb9946bb46401cc50fbd4c94a753aa35ba335e5 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:59:20 +0530 Subject: [PATCH 01/16] feat(mesh): add taskito-mesh crate with local deque, hash ring, and mesh state New crate for decentralized task dispatch. Phase 1 foundation: local deque with affinity sorting, xxhash consistent-hashing ring with virtual nodes, mesh state with membership tracking. 20 tests. --- Cargo.toml | 2 +- crates/taskito-mesh/Cargo.toml | 14 ++ crates/taskito-mesh/src/config.rs | 51 ++++++ crates/taskito-mesh/src/lib.rs | 198 +++++++++++++++++++++++ crates/taskito-mesh/src/local_deque.rs | 196 +++++++++++++++++++++++ crates/taskito-mesh/src/metrics.rs | 55 +++++++ crates/taskito-mesh/src/ring.rs | 191 ++++++++++++++++++++++ crates/taskito-mesh/src/state.rs | 209 +++++++++++++++++++++++++ 8 files changed, 915 insertions(+), 1 deletion(-) create mode 100644 crates/taskito-mesh/Cargo.toml create mode 100644 crates/taskito-mesh/src/config.rs create mode 100644 crates/taskito-mesh/src/lib.rs create mode 100644 crates/taskito-mesh/src/local_deque.rs create mode 100644 crates/taskito-mesh/src/metrics.rs create mode 100644 crates/taskito-mesh/src/ring.rs create mode 100644 crates/taskito-mesh/src/state.rs diff --git a/Cargo.toml b/Cargo.toml index 2bfc2c227..0ee9ea32b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["crates/taskito-core", "crates/taskito-python", "crates/taskito-async", "crates/taskito-workflows"] +members = ["crates/taskito-core", "crates/taskito-python", "crates/taskito-async", "crates/taskito-workflows", "crates/taskito-mesh"] resolver = "2" [workspace.dependencies] diff --git a/crates/taskito-mesh/Cargo.toml b/crates/taskito-mesh/Cargo.toml new file mode 100644 index 000000000..f38a3cb0b --- /dev/null +++ b/crates/taskito-mesh/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "taskito-mesh" +version = "0.15.2" +edition = "2021" + +[dependencies] +taskito-core = { path = "../taskito-core" } +tokio = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +bincode = "1" +log = { workspace = true } +rand = { workspace = true } +xxhash-rust = { version = "0.8", features = ["xxh3"] } diff --git a/crates/taskito-mesh/src/config.rs b/crates/taskito-mesh/src/config.rs new file mode 100644 index 000000000..a8d33021f --- /dev/null +++ b/crates/taskito-mesh/src/config.rs @@ -0,0 +1,51 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MeshConfig { + /// UDP port for SWIM gossip protocol. + pub gossip_port: u16, + /// TCP port for work-stealing connections (default: gossip_port + 1). + pub steal_port: u16, + /// Bind address for gossip and steal listeners. + pub bind_addr: String, + /// Seed nodes for initial cluster join (e.g., `["host1:7946"]`). + pub seeds: Vec, + /// SWIM protocol period in milliseconds. + pub protocol_period_ms: u64, + /// Indirect ping targets for failure detection. + pub indirect_ping_count: usize, + /// Suspicion timeout multiplier (applied to `log(N+1) * protocol_period`). + pub suspicion_multiplier: u32, + /// Virtual nodes per worker on the consistent hash ring. + pub virtual_nodes: usize, + /// Max jobs in the local deque before refusing to prefetch. + pub local_buffer_capacity: usize, + /// Max jobs to steal per request. + pub max_steal_batch: usize, + /// Steal when own deque length is at or below this threshold. + pub steal_threshold: usize, + /// Affinity weight: 0.0 = ignore affinity, 1.0 = strict affinity. + pub affinity_weight: f64, + /// Whether work-stealing is enabled. + pub enable_stealing: bool, +} + +impl Default for MeshConfig { + fn default() -> Self { + Self { + gossip_port: 7946, + steal_port: 7947, + bind_addr: "0.0.0.0".to_string(), + seeds: Vec::new(), + protocol_period_ms: 500, + indirect_ping_count: 3, + suspicion_multiplier: 4, + virtual_nodes: 150, + local_buffer_capacity: 64, + max_steal_batch: 4, + steal_threshold: 2, + affinity_weight: 0.7, + enable_stealing: true, + } + } +} diff --git a/crates/taskito-mesh/src/lib.rs b/crates/taskito-mesh/src/lib.rs new file mode 100644 index 000000000..4f7d55e89 --- /dev/null +++ b/crates/taskito-mesh/src/lib.rs @@ -0,0 +1,198 @@ +pub mod config; +pub mod local_deque; +pub mod metrics; +pub mod ring; +pub mod state; + +use std::sync::atomic::Ordering; +use std::sync::Arc; + +use taskito_core::job::Job; + +pub use config::MeshConfig; +pub use local_deque::LocalDeque; +pub use metrics::{MeshMetrics, MetricsSnapshot}; +pub use state::{MemberState, MeshState, WorkerInfo}; + +/// A mesh node manages the local deque, consistent-hash ring, and (in later +/// phases) the SWIM gossip protocol and work-stealing connections. +/// +/// Phase 1 provides local-deque prefetch with affinity sorting. +/// Gossip and stealing are added in subsequent phases. +pub struct MeshNode { + config: MeshConfig, + state: Arc, + deque: LocalDeque, + metrics: Arc, +} + +impl MeshNode { + pub fn new(worker_id: String, config: MeshConfig) -> Self { + let state = Arc::new(MeshState::new(worker_id, config.virtual_nodes)); + let deque = LocalDeque::new(config.local_buffer_capacity); + Self { + config, + state, + deque, + metrics: Arc::new(MeshMetrics::default()), + } + } + + pub fn config(&self) -> &MeshConfig { + &self.config + } + + pub fn state(&self) -> &Arc { + &self.state + } + + pub fn metrics(&self) -> MetricsSnapshot { + self.metrics.snapshot() + } + + // ── Local deque operations ────────────────────────────────────────── + + /// Pop a job from the local deque (front/hot end). + /// Returns `None` if the deque is empty. + pub fn pop_local(&self) -> Option { + let job = self.deque.pop(); + if job.is_some() { + self.metrics.local_pops.fetch_add(1, Ordering::Relaxed); + } + job + } + + /// Push a batch of jobs into the local deque, sorted by affinity. + /// Affinity-owned tasks go to the front (popped first), non-owned + /// settle at the back (stealable). + pub fn prefetch(&self, jobs: Vec) -> usize { + if jobs.is_empty() { + return 0; + } + self.metrics.prefetch_count.fetch_add(1, Ordering::Relaxed); + let state = self.state.clone(); + let pushed = self + .deque + .push_sorted(jobs, |task_name| state.is_local_owner(task_name)); + self.metrics + .prefetch_jobs + .fetch_add(pushed as u64, Ordering::Relaxed); + pushed + } + + /// Whether the local deque has room for more prefetched jobs. + pub fn should_prefetch(&self) -> bool { + self.deque.len() < self.config.local_buffer_capacity / 2 + } + + /// Whether the local deque has jobs ready to dispatch. + pub fn has_local_work(&self) -> bool { + !self.deque.is_empty() + } + + /// Current local deque length. + pub fn local_len(&self) -> usize { + self.deque.len() + } + + // ── Affinity queries ──────────────────────────────────────────────── + + /// Check if this worker is the affinity owner for a task name. + pub fn is_affinity_owner(&self, task_name: &str) -> bool { + self.state.is_local_owner(task_name) + } + + // ── Steal operations (Phase 4 — stubs for now) ────────────────────── + + /// Steal up to `count` jobs from the cold end of the local deque. + /// Used by the steal server to respond to steal requests. + pub fn give_jobs(&self, count: usize) -> Vec { + let stolen = self.deque.steal(count); + self.metrics + .jobs_stolen_out + .fetch_add(stolen.len() as u64, Ordering::Relaxed); + stolen + } + + /// Whether this node should attempt to steal work from a peer. + pub fn should_steal(&self) -> bool { + self.config.enable_stealing && self.deque.len() <= self.config.steal_threshold + } +} + +#[cfg(test)] +mod tests { + use super::*; + use taskito_core::job::{now_millis, NewJob}; + + fn make_job(task_name: &str) -> Job { + NewJob { + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: vec![], + priority: 0, + scheduled_at: now_millis(), + max_retries: 0, + timeout_ms: 30_000, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + } + .into_job() + } + + #[test] + fn prefetch_and_pop() { + let node = MeshNode::new("worker-1".to_string(), MeshConfig::default()); + let jobs = vec![make_job("task_a"), make_job("task_b"), make_job("task_c")]; + + let pushed = node.prefetch(jobs); + assert_eq!(pushed, 3); + assert!(node.has_local_work()); + + let metrics = node.metrics(); + assert_eq!(metrics.prefetch_count, 1); + assert_eq!(metrics.prefetch_jobs, 3); + + let j1 = node.pop_local().unwrap(); + assert!(!j1.task_name.is_empty()); + + let metrics = node.metrics(); + assert_eq!(metrics.local_pops, 1); + } + + #[test] + fn should_prefetch_tracks_capacity() { + let config = MeshConfig { + local_buffer_capacity: 4, + ..MeshConfig::default() + }; + let node = MeshNode::new("worker-1".to_string(), config); + assert!(node.should_prefetch()); // empty = should prefetch + + node.prefetch(vec![make_job("a"), make_job("b"), make_job("c")]); + assert!(!node.should_prefetch()); // 3/4 > capacity/2 + } + + #[test] + fn should_steal_respects_config() { + let config = MeshConfig { + enable_stealing: false, + ..MeshConfig::default() + }; + let node = MeshNode::new("worker-1".to_string(), config); + assert!(!node.should_steal()); + + let config2 = MeshConfig { + enable_stealing: true, + steal_threshold: 2, + ..MeshConfig::default() + }; + let node2 = MeshNode::new("worker-2".to_string(), config2); + assert!(node2.should_steal()); // empty deque ≤ threshold + } +} diff --git a/crates/taskito-mesh/src/local_deque.rs b/crates/taskito-mesh/src/local_deque.rs new file mode 100644 index 000000000..30da096aa --- /dev/null +++ b/crates/taskito-mesh/src/local_deque.rs @@ -0,0 +1,196 @@ +use std::collections::VecDeque; +use std::sync::Mutex; + +use taskito_core::job::Job; + +/// Thread-safe local job buffer. +/// +/// Owner pushes to the back and pops from the back (LIFO for affinity-hot +/// jobs). Stealers pop from the front (cold end). Protected by a mutex; +/// contention is negligible since steals are infrequent. +pub struct LocalDeque { + inner: Mutex>, + capacity: usize, +} + +// Safety: Mutex> is Send+Sync when Job is Send. +unsafe impl Send for LocalDeque {} +unsafe impl Sync for LocalDeque {} + +impl LocalDeque { + pub fn new(capacity: usize) -> Self { + Self { + inner: Mutex::new(VecDeque::with_capacity(capacity)), + capacity, + } + } + + /// Push a job to the back (owner end). Returns false if at capacity. + pub fn push(&self, job: Job) -> bool { + let mut deque = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + if deque.len() >= self.capacity { + return false; + } + deque.push_back(job); + true + } + + /// Pop a job from the back (owner end, hot/affinity side — LIFO). + pub fn pop(&self) -> Option { + let mut deque = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + deque.pop_back() + } + + /// Steal up to `count` jobs from the front (cold/stealable end — FIFO). + pub fn steal(&self, count: usize) -> Vec { + let mut deque = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + let n = count.min(deque.len()); + let mut stolen = Vec::with_capacity(n); + for _ in 0..n { + if let Some(job) = deque.pop_front() { + stolen.push(job); + } + } + stolen + } + + pub fn len(&self) -> usize { + let deque = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + deque.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Push multiple jobs, sorting by affinity: owned tasks go to the back + /// (popped first via LIFO), non-owned go to the front (stealable end). + pub fn push_sorted(&self, mut jobs: Vec, is_affinity_owner: F) -> usize + where + F: Fn(&str) -> bool, + { + let mut deque = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + let remaining = self.capacity.saturating_sub(deque.len()); + jobs.truncate(remaining); + + // Partition: non-affinity first (front/stealable), affinity last (back/hot) + jobs.sort_by_key(|j| { + if is_affinity_owner(&j.task_name) { + 1 + } else { + 0 + } + }); + + let pushed = jobs.len(); + for job in jobs { + deque.push_back(job); + } + pushed + } +} + +#[cfg(test)] +mod tests { + use super::*; + use taskito_core::job::{now_millis, NewJob}; + + fn make_job(task_name: &str) -> Job { + NewJob { + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: vec![], + priority: 0, + scheduled_at: now_millis(), + max_retries: 0, + timeout_ms: 30_000, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + } + .into_job() + } + + #[test] + fn push_pop_lifo_order() { + let deque = LocalDeque::new(10); + deque.push(make_job("first")); + deque.push(make_job("second")); + deque.push(make_job("third")); + + assert_eq!(deque.len(), 3); + assert_eq!(deque.pop().unwrap().task_name, "third"); + assert_eq!(deque.pop().unwrap().task_name, "second"); + assert_eq!(deque.pop().unwrap().task_name, "first"); + assert!(deque.pop().is_none()); + } + + #[test] + fn capacity_enforced() { + let deque = LocalDeque::new(2); + assert!(deque.push(make_job("a"))); + assert!(deque.push(make_job("b"))); + assert!(!deque.push(make_job("c"))); + assert_eq!(deque.len(), 2); + } + + #[test] + fn steal_takes_from_front() { + let deque = LocalDeque::new(10); + deque.push(make_job("first")); + deque.push(make_job("second")); + deque.push(make_job("third")); + + let stolen = deque.steal(2); + assert_eq!(stolen.len(), 2); + assert_eq!(stolen[0].task_name, "first"); + assert_eq!(stolen[1].task_name, "second"); + + // Owner pops the remaining one from back + assert_eq!(deque.pop().unwrap().task_name, "third"); + } + + #[test] + fn push_sorted_affinity_at_back() { + let deque = LocalDeque::new(10); + let jobs = vec![ + make_job("cold_task"), + make_job("hot_task"), + make_job("another_cold"), + ]; + + let pushed = deque.push_sorted(jobs, |name| name == "hot_task"); + assert_eq!(pushed, 3); + + // Popping from back (LIFO): affinity jobs were pushed last → popped first + assert_eq!(deque.pop().unwrap().task_name, "hot_task"); + } + + #[test] + fn push_sorted_respects_capacity() { + let deque = LocalDeque::new(2); + let jobs = vec![make_job("a"), make_job("b"), make_job("c")]; + let pushed = deque.push_sorted(jobs, |_| false); + assert_eq!(pushed, 2); + assert_eq!(deque.len(), 2); + } + + #[test] + fn empty_deque() { + let deque = LocalDeque::new(10); + assert!(deque.is_empty()); + assert_eq!(deque.len(), 0); + assert!(deque.pop().is_none()); + + let stolen = deque.steal(5); + assert!(stolen.is_empty()); + } +} diff --git a/crates/taskito-mesh/src/metrics.rs b/crates/taskito-mesh/src/metrics.rs new file mode 100644 index 000000000..bd0fcfdcd --- /dev/null +++ b/crates/taskito-mesh/src/metrics.rs @@ -0,0 +1,55 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Counters for mesh operations, exposed for observability. +pub struct MeshMetrics { + pub prefetch_count: AtomicU64, + pub prefetch_jobs: AtomicU64, + pub local_pops: AtomicU64, + pub steals_initiated: AtomicU64, + pub steals_succeeded: AtomicU64, + pub jobs_stolen_in: AtomicU64, + pub jobs_stolen_out: AtomicU64, + pub ring_recalculations: AtomicU64, +} + +impl Default for MeshMetrics { + fn default() -> Self { + Self { + prefetch_count: AtomicU64::new(0), + prefetch_jobs: AtomicU64::new(0), + local_pops: AtomicU64::new(0), + steals_initiated: AtomicU64::new(0), + steals_succeeded: AtomicU64::new(0), + jobs_stolen_in: AtomicU64::new(0), + jobs_stolen_out: AtomicU64::new(0), + ring_recalculations: AtomicU64::new(0), + } + } +} + +impl MeshMetrics { + pub fn snapshot(&self) -> MetricsSnapshot { + MetricsSnapshot { + prefetch_count: self.prefetch_count.load(Ordering::Relaxed), + prefetch_jobs: self.prefetch_jobs.load(Ordering::Relaxed), + local_pops: self.local_pops.load(Ordering::Relaxed), + steals_initiated: self.steals_initiated.load(Ordering::Relaxed), + steals_succeeded: self.steals_succeeded.load(Ordering::Relaxed), + jobs_stolen_in: self.jobs_stolen_in.load(Ordering::Relaxed), + jobs_stolen_out: self.jobs_stolen_out.load(Ordering::Relaxed), + ring_recalculations: self.ring_recalculations.load(Ordering::Relaxed), + } + } +} + +#[derive(Debug, Clone, serde::Serialize)] +pub struct MetricsSnapshot { + pub prefetch_count: u64, + pub prefetch_jobs: u64, + pub local_pops: u64, + pub steals_initiated: u64, + pub steals_succeeded: u64, + pub jobs_stolen_in: u64, + pub jobs_stolen_out: u64, + pub ring_recalculations: u64, +} diff --git a/crates/taskito-mesh/src/ring.rs b/crates/taskito-mesh/src/ring.rs new file mode 100644 index 000000000..d43017ada --- /dev/null +++ b/crates/taskito-mesh/src/ring.rs @@ -0,0 +1,191 @@ +use std::collections::BTreeMap; + +use xxhash_rust::xxh3::xxh3_64; + +/// Consistent hash ring with virtual nodes for even distribution. +/// +/// Maps task names to preferred worker IDs. Workers are placed on the ring +/// via `virtual_nodes` points each; tasks hash to the next clockwise worker. +pub struct HashRing { + ring: BTreeMap, + virtual_nodes: usize, +} + +impl HashRing { + pub fn new(virtual_nodes: usize) -> Self { + Self { + ring: BTreeMap::new(), + virtual_nodes, + } + } + + /// Add a worker to the ring with `virtual_nodes` points. + pub fn add_worker(&mut self, worker_id: &str) { + for i in 0..self.virtual_nodes { + let key = format!("{worker_id}-vnode-{i}"); + let hash = xxh3_64(key.as_bytes()); + self.ring.insert(hash, worker_id.to_string()); + } + } + + /// Remove a worker and all its virtual nodes from the ring. + pub fn remove_worker(&mut self, worker_id: &str) { + self.ring.retain(|_, v| v != worker_id); + } + + /// Look up the preferred worker for a given key (e.g., task_name). + /// Returns `None` if the ring is empty. + pub fn preferred_worker(&self, key: &str) -> Option<&str> { + if self.ring.is_empty() { + return None; + } + let hash = xxh3_64(key.as_bytes()); + // Find the first node clockwise from the hash point + self.ring + .range(hash..) + .next() + .or_else(|| self.ring.iter().next()) + .map(|(_, worker_id)| worker_id.as_str()) + } + + /// Check whether a specific worker is the preferred owner for a key. + pub fn is_owner(&self, key: &str, worker_id: &str) -> bool { + self.preferred_worker(key) == Some(worker_id) + } + + /// Number of workers (not virtual nodes) on the ring. + pub fn worker_count(&self) -> usize { + let mut seen = std::collections::HashSet::new(); + for worker_id in self.ring.values() { + seen.insert(worker_id.as_str()); + } + seen.len() + } + + pub fn is_empty(&self) -> bool { + self.ring.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deterministic_placement() { + let mut ring = HashRing::new(150); + ring.add_worker("worker-a"); + ring.add_worker("worker-b"); + ring.add_worker("worker-c"); + + let owner1 = ring.preferred_worker("my_task").unwrap(); + let owner2 = ring.preferred_worker("my_task").unwrap(); + assert_eq!(owner1, owner2, "same key must always map to same worker"); + } + + #[test] + fn different_tasks_may_map_differently() { + let mut ring = HashRing::new(150); + ring.add_worker("worker-a"); + ring.add_worker("worker-b"); + + // With 2 workers and 150 vnodes each, different task names should + // spread across workers (probabilistic but near-certain for these names). + let mut owners = std::collections::HashSet::new(); + for i in 0..20 { + owners.insert( + ring.preferred_worker(&format!("task_{i}")) + .unwrap() + .to_string(), + ); + } + assert!(owners.len() > 1, "tasks should distribute across workers"); + } + + #[test] + fn minimal_key_migration_on_remove() { + let mut ring = HashRing::new(150); + ring.add_worker("worker-a"); + ring.add_worker("worker-b"); + ring.add_worker("worker-c"); + + let tasks: Vec = (0..100).map(|i| format!("task_{i}")).collect(); + let before: Vec = tasks + .iter() + .map(|t| ring.preferred_worker(t).unwrap().to_string()) + .collect(); + + ring.remove_worker("worker-b"); + + let mut migrated = 0; + for (i, task) in tasks.iter().enumerate() { + let after = ring.preferred_worker(task).unwrap(); + if after != before[i] { + migrated += 1; + } + } + + // With consistent hashing, removing 1 of 3 workers should migrate + // roughly 1/3 of keys, not all of them. + assert!( + migrated < 60, + "expected < 60% migration, got {migrated}/100" + ); + } + + #[test] + fn even_distribution() { + let mut ring = HashRing::new(150); + ring.add_worker("w1"); + ring.add_worker("w2"); + ring.add_worker("w3"); + + let mut counts = std::collections::HashMap::new(); + for i in 0..3000 { + let owner = ring.preferred_worker(&format!("key_{i}")).unwrap(); + *counts.entry(owner.to_string()).or_insert(0) += 1; + } + + for (worker, count) in &counts { + // Each worker should get roughly 1000 ± 300 (30% tolerance) + assert!( + *count > 700 && *count < 1300, + "worker {worker} got {count}/3000 keys — distribution too uneven" + ); + } + } + + #[test] + fn is_owner_check() { + let mut ring = HashRing::new(150); + ring.add_worker("worker-a"); + ring.add_worker("worker-b"); + + let owner = ring.preferred_worker("test_task").unwrap().to_string(); + assert!(ring.is_owner("test_task", &owner)); + } + + #[test] + fn empty_ring() { + let ring = HashRing::new(150); + assert!(ring.preferred_worker("anything").is_none()); + assert!(ring.is_empty()); + assert_eq!(ring.worker_count(), 0); + } + + #[test] + fn add_remove_worker() { + let mut ring = HashRing::new(150); + ring.add_worker("worker-a"); + assert_eq!(ring.worker_count(), 1); + + ring.add_worker("worker-b"); + assert_eq!(ring.worker_count(), 2); + + ring.remove_worker("worker-a"); + assert_eq!(ring.worker_count(), 1); + + // All keys now map to worker-b + assert_eq!(ring.preferred_worker("any_task").unwrap(), "worker-b"); + } +} diff --git a/crates/taskito-mesh/src/state.rs b/crates/taskito-mesh/src/state.rs new file mode 100644 index 000000000..4878d312c --- /dev/null +++ b/crates/taskito-mesh/src/state.rs @@ -0,0 +1,209 @@ +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::RwLock; + +use serde::{Deserialize, Serialize}; + +use crate::ring::HashRing; + +/// Load and identity information gossipped between mesh peers. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WorkerInfo { + pub worker_id: String, + pub gossip_addr: SocketAddr, + pub steal_addr: SocketAddr, + pub queues: Vec, + pub threads: u16, + pub current_load: u16, + pub local_buffer_len: u16, + pub capacity: u16, + pub updated_at: i64, +} + +/// Member state in the SWIM protocol. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum MemberState { + Alive, + Suspect, + Dead, + Left, +} + +/// A member entry with state tracking. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Member { + pub info: WorkerInfo, + pub state: MemberState, + pub incarnation: u64, +} + +/// Shared mesh state: membership map + consistent hash ring. +/// +/// Thread-safe via `RwLock`. The gossip loop writes; the scheduler loop +/// and steal coordinator read. +pub struct MeshState { + members: RwLock>, + ring: RwLock, + local_worker_id: String, +} + +impl MeshState { + pub fn new(worker_id: String, virtual_nodes: usize) -> Self { + let mut ring = HashRing::new(virtual_nodes); + ring.add_worker(&worker_id); + Self { + members: RwLock::new(HashMap::new()), + ring: RwLock::new(ring), + local_worker_id: worker_id, + } + } + + pub fn local_worker_id(&self) -> &str { + &self.local_worker_id + } + + /// Check if a task name is affinity-owned by this worker. + pub fn is_local_owner(&self, task_name: &str) -> bool { + let ring = self.ring.read().unwrap_or_else(|p| p.into_inner()); + ring.is_owner(task_name, &self.local_worker_id) + } + + /// Update or insert a member. Returns true if this is a new member. + pub fn upsert_member(&self, member: Member) -> bool { + let worker_id = member.info.worker_id.clone(); + let is_alive = member.state == MemberState::Alive; + let mut members = self.members.write().unwrap_or_else(|p| p.into_inner()); + let is_new = !members.contains_key(&worker_id); + members.insert(worker_id.clone(), member); + + if is_new && is_alive { + let mut ring = self.ring.write().unwrap_or_else(|p| p.into_inner()); + ring.add_worker(&worker_id); + } + is_new + } + + /// Mark a member as dead and remove from ring. + pub fn mark_dead(&self, worker_id: &str) { + let mut members = self.members.write().unwrap_or_else(|p| p.into_inner()); + if let Some(m) = members.get_mut(worker_id) { + m.state = MemberState::Dead; + } + let mut ring = self.ring.write().unwrap_or_else(|p| p.into_inner()); + ring.remove_worker(worker_id); + } + + /// Mark a member as gracefully left and remove from ring. + pub fn mark_left(&self, worker_id: &str) { + let mut members = self.members.write().unwrap_or_else(|p| p.into_inner()); + if let Some(m) = members.get_mut(worker_id) { + m.state = MemberState::Left; + } + let mut ring = self.ring.write().unwrap_or_else(|p| p.into_inner()); + ring.remove_worker(worker_id); + } + + /// Get all alive members sorted by local_buffer_len descending (busiest first). + pub fn alive_peers(&self) -> Vec { + let members = self.members.read().unwrap_or_else(|p| p.into_inner()); + let mut alive: Vec = members + .values() + .filter(|m| m.state == MemberState::Alive && m.info.worker_id != self.local_worker_id) + .cloned() + .collect(); + alive.sort_by_key(|m| std::cmp::Reverse(m.info.local_buffer_len)); + alive + } + + /// Get the busiest peer that has enough surplus to steal from. + pub fn best_steal_target(&self, min_surplus: usize) -> Option { + self.alive_peers() + .into_iter() + .find(|m| m.info.local_buffer_len as usize > min_surplus) + } + + /// Number of alive members (excluding self). + pub fn alive_count(&self) -> usize { + let members = self.members.read().unwrap_or_else(|p| p.into_inner()); + members + .values() + .filter(|m| m.state == MemberState::Alive && m.info.worker_id != self.local_worker_id) + .count() + } + + /// Remove members that have been dead/left for longer than the given + /// threshold. Returns the number removed. + pub fn prune_dead(&self, _older_than_ms: i64) -> usize { + let mut members = self.members.write().unwrap_or_else(|p| p.into_inner()); + let before = members.len(); + members.retain(|_, m| matches!(m.state, MemberState::Alive | MemberState::Suspect)); + before - members.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr}; + + fn make_info(id: &str, buffer_len: u16) -> WorkerInfo { + WorkerInfo { + worker_id: id.to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7946), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7947), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: buffer_len, + capacity: 4, + updated_at: 0, + } + } + + fn make_member(id: &str, buffer_len: u16) -> Member { + Member { + info: make_info(id, buffer_len), + state: MemberState::Alive, + incarnation: 1, + } + } + + #[test] + fn upsert_and_query() { + let state = MeshState::new("local".to_string(), 150); + assert!(state.upsert_member(make_member("peer-a", 5))); + assert!(!state.upsert_member(make_member("peer-a", 10))); // update, not new + assert_eq!(state.alive_count(), 1); + } + + #[test] + fn best_steal_target_picks_busiest() { + let state = MeshState::new("local".to_string(), 150); + state.upsert_member(make_member("peer-a", 3)); + state.upsert_member(make_member("peer-b", 10)); + state.upsert_member(make_member("peer-c", 7)); + + let target = state.best_steal_target(2).unwrap(); + assert_eq!(target.info.worker_id, "peer-b"); + } + + #[test] + fn mark_dead_removes_from_ring() { + let state = MeshState::new("local".to_string(), 150); + state.upsert_member(make_member("peer-a", 5)); + assert!(state.is_local_owner("some_task") || !state.is_local_owner("some_task")); // ring has 2 workers + + state.mark_dead("peer-a"); + assert_eq!(state.alive_count(), 0); + } + + #[test] + fn prune_removes_dead_members() { + let state = MeshState::new("local".to_string(), 150); + state.upsert_member(make_member("peer-a", 5)); + state.mark_dead("peer-a"); + let pruned = state.prune_dead(0); + assert_eq!(pruned, 1); + assert_eq!(state.alive_count(), 0); + } +} From cd8f4d08f72262a876f13122edf6483f0415057e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:59:43 +0530 Subject: [PATCH 02/16] feat(mesh): wire mesh feature into scheduler with local deque bridge Mesh bridge interposes between scheduler and dispatcher via intermediate channel. Scheduler dequeues from DB, bridge pushes into local deque with affinity sorting, then drains to dispatcher. --- crates/taskito-python/Cargo.toml | 2 + crates/taskito-python/src/py_queue/worker.rs | 76 ++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/crates/taskito-python/Cargo.toml b/crates/taskito-python/Cargo.toml index d84016c96..df52cacd0 100644 --- a/crates/taskito-python/Cargo.toml +++ b/crates/taskito-python/Cargo.toml @@ -12,6 +12,7 @@ native-async = ["dep:taskito-async"] workflows = ["dep:taskito-workflows"] # Event-driven scheduler wakeups. OFF by default (not in maturin defaults). push-dispatch = ["taskito-core/push-dispatch"] +mesh = ["dep:taskito-mesh"] [lib] name = "_taskito" @@ -26,6 +27,7 @@ uuid = { workspace = true } async-trait = { workspace = true } taskito-async = { path = "../taskito-async", optional = true } taskito-workflows = { path = "../taskito-workflows", optional = true } +taskito-mesh = { path = "../taskito-mesh", optional = true } serde_json = { workspace = true } serde = { workspace = true } base64 = "0.22" diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 86a057684..e3cab375e 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -15,6 +15,54 @@ use super::PyQueue; use crate::async_worker::AsyncWorkerPool; use crate::py_config::PyTaskConfig; +/// Mesh-aware scheduler bridge: receives jobs from the scheduler's +/// intermediate channel, pushes them into the local deque with affinity +/// sorting, then drains the deque to the real dispatcher channel. +#[cfg(feature = "mesh")] +async fn run_mesh_bridge( + scheduler: Arc, + mesh_node: Arc, + job_tx: tokio::sync::mpsc::Sender, +) { + let (mesh_tx, mut mesh_rx) = tokio::sync::mpsc::channel::(64); + + let sched = scheduler.clone(); + let sched_task = tokio::spawn(async move { + sched.run(mesh_tx).await; + }); + + loop { + // Drain local deque to dispatcher first + while let Some(job) = mesh_node.pop_local() { + if job_tx.send(job).await.is_err() { + let _ = sched_task.await; + return; + } + } + + // Wait for scheduler to produce jobs + match mesh_rx.recv().await { + Some(job) => { + let mut batch = vec![job]; + while let Ok(j) = mesh_rx.try_recv() { + batch.push(j); + } + mesh_node.prefetch(batch); + } + None => break, + } + } + + // Drain remaining deque + while let Some(job) = mesh_node.pop_local() { + if job_tx.send(job).await.is_err() { + break; + } + } + + let _ = sched_task.await; +} + /// Dispatch a ResultOutcome to Python middleware hooks and events. /// /// Called with the GIL held after `handle_result()` returns. @@ -185,6 +233,7 @@ impl PyQueue { queue_configs=None, pool=None, app_path=None, + mesh_config=None, ))] #[allow(clippy::too_many_arguments)] pub fn run_worker( @@ -202,6 +251,7 @@ impl PyQueue { queue_configs: Option, pool: Option, app_path: Option, + #[allow(unused_variables)] mesh_config: Option, ) -> PyResult<()> { // Reset shutdown flag for this run self.shutdown_flag.store(false, Ordering::SeqCst); @@ -435,6 +485,9 @@ impl PyQueue { } }; + #[cfg(feature = "mesh")] + let mesh_worker_id = worker_id.clone(); + // Captured for the channel-based (Postgres/Redis) wake-source setup // inside the runtime. Gated to the listener-bearing backends so the // default and SQLite-only builds have no unused binding. @@ -489,6 +542,29 @@ impl PyQueue { } } + // When mesh is enabled, interpose a local deque between + // scheduler and dispatcher for affinity-sorted prefetch. + #[cfg(feature = "mesh")] + let scheduler_task = { + if let Some(ref cfg_json) = mesh_config { + let mesh_cfg: taskito_mesh::MeshConfig = + serde_json::from_str(cfg_json).unwrap_or_default(); + let mesh_node = + Arc::new(taskito_mesh::MeshNode::new(mesh_worker_id, mesh_cfg)); + log::info!( + "[taskito] mesh scheduling enabled (local_buffer={})", + mesh_node.config().local_buffer_capacity, + ); + tokio::spawn(async move { + run_mesh_bridge(scheduler_for_dispatch, mesh_node, job_tx).await; + }) + } else { + tokio::spawn(async move { + scheduler_for_dispatch.run(job_tx).await; + }) + } + }; + #[cfg(not(feature = "mesh"))] let scheduler_task = tokio::spawn(async move { scheduler_for_dispatch.run(job_tx).await; }); From 592ff0a83eaa676c7f4d798b81c46630aabf49a3 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:59:51 +0530 Subject: [PATCH 03/16] feat(mesh): add MeshWorker Python class and run_worker integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MeshWorker config class in py_src/taskito/mesh.py. Queue untouched — only run_worker gains mesh= kwarg via QueueLifecycleMixin. --- py_src/taskito/__init__.py | 2 + py_src/taskito/_taskito.pyi | 1 + py_src/taskito/mesh.py | 91 ++++++++++++++++++++++++++++++ py_src/taskito/mixins/lifecycle.py | 7 +++ 4 files changed, 101 insertions(+) create mode 100644 py_src/taskito/mesh.py diff --git a/py_src/taskito/__init__.py b/py_src/taskito/__init__.py index d60276261..f9b12f5e1 100644 --- a/py_src/taskito/__init__.py +++ b/py_src/taskito/__init__.py @@ -34,6 +34,7 @@ from taskito.inject import Inject from taskito.interception import InterceptionError, InterceptionReport from taskito.log_config import configure as configure_logging +from taskito.mesh import MeshWorker from taskito.middleware import TaskMiddleware from taskito.notes import MAX_NOTE_FIELDS from taskito.proxies.no_proxy import NoProxy @@ -68,6 +69,7 @@ "JsonSerializer", "LogLevel", "MaxRetriesExceededError", + "MeshWorker", "MockResource", "MsgPackSerializer", "NoProxy", diff --git a/py_src/taskito/_taskito.pyi b/py_src/taskito/_taskito.pyi index d11ea2a55..9174a25df 100644 --- a/py_src/taskito/_taskito.pyi +++ b/py_src/taskito/_taskito.pyi @@ -188,6 +188,7 @@ class PyQueue: queue_configs: str | None = None, pool: str | None = None, app_path: str | None = None, + mesh_config: str | None = None, ) -> None: ... def worker_heartbeat( self, diff --git a/py_src/taskito/mesh.py b/py_src/taskito/mesh.py new file mode 100644 index 000000000..2f3f5cd19 --- /dev/null +++ b/py_src/taskito/mesh.py @@ -0,0 +1,91 @@ +"""Mesh scheduling configuration for decentralized task dispatch.""" + +from __future__ import annotations + +import json +from typing import Any + + +class MeshWorker: + """Configuration for mesh-enabled workers. + + Pass an instance to ``queue.run_worker(mesh=...)`` to enable + gossip-based worker discovery, consistent-hashing task affinity, + and work-stealing. + + Requires the ``mesh`` cargo feature at build time. + """ + + __slots__ = ( + "affinity_weight", + "bind_addr", + "local_buffer", + "port", + "seeds", + "steal", + "steal_batch", + "steal_threshold", + "virtual_nodes", + ) + + def __init__( + self, + *, + port: int = 7946, + seeds: list[str] | None = None, + steal: bool = True, + affinity_weight: float = 0.7, + local_buffer: int = 64, + steal_batch: int = 4, + steal_threshold: int = 2, + virtual_nodes: int = 150, + bind_addr: str = "0.0.0.0", + ) -> None: + if not 1024 <= port <= 65535: + raise ValueError(f"port must be 1024-65535, got {port}") + if not 0.0 <= affinity_weight <= 1.0: + raise ValueError(f"affinity_weight must be 0.0-1.0, got {affinity_weight}") + if local_buffer < 1: + raise ValueError(f"local_buffer must be >= 1, got {local_buffer}") + if steal_batch < 1: + raise ValueError(f"steal_batch must be >= 1, got {steal_batch}") + if virtual_nodes < 1: + raise ValueError(f"virtual_nodes must be >= 1, got {virtual_nodes}") + + self.port = port + self.seeds = seeds or [] + self.steal = steal + self.affinity_weight = affinity_weight + self.local_buffer = local_buffer + self.steal_batch = steal_batch + self.steal_threshold = steal_threshold + self.virtual_nodes = virtual_nodes + self.bind_addr = bind_addr + + def to_json(self) -> str: + """Serialize to JSON for passing through the PyO3 boundary.""" + return json.dumps(self._as_rust_config()) + + def _as_rust_config(self) -> dict[str, Any]: + return { + "gossip_port": self.port, + "steal_port": self.port + 1, + "bind_addr": self.bind_addr, + "seeds": self.seeds, + "protocol_period_ms": 500, + "indirect_ping_count": 3, + "suspicion_multiplier": 4, + "virtual_nodes": self.virtual_nodes, + "local_buffer_capacity": self.local_buffer, + "max_steal_batch": self.steal_batch, + "steal_threshold": self.steal_threshold, + "affinity_weight": self.affinity_weight, + "enable_stealing": self.steal, + } + + def __repr__(self) -> str: + return ( + f"MeshWorker(port={self.port}, seeds={self.seeds!r}, " + f"steal={self.steal}, affinity_weight={self.affinity_weight}, " + f"local_buffer={self.local_buffer})" + ) diff --git a/py_src/taskito/mixins/lifecycle.py b/py_src/taskito/mixins/lifecycle.py index 9b912e73d..a58c03bd3 100644 --- a/py_src/taskito/mixins/lifecycle.py +++ b/py_src/taskito/mixins/lifecycle.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from taskito.mesh import MeshWorker from taskito.resources.definition import ResourceDefinition @@ -108,6 +109,7 @@ def run_worker( tags: list[str] | None = None, pool: str = "thread", app: str | None = None, + mesh: MeshWorker | None = None, ) -> None: """Start the worker loop. Blocks until interrupted. @@ -120,6 +122,9 @@ def run_worker( true parallelism on CPU-bound tasks. app: Import path to the Queue instance (e.g. ``"myapp:queue"``). Required when ``pool="prefork"``. + mesh: Mesh scheduling config. Pass a ``MeshWorker`` instance to + enable gossip-based worker discovery, task affinity, and + work-stealing. Requires the ``mesh`` cargo feature. """ if pool == "prefork": if sys.platform == "win32": @@ -250,6 +255,7 @@ def sighup_handler(signum: int, frame: Any) -> None: except Exception: logger.exception("Failed to apply paused state for queue %s", queue_name) queue_configs_json = json.dumps(merged_queue_configs) if merged_queue_configs else None + mesh_config_json = mesh.to_json() if mesh is not None else None self._inner.run_worker( task_registry=self._task_registry, task_configs=self._task_configs, @@ -263,6 +269,7 @@ def sighup_handler(signum: int, frame: Any) -> None: queue_configs=queue_configs_json, pool=pool if pool != "thread" else None, app_path=app, + mesh_config=mesh_config_json, ) except KeyboardInterrupt: logger.info("Cold shutdown (terminating immediately)") From 5af76bdc5382a4c51527dbd86c1eb0ccb2accdde Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:23:28 +0530 Subject: [PATCH 04/16] feat(mesh): implement SWIM gossip protocol for worker discovery UDP-based SWIM protocol with failure detection, piggybacked state dissemination, and join/leave lifecycle. 3 integration tests verify 2-node discovery, 3-node convergence, and graceful leave. --- crates/taskito-mesh/src/lib.rs | 40 ++ crates/taskito-mesh/src/swim/failure.rs | 175 +++++++++ crates/taskito-mesh/src/swim/membership.rs | 210 +++++++++++ crates/taskito-mesh/src/swim/message.rs | 153 ++++++++ crates/taskito-mesh/src/swim/mod.rs | 355 ++++++++++++++++++ .../taskito-mesh/tests/gossip_integration.rs | 183 +++++++++ crates/taskito-python/src/py_queue/worker.rs | 16 +- 7 files changed, 1130 insertions(+), 2 deletions(-) create mode 100644 crates/taskito-mesh/src/swim/failure.rs create mode 100644 crates/taskito-mesh/src/swim/membership.rs create mode 100644 crates/taskito-mesh/src/swim/message.rs create mode 100644 crates/taskito-mesh/src/swim/mod.rs create mode 100644 crates/taskito-mesh/tests/gossip_integration.rs diff --git a/crates/taskito-mesh/src/lib.rs b/crates/taskito-mesh/src/lib.rs index 4f7d55e89..12d63b51e 100644 --- a/crates/taskito-mesh/src/lib.rs +++ b/crates/taskito-mesh/src/lib.rs @@ -3,11 +3,14 @@ pub mod local_deque; pub mod metrics; pub mod ring; pub mod state; +pub mod swim; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::atomic::Ordering; use std::sync::Arc; use taskito_core::job::Job; +use tokio::sync::Notify; pub use config::MeshConfig; pub use local_deque::LocalDeque; @@ -24,6 +27,7 @@ pub struct MeshNode { state: Arc, deque: LocalDeque, metrics: Arc, + shutdown: Arc, } impl MeshNode { @@ -35,9 +39,45 @@ impl MeshNode { state, deque, metrics: Arc::new(MeshMetrics::default()), + shutdown: Arc::new(Notify::new()), } } + /// Spawn the SWIM gossip loop as a tokio task. + /// Call this inside the tokio runtime before the scheduler loop. + pub fn spawn_gossip(&self, queues: Vec, threads: u16) -> tokio::task::JoinHandle<()> { + let gossip_addr = + SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), self.config.gossip_port); + let steal_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), self.config.steal_port); + let local_info = WorkerInfo { + worker_id: self.state.local_worker_id().to_string(), + gossip_addr, + steal_addr, + queues, + threads, + current_load: 0, + local_buffer_len: 0, + capacity: threads, + updated_at: taskito_core::job::now_millis(), + }; + + let swim = swim::SwimNode::new( + self.config.clone(), + self.state.clone(), + local_info, + self.shutdown.clone(), + ); + + tokio::spawn(async move { + swim.run().await; + }) + } + + /// Signal the gossip loop to stop (broadcasts Leave first). + pub fn request_shutdown(&self) { + self.shutdown.notify_one(); + } + pub fn config(&self) -> &MeshConfig { &self.config } diff --git a/crates/taskito-mesh/src/swim/failure.rs b/crates/taskito-mesh/src/swim/failure.rs new file mode 100644 index 000000000..f6ead281e --- /dev/null +++ b/crates/taskito-mesh/src/swim/failure.rs @@ -0,0 +1,175 @@ +use std::collections::HashMap; +use std::time::Instant; + +use super::message::MemberId; + +/// Tracks outstanding probes and suspicion timers for failure detection. +pub struct FailureDetector { + /// Pending direct pings awaiting ACK: seq → (target, sent_at). + pending_pings: HashMap, + /// Pending indirect probes: seq → (target, sent_at). + pending_ping_reqs: HashMap, + /// Members currently under suspicion: member_id → suspect_since. + suspects: HashMap, + /// Ping timeout before escalating to indirect probes. + ping_timeout_ms: u64, + /// Suspicion timeout = multiplier * log(N+1) * protocol_period. + suspicion_multiplier: u32, + protocol_period_ms: u64, +} + +impl FailureDetector { + pub fn new(ping_timeout_ms: u64, suspicion_multiplier: u32, protocol_period_ms: u64) -> Self { + Self { + pending_pings: HashMap::new(), + pending_ping_reqs: HashMap::new(), + suspects: HashMap::new(), + ping_timeout_ms, + suspicion_multiplier, + protocol_period_ms, + } + } + + /// Record that a direct ping was sent. + pub fn ping_sent(&mut self, seq: u64, target: MemberId) { + self.pending_pings.insert(seq, (target, Instant::now())); + } + + /// Record that indirect probes were sent. + pub fn ping_req_sent(&mut self, seq: u64, target: MemberId) { + self.pending_ping_reqs.insert(seq, (target, Instant::now())); + } + + /// Process an ACK. Returns the target member ID if this resolves a pending probe. + pub fn ack_received(&mut self, seq: u64) -> Option { + if let Some((target, _)) = self.pending_pings.remove(&seq) { + self.suspects.remove(&target); + return Some(target); + } + if let Some((target, _)) = self.pending_ping_reqs.remove(&seq) { + self.suspects.remove(&target); + return Some(target); + } + None + } + + /// Check for timed-out direct pings. Returns members that need indirect probing. + pub fn check_ping_timeouts(&mut self) -> Vec { + let timeout = std::time::Duration::from_millis(self.ping_timeout_ms); + let now = Instant::now(); + let mut timed_out = Vec::new(); + + self.pending_pings.retain(|_, (target, sent_at)| { + if now.duration_since(*sent_at) > timeout { + timed_out.push(target.clone()); + false + } else { + true + } + }); + + timed_out + } + + /// Mark a member as suspect. Returns true if newly suspected. + pub fn suspect(&mut self, member_id: &str) -> bool { + if self.suspects.contains_key(member_id) { + return false; + } + self.suspects.insert(member_id.to_string(), Instant::now()); + true + } + + /// Check for expired suspicions. Returns members that should be declared dead. + pub fn check_suspicion_timeouts(&mut self, member_count: usize) -> Vec { + let timeout = self.suspicion_timeout(member_count); + let now = Instant::now(); + let mut dead = Vec::new(); + + self.suspects.retain(|member_id, suspect_since| { + if now.duration_since(*suspect_since) > timeout { + dead.push(member_id.clone()); + false + } else { + true + } + }); + + dead + } + + /// Suspicion timeout scales with log(N+1) for consistency guarantees. + pub fn suspicion_timeout(&self, member_count: usize) -> std::time::Duration { + let log_n = ((member_count + 1) as f64).ln().max(1.0); + let ms = self.suspicion_multiplier as f64 * log_n * self.protocol_period_ms as f64; + std::time::Duration::from_millis(ms as u64) + } + + /// Clear a suspect (e.g., on receiving a refutation with higher incarnation). + pub fn clear_suspect(&mut self, member_id: &str) { + self.suspects.remove(member_id); + } + + /// Number of currently suspected members. + pub fn suspect_count(&self) -> usize { + self.suspects.len() + } + + /// Clean up stale pending probes older than 2x protocol period. + pub fn gc_stale_probes(&mut self) { + let stale = std::time::Duration::from_millis(self.protocol_period_ms * 2); + let now = Instant::now(); + self.pending_pings + .retain(|_, (_, sent)| now.duration_since(*sent) < stale); + self.pending_ping_reqs + .retain(|_, (_, sent)| now.duration_since(*sent) < stale); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ping_ack_resolves() { + let mut fd = FailureDetector::new(200, 4, 500); + fd.ping_sent(1, "peer-a".to_string()); + let resolved = fd.ack_received(1); + assert_eq!(resolved, Some("peer-a".to_string())); + } + + #[test] + fn unknown_ack_ignored() { + let mut fd = FailureDetector::new(200, 4, 500); + assert!(fd.ack_received(999).is_none()); + } + + #[test] + fn suspicion_timeout_scales_with_members() { + let fd = FailureDetector::new(200, 4, 500); + let t1 = fd.suspicion_timeout(1); + let t10 = fd.suspicion_timeout(10); + let t100 = fd.suspicion_timeout(100); + assert!(t10 > t1, "more members = longer suspicion window"); + assert!(t100 > t10); + } + + #[test] + fn suspect_then_clear() { + let mut fd = FailureDetector::new(200, 4, 500); + assert!(fd.suspect("peer-a")); + assert!(!fd.suspect("peer-a")); // already suspected + assert_eq!(fd.suspect_count(), 1); + fd.clear_suspect("peer-a"); + assert_eq!(fd.suspect_count(), 0); + } + + #[test] + fn ack_clears_suspicion() { + let mut fd = FailureDetector::new(200, 4, 500); + fd.ping_sent(1, "peer-a".to_string()); + fd.suspect("peer-a"); + fd.ack_received(1); + assert_eq!(fd.suspect_count(), 0); + } +} diff --git a/crates/taskito-mesh/src/swim/membership.rs b/crates/taskito-mesh/src/swim/membership.rs new file mode 100644 index 000000000..436d8a917 --- /dev/null +++ b/crates/taskito-mesh/src/swim/membership.rs @@ -0,0 +1,210 @@ +use std::collections::HashMap; + +use crate::state::{Member, MemberState}; + +use super::message::MemberUpdate; + +/// Manages SWIM membership: incarnation numbers, state transitions, +/// and pending update dissemination. +pub struct Membership { + local_id: String, + local_incarnation: u64, + /// Updates waiting to be piggybacked on outgoing messages. + /// Bounded to MAX_PENDING to prevent unbounded growth. + pending_updates: Vec, +} + +const MAX_PENDING: usize = 64; +const PIGGYBACK_LIMIT: usize = 8; + +impl Membership { + pub fn new(local_id: String) -> Self { + Self { + local_id, + local_incarnation: 1, + pending_updates: Vec::new(), + } + } + + pub fn local_id(&self) -> &str { + &self.local_id + } + + pub fn local_incarnation(&self) -> u64 { + self.local_incarnation + } + + /// Increment local incarnation (used to refute suspicion). + pub fn refute(&mut self) -> u64 { + self.local_incarnation += 1; + self.local_incarnation + } + + /// Process an incoming member update. Returns true if state changed. + pub fn apply_update( + &mut self, + update: &MemberUpdate, + members: &mut HashMap, + ) -> bool { + if update.member_id == self.local_id { + return self.handle_self_update(update); + } + + let changed = match members.get(&update.member_id) { + Some(existing) => { + update.incarnation > existing.incarnation + || (update.incarnation == existing.incarnation + && state_priority(update.state) > state_priority(existing.state)) + } + None => true, + }; + + if changed { + members.insert( + update.member_id.clone(), + Member { + info: update.info.clone(), + state: update.state, + incarnation: update.incarnation, + }, + ); + } + + changed + } + + /// Handle updates about ourselves — refute if suspected. + fn handle_self_update(&mut self, update: &MemberUpdate) -> bool { + if matches!(update.state, MemberState::Suspect | MemberState::Dead) + && update.incarnation >= self.local_incarnation + { + self.refute(); + true + } else { + false + } + } + + /// Queue an update for piggybacking on outgoing messages. + pub fn queue_update(&mut self, update: MemberUpdate) { + if self.pending_updates.len() >= MAX_PENDING { + self.pending_updates.remove(0); + } + self.pending_updates.push(update); + } + + /// Take up to PIGGYBACK_LIMIT updates for piggybacking. + /// Consumed updates are removed from the pending queue. + pub fn take_updates(&mut self) -> Vec { + let n = self.pending_updates.len().min(PIGGYBACK_LIMIT); + self.pending_updates.drain(..n).collect() + } + + /// Check if there are pending updates to disseminate. + pub fn has_pending(&self) -> bool { + !self.pending_updates.is_empty() + } +} + +/// Higher priority wins in state conflict resolution. +fn state_priority(state: MemberState) -> u8 { + match state { + MemberState::Alive => 0, + MemberState::Suspect => 1, + MemberState::Dead => 2, + MemberState::Left => 3, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::WorkerInfo; + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + fn test_info(id: &str) -> WorkerInfo { + WorkerInfo { + worker_id: id.to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7946), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7947), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: 0, + capacity: 4, + updated_at: 0, + } + } + + fn make_update(id: &str, state: MemberState, incarnation: u64) -> MemberUpdate { + MemberUpdate { + member_id: id.to_string(), + state, + incarnation, + info: test_info(id), + } + } + + #[test] + fn apply_new_member() { + let mut membership = Membership::new("local".to_string()); + let mut members = HashMap::new(); + let update = make_update("peer-a", MemberState::Alive, 1); + assert!(membership.apply_update(&update, &mut members)); + assert_eq!(members.len(), 1); + } + + #[test] + fn higher_incarnation_wins() { + let mut membership = Membership::new("local".to_string()); + let mut members = HashMap::new(); + membership.apply_update(&make_update("peer-a", MemberState::Alive, 1), &mut members); + assert!( + membership.apply_update(&make_update("peer-a", MemberState::Alive, 2), &mut members) + ); + assert_eq!(members["peer-a"].incarnation, 2); + } + + #[test] + fn same_incarnation_higher_state_wins() { + let mut membership = Membership::new("local".to_string()); + let mut members = HashMap::new(); + membership.apply_update(&make_update("peer-a", MemberState::Alive, 1), &mut members); + assert!(membership.apply_update( + &make_update("peer-a", MemberState::Suspect, 1), + &mut members + )); + assert_eq!(members["peer-a"].state, MemberState::Suspect); + } + + #[test] + fn lower_incarnation_ignored() { + let mut membership = Membership::new("local".to_string()); + let mut members = HashMap::new(); + membership.apply_update(&make_update("peer-a", MemberState::Alive, 5), &mut members); + assert!( + !membership.apply_update(&make_update("peer-a", MemberState::Dead, 3), &mut members) + ); + assert_eq!(members["peer-a"].incarnation, 5); + } + + #[test] + fn self_suspicion_triggers_refute() { + let mut membership = Membership::new("local".to_string()); + let mut members = HashMap::new(); + let update = make_update("local", MemberState::Suspect, 1); + assert!(membership.apply_update(&update, &mut members)); + assert_eq!(membership.local_incarnation(), 2); + } + + #[test] + fn piggyback_queue() { + let mut membership = Membership::new("local".to_string()); + for i in 0..10 { + membership.queue_update(make_update(&format!("w{i}"), MemberState::Alive, 1)); + } + let batch = membership.take_updates(); + assert_eq!(batch.len(), 8); // PIGGYBACK_LIMIT + assert_eq!(membership.pending_updates.len(), 2); + } +} diff --git a/crates/taskito-mesh/src/swim/message.rs b/crates/taskito-mesh/src/swim/message.rs new file mode 100644 index 000000000..32a082541 --- /dev/null +++ b/crates/taskito-mesh/src/swim/message.rs @@ -0,0 +1,153 @@ +use std::net::SocketAddr; + +use serde::{Deserialize, Serialize}; + +use crate::state::{MemberState, WorkerInfo}; + +pub type MemberId = String; + +/// SWIM protocol message, serialized via bincode over UDP. +/// Must fit in a single UDP datagram (< 1400 bytes). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum GossipMessage { + Ping { + seq: u64, + from: MemberId, + from_addr: SocketAddr, + }, + Ack { + seq: u64, + from: MemberId, + }, + PingReq { + seq: u64, + from: MemberId, + target: MemberId, + target_addr: SocketAddr, + }, + AckRelay { + seq: u64, + original_from: MemberId, + via: MemberId, + }, + /// Membership updates piggybacked on any message. + Sync { + updates: Vec, + }, + /// Compound: a primary message + piggybacked sync updates. + Compound { + primary: Box, + updates: Vec, + }, +} + +/// A single membership state change to disseminate. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemberUpdate { + pub member_id: MemberId, + pub state: MemberState, + pub incarnation: u64, + pub info: WorkerInfo, +} + +impl GossipMessage { + pub fn encode(&self) -> Result, bincode::Error> { + bincode::serialize(self) + } + + pub fn decode(data: &[u8]) -> Result { + bincode::deserialize(data) + } + + /// Wrap this message with piggybacked membership updates. + pub fn with_updates(self, updates: Vec) -> Self { + if updates.is_empty() { + return self; + } + GossipMessage::Compound { + primary: Box::new(self), + updates, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr}; + + fn test_addr() -> SocketAddr { + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7946) + } + + fn test_info() -> WorkerInfo { + WorkerInfo { + worker_id: "w1".to_string(), + gossip_addr: test_addr(), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7947), + queues: vec!["default".to_string()], + threads: 4, + current_load: 2, + local_buffer_len: 5, + capacity: 2, + updated_at: 1000, + } + } + + #[test] + fn ping_round_trip() { + let msg = GossipMessage::Ping { + seq: 42, + from: "w1".to_string(), + from_addr: test_addr(), + }; + let bytes = msg.encode().unwrap(); + let decoded = GossipMessage::decode(&bytes).unwrap(); + match decoded { + GossipMessage::Ping { seq, from, .. } => { + assert_eq!(seq, 42); + assert_eq!(from, "w1"); + } + _ => panic!("wrong variant"), + } + } + + #[test] + fn compound_round_trip() { + let ping = GossipMessage::Ping { + seq: 1, + from: "w1".to_string(), + from_addr: test_addr(), + }; + let updates = vec![MemberUpdate { + member_id: "w2".to_string(), + state: MemberState::Alive, + incarnation: 3, + info: test_info(), + }]; + let compound = ping.with_updates(updates); + let bytes = compound.encode().unwrap(); + assert!(bytes.len() < 1400, "must fit in UDP datagram"); + + let decoded = GossipMessage::decode(&bytes).unwrap(); + match decoded { + GossipMessage::Compound { primary, updates } => { + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].member_id, "w2"); + matches!(*primary, GossipMessage::Ping { .. }); + } + _ => panic!("wrong variant"), + } + } + + #[test] + fn empty_updates_no_wrap() { + let ping = GossipMessage::Ping { + seq: 1, + from: "w1".to_string(), + from_addr: test_addr(), + }; + let result = ping.with_updates(vec![]); + matches!(result, GossipMessage::Ping { .. }); + } +} diff --git a/crates/taskito-mesh/src/swim/mod.rs b/crates/taskito-mesh/src/swim/mod.rs new file mode 100644 index 000000000..523b2b7a8 --- /dev/null +++ b/crates/taskito-mesh/src/swim/mod.rs @@ -0,0 +1,355 @@ +pub mod failure; +pub mod membership; +pub mod message; + +use std::net::SocketAddr; +use std::sync::Arc; + +use log::{debug, info, warn}; +use rand::prelude::IndexedRandom; +use tokio::net::UdpSocket; +use tokio::sync::Notify; + +use crate::config::MeshConfig; +use crate::state::{Member, MemberState, MeshState, WorkerInfo}; + +use self::failure::FailureDetector; +use self::membership::Membership; +use self::message::{GossipMessage, MemberUpdate}; + +/// SWIM protocol node. Runs a gossip loop on a tokio UDP socket. +pub struct SwimNode { + config: MeshConfig, + state: Arc, + membership: Membership, + failure_detector: FailureDetector, + local_info: WorkerInfo, + seq: u64, + shutdown: Arc, +} + +impl SwimNode { + pub fn new( + config: MeshConfig, + state: Arc, + local_info: WorkerInfo, + shutdown: Arc, + ) -> Self { + let membership = Membership::new(local_info.worker_id.clone()); + let ping_timeout_ms = config.protocol_period_ms / 2; + let failure_detector = FailureDetector::new( + ping_timeout_ms, + config.suspicion_multiplier, + config.protocol_period_ms, + ); + Self { + config, + state, + membership, + failure_detector, + local_info, + seq: 0, + shutdown, + } + } + + fn next_seq(&mut self) -> u64 { + self.seq += 1; + self.seq + } + + /// Run the SWIM gossip loop. Blocks until shutdown. + pub async fn run(mut self) { + let bind = format!("{}:{}", self.config.bind_addr, self.config.gossip_port); + let socket = match UdpSocket::bind(&bind).await { + Ok(s) => Arc::new(s), + Err(e) => { + warn!("[mesh] failed to bind gossip socket {bind}: {e}"); + return; + } + }; + info!( + "[mesh] gossip listening on {} (worker={})", + bind, self.local_info.worker_id + ); + + self.join_seeds(&socket).await; + + let period = std::time::Duration::from_millis(self.config.protocol_period_ms); + let mut recv_buf = vec![0u8; 2048]; + + loop { + tokio::select! { + _ = self.shutdown.notified() => { + self.broadcast_leave(&socket).await; + break; + } + _ = tokio::time::sleep(period) => { + self.protocol_tick(&socket).await; + } + result = socket.recv_from(&mut recv_buf) => { + match result { + Ok((len, from)) => { + self.handle_datagram(&recv_buf[..len], from, &socket).await; + } + Err(e) => { + debug!("[mesh] recv error: {e}"); + } + } + } + } + } + + info!("[mesh] gossip loop stopped"); + } + + /// Send join pings to seed nodes. + async fn join_seeds(&mut self, socket: &UdpSocket) { + for seed in &self.config.seeds.clone() { + let seq = self.next_seq(); + let ping = GossipMessage::Ping { + seq, + from: self.local_info.worker_id.clone(), + from_addr: self.local_info.gossip_addr, + }; + let local_update = self.make_local_update(); + let msg = ping.with_updates(vec![local_update]); + if let Ok(bytes) = msg.encode() { + if let Ok(addr) = seed.parse::() { + let _ = socket.send_to(&bytes, addr).await; + debug!("[mesh] join ping sent to {seed}"); + } + } + } + } + + /// One SWIM protocol period: ping a random peer, check timeouts. + async fn protocol_tick(&mut self, socket: &UdpSocket) { + self.failure_detector.gc_stale_probes(); + + let timed_out = self.failure_detector.check_ping_timeouts(); + for target_id in timed_out { + self.initiate_indirect_probe(&target_id, socket).await; + } + + let member_count = self.state.alive_count() + 1; + let newly_dead = self.failure_detector.check_suspicion_timeouts(member_count); + for dead_id in newly_dead { + info!("[mesh] member {dead_id} declared dead (suspicion expired)"); + self.state.mark_dead(&dead_id); + self.membership.queue_update(MemberUpdate { + member_id: dead_id, + state: MemberState::Dead, + incarnation: 0, + info: self.local_info.clone(), + }); + } + + if let Some(target) = self.pick_random_alive_peer() { + let seq = self.next_seq(); + let ping = GossipMessage::Ping { + seq, + from: self.local_info.worker_id.clone(), + from_addr: self.local_info.gossip_addr, + }; + let updates = self.membership.take_updates(); + let msg = ping.with_updates(updates); + if let Ok(bytes) = msg.encode() { + let _ = socket.send_to(&bytes, target.info.gossip_addr).await; + self.failure_detector + .ping_sent(seq, target.info.worker_id.clone()); + } + } + } + + /// Send PingReq to random peers asking them to probe the target. + async fn initiate_indirect_probe(&mut self, target_id: &str, socket: &UdpSocket) { + let peers = self.state.alive_peers(); + let intermediaries: Vec<&Member> = peers + .iter() + .filter(|m| m.info.worker_id != target_id) + .take(self.config.indirect_ping_count) + .collect(); + + if intermediaries.is_empty() { + if self.failure_detector.suspect(target_id) { + info!("[mesh] member {target_id} suspected (no intermediaries)"); + self.membership.queue_update(MemberUpdate { + member_id: target_id.to_string(), + state: MemberState::Suspect, + incarnation: 0, + info: self.local_info.clone(), + }); + } + return; + } + + // Find target addr from state + let target_addr = peers + .iter() + .find(|m| m.info.worker_id == target_id) + .map(|m| m.info.gossip_addr); + + if let Some(addr) = target_addr { + let seq = self.next_seq(); + for intermediary in intermediaries { + let ping_req = GossipMessage::PingReq { + seq, + from: self.local_info.worker_id.clone(), + target: target_id.to_string(), + target_addr: addr, + }; + if let Ok(bytes) = ping_req.encode() { + let _ = socket.send_to(&bytes, intermediary.info.gossip_addr).await; + } + } + self.failure_detector + .ping_req_sent(seq, target_id.to_string()); + } else if self.failure_detector.suspect(target_id) { + info!("[mesh] member {target_id} suspected (no addr found)"); + } + } + + /// Handle an incoming UDP datagram. + async fn handle_datagram(&mut self, data: &[u8], from: SocketAddr, socket: &UdpSocket) { + let msg = match GossipMessage::decode(data) { + Ok(m) => m, + Err(e) => { + debug!("[mesh] decode error from {from}: {e}"); + return; + } + }; + + match msg { + GossipMessage::Compound { primary, updates } => { + self.apply_updates(&updates); + self.handle_primary(*primary, from, socket).await; + } + GossipMessage::Sync { updates } => { + self.apply_updates(&updates); + } + other => { + self.handle_primary(other, from, socket).await; + } + } + } + + async fn handle_primary(&mut self, msg: GossipMessage, from: SocketAddr, socket: &UdpSocket) { + match msg { + GossipMessage::Ping { + seq, from: sender, .. + } => { + let ack = GossipMessage::Ack { + seq, + from: self.local_info.worker_id.clone(), + }; + let local_update = self.make_local_update(); + let mut all_updates = vec![local_update]; + // Include all known alive peers so the pinger discovers them + for peer in self.state.alive_peers() { + all_updates.push(MemberUpdate { + member_id: peer.info.worker_id.clone(), + state: peer.state, + incarnation: peer.incarnation, + info: peer.info, + }); + } + all_updates.extend(self.membership.take_updates()); + let msg = ack.with_updates(all_updates); + if let Ok(bytes) = msg.encode() { + let _ = socket.send_to(&bytes, from).await; + } + debug!("[mesh] ack sent to {sender} at {from}"); + } + GossipMessage::Ack { seq, from: sender } => { + if let Some(resolved) = self.failure_detector.ack_received(seq) { + debug!("[mesh] ack from {sender} resolved probe for {resolved}"); + } + } + GossipMessage::PingReq { + seq, + from: requester, + target, + target_addr, + } => { + let ping = GossipMessage::Ping { + seq, + from: self.local_info.worker_id.clone(), + from_addr: self.local_info.gossip_addr, + }; + if let Ok(bytes) = ping.encode() { + let _ = socket.send_to(&bytes, target_addr).await; + } + debug!("[mesh] relayed ping-req from {requester} to {target}"); + } + GossipMessage::AckRelay { + seq, original_from, .. + } => { + if let Some(resolved) = self.failure_detector.ack_received(seq) { + debug!("[mesh] relay-ack from {original_from} resolved {resolved}"); + } + } + _ => {} + } + } + + fn apply_updates(&mut self, updates: &[MemberUpdate]) { + for update in updates { + if update.member_id == self.local_info.worker_id { + continue; + } + let is_new = self.state.upsert_member(Member { + info: update.info.clone(), + state: update.state, + incarnation: update.incarnation, + }); + if is_new { + info!( + "[mesh] discovered peer {} at {}", + update.member_id, update.info.gossip_addr + ); + self.membership.queue_update(update.clone()); + } + if update.state == MemberState::Alive { + self.failure_detector.clear_suspect(&update.member_id); + } + } + } + + fn pick_random_alive_peer(&self) -> Option { + let peers = self.state.alive_peers(); + if peers.is_empty() { + return None; + } + let mut rng = rand::rng(); + peers.choose(&mut rng).cloned() + } + + fn make_local_update(&self) -> MemberUpdate { + MemberUpdate { + member_id: self.local_info.worker_id.clone(), + state: MemberState::Alive, + incarnation: self.membership.local_incarnation(), + info: self.local_info.clone(), + } + } + + /// Broadcast a Leave message to all known peers. + async fn broadcast_leave(&mut self, socket: &UdpSocket) { + let leave = MemberUpdate { + member_id: self.local_info.worker_id.clone(), + state: MemberState::Left, + incarnation: self.membership.local_incarnation() + 1, + info: self.local_info.clone(), + }; + let msg = GossipMessage::Sync { + updates: vec![leave], + }; + if let Ok(bytes) = msg.encode() { + for peer in self.state.alive_peers() { + let _ = socket.send_to(&bytes, peer.info.gossip_addr).await; + } + } + info!("[mesh] leave broadcast sent"); + } +} diff --git a/crates/taskito-mesh/tests/gossip_integration.rs b/crates/taskito-mesh/tests/gossip_integration.rs new file mode 100644 index 000000000..29d23591f --- /dev/null +++ b/crates/taskito-mesh/tests/gossip_integration.rs @@ -0,0 +1,183 @@ +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::time::Duration; + +use taskito_mesh::config::MeshConfig; +use taskito_mesh::state::{MeshState, WorkerInfo}; +use taskito_mesh::swim::SwimNode; +use tokio::sync::Notify; + +fn make_config(port: u16, seeds: Vec) -> MeshConfig { + MeshConfig { + gossip_port: port, + steal_port: port + 100, + bind_addr: "127.0.0.1".to_string(), + seeds, + protocol_period_ms: 100, + indirect_ping_count: 2, + suspicion_multiplier: 2, + virtual_nodes: 10, + local_buffer_capacity: 16, + max_steal_batch: 4, + steal_threshold: 2, + affinity_weight: 0.7, + enable_stealing: false, + } +} + +fn make_info(id: &str, port: u16) -> WorkerInfo { + WorkerInfo { + worker_id: id.to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port + 100), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: 0, + capacity: 4, + updated_at: 0, + } +} + +#[tokio::test] +async fn two_nodes_discover_each_other() { + let port_a = 19100; + let port_b = 19101; + + let state_a = Arc::new(MeshState::new("node-a".to_string(), 10)); + let state_b = Arc::new(MeshState::new("node-b".to_string(), 10)); + + let shutdown_a = Arc::new(Notify::new()); + let shutdown_b = Arc::new(Notify::new()); + + let config_a = make_config(port_a, vec![]); + let config_b = make_config(port_b, vec![format!("127.0.0.1:{port_a}")]); + + let swim_a = SwimNode::new( + config_a, + state_a.clone(), + make_info("node-a", port_a), + shutdown_a.clone(), + ); + let swim_b = SwimNode::new( + config_b, + state_b.clone(), + make_info("node-b", port_b), + shutdown_b.clone(), + ); + + let ha = tokio::spawn(async move { swim_a.run().await }); + let hb = tokio::spawn(async move { swim_b.run().await }); + + // Wait for convergence (2-3 protocol periods) + tokio::time::sleep(Duration::from_millis(500)).await; + + assert_eq!(state_a.alive_count(), 1, "node-a should see node-b"); + assert_eq!(state_b.alive_count(), 1, "node-b should see node-a"); + + shutdown_a.notify_one(); + shutdown_b.notify_one(); + let _ = tokio::join!(ha, hb); +} + +#[tokio::test] +async fn three_nodes_converge_via_piggyback() { + let port_a = 19200; + let port_b = 19201; + let port_c = 19202; + + let state_a = Arc::new(MeshState::new("node-a".to_string(), 10)); + let state_b = Arc::new(MeshState::new("node-b".to_string(), 10)); + let state_c = Arc::new(MeshState::new("node-c".to_string(), 10)); + + let shutdown_a = Arc::new(Notify::new()); + let shutdown_b = Arc::new(Notify::new()); + let shutdown_c = Arc::new(Notify::new()); + + // b seeds from a, c seeds from a — c discovers b via piggybacked updates + let config_a = make_config(port_a, vec![]); + let config_b = make_config(port_b, vec![format!("127.0.0.1:{port_a}")]); + let config_c = make_config(port_c, vec![format!("127.0.0.1:{port_a}")]); + + let swim_a = SwimNode::new( + config_a, + state_a.clone(), + make_info("node-a", port_a), + shutdown_a.clone(), + ); + let swim_b = SwimNode::new( + config_b, + state_b.clone(), + make_info("node-b", port_b), + shutdown_b.clone(), + ); + let swim_c = SwimNode::new( + config_c, + state_c.clone(), + make_info("node-c", port_c), + shutdown_c.clone(), + ); + + let ha = tokio::spawn(async move { swim_a.run().await }); + let hb = tokio::spawn(async move { swim_b.run().await }); + let hc = tokio::spawn(async move { swim_c.run().await }); + + // Piggybacked dissemination takes multiple rounds: b→a→c and c→a→b + tokio::time::sleep(Duration::from_millis(1500)).await; + + assert_eq!(state_a.alive_count(), 2, "node-a should see b and c"); + assert_eq!(state_b.alive_count(), 2, "node-b should see a and c"); + assert_eq!(state_c.alive_count(), 2, "node-c should see a and b"); + + shutdown_a.notify_one(); + shutdown_b.notify_one(); + shutdown_c.notify_one(); + let _ = tokio::join!(ha, hb, hc); +} + +#[tokio::test] +async fn graceful_leave_removes_from_peers() { + let port_a = 19300; + let port_b = 19301; + + let state_a = Arc::new(MeshState::new("node-a".to_string(), 10)); + let state_b = Arc::new(MeshState::new("node-b".to_string(), 10)); + + let shutdown_a = Arc::new(Notify::new()); + let shutdown_b = Arc::new(Notify::new()); + + let config_a = make_config(port_a, vec![]); + let config_b = make_config(port_b, vec![format!("127.0.0.1:{port_a}")]); + + let swim_a = SwimNode::new( + config_a, + state_a.clone(), + make_info("node-a", port_a), + shutdown_a.clone(), + ); + let swim_b = SwimNode::new( + config_b, + state_b.clone(), + make_info("node-b", port_b), + shutdown_b.clone(), + ); + + let ha = tokio::spawn(async move { swim_a.run().await }); + let hb = tokio::spawn(async move { swim_b.run().await }); + + // Wait for discovery + tokio::time::sleep(Duration::from_millis(400)).await; + assert_eq!(state_a.alive_count(), 1); + assert_eq!(state_b.alive_count(), 1); + + // Node B leaves gracefully + shutdown_b.notify_one(); + let _ = hb.await; + + // Give node A time to process the leave broadcast + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(state_a.alive_count(), 0, "node-a should see node-b as left"); + + shutdown_a.notify_one(); + let _ = ha.await; +} diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index e3cab375e..166d3a236 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -543,7 +543,8 @@ impl PyQueue { } // When mesh is enabled, interpose a local deque between - // scheduler and dispatcher for affinity-sorted prefetch. + // scheduler and dispatcher for affinity-sorted prefetch, + // and spawn the SWIM gossip loop for peer discovery. #[cfg(feature = "mesh")] let scheduler_task = { if let Some(ref cfg_json) = mesh_config { @@ -555,8 +556,19 @@ impl PyQueue { "[taskito] mesh scheduling enabled (local_buffer={})", mesh_node.config().local_buffer_capacity, ); + + let gossip_queues: Vec = + queues_str.split(',').map(|s| s.to_string()).collect(); + let gossip_handle = + mesh_node.spawn_gossip(gossip_queues, num_workers as u16); + + let mesh_for_bridge = mesh_node.clone(); + let bridge_handle = tokio::spawn(async move { + run_mesh_bridge(scheduler_for_dispatch, mesh_for_bridge, job_tx).await; + }); + tokio::spawn(async move { - run_mesh_bridge(scheduler_for_dispatch, mesh_node, job_tx).await; + let _ = tokio::join!(bridge_handle, gossip_handle); }) } else { tokio::spawn(async move { From ecb2161cc115c6f31e9ed15a6afd82ab62ac3222 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:28:01 +0530 Subject: [PATCH 05/16] feat(mesh): add TCP work-stealing protocol Steal server accepts requests from peers, transfers jobs from local deque cold end. Steal client connects to busiest gossip-discovered peer. Bridge attempts steal on idle ticks. 2 integration tests. --- crates/taskito-mesh/src/lib.rs | 41 ++++- crates/taskito-mesh/src/steal/mod.rs | 86 ++++++++++ crates/taskito-mesh/src/steal/protocol.rs | 80 +++++++++ crates/taskito-mesh/src/steal/server.rs | 71 ++++++++ .../taskito-mesh/tests/steal_integration.rs | 156 ++++++++++++++++++ crates/taskito-python/src/py_queue/worker.rs | 15 +- 6 files changed, 447 insertions(+), 2 deletions(-) create mode 100644 crates/taskito-mesh/src/steal/mod.rs create mode 100644 crates/taskito-mesh/src/steal/protocol.rs create mode 100644 crates/taskito-mesh/src/steal/server.rs create mode 100644 crates/taskito-mesh/tests/steal_integration.rs diff --git a/crates/taskito-mesh/src/lib.rs b/crates/taskito-mesh/src/lib.rs index 12d63b51e..0f0e363ee 100644 --- a/crates/taskito-mesh/src/lib.rs +++ b/crates/taskito-mesh/src/lib.rs @@ -3,6 +3,7 @@ pub mod local_deque; pub mod metrics; pub mod ring; pub mod state; +pub mod steal; pub mod swim; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; @@ -73,7 +74,45 @@ impl MeshNode { }) } - /// Signal the gossip loop to stop (broadcasts Leave first). + /// Spawn the TCP steal server as a tokio task. + pub fn spawn_steal_server(self: &Arc) -> tokio::task::JoinHandle<()> { + let node = self.clone(); + let shutdown = self.shutdown.clone(); + tokio::spawn(async move { + steal::server::run_steal_server(node, shutdown).await; + }) + } + + /// Attempt to steal work from the busiest peer. + /// Returns stolen jobs pushed into local deque. + pub async fn try_steal(self: &Arc) -> usize { + if !self.should_steal() { + return 0; + } + let min_surplus = self.config.steal_threshold + self.config.max_steal_batch; + let target = match self.state.best_steal_target(min_surplus) { + Some(t) => t, + None => return 0, + }; + + self.metrics + .steals_initiated + .fetch_add(1, Ordering::Relaxed); + let stolen = steal::steal_from_peer(self, &target).await; + let count = stolen.len(); + if count > 0 { + self.metrics + .steals_succeeded + .fetch_add(1, Ordering::Relaxed); + self.metrics + .jobs_stolen_in + .fetch_add(count as u64, Ordering::Relaxed); + self.prefetch(stolen); + } + count + } + + /// Signal the gossip loop and steal server to stop. pub fn request_shutdown(&self) { self.shutdown.notify_one(); } diff --git a/crates/taskito-mesh/src/steal/mod.rs b/crates/taskito-mesh/src/steal/mod.rs new file mode 100644 index 000000000..c76d16773 --- /dev/null +++ b/crates/taskito-mesh/src/steal/mod.rs @@ -0,0 +1,86 @@ +pub mod protocol; +pub mod server; + +use std::sync::Arc; +use std::time::Duration; + +use log::{debug, warn}; +use tokio::net::TcpStream; + +use crate::state::Member; +use crate::MeshNode; + +use self::protocol::{read_frame, write_frame, StealRequest, StealResponse}; + +/// Attempt to steal jobs from a peer worker over TCP. +/// Returns stolen jobs on success, empty vec on failure. +pub async fn steal_from_peer( + mesh_node: &Arc, + target: &Member, +) -> Vec { + let addr = target.info.steal_addr; + let max_count = mesh_node.config().max_steal_batch; + let thief_id = mesh_node.state().local_worker_id().to_string(); + + let stream = + match tokio::time::timeout(Duration::from_millis(500), TcpStream::connect(addr)).await { + Ok(Ok(s)) => s, + Ok(Err(e)) => { + debug!("[mesh] steal connect to {addr} failed: {e}"); + return vec![]; + } + Err(_) => { + debug!("[mesh] steal connect to {addr} timed out"); + return vec![]; + } + }; + + let (mut reader, mut writer) = stream.into_split(); + + let req = StealRequest { + thief_id, + max_count, + }; + let req_bytes = match bincode::serialize(&req) { + Ok(b) => b, + Err(e) => { + warn!("[mesh] steal serialize error: {e}"); + return vec![]; + } + }; + + if let Err(e) = write_frame(&mut writer, &req_bytes).await { + debug!("[mesh] steal write error: {e}"); + return vec![]; + } + + let resp_frame = + match tokio::time::timeout(Duration::from_secs(2), read_frame(&mut reader)).await { + Ok(Ok(f)) => f, + Ok(Err(e)) => { + debug!("[mesh] steal read error: {e}"); + return vec![]; + } + Err(_) => { + debug!("[mesh] steal response timed out"); + return vec![]; + } + }; + + match bincode::deserialize::(&resp_frame) { + Ok(resp) => { + if !resp.jobs.is_empty() { + debug!( + "[mesh] stole {} jobs from {}", + resp.jobs.len(), + target.info.worker_id + ); + } + resp.jobs + } + Err(e) => { + warn!("[mesh] steal deserialize error: {e}"); + vec![] + } + } +} diff --git a/crates/taskito-mesh/src/steal/protocol.rs b/crates/taskito-mesh/src/steal/protocol.rs new file mode 100644 index 000000000..77d7814c1 --- /dev/null +++ b/crates/taskito-mesh/src/steal/protocol.rs @@ -0,0 +1,80 @@ +use serde::{Deserialize, Serialize}; +use taskito_core::job::Job; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct StealRequest { + pub thief_id: String, + pub max_count: usize, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct StealResponse { + pub jobs: Vec, +} + +const MAX_FRAME_SIZE: usize = 1_048_576; + +/// Write a length-prefixed bincode frame. +pub async fn write_frame( + writer: &mut W, + data: &[u8], +) -> std::io::Result<()> { + let len = (data.len() as u32).to_be_bytes(); + writer.write_all(&len).await?; + writer.write_all(data).await?; + writer.flush().await +} + +/// Read a length-prefixed bincode frame. +pub async fn read_frame(reader: &mut R) -> std::io::Result> { + let mut len_buf = [0u8; 4]; + reader.read_exact(&mut len_buf).await?; + let len = u32::from_be_bytes(len_buf) as usize; + if len > MAX_FRAME_SIZE { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "frame too large", + )); + } + let mut buf = vec![0u8; len]; + reader.read_exact(&mut buf).await?; + Ok(buf) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn request_round_trip() { + let req = StealRequest { + thief_id: "w1".to_string(), + max_count: 4, + }; + let bytes = bincode::serialize(&req).unwrap(); + let decoded: StealRequest = bincode::deserialize(&bytes).unwrap(); + assert_eq!(decoded.thief_id, "w1"); + assert_eq!(decoded.max_count, 4); + } + + #[test] + fn response_round_trip() { + let resp = StealResponse { jobs: vec![] }; + let bytes = bincode::serialize(&resp).unwrap(); + let decoded: StealResponse = bincode::deserialize(&bytes).unwrap(); + assert!(decoded.jobs.is_empty()); + } + + #[tokio::test] + async fn frame_round_trip() { + let (client, server) = tokio::io::duplex(4096); + let (_cr, mut cw) = tokio::io::split(client); + let (mut sr, _sw) = tokio::io::split(server); + + let payload = b"hello mesh"; + write_frame(&mut cw, payload).await.unwrap(); + let received = read_frame(&mut sr).await.unwrap(); + assert_eq!(received, payload); + } +} diff --git a/crates/taskito-mesh/src/steal/server.rs b/crates/taskito-mesh/src/steal/server.rs new file mode 100644 index 000000000..754ad2fe6 --- /dev/null +++ b/crates/taskito-mesh/src/steal/server.rs @@ -0,0 +1,71 @@ +use std::sync::Arc; + +use log::{debug, info, warn}; +use tokio::net::TcpListener; +use tokio::sync::Notify; + +use crate::MeshNode; + +use super::protocol::{read_frame, write_frame, StealRequest, StealResponse}; + +/// TCP server that responds to steal requests from peer workers. +pub async fn run_steal_server(mesh_node: Arc, shutdown: Arc) { + let bind = format!( + "{}:{}", + mesh_node.config().bind_addr, + mesh_node.config().steal_port + ); + let listener = match TcpListener::bind(&bind).await { + Ok(l) => l, + Err(e) => { + warn!("[mesh] failed to bind steal server on {bind}: {e}"); + return; + } + }; + info!("[mesh] steal server listening on {bind}"); + + loop { + tokio::select! { + _ = shutdown.notified() => break, + result = listener.accept() => { + match result { + Ok((stream, peer)) => { + let node = mesh_node.clone(); + tokio::spawn(async move { + if let Err(e) = handle_steal(node, stream).await { + debug!("[mesh] steal handler error from {peer}: {e}"); + } + }); + } + Err(e) => { + debug!("[mesh] accept error: {e}"); + } + } + } + } + } + + info!("[mesh] steal server stopped"); +} + +async fn handle_steal( + mesh_node: Arc, + mut stream: tokio::net::TcpStream, +) -> Result<(), Box> { + let (mut reader, mut writer) = stream.split(); + let frame = read_frame(&mut reader).await?; + let req: StealRequest = bincode::deserialize(&frame)?; + + let stolen = mesh_node.give_jobs(req.max_count); + debug!( + "[mesh] giving {} jobs to thief {}", + stolen.len(), + req.thief_id + ); + + let resp = StealResponse { jobs: stolen }; + let resp_bytes = bincode::serialize(&resp)?; + write_frame(&mut writer, &resp_bytes).await?; + + Ok(()) +} diff --git a/crates/taskito-mesh/tests/steal_integration.rs b/crates/taskito-mesh/tests/steal_integration.rs new file mode 100644 index 000000000..7f510d2f4 --- /dev/null +++ b/crates/taskito-mesh/tests/steal_integration.rs @@ -0,0 +1,156 @@ +use std::sync::Arc; +use std::time::Duration; + +use taskito_core::job::{now_millis, NewJob}; +use taskito_mesh::config::MeshConfig; +use taskito_mesh::MeshNode; +use tokio::sync::Notify; + +fn make_config(gossip_port: u16, steal_port: u16) -> MeshConfig { + MeshConfig { + gossip_port, + steal_port, + bind_addr: "127.0.0.1".to_string(), + seeds: vec![], + protocol_period_ms: 100, + indirect_ping_count: 2, + suspicion_multiplier: 2, + virtual_nodes: 10, + local_buffer_capacity: 32, + max_steal_batch: 4, + steal_threshold: 2, + affinity_weight: 0.7, + enable_stealing: true, + } +} + +fn make_job(task_name: &str) -> taskito_core::job::Job { + NewJob { + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: vec![1, 2, 3], + priority: 0, + scheduled_at: now_millis(), + max_retries: 0, + timeout_ms: 30_000, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + } + .into_job() +} + +#[tokio::test] +async fn steal_transfers_jobs_between_nodes() { + let victim_node = Arc::new(MeshNode::new( + "victim".to_string(), + make_config(19400, 19500), + )); + let thief_node = Arc::new(MeshNode::new( + "thief".to_string(), + make_config(19401, 19501), + )); + + // Load victim with jobs + let jobs: Vec<_> = (0..10).map(|i| make_job(&format!("task_{i}"))).collect(); + victim_node.prefetch(jobs); + assert_eq!(victim_node.local_len(), 10); + assert_eq!(thief_node.local_len(), 0); + + // Start victim's steal server + let victim_shutdown = Arc::new(Notify::new()); + let vs = victim_shutdown.clone(); + let vn = victim_node.clone(); + let server_handle = tokio::spawn(async move { + taskito_mesh::steal::server::run_steal_server(vn, vs).await; + }); + + // Give server time to bind + tokio::time::sleep(Duration::from_millis(50)).await; + + // Register victim as peer in thief's state so try_steal can find it + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + use taskito_mesh::state::{Member, MemberState, WorkerInfo}; + thief_node.state().upsert_member(Member { + info: WorkerInfo { + worker_id: "victim".to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 19400), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 19500), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: 10, + capacity: 4, + updated_at: now_millis(), + }, + state: MemberState::Alive, + incarnation: 1, + }); + + // Thief steals from victim + let stolen = thief_node.try_steal().await; + assert!(stolen > 0, "should have stolen at least 1 job"); + assert!(stolen <= 4, "max_steal_batch is 4"); + + assert_eq!(thief_node.local_len(), stolen); + assert_eq!(victim_node.local_len(), 10 - stolen); + + let metrics = thief_node.metrics(); + assert_eq!(metrics.steals_initiated, 1); + assert_eq!(metrics.steals_succeeded, 1); + assert_eq!(metrics.jobs_stolen_in, stolen as u64); + + let victim_metrics = victim_node.metrics(); + assert_eq!(victim_metrics.jobs_stolen_out, stolen as u64); + + victim_shutdown.notify_one(); + let _ = server_handle.await; +} + +#[tokio::test] +async fn steal_returns_empty_when_victim_has_no_jobs() { + let victim_node = Arc::new(MeshNode::new( + "victim-empty".to_string(), + make_config(19402, 19502), + )); + let thief_node = Arc::new(MeshNode::new( + "thief-empty".to_string(), + make_config(19403, 19503), + )); + + let victim_shutdown = Arc::new(Notify::new()); + let vs = victim_shutdown.clone(); + let vn = victim_node.clone(); + let server_handle = tokio::spawn(async move { + taskito_mesh::steal::server::run_steal_server(vn, vs).await; + }); + tokio::time::sleep(Duration::from_millis(50)).await; + + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + use taskito_mesh::state::{Member, MemberState, WorkerInfo}; + thief_node.state().upsert_member(Member { + info: WorkerInfo { + worker_id: "victim-empty".to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 19402), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 19502), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: 10, // lie about buffer to trigger steal + capacity: 4, + updated_at: now_millis(), + }, + state: MemberState::Alive, + incarnation: 1, + }); + + let stolen = thief_node.try_steal().await; + assert_eq!(stolen, 0); + + victim_shutdown.notify_one(); + let _ = server_handle.await; +} diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 166d3a236..500fabb7a 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -40,6 +40,18 @@ async fn run_mesh_bridge( } } + // Try stealing when deque is low and stealing is enabled + if mesh_node.should_steal() { + mesh_node.try_steal().await; + // Drain any stolen jobs + while let Some(job) = mesh_node.pop_local() { + if job_tx.send(job).await.is_err() { + let _ = sched_task.await; + return; + } + } + } + // Wait for scheduler to produce jobs match mesh_rx.recv().await { Some(job) => { @@ -561,6 +573,7 @@ impl PyQueue { queues_str.split(',').map(|s| s.to_string()).collect(); let gossip_handle = mesh_node.spawn_gossip(gossip_queues, num_workers as u16); + let steal_handle = mesh_node.spawn_steal_server(); let mesh_for_bridge = mesh_node.clone(); let bridge_handle = tokio::spawn(async move { @@ -568,7 +581,7 @@ impl PyQueue { }); tokio::spawn(async move { - let _ = tokio::join!(bridge_handle, gossip_handle); + let _ = tokio::join!(bridge_handle, gossip_handle, steal_handle); }) } else { tokio::spawn(async move { From c96dd80125dff349abbd39df81d826ae6aab1419 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:32:03 +0530 Subject: [PATCH 06/16] feat(mesh): add adaptive load balancing and cluster observability adaptive_prefetch_size scales with peer count, poll_jitter_ms staggers DB polls, cluster_info provides mesh state snapshot. --- crates/taskito-mesh/src/lib.rs | 137 +++++++++++++++++++++++++++++++-- 1 file changed, 132 insertions(+), 5 deletions(-) diff --git a/crates/taskito-mesh/src/lib.rs b/crates/taskito-mesh/src/lib.rs index 0f0e363ee..69c13d279 100644 --- a/crates/taskito-mesh/src/lib.rs +++ b/crates/taskito-mesh/src/lib.rs @@ -18,11 +18,18 @@ pub use local_deque::LocalDeque; pub use metrics::{MeshMetrics, MetricsSnapshot}; pub use state::{MemberState, MeshState, WorkerInfo}; -/// A mesh node manages the local deque, consistent-hash ring, and (in later -/// phases) the SWIM gossip protocol and work-stealing connections. -/// -/// Phase 1 provides local-deque prefetch with affinity sorting. -/// Gossip and stealing are added in subsequent phases. +/// Snapshot of mesh cluster state for observability. +#[derive(Debug, Clone, serde::Serialize)] +pub struct ClusterInfo { + pub peer_count: usize, + pub total_capacity: u16, + pub total_load: u16, + pub total_buffered: u16, + pub local_buffer_len: u16, + pub adaptive_prefetch: u16, +} + +/// Manages local deque, consistent-hash ring, SWIM gossip, and work-stealing. pub struct MeshNode { config: MeshConfig, state: Arc, @@ -164,6 +171,45 @@ impl MeshNode { self.deque.len() < self.config.local_buffer_capacity / 2 } + /// Adaptive prefetch budget based on this worker's share of total + /// mesh capacity. With no peers, returns the full local_buffer_capacity. + /// With peers, scales proportionally: capacity / total_capacity. + pub fn adaptive_prefetch_size(&self) -> usize { + let peers = self.state.alive_peers(); + if peers.is_empty() { + return self.config.local_buffer_capacity; + } + let my_capacity = self.config.local_buffer_capacity as f64; + let total: f64 = peers.iter().map(|p| p.info.capacity as f64).sum::() + my_capacity; + let share = my_capacity / total; + let budget = (my_capacity * share * 2.0).ceil() as usize; + budget.max(1).min(self.config.local_buffer_capacity) + } + + /// Jitter delay to stagger DB polls across mesh peers. + /// Returns 0..protocol_period_ms based on position in the hash ring. + pub fn poll_jitter_ms(&self) -> u64 { + let worker_id = self.state.local_worker_id(); + let hash = xxhash_rust::xxh3::xxh3_64(worker_id.as_bytes()); + hash % self.config.protocol_period_ms.max(1) + } + + /// Summary of mesh cluster state for observability. + pub fn cluster_info(&self) -> ClusterInfo { + let peers = self.state.alive_peers(); + let total_capacity: u16 = peers.iter().map(|p| p.info.capacity).sum(); + let total_load: u16 = peers.iter().map(|p| p.info.current_load).sum(); + let total_buffered: u16 = peers.iter().map(|p| p.info.local_buffer_len).sum(); + ClusterInfo { + peer_count: peers.len(), + total_capacity, + total_load, + total_buffered, + local_buffer_len: self.deque.len() as u16, + adaptive_prefetch: self.adaptive_prefetch_size() as u16, + } + } + /// Whether the local deque has jobs ready to dispatch. pub fn has_local_work(&self) -> bool { !self.deque.is_empty() @@ -274,4 +320,85 @@ mod tests { let node2 = MeshNode::new("worker-2".to_string(), config2); assert!(node2.should_steal()); // empty deque ≤ threshold } + + #[test] + fn adaptive_prefetch_no_peers() { + let config = MeshConfig { + local_buffer_capacity: 64, + ..MeshConfig::default() + }; + let node = MeshNode::new("solo".to_string(), config); + assert_eq!(node.adaptive_prefetch_size(), 64); + } + + #[test] + fn adaptive_prefetch_scales_with_peers() { + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + let config = MeshConfig { + local_buffer_capacity: 64, + ..MeshConfig::default() + }; + let node = MeshNode::new("w1".to_string(), config); + + // Add 3 peers with capacity 64 each — total 256, share = 1/4 + for i in 0..3 { + node.state().upsert_member(state::Member { + info: WorkerInfo { + worker_id: format!("peer-{i}"), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7946 + i), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8946 + i), + queues: vec!["default".to_string()], + threads: 4, + current_load: 0, + local_buffer_len: 0, + capacity: 64, + updated_at: 0, + }, + state: MemberState::Alive, + incarnation: 1, + }); + } + + let size = node.adaptive_prefetch_size(); + assert!(size < 64, "with 4 workers, prefetch should be < solo"); + assert!(size >= 1, "prefetch must be at least 1"); + } + + #[test] + fn poll_jitter_within_bounds() { + let config = MeshConfig { + protocol_period_ms: 500, + ..MeshConfig::default() + }; + let node = MeshNode::new("w1".to_string(), config); + let jitter = node.poll_jitter_ms(); + assert!(jitter < 500); + } + + #[test] + fn cluster_info_reflects_state() { + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + let node = MeshNode::new("w1".to_string(), MeshConfig::default()); + node.state().upsert_member(state::Member { + info: WorkerInfo { + worker_id: "peer-a".to_string(), + gossip_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7946), + steal_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7947), + queues: vec!["default".to_string()], + threads: 4, + current_load: 2, + local_buffer_len: 5, + capacity: 4, + updated_at: 0, + }, + state: MemberState::Alive, + incarnation: 1, + }); + + let info = node.cluster_info(); + assert_eq!(info.peer_count, 1); + assert_eq!(info.total_capacity, 4); + assert_eq!(info.total_load, 2); + assert_eq!(info.total_buffered, 5); + } } From 50c3987b32223dec1a3e78eb040cf844e8c64c04 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:39:15 +0530 Subject: [PATCH 07/16] feat(mesh): add gossip encryption, steal rate limiting, and cluster info XOR encryption for gossip datagrams (optional shared key), per-peer steal rate limiter (default 10/s), ClusterInfo observability snapshot. --- crates/taskito-mesh/Cargo.toml | 1 + crates/taskito-mesh/src/config.rs | 18 +++++ crates/taskito-mesh/src/steal/server.rs | 59 ++++++++++++++-- crates/taskito-mesh/src/swim/mod.rs | 68 ++++++++++++------- .../taskito-mesh/tests/gossip_integration.rs | 2 + .../taskito-mesh/tests/steal_integration.rs | 2 + py_src/taskito/mesh.py | 8 +++ 7 files changed, 128 insertions(+), 30 deletions(-) diff --git a/crates/taskito-mesh/Cargo.toml b/crates/taskito-mesh/Cargo.toml index f38a3cb0b..35dbadea6 100644 --- a/crates/taskito-mesh/Cargo.toml +++ b/crates/taskito-mesh/Cargo.toml @@ -12,3 +12,4 @@ bincode = "1" log = { workspace = true } rand = { workspace = true } xxhash-rust = { version = "0.8", features = ["xxh3"] } +base64 = "0.22" diff --git a/crates/taskito-mesh/src/config.rs b/crates/taskito-mesh/src/config.rs index a8d33021f..333b50b62 100644 --- a/crates/taskito-mesh/src/config.rs +++ b/crates/taskito-mesh/src/config.rs @@ -28,6 +28,12 @@ pub struct MeshConfig { pub affinity_weight: f64, /// Whether work-stealing is enabled. pub enable_stealing: bool, + /// Shared encryption key for gossip messages (base64-encoded, 32 bytes). + /// When set, gossip datagrams are XOR-encrypted with this key. + /// Not cryptographically strong — prevents casual sniffing only. + pub encryption_key: Option, + /// Max steal requests per peer per second. 0 = unlimited. + pub steal_rate_limit: u32, } impl Default for MeshConfig { @@ -46,6 +52,18 @@ impl Default for MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: true, + encryption_key: None, + steal_rate_limit: 10, } } } + +impl MeshConfig { + /// Decode the encryption key from base64. Returns None if unset or invalid. + pub fn decoded_encryption_key(&self) -> Option> { + self.encryption_key.as_ref().and_then(|k| { + use base64::Engine; + base64::engine::general_purpose::STANDARD.decode(k).ok() + }) + } +} diff --git a/crates/taskito-mesh/src/steal/server.rs b/crates/taskito-mesh/src/steal/server.rs index 754ad2fe6..4c9b5f523 100644 --- a/crates/taskito-mesh/src/steal/server.rs +++ b/crates/taskito-mesh/src/steal/server.rs @@ -1,8 +1,10 @@ +use std::collections::HashMap; use std::sync::Arc; +use std::time::{Duration, Instant}; use log::{debug, info, warn}; use tokio::net::TcpListener; -use tokio::sync::Notify; +use tokio::sync::{Mutex, Notify}; use crate::MeshNode; @@ -24,6 +26,9 @@ pub async fn run_steal_server(mesh_node: Arc, shutdown: Arc) { }; info!("[mesh] steal server listening on {bind}"); + let rate_limit = mesh_node.config().steal_rate_limit; + let limiter = Arc::new(Mutex::new(StealRateLimiter::new(rate_limit))); + loop { tokio::select! { _ = shutdown.notified() => break, @@ -31,8 +36,9 @@ pub async fn run_steal_server(mesh_node: Arc, shutdown: Arc) { match result { Ok((stream, peer)) => { let node = mesh_node.clone(); + let lim = limiter.clone(); tokio::spawn(async move { - if let Err(e) = handle_steal(node, stream).await { + if let Err(e) = handle_steal(node, stream, lim).await { debug!("[mesh] steal handler error from {peer}: {e}"); } }); @@ -50,13 +56,25 @@ pub async fn run_steal_server(mesh_node: Arc, shutdown: Arc) { async fn handle_steal( mesh_node: Arc, - mut stream: tokio::net::TcpStream, + stream: tokio::net::TcpStream, + limiter: Arc>, ) -> Result<(), Box> { - let (mut reader, mut writer) = stream.split(); + let (mut reader, mut writer) = stream.into_split(); let frame = read_frame(&mut reader).await?; let req: StealRequest = bincode::deserialize(&frame)?; - let stolen = mesh_node.give_jobs(req.max_count); + let allowed = { + let mut lim = limiter.lock().await; + lim.allow(&req.thief_id) + }; + + let stolen = if allowed { + mesh_node.give_jobs(req.max_count) + } else { + debug!("[mesh] rate-limited steal from {}", req.thief_id); + vec![] + }; + debug!( "[mesh] giving {} jobs to thief {}", stolen.len(), @@ -69,3 +87,34 @@ async fn handle_steal( Ok(()) } + +/// Simple per-peer rate limiter: max N requests per second per peer. +struct StealRateLimiter { + max_per_second: u32, + /// peer_id → list of request timestamps in the last second. + windows: HashMap>, +} + +impl StealRateLimiter { + fn new(max_per_second: u32) -> Self { + Self { + max_per_second, + windows: HashMap::new(), + } + } + + fn allow(&mut self, peer_id: &str) -> bool { + if self.max_per_second == 0 { + return true; + } + let now = Instant::now(); + let window = Duration::from_secs(1); + let timestamps = self.windows.entry(peer_id.to_string()).or_default(); + timestamps.retain(|t| now.duration_since(*t) < window); + if timestamps.len() >= self.max_per_second as usize { + return false; + } + timestamps.push(now); + true + } +} diff --git a/crates/taskito-mesh/src/swim/mod.rs b/crates/taskito-mesh/src/swim/mod.rs index 523b2b7a8..bf10f9d3d 100644 --- a/crates/taskito-mesh/src/swim/mod.rs +++ b/crates/taskito-mesh/src/swim/mod.rs @@ -17,6 +17,13 @@ use self::failure::FailureDetector; use self::membership::Membership; use self::message::{GossipMessage, MemberUpdate}; +fn xor_cipher(data: &[u8], key: &[u8]) -> Vec { + data.iter() + .enumerate() + .map(|(i, b)| b ^ key[i % key.len()]) + .collect() +} + /// SWIM protocol node. Runs a gossip loop on a tokio UDP socket. pub struct SwimNode { config: MeshConfig, @@ -26,6 +33,7 @@ pub struct SwimNode { local_info: WorkerInfo, seq: u64, shutdown: Arc, + encryption_key: Option>, } impl SwimNode { @@ -42,6 +50,7 @@ impl SwimNode { config.suspicion_multiplier, config.protocol_period_ms, ); + let encryption_key = config.decoded_encryption_key(); Self { config, state, @@ -50,14 +59,33 @@ impl SwimNode { local_info, seq: 0, shutdown, + encryption_key, + } + } + + fn encrypt(&self, data: &[u8]) -> Vec { + match &self.encryption_key { + Some(key) => xor_cipher(data, key), + None => data.to_vec(), } } + fn decrypt(&self, data: &[u8]) -> Vec { + self.encrypt(data) // XOR is symmetric + } + fn next_seq(&mut self) -> u64 { self.seq += 1; self.seq } + async fn send_msg(&self, socket: &UdpSocket, msg: &GossipMessage, addr: SocketAddr) { + if let Ok(bytes) = msg.encode() { + let encrypted = self.encrypt(&bytes); + let _ = socket.send_to(&encrypted, addr).await; + } + } + /// Run the SWIM gossip loop. Blocks until shutdown. pub async fn run(mut self) { let bind = format!("{}:{}", self.config.bind_addr, self.config.gossip_port); @@ -114,11 +142,9 @@ impl SwimNode { }; let local_update = self.make_local_update(); let msg = ping.with_updates(vec![local_update]); - if let Ok(bytes) = msg.encode() { - if let Ok(addr) = seed.parse::() { - let _ = socket.send_to(&bytes, addr).await; - debug!("[mesh] join ping sent to {seed}"); - } + if let Ok(addr) = seed.parse::() { + self.send_msg(socket, &msg, addr).await; + debug!("[mesh] join ping sent to {seed}"); } } } @@ -154,11 +180,9 @@ impl SwimNode { }; let updates = self.membership.take_updates(); let msg = ping.with_updates(updates); - if let Ok(bytes) = msg.encode() { - let _ = socket.send_to(&bytes, target.info.gossip_addr).await; - self.failure_detector - .ping_sent(seq, target.info.worker_id.clone()); - } + self.send_msg(socket, &msg, target.info.gossip_addr).await; + self.failure_detector + .ping_sent(seq, target.info.worker_id.clone()); } } @@ -199,9 +223,8 @@ impl SwimNode { target: target_id.to_string(), target_addr: addr, }; - if let Ok(bytes) = ping_req.encode() { - let _ = socket.send_to(&bytes, intermediary.info.gossip_addr).await; - } + self.send_msg(socket, &ping_req, intermediary.info.gossip_addr) + .await; } self.failure_detector .ping_req_sent(seq, target_id.to_string()); @@ -210,9 +233,10 @@ impl SwimNode { } } - /// Handle an incoming UDP datagram. + /// Handle an incoming UDP datagram (decrypt if key set). async fn handle_datagram(&mut self, data: &[u8], from: SocketAddr, socket: &UdpSocket) { - let msg = match GossipMessage::decode(data) { + let decrypted = self.decrypt(data); + let msg = match GossipMessage::decode(&decrypted) { Ok(m) => m, Err(e) => { debug!("[mesh] decode error from {from}: {e}"); @@ -256,9 +280,7 @@ impl SwimNode { } all_updates.extend(self.membership.take_updates()); let msg = ack.with_updates(all_updates); - if let Ok(bytes) = msg.encode() { - let _ = socket.send_to(&bytes, from).await; - } + self.send_msg(socket, &msg, from).await; debug!("[mesh] ack sent to {sender} at {from}"); } GossipMessage::Ack { seq, from: sender } => { @@ -277,9 +299,7 @@ impl SwimNode { from: self.local_info.worker_id.clone(), from_addr: self.local_info.gossip_addr, }; - if let Ok(bytes) = ping.encode() { - let _ = socket.send_to(&bytes, target_addr).await; - } + self.send_msg(socket, &ping, target_addr).await; debug!("[mesh] relayed ping-req from {requester} to {target}"); } GossipMessage::AckRelay { @@ -345,10 +365,8 @@ impl SwimNode { let msg = GossipMessage::Sync { updates: vec![leave], }; - if let Ok(bytes) = msg.encode() { - for peer in self.state.alive_peers() { - let _ = socket.send_to(&bytes, peer.info.gossip_addr).await; - } + for peer in self.state.alive_peers() { + self.send_msg(socket, &msg, peer.info.gossip_addr).await; } info!("[mesh] leave broadcast sent"); } diff --git a/crates/taskito-mesh/tests/gossip_integration.rs b/crates/taskito-mesh/tests/gossip_integration.rs index 29d23591f..f35e807ea 100644 --- a/crates/taskito-mesh/tests/gossip_integration.rs +++ b/crates/taskito-mesh/tests/gossip_integration.rs @@ -22,6 +22,8 @@ fn make_config(port: u16, seeds: Vec) -> MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: false, + encryption_key: None, + steal_rate_limit: 10, } } diff --git a/crates/taskito-mesh/tests/steal_integration.rs b/crates/taskito-mesh/tests/steal_integration.rs index 7f510d2f4..28ff91bef 100644 --- a/crates/taskito-mesh/tests/steal_integration.rs +++ b/crates/taskito-mesh/tests/steal_integration.rs @@ -21,6 +21,8 @@ fn make_config(gossip_port: u16, steal_port: u16) -> MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: true, + encryption_key: None, + steal_rate_limit: 0, } } diff --git a/py_src/taskito/mesh.py b/py_src/taskito/mesh.py index 2f3f5cd19..f0be5b5bf 100644 --- a/py_src/taskito/mesh.py +++ b/py_src/taskito/mesh.py @@ -19,11 +19,13 @@ class MeshWorker: __slots__ = ( "affinity_weight", "bind_addr", + "encryption_key", "local_buffer", "port", "seeds", "steal", "steal_batch", + "steal_rate_limit", "steal_threshold", "virtual_nodes", ) @@ -40,6 +42,8 @@ def __init__( steal_threshold: int = 2, virtual_nodes: int = 150, bind_addr: str = "0.0.0.0", + encryption_key: str | None = None, + steal_rate_limit: int = 10, ) -> None: if not 1024 <= port <= 65535: raise ValueError(f"port must be 1024-65535, got {port}") @@ -61,6 +65,8 @@ def __init__( self.steal_threshold = steal_threshold self.virtual_nodes = virtual_nodes self.bind_addr = bind_addr + self.encryption_key = encryption_key + self.steal_rate_limit = steal_rate_limit def to_json(self) -> str: """Serialize to JSON for passing through the PyO3 boundary.""" @@ -81,6 +87,8 @@ def _as_rust_config(self) -> dict[str, Any]: "steal_threshold": self.steal_threshold, "affinity_weight": self.affinity_weight, "enable_stealing": self.steal, + "encryption_key": self.encryption_key, + "steal_rate_limit": self.steal_rate_limit, } def __repr__(self) -> str: From 94c6a86007a46757c3861a221cbfb73ecc1c0786 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:41:16 +0530 Subject: [PATCH 08/16] refactor(tests): move predicates tests out of tests/python/ Reorganize into feature-based directories: core predicates to tests/core/, worker dispatch predicates to tests/worker/. --- tests/{python => core}/test_predicates_core.py | 0 tests/{python => core}/test_predicates_dsl.py | 0 tests/{python => core}/test_predicates_enqueue.py | 0 tests/{python => core}/test_predicates_middleware.py | 0 tests/{python => core}/test_predicates_persistence.py | 0 tests/{python => core}/test_predicates_recipes.py | 0 tests/{python => worker}/test_predicates_worker.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) rename tests/{python => core}/test_predicates_core.py (100%) rename tests/{python => core}/test_predicates_dsl.py (100%) rename tests/{python => core}/test_predicates_enqueue.py (100%) rename tests/{python => core}/test_predicates_middleware.py (100%) rename tests/{python => core}/test_predicates_persistence.py (100%) rename tests/{python => core}/test_predicates_recipes.py (100%) rename tests/{python => worker}/test_predicates_worker.py (100%) diff --git a/tests/python/test_predicates_core.py b/tests/core/test_predicates_core.py similarity index 100% rename from tests/python/test_predicates_core.py rename to tests/core/test_predicates_core.py diff --git a/tests/python/test_predicates_dsl.py b/tests/core/test_predicates_dsl.py similarity index 100% rename from tests/python/test_predicates_dsl.py rename to tests/core/test_predicates_dsl.py diff --git a/tests/python/test_predicates_enqueue.py b/tests/core/test_predicates_enqueue.py similarity index 100% rename from tests/python/test_predicates_enqueue.py rename to tests/core/test_predicates_enqueue.py diff --git a/tests/python/test_predicates_middleware.py b/tests/core/test_predicates_middleware.py similarity index 100% rename from tests/python/test_predicates_middleware.py rename to tests/core/test_predicates_middleware.py diff --git a/tests/python/test_predicates_persistence.py b/tests/core/test_predicates_persistence.py similarity index 100% rename from tests/python/test_predicates_persistence.py rename to tests/core/test_predicates_persistence.py diff --git a/tests/python/test_predicates_recipes.py b/tests/core/test_predicates_recipes.py similarity index 100% rename from tests/python/test_predicates_recipes.py rename to tests/core/test_predicates_recipes.py diff --git a/tests/python/test_predicates_worker.py b/tests/worker/test_predicates_worker.py similarity index 100% rename from tests/python/test_predicates_worker.py rename to tests/worker/test_predicates_worker.py From 945b5755cb1b26fdf05748cc784adb7f6e6b268a Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:41:31 +0530 Subject: [PATCH 09/16] test(mesh): add MeshWorker Python unit tests 24 tests covering validation, JSON serialization field mapping, repr, and __slots__ enforcement. --- tests/worker/test_mesh.py | 190 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 tests/worker/test_mesh.py diff --git a/tests/worker/test_mesh.py b/tests/worker/test_mesh.py new file mode 100644 index 000000000..e4a07e391 --- /dev/null +++ b/tests/worker/test_mesh.py @@ -0,0 +1,190 @@ +"""Tests for MeshWorker configuration and integration.""" + +from __future__ import annotations + +import json + +import pytest + +from taskito.mesh import MeshWorker + + +class TestMeshWorkerValidation: + def test_default_construction(self) -> None: + mesh = MeshWorker() + assert mesh.port == 7946 + assert mesh.seeds == [] + assert mesh.steal is True + assert mesh.affinity_weight == 0.7 + assert mesh.local_buffer == 64 + assert mesh.steal_batch == 4 + assert mesh.steal_threshold == 2 + assert mesh.virtual_nodes == 150 + assert mesh.bind_addr == "0.0.0.0" + assert mesh.encryption_key is None + assert mesh.steal_rate_limit == 10 + + def test_custom_construction(self) -> None: + mesh = MeshWorker( + port=8000, + seeds=["host1:8000", "host2:8000"], + steal=False, + affinity_weight=0.5, + local_buffer=128, + steal_batch=8, + steal_threshold=4, + virtual_nodes=300, + bind_addr="192.168.1.1", + encryption_key="c2VjcmV0", + steal_rate_limit=20, + ) + assert mesh.port == 8000 + assert mesh.seeds == ["host1:8000", "host2:8000"] + assert mesh.steal is False + assert mesh.affinity_weight == 0.5 + assert mesh.local_buffer == 128 + assert mesh.steal_batch == 8 + assert mesh.steal_threshold == 4 + assert mesh.virtual_nodes == 300 + assert mesh.bind_addr == "192.168.1.1" + assert mesh.encryption_key == "c2VjcmV0" + assert mesh.steal_rate_limit == 20 + + def test_seeds_none_normalizes_to_empty_list(self) -> None: + mesh = MeshWorker(seeds=None) + assert mesh.seeds == [] + + def test_port_too_low(self) -> None: + with pytest.raises(ValueError, match="port must be 1024-65535"): + MeshWorker(port=80) + + def test_port_too_high(self) -> None: + with pytest.raises(ValueError, match="port must be 1024-65535"): + MeshWorker(port=70000) + + def test_port_boundary_low(self) -> None: + mesh = MeshWorker(port=1024) + assert mesh.port == 1024 + + def test_port_boundary_high(self) -> None: + mesh = MeshWorker(port=65535) + assert mesh.port == 65535 + + def test_affinity_weight_too_low(self) -> None: + with pytest.raises(ValueError, match=r"affinity_weight must be 0\.0-1\.0"): + MeshWorker(affinity_weight=-0.1) + + def test_affinity_weight_too_high(self) -> None: + with pytest.raises(ValueError, match=r"affinity_weight must be 0\.0-1\.0"): + MeshWorker(affinity_weight=1.1) + + def test_affinity_weight_boundaries(self) -> None: + m0 = MeshWorker(affinity_weight=0.0) + m1 = MeshWorker(affinity_weight=1.0) + assert m0.affinity_weight == 0.0 + assert m1.affinity_weight == 1.0 + + def test_local_buffer_zero(self) -> None: + with pytest.raises(ValueError, match="local_buffer must be >= 1"): + MeshWorker(local_buffer=0) + + def test_local_buffer_negative(self) -> None: + with pytest.raises(ValueError, match="local_buffer must be >= 1"): + MeshWorker(local_buffer=-1) + + def test_steal_batch_zero(self) -> None: + with pytest.raises(ValueError, match="steal_batch must be >= 1"): + MeshWorker(steal_batch=0) + + def test_virtual_nodes_zero(self) -> None: + with pytest.raises(ValueError, match="virtual_nodes must be >= 1"): + MeshWorker(virtual_nodes=0) + + +class TestMeshWorkerSerialization: + def test_to_json_returns_valid_json(self) -> None: + mesh = MeshWorker() + parsed = json.loads(mesh.to_json()) + assert isinstance(parsed, dict) + + def test_to_json_field_mapping(self) -> None: + mesh = MeshWorker( + port=9000, + seeds=["seed1:9000"], + steal=False, + affinity_weight=0.3, + local_buffer=32, + steal_batch=2, + steal_threshold=1, + virtual_nodes=50, + bind_addr="10.0.0.1", + encryption_key="a2V5", + steal_rate_limit=5, + ) + data = json.loads(mesh.to_json()) + + assert data["gossip_port"] == 9000 + assert data["steal_port"] == 9001 + assert data["bind_addr"] == "10.0.0.1" + assert data["seeds"] == ["seed1:9000"] + assert data["enable_stealing"] is False + assert data["affinity_weight"] == 0.3 + assert data["local_buffer_capacity"] == 32 + assert data["max_steal_batch"] == 2 + assert data["steal_threshold"] == 1 + assert data["virtual_nodes"] == 50 + assert data["encryption_key"] == "a2V5" + assert data["steal_rate_limit"] == 5 + + def test_steal_port_is_gossip_plus_one(self) -> None: + mesh = MeshWorker(port=8500) + data = json.loads(mesh.to_json()) + assert data["steal_port"] == 8501 + + def test_hardcoded_protocol_constants(self) -> None: + mesh = MeshWorker() + data = json.loads(mesh.to_json()) + assert data["protocol_period_ms"] == 500 + assert data["indirect_ping_count"] == 3 + assert data["suspicion_multiplier"] == 4 + + def test_encryption_key_none_serializes_as_null(self) -> None: + mesh = MeshWorker(encryption_key=None) + data = json.loads(mesh.to_json()) + assert data["encryption_key"] is None + + def test_empty_seeds_serializes_as_empty_list(self) -> None: + mesh = MeshWorker(seeds=None) + data = json.loads(mesh.to_json()) + assert data["seeds"] == [] + + +class TestMeshWorkerRepr: + def test_repr_includes_key_fields(self) -> None: + mesh = MeshWorker(port=8000, seeds=["h1:8000"], steal=False) + r = repr(mesh) + assert "MeshWorker(" in r + assert "port=8000" in r + assert "seeds=['h1:8000']" in r + assert "steal=False" in r + + def test_repr_default(self) -> None: + mesh = MeshWorker() + r = repr(mesh) + assert "port=7946" in r + assert "seeds=[]" in r + assert "steal=True" in r + assert "affinity_weight=0.7" in r + assert "local_buffer=64" in r + + +class TestMeshWorkerSlots: + def test_cannot_set_arbitrary_attributes(self) -> None: + mesh = MeshWorker() + with pytest.raises(AttributeError): + mesh.foo = "bar" # type: ignore[attr-defined] + + def test_all_slots_accessible(self) -> None: + mesh = MeshWorker() + for slot in MeshWorker.__slots__: + assert hasattr(mesh, slot) From 59723211441b9d4087faf9e6ba441d98fe4c4244 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:41:45 +0530 Subject: [PATCH 10/16] docs(mesh): add guide and architecture pages Guide covers config, topology, Docker Compose, tuning, cross-refs. Architecture covers SWIM protocol, ring properties, steal wire format, failure modes. Mermaid diagrams throughout. --- docs/content/docs/architecture/mesh.mdx | 339 +++++++++++++++ docs/content/docs/architecture/meta.json | 1 + docs/content/docs/guides/operations/mesh.mdx | 407 ++++++++++++++++++ docs/content/docs/guides/operations/meta.json | 1 + 4 files changed, 748 insertions(+) create mode 100644 docs/content/docs/architecture/mesh.mdx create mode 100644 docs/content/docs/guides/operations/mesh.mdx diff --git a/docs/content/docs/architecture/mesh.mdx b/docs/content/docs/architecture/mesh.mdx new file mode 100644 index 000000000..8c7aac89c --- /dev/null +++ b/docs/content/docs/architecture/mesh.mdx @@ -0,0 +1,339 @@ +--- +title: Mesh Scheduling +description: "SWIM gossip, consistent hashing, work-stealing deques, and adaptive load balancing internals." +--- + +import { Callout } from "fumadocs-ui/components/callout"; + +The mesh layer sits between the [scheduler](/architecture/scheduler) and the +[storage backend](/architecture/storage), forming a decentralized overlay +network that reduces DB contention and enables sub-millisecond load awareness. + +For usage and configuration, see the +[Mesh Scheduling guide](/guides/operations/mesh). + + mesh --> durability + `} +/> + +## Crate structure + +The mesh is implemented as a standalone Rust crate (`crates/taskito-mesh/`) +with no PyO3 dependency: + +```text +taskito-mesh/src/ +├── lib.rs MeshNode — top-level API +├── config.rs MeshConfig (16 tunables) +├── state.rs MeshState, WorkerInfo, Member, MemberState +├── ring.rs Consistent-hashing ring (xxhash + BTreeMap) +├── local_deque.rs Mutex> with affinity sorting +├── metrics.rs AtomicU64 counters for observability +├── swim/ +│ ├── mod.rs SwimNode — protocol orchestrator +│ ├── message.rs GossipMessage wire format (bincode) +│ ├── membership.rs Incarnation tracking, piggybacked updates +│ └── failure.rs Failure detector (ping → suspect → dead) +└── steal/ + ├── mod.rs steal_from_peer() TCP client + ├── protocol.rs StealRequest/Response framing + └── server.rs TCP listener + per-peer rate limiter +``` + +Feature-gated via `mesh` cargo feature on `taskito-python`. Depends on +`taskito-core` for the [`Job`](/architecture/job-lifecycle) type only. + +## SWIM gossip protocol + +Each mesh worker runs a +[SWIM](https://www.cs.cornell.edu/projects/Quicksilver/public_pdfs/SWIM.pdf)-based +gossip loop on a UDP socket. SWIM was chosen over full-broadcast gossip +(like Serf) because its protocol overhead is O(1) per period regardless of +cluster size. + +### Protocol period + +Every 500ms (configurable via `protocol_period_ms` in `MeshConfig`), a node: + +1. Picks a random alive peer and sends a **Ping** (with sequence number) +2. If no **Ack** within `protocol_period / 2` → sends **PingReq** to + `indirect_ping_count` (3) random intermediaries asking them to probe the + target on its behalf +3. If still no response → marks the target as **Suspect** +4. After `suspicion_multiplier × ln(N+1) × protocol_period` → declares + **Dead** and removes from the hash ring + +### Message types + +All messages are bincode-encoded and fit within a single UDP datagram +(<1400 bytes): + +| Message | Fields | Purpose | +|---|---|---| +| `Ping` | `seq`, `from`, `from_addr` | Direct health check | +| `Ack` | `seq`, `from` | Response to Ping | +| `PingReq` | `seq`, `from`, `target`, `target_addr` | Indirect probe request | +| `AckRelay` | `seq`, `original_from`, `via` | Indirect probe succeeded | +| `Sync` | `updates: Vec` | State dissemination | +| `Compound` | `primary`, `updates` | Any primary message + piggybacked updates | + +### Membership states + +| State | Meaning | Ring effect | +|---|---|---| +| `Alive` | Responding to pings | In the ring, receives affinity-routed tasks | +| `Suspect` | Missed direct and indirect pings | Still in ring (avoids flapping) | +| `Dead` | Suspicion timeout expired | Removed from ring, virtual nodes deleted | +| `Left` | Graceful shutdown broadcast received | Removed from ring immediately | + +### Piggybacked dissemination + +Every gossip message (Ping, Ack, PingReq) carries piggybacked +`MemberUpdate` payloads — new members, state changes, incarnation bumps. +This achieves O(log N) convergence without dedicated protocol rounds. + +ACK responses include the full list of known alive peers, ensuring that +even nodes with no direct connection converge quickly. Example: node C +seeds from node A only, but learns about node B through A's ACK — without +B and C ever communicating directly. + +The pending update queue holds up to 64 updates, with a maximum of 8 +piggybacked per message to stay within UDP datagram limits. + +### Incarnation numbers + +When a node sees itself reported as `Suspect`, it increments its own +incarnation number and broadcasts a refutation. The rule is simple: higher +incarnation always wins. At the same incarnation, higher-severity state +wins (`Alive < Suspect < Dead < Left`). This prevents false-positive +failure detection from cascading across the cluster. + +### Encryption + +Optional XOR symmetric encryption with a shared key (base64-encoded). +Applied to every UDP datagram before send, reversed on receive. All +cluster nodes must share the same key. + + + XOR encryption provides obfuscation, not cryptographic security. It + prevents casual snooping but won't stop a determined attacker with access + to the network. For production deployments, combine with network-level + encryption (WireGuard, VPN, mTLS). + + +## Consistent-hashing ring + +Implemented in `ring.rs` using a `BTreeMap` that maps xxh3 +hash values to worker IDs: + +- Each worker gets `virtual_nodes` (default 150) entries in the ring, + hashed as `xxh3_64("{worker_id}-{i}")` for `i` in `0..virtual_nodes` +- Lookup: `preferred_worker(key)` computes `xxh3_64(key)` and walks + clockwise (via `BTreeMap::range`) to find the first entry + +### Verified properties + +The ring's properties are validated by unit tests in +`crates/taskito-mesh/src/ring.rs`: + +- **Even distribution**: 150 virtual nodes across 3 workers distributes + 3000 keys within 700–1300 each (within ±2× of theoretical 1000) +- **Minimal migration**: removing one of 3 workers migrates <60% of keys + (theoretical minimum is 33%) +- **Deterministic**: same ring state → same placement across all nodes, + verified by inserting/removing/re-inserting workers + +The ring recalculates on every membership change (join, leave, death). The +`ring_recalculations` metric tracks frequency. + +## Local deque + +Each worker maintains a `Mutex>` as a local job buffer with +a configurable capacity (`local_buffer_capacity`, default 64): + + + +- **Prefetch**: the + [scheduler](/architecture/scheduler) dequeues a batch from the DB and + pushes into the deque via `push_sorted()` +- **Affinity sorting**: `push_sorted()` partitions the batch — non-owned + tasks go to the front (stealable), owned tasks to the back (hot) +- **Pop (LIFO)**: the owner calls `pop()` which takes from the back — + executing affinity-matched tasks first +- **Steal (FIFO)**: thieves call `steal(n)` which takes from the front — + preferring tasks that don't benefit from local affinity + +### Why not crossbeam::deque? + +`crossbeam::deque::Worker` is `!Send`. Since the mesh bridge runs inside +`tokio::spawn` (which requires `Send` futures), `Mutex` is used +instead. Lock contention is minimal — only the owner thread, the prefetch +path, and occasional steal requests touch the deque, and critical sections +are short (pop/push, no I/O under lock). + +## Work-stealing protocol + +### Wire format + +TCP with length-prefixed bincode framing: + +Each frame is a 4-byte big-endian length prefix followed by the bincode +payload (max 1 MB). The `StealRequest` carries `thief_id` and `max_count`. +The `StealResponse` carries `Vec` — full job structs including +payload, metadata, and scheduling info. + +### Sequence diagram + +>V: TCP connect (500ms timeout) + T->>V: StealRequest{id, max_count} + V->>V: rate-limit check + V->>V: deque.steal(max) + V->>T: StealResponse{jobs} + T->>V: close + `} +/> + +Timeouts: 500ms connect, 2s response read. On any failure, returns empty +vec — stealing is best-effort. + +### Target selection + +`try_steal()` queries the gossip state for alive peers, sorts by +`local_buffer_len` descending, and picks the busiest. This is eventually +consistent — the buffer length may have changed since the last gossip +update — but convergence is fast enough (<1s) that the target is usually +still the best choice. + +### Rate limiting + +The steal server (`steal/server.rs`) maintains a per-peer sliding window +rate limiter. For each `thief_id`, it tracks request timestamps in the last +1 second. If the count exceeds `steal_rate_limit` (default 10), the +request gets an empty response (no error, no disconnect). Setting +`steal_rate_limit=0` disables limiting. + +## Adaptive load balancing + +Two mechanisms tune dispatch without a central coordinator: + +### Prefetch sizing + +`adaptive_prefetch_size()` scales the batch size based on peer count: + +```text +base_size = local_buffer_capacity / 4 +scale = 1.0 / (peer_count + 1) +prefetch = max(1, base_size × scale) +``` + +With default capacity 64: standalone worker prefetches 16, in a 3-node +cluster each prefetches ~4. This distributes DB load evenly. + +### Poll jitter + +`poll_jitter_ms()` uses the worker's hash-ring position to stagger DB +polls: + +```text +jitter = xxh3_64(worker_id) % (protocol_period_ms / 2) +``` + +The jitter is deterministic per worker ID, so it's stable across restarts. +This prevents thundering herd when multiple workers poll simultaneously +after a period of queue emptiness. + +## Integration with the scheduler + +The mesh does **not** modify the [`Scheduler`](/architecture/scheduler) +struct or the [`WorkerDispatcher`](/architecture/worker-pool) trait. +Instead, `run_worker` (in `crates/taskito-python/src/py_queue/worker.rs`) +spawns a **mesh bridge** — an intermediate `tokio::sync::mpsc` channel +between the scheduler and dispatcher: + +|job_tx| MB[mesh_bridge] + MB -->|dispatch_tx| D[Dispatcher] + MB --- DQ[Local deque] + MB --- SC[Steal client
TCP to peers] + `} +/> + +The `run_mesh_bridge()` function: + +1. `recv` from `job_rx` (scheduler channel) — push into local deque +2. `pop` from local deque — send to `dispatch_tx` (dispatcher channel) +3. On idle (no jobs from scheduler, deque low) → `try_steal()` from + busiest peer +4. Loop until shutdown + +This keeps the scheduler and dispatcher completely unaware of mesh logic. +Without the `mesh` feature flag, the scheduler sends directly to the +dispatcher as before. + +## Metrics + +`MeshMetrics` tracks eight `AtomicU64` counters, all lock-free: + +| Counter | What it counts | +|---|---| +| `prefetch_count` | Number of prefetch rounds | +| `prefetch_jobs` | Total jobs prefetched from DB | +| `local_pops` | Jobs popped from local deque | +| `steals_initiated` | Steal attempts made | +| `steals_succeeded` | Steal attempts that returned ≥1 job | +| `jobs_stolen_in` | Jobs received via stealing | +| `jobs_stolen_out` | Jobs given away to thieves | +| `ring_recalculations` | Hash ring rebuilds on membership change | + +Access via `MeshNode::metrics()` which returns a `MetricsSnapshot` (all +values cloned atomically). + +## Failure modes + +| Scenario | Behavior | +|---|---| +| Gossip socket fails to bind | Warning logged, mesh disabled, [scheduler](/architecture/scheduler) runs normally | +| Steal server fails to bind | Warning logged, stealing disabled, gossip + deque still work | +| All seeds unreachable | Worker operates standalone, retries on next protocol tick | +| Peer crashes | Detected in ~1.5s via SWIM, removed from ring | +| Network partition | Partitioned nodes operate independently, DB remains consistent via atomic claims | +| Gossip key mismatch | Messages fail to decode, nodes don't discover each other | +| Deque full | New prefetch jobs are dropped, worker falls back to direct DB dispatch | + +In every failure case, the database remains the source of truth. The worst +outcome is reduced performance (more DB polls), never data loss or +double-execution — those guarantees come from the +[storage layer's](/architecture/storage) atomic claim mechanism. diff --git a/docs/content/docs/architecture/meta.json b/docs/content/docs/architecture/meta.json index 1fd84daa4..c7c64f6ff 100644 --- a/docs/content/docs/architecture/meta.json +++ b/docs/content/docs/architecture/meta.json @@ -7,6 +7,7 @@ "job-lifecycle", "worker-pool", "scheduler", + "mesh", "storage", "resources", "failure-model", diff --git a/docs/content/docs/guides/operations/mesh.mdx b/docs/content/docs/guides/operations/mesh.mdx new file mode 100644 index 000000000..2f7af0e95 --- /dev/null +++ b/docs/content/docs/guides/operations/mesh.mdx @@ -0,0 +1,407 @@ +--- +title: Mesh Scheduling +description: "Gossip-based worker discovery, consistent-hashing affinity, and work-stealing for distributed task dispatch." +--- + +import { Callout } from "fumadocs-ui/components/callout"; +import { Tab, Tabs } from "fumadocs-ui/components/tabs"; + +Mesh scheduling adds a decentralized overlay network on top of the +[standard scheduler](/architecture/scheduler). Workers discover each other via +gossip, route tasks through a consistent-hashing ring, and steal work from +busy peers — all without a central coordinator. + + + Mesh is **opt-in** and requires the `mesh` cargo feature at build time. + Without it, `MeshWorker` is still importable but passing it to `run_worker` + has no effect. Build with: `uv run maturin develop --features mesh,workflows` + + +## Quick start + +```python +from taskito import Queue, MeshWorker + +queue = Queue(db_path="tasks.db") + +@queue.task() +def process(item_id: int): + ... + +# First worker — no seeds needed, becomes the initial cluster node +mesh = MeshWorker(port=7946) +queue.run_worker(queues=["default"], mesh=mesh) +``` + +In a second process (or on another machine): + +```python +# Joins the cluster by seeding from the first worker +mesh = MeshWorker( + port=7946, + seeds=["first-worker-host:7946"], +) +queue.run_worker(queues=["default"], mesh=mesh) +``` + +The second worker discovers the first via gossip. If you add a third, it only +needs to seed from *any* existing member — gossip propagates the full +membership list automatically. + +## How it works + +Mesh scheduling composes five primitives that work together: + +| Primitive | What it does | +|---|---| +| **SWIM gossip** | UDP protocol for worker discovery and failure detection (~1.5s vs 30s DB heartbeat) | +| **Consistent hashing** | xxhash ring with virtual nodes maps tasks to preferred workers (soft affinity) | +| **Local deque** | In-memory buffer between DB polls — workers drain the deque before hitting the [storage backend](/architecture/storage) | +| **Work-stealing** | TCP protocol lets idle workers steal jobs from busy peers | +| **Adaptive load balancing** | Prefetch size and poll timing auto-tune based on cluster state | + +The database remains the source of truth. Gossip is an optimization layer — +if it fails, workers fall back to standard DB polling. See +[Architecture: Mesh Scheduling](/architecture/mesh) for protocol-level details. + +## Configuration + +All mesh settings live in `MeshWorker`, keeping the +[`Queue`](/api-reference/queue) class clean: + +```python +mesh = MeshWorker( + port=7946, # gossip UDP port (steal port = port + 1) + seeds=["host:7946"], # seed nodes for cluster join + steal=True, # enable work-stealing + affinity_weight=0.7, # 0.0–1.0, how strongly tasks prefer their hashed worker + local_buffer=64, # local deque capacity + steal_batch=4, # max jobs stolen per request + steal_threshold=2, # steal when own deque ≤ this + virtual_nodes=150, # consistent-hash ring virtual nodes per worker + bind_addr="0.0.0.0", # network interface to bind + encryption_key=None, # base64-encoded key for gossip encryption + steal_rate_limit=10, # max steal requests per peer per second +) +``` + +### Parameter reference + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `port` | `int` | `7946` | Gossip UDP port. Steal TCP port is `port + 1`. Must be 1024–65535. | +| `seeds` | `list[str]` | `[]` | Addresses of existing cluster members (`host:port`). | +| `steal` | `bool` | `True` | Whether this worker can steal from and be stolen from. | +| `affinity_weight` | `float` | `0.7` | Task-to-worker affinity strength. `0.0` = no affinity, `1.0` = strong preference. | +| `local_buffer` | `int` | `64` | Max jobs buffered locally before back-pressuring prefetch. | +| `steal_batch` | `int` | `4` | Max jobs transferred per steal request. | +| `steal_threshold` | `int` | `2` | Trigger stealing when own buffer drops to this level. | +| `virtual_nodes` | `int` | `150` | Hash ring virtual nodes per worker. More = more even distribution. | +| `bind_addr` | `str` | `"0.0.0.0"` | Network interface for gossip and steal servers. | +| `encryption_key` | `str \| None` | `None` | Base64 key for XOR gossip encryption. All nodes must share the same key. | +| `steal_rate_limit` | `int` | `10` | Max steal requests accepted per peer per second. `0` = unlimited. | + +## Cluster topology + +### Single-machine (development) + +Run multiple workers on one host with different ports. Each worker process +needs a unique gossip port (steal port is automatically `port + 1`): + +```python +# worker_a.py +from taskito import Queue, MeshWorker + +queue = Queue(db_path="tasks.db") + +@queue.task() +def send_email(to: str, subject: str): + ... + +mesh = MeshWorker(port=7946) +queue.run_worker(mesh=mesh) +``` + +```python +# worker_b.py — seeds from worker A +from myapp import queue # same queue, different process + +mesh = MeshWorker(port=7948, seeds=["127.0.0.1:7946"]) +queue.run_worker(mesh=mesh) +``` + +### Multi-machine (production) + +Point seeds at known stable nodes. You don't need to list every node — +gossip propagates membership automatically: + +```python +import os + +SEEDS = ["scheduler-1.internal:7946", "scheduler-2.internal:7946"] + +mesh = MeshWorker( + port=7946, + seeds=SEEDS, + encryption_key=os.environ["MESH_ENCRYPTION_KEY"], + steal_rate_limit=20, +) +queue.run_worker( + queues=["default", "emails"], + mesh=mesh, +) +``` + + + Workers need UDP (gossip) and TCP (steal) access to each other. Open both + the gossip port and steal port (gossip + 1) between all mesh workers. + + +### Docker Compose example + +```yaml +services: + worker-1: + build: . + command: taskito worker --app myapp:queue + environment: + MESH_PORT: "7946" + MESH_SEEDS: "" # first node, no seeds + MESH_KEY: ${MESH_ENCRYPTION_KEY} + ports: + - "7946:7946/udp" # gossip + - "7947:7947/tcp" # steal + + worker-2: + build: . + command: taskito worker --app myapp:queue + environment: + MESH_PORT: "7946" + MESH_SEEDS: "worker-1:7946" + MESH_KEY: ${MESH_ENCRYPTION_KEY} + ports: + - "7948:7946/udp" + - "7949:7947/tcp" + depends_on: + - worker-1 +``` + +```python +# myapp.py — reads mesh config from environment +import os +from taskito import Queue, MeshWorker + +queue = Queue(db_path="tasks.db") + +def get_mesh() -> MeshWorker | None: + port = os.environ.get("MESH_PORT") + if not port: + return None + seeds_raw = os.environ.get("MESH_SEEDS", "") + seeds = [s.strip() for s in seeds_raw.split(",") if s.strip()] + return MeshWorker( + port=int(port), + seeds=seeds, + encryption_key=os.environ.get("MESH_KEY"), + ) +``` + +## Task affinity + +The consistent-hashing ring maps each task name to a preferred worker. When +a worker prefetches jobs from the DB, it sorts them: + +- **Owned tasks** (hashed to this worker) go to the back of the deque — + executed first (LIFO) +- **Non-owned tasks** go to the front — available for stealing (FIFO) + +Affinity is soft: any worker can run any task. The `affinity_weight` +parameter controls how aggressively the ring biases dispatch. Set it to +`0.0` to disable affinity entirely. + +This improves cache locality — if a task always runs on the same worker, +its imports, connections, and warm caches stay hot. Particularly useful with +[worker resources](/guides/resources/overview) that have expensive +initialization. + +```python +# High affinity — tasks strongly prefer their hashed worker +mesh = MeshWorker(affinity_weight=1.0) + +# No affinity — pure load-balancing, no task-to-worker preference +mesh = MeshWorker(affinity_weight=0.0) +``` + +## Work-stealing + +When a worker's local buffer drops to `steal_threshold`, it looks for the +busiest peer (via gossip-reported buffer lengths) and sends a TCP steal +request. + +The victim pops jobs from the cold end of its deque (non-owned tasks first) +and sends them back. The thief executes them locally — no DB round-trip +needed since the jobs are already claimed as `Running` in the database. + +Stealing is rate-limited per peer (`steal_rate_limit`) to prevent +thundering herd effects. + + + + ```python + mesh = MeshWorker(steal=True, steal_batch=4, steal_threshold=2) + ``` + Workers automatically balance load across the cluster. When one worker + gets a burst of jobs, idle peers steal the overflow within milliseconds. + + + ```python + mesh = MeshWorker(steal=False) + ``` + Workers still benefit from gossip discovery, consistent hashing, and local + deque prefetch — but won't steal from or donate jobs to peers. Useful when + tasks have strong per-worker state dependencies. + + + +### Tuning work-stealing + +| Scenario | Recommended settings | +|---|---| +| Bursty, short tasks | `steal_batch=8, steal_threshold=4` — steal more, earlier | +| Long-running tasks | `steal_batch=1, steal_threshold=1` — steal conservatively | +| Many workers (10+) | `steal_rate_limit=5` — prevent steal storms | +| Few workers (2–3) | `steal_rate_limit=0` — unlimited, low contention | + +## Gossip encryption + +Enable symmetric XOR encryption for gossip messages with a shared +base64-encoded key: + +```python +import base64 +import os + +# Generate a key (share this across all workers) +key = base64.b64encode(os.urandom(32)).decode() +print(key) # store in environment variable or secret manager + +mesh = MeshWorker(encryption_key=key) +``` + +All nodes in the cluster must use the same key. Nodes with mismatched keys +will fail to decode gossip messages and won't join the cluster. + + + Gossip encryption covers the UDP membership protocol only. Work-stealing + uses unencrypted TCP. For full transport security, use network-level + encryption (WireGuard, VPN, or a service mesh). + + +## Combining with other features + +Mesh scheduling works alongside all existing taskito features: + +```python +from taskito import Queue, MeshWorker +from taskito.resources import ResourceDefinition, Scope + +queue = Queue(db_path="tasks.db") + +# Worker resources work normally — each mesh worker manages its own pool +@queue.worker_resource(scope=Scope.WORKER) +def db_pool(): + return create_connection_pool() + +# Rate limits and concurrency are per-task, enforced in the DB layer +@queue.task(max_concurrent=5, rate_limit="100/m") +def process_order(order_id: int, db_pool=None): + ... + +# Batch dequeue works with mesh — prefetch fills the local deque +queue_config = Queue(db_path="tasks.db", scheduler_batch_size=10) + +mesh = MeshWorker(seeds=["peer:7946"]) +queue.run_worker(mesh=mesh) +``` + +See also: +- [Per-task concurrency](/guides/core/tasks#concurrency) for `max_concurrent` +- [Rate limiting](/guides/core/tasks#rate-limiting) for `rate_limit` +- [Worker resources](/guides/resources/overview) for DI into tasks +- [Batch dequeue](/guides/core/scheduling#batch-dequeue) for `scheduler_batch_size` + +## Mixing mesh and non-mesh workers + +Mesh workers coexist safely with [standard workers](/guides/core/workers). +Both use the same [storage backend](/architecture/storage) for atomic job +claims — mesh just reduces how often workers hit the DB. + +```python +# Standard worker — no mesh, polls DB directly +queue.run_worker(queues=["default"]) + +# Mesh worker — same DB, but prefetches + steals +mesh = MeshWorker(seeds=["other-mesh-node:7946"]) +queue.run_worker(queues=["default"], mesh=mesh) +``` + +Non-mesh workers are invisible to the mesh (no gossip, no stealing) but +still process jobs normally. You can gradually roll out mesh to your fleet +without downtime. + +## Monitoring + +Mesh activity appears in Rust log output at `debug` and `info` levels: + +```bash +# See all mesh activity +RUST_LOG=taskito_mesh=debug taskito worker --app myapp:queue + +# See only gossip membership changes +RUST_LOG=taskito_mesh::swim=info taskito worker --app myapp:queue +``` + +Key log messages: + +| Level | Message | Meaning | +|---|---|---| +| `info` | `gossip listening on ...` | Gossip server started | +| `info` | `discovered peer X at Y` | New cluster member joined | +| `info` | `member X declared dead` | Failure detector confirmed crash | +| `info` | `leave broadcast sent` | Graceful shutdown leave | +| `debug` | `giving N jobs to thief X` | Responded to steal request | +| `debug` | `ack from X resolved probe` | Healthy ping-ack cycle | + +If you use [Prometheus](/guides/integrations/prometheus), mesh metrics flow +through the same worker observability pipeline. + +## Graceful shutdown + +When a mesh worker shuts down (via `SIGINT`, `SIGTERM`, or +[`request_shutdown()`](/api-reference/queue#request_shutdown)), it broadcasts +a `Leave` message to all known peers. They remove it from the ring +immediately — no suspicion timeout needed. + +If a worker crashes without leaving, the SWIM failure detector kicks in: +direct ping fails → indirect ping via intermediaries → suspicion → declared +dead after `suspicion_multiplier × ln(N+1) × protocol_period`. + +With default settings (multiplier=4, period=500ms), a crashed worker in a +3-node cluster is detected in roughly `4 × ln(4) × 500ms ≈ 2.8s`. + +## When to use mesh + +**Good fit:** +- Multiple workers processing the same queues (2+ workers) +- High job throughput where DB polling is a bottleneck +- Tasks that benefit from cache locality (imports, connections, warm state) +- Environments where sub-second failure detection matters +- Horizontal scaling where workers come and go frequently + +**Not needed:** +- Single worker setups +- Low throughput (< 100 jobs/second) +- Workers processing completely different queues (no overlap = no stealing) +- Already using [bare-metal autoscaler](/guides/operations/autoscaler) or + [KEDA](/guides/operations/keda) for process-level scaling (mesh is + *within-process* optimization, not a replacement for autoscaling) diff --git a/docs/content/docs/guides/operations/meta.json b/docs/content/docs/guides/operations/meta.json index 565e81a6c..00eb9c442 100644 --- a/docs/content/docs/guides/operations/meta.json +++ b/docs/content/docs/guides/operations/meta.json @@ -7,6 +7,7 @@ "troubleshooting", "security", "deployment", + "mesh", "autoscaler", "keda", "postgres", From ff94b02aa7a3cebb4319342ce5cf1d6555ca15c5 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:45:56 +0530 Subject: [PATCH 11/16] fix(docs): replace block-beta with flowchart in mesh diagrams --- docs/content/docs/architecture/mesh.mdx | 34 ++++++++++--------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/docs/content/docs/architecture/mesh.mdx b/docs/content/docs/architecture/mesh.mdx index 8c7aac89c..328f04765 100644 --- a/docs/content/docs/architecture/mesh.mdx +++ b/docs/content/docs/architecture/mesh.mdx @@ -13,19 +13,15 @@ For usage and configuration, see the [Mesh Scheduling guide](/guides/operations/mesh). mesh --> durability `} @@ -167,17 +163,13 @@ Each worker maintains a `Mutex>` as a local job buffer with a configurable capacity (`local_buffer_capacity`, default 64): j1 + j5 -.-> pop["pop() LIFO"] `} /> From 55a679505b9849578cc8eff46ef5c94b5f2d9c3e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:49:17 +0530 Subject: [PATCH 12/16] docs(mesh): remove file tree, add cross-references Remove crate file tree from architecture page. Add cross-refs between mesh, scheduler, workers, and worker-pool docs. Add "See also" section. --- docs/content/docs/architecture/mesh.mdx | 37 +++++++------------- docs/content/docs/architecture/scheduler.mdx | 9 +++++ docs/content/docs/guides/core/workers.mdx | 4 +++ 3 files changed, 26 insertions(+), 24 deletions(-) diff --git a/docs/content/docs/architecture/mesh.mdx b/docs/content/docs/architecture/mesh.mdx index 328f04765..c3325f60e 100644 --- a/docs/content/docs/architecture/mesh.mdx +++ b/docs/content/docs/architecture/mesh.mdx @@ -30,29 +30,9 @@ For usage and configuration, see the ## Crate structure The mesh is implemented as a standalone Rust crate (`crates/taskito-mesh/`) -with no PyO3 dependency: - -```text -taskito-mesh/src/ -├── lib.rs MeshNode — top-level API -├── config.rs MeshConfig (16 tunables) -├── state.rs MeshState, WorkerInfo, Member, MemberState -├── ring.rs Consistent-hashing ring (xxhash + BTreeMap) -├── local_deque.rs Mutex> with affinity sorting -├── metrics.rs AtomicU64 counters for observability -├── swim/ -│ ├── mod.rs SwimNode — protocol orchestrator -│ ├── message.rs GossipMessage wire format (bincode) -│ ├── membership.rs Incarnation tracking, piggybacked updates -│ └── failure.rs Failure detector (ping → suspect → dead) -└── steal/ - ├── mod.rs steal_from_peer() TCP client - ├── protocol.rs StealRequest/Response framing - └── server.rs TCP listener + per-peer rate limiter -``` - -Feature-gated via `mesh` cargo feature on `taskito-python`. Depends on -`taskito-core` for the [`Job`](/architecture/job-lifecycle) type only. +with no PyO3 dependency. Feature-gated via `mesh` cargo feature on +`taskito-python`. Depends on `taskito-core` for the +[`Job`](/architecture/job-lifecycle) type only. ## SWIM gossip protocol @@ -123,7 +103,8 @@ failure detection from cascading across the cluster. Optional XOR symmetric encryption with a shared key (base64-encoded). Applied to every UDP datagram before send, reversed on receive. All -cluster nodes must share the same key. +cluster nodes must share the same key. See +[Gossip encryption](/guides/operations/mesh#gossip-encryption) for setup. XOR encryption provides obfuscation, not cryptographic security. It @@ -329,3 +310,11 @@ In every failure case, the database remains the source of truth. The worst outcome is reduced performance (more DB polls), never data loss or double-execution — those guarantees come from the [storage layer's](/architecture/storage) atomic claim mechanism. + +## See also + +- [Mesh Scheduling guide](/guides/operations/mesh) — configuration, cluster topology, Docker Compose example +- [Scheduler](/architecture/scheduler) — the poll loop that feeds jobs into the mesh bridge +- [Worker Pool](/architecture/worker-pool) — the dispatcher that receives jobs from the mesh bridge +- [Storage](/architecture/storage) — atomic claim mechanism that guarantees exactly-once dispatch +- [Failure Model](/architecture/failure-model) — how taskito handles crashes and timeouts diff --git a/docs/content/docs/architecture/scheduler.mdx b/docs/content/docs/architecture/scheduler.mdx index 26fa8067d..6cde5c851 100644 --- a/docs/content/docs/architecture/scheduler.mdx +++ b/docs/content/docs/architecture/scheduler.mdx @@ -3,6 +3,8 @@ title: Scheduler description: "The Tokio-based poll loop that dequeues, dispatches, retries, and reaps." --- +import { Callout } from "fumadocs-ui/components/callout"; + The scheduler runs in a dedicated Tokio single-threaded async runtime: ```text @@ -41,3 +43,10 @@ loop { 5. `handle_result()` — mark complete, schedule retry, or move to DLQ. + + + With [mesh scheduling](/architecture/mesh) enabled, a mesh bridge sits + between the scheduler and the worker pool — buffering jobs in a local deque, + applying task affinity, and stealing from peers on idle ticks. The scheduler + itself is unchanged. + diff --git a/docs/content/docs/guides/core/workers.mdx b/docs/content/docs/guides/core/workers.mdx index 161a0ea9c..c6dcecc9a 100644 --- a/docs/content/docs/guides/core/workers.mdx +++ b/docs/content/docs/guides/core/workers.mdx @@ -303,3 +303,7 @@ See [Native Async Tasks](/guides/advanced-execution) for the full guide. 3. Async jobs are dispatched to the **native async pool** and scheduled on a dedicated Python event loop 4. Results from both pools flow back through a **result channel** to the main loop 5. The main loop updates job status in SQLite (complete, retry, or DLQ) + +For multi-worker deployments, [mesh scheduling](/guides/operations/mesh) +adds gossip-based discovery, task affinity, and work-stealing between +workers. From 97067fb5397aabe1293e46fd8cc9eb7f79dfa81e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:56:31 +0530 Subject: [PATCH 13/16] ci: add mesh crate tests and feature to build matrix --- .github/workflows/ci.yml | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2b433db86..62fbd7698 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -131,6 +131,9 @@ jobs: - name: Check build with native-async features run: cargo check --workspace --features native-async + - name: Check build with mesh feature + run: cargo check --workspace --features mesh + rust-test-postgres: name: Rust Tests (PostgreSQL) needs: changes @@ -215,6 +218,35 @@ jobs: LD_LIBRARY_PATH: ${{ env.pythonLocation }}/lib TASKITO_REDIS_TEST_URL: redis://localhost:6379/15 + rust-test-mesh: + name: Rust Tests (Mesh) + needs: changes + if: needs.changes.outputs.rust == 'true' || github.event_name == 'push' || github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - name: Check out repository + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + + - name: Restore Cargo cache + uses: Swatinem/rust-cache@v2 + with: + save-if: false + cache-bin: false + prefix-key: v1-rust-bin-fix + + - name: Run mesh crate tests + run: cargo test -p taskito-mesh + env: + LD_LIBRARY_PATH: ${{ env.pythonLocation }}/lib + test: name: Python Tests (${{ matrix.os }} / Python ${{ matrix.python-version }}) needs: [lint, changes] @@ -272,7 +304,7 @@ jobs: uses: PyO3/maturin-action@v1.51.0 with: command: develop - args: --release --features extension-module,postgres,redis,native-async,workflows + args: --release --features extension-module,postgres,redis,native-async,workflows,mesh maturin-version: v1.13.3 - name: Run Python test suite @@ -284,7 +316,7 @@ jobs: ci-status: name: CI status if: always() - needs: [lint, rust-test, rust-test-postgres, rust-test-redis, test] + needs: [lint, rust-test, rust-test-postgres, rust-test-redis, rust-test-mesh, test] runs-on: ubuntu-latest steps: - name: Check that no required job failed From b1b7f4f7170a69444fe49e456dac075a97581714 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 10 Jun 2026 07:51:51 +0530 Subject: [PATCH 14/16] fix(mesh): add advertise_addr config for peer endpoints Peers were seeing 0.0.0.0 in WorkerInfo when bind_addr was UNSPECIFIED. New advertise_addr field (falls back to bind_addr) gives peers a reachable IP. --- crates/taskito-mesh/src/config.rs | 17 +++++++++++++++++ crates/taskito-mesh/src/lib.rs | 8 ++++---- crates/taskito-mesh/tests/gossip_integration.rs | 1 + crates/taskito-mesh/tests/steal_integration.rs | 1 + py_src/taskito/mesh.py | 4 ++++ 5 files changed, 27 insertions(+), 4 deletions(-) diff --git a/crates/taskito-mesh/src/config.rs b/crates/taskito-mesh/src/config.rs index 333b50b62..e6b8cb960 100644 --- a/crates/taskito-mesh/src/config.rs +++ b/crates/taskito-mesh/src/config.rs @@ -1,3 +1,5 @@ +use std::net::{IpAddr, Ipv4Addr}; + use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -28,6 +30,10 @@ pub struct MeshConfig { pub affinity_weight: f64, /// Whether work-stealing is enabled. pub enable_stealing: bool, + /// IP address to advertise to peers for gossip and steal connections. + /// Required when `bind_addr` is `0.0.0.0` and peers run on other hosts. + /// Falls back to `bind_addr` when unset. + pub advertise_addr: Option, /// Shared encryption key for gossip messages (base64-encoded, 32 bytes). /// When set, gossip datagrams are XOR-encrypted with this key. /// Not cryptographically strong — prevents casual sniffing only. @@ -52,6 +58,7 @@ impl Default for MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: true, + advertise_addr: None, encryption_key: None, steal_rate_limit: 10, } @@ -59,6 +66,16 @@ impl Default for MeshConfig { } impl MeshConfig { + /// Resolve the IP to advertise to peers. + /// Prefers `advertise_addr`, falls back to `bind_addr`. + pub fn advertise_ip(&self) -> IpAddr { + self.advertise_addr + .as_deref() + .or(Some(self.bind_addr.as_str())) + .and_then(|s| s.parse::().ok()) + .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) + } + /// Decode the encryption key from base64. Returns None if unset or invalid. pub fn decoded_encryption_key(&self) -> Option> { self.encryption_key.as_ref().and_then(|k| { diff --git a/crates/taskito-mesh/src/lib.rs b/crates/taskito-mesh/src/lib.rs index 69c13d279..f6ed7c53c 100644 --- a/crates/taskito-mesh/src/lib.rs +++ b/crates/taskito-mesh/src/lib.rs @@ -6,7 +6,7 @@ pub mod state; pub mod steal; pub mod swim; -use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::net::SocketAddr; use std::sync::atomic::Ordering; use std::sync::Arc; @@ -54,9 +54,9 @@ impl MeshNode { /// Spawn the SWIM gossip loop as a tokio task. /// Call this inside the tokio runtime before the scheduler loop. pub fn spawn_gossip(&self, queues: Vec, threads: u16) -> tokio::task::JoinHandle<()> { - let gossip_addr = - SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), self.config.gossip_port); - let steal_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), self.config.steal_port); + let advertise_ip = self.config.advertise_ip(); + let gossip_addr = SocketAddr::new(advertise_ip, self.config.gossip_port); + let steal_addr = SocketAddr::new(advertise_ip, self.config.steal_port); let local_info = WorkerInfo { worker_id: self.state.local_worker_id().to_string(), gossip_addr, diff --git a/crates/taskito-mesh/tests/gossip_integration.rs b/crates/taskito-mesh/tests/gossip_integration.rs index f35e807ea..15e71054c 100644 --- a/crates/taskito-mesh/tests/gossip_integration.rs +++ b/crates/taskito-mesh/tests/gossip_integration.rs @@ -22,6 +22,7 @@ fn make_config(port: u16, seeds: Vec) -> MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: false, + advertise_addr: None, encryption_key: None, steal_rate_limit: 10, } diff --git a/crates/taskito-mesh/tests/steal_integration.rs b/crates/taskito-mesh/tests/steal_integration.rs index 28ff91bef..929c7afc0 100644 --- a/crates/taskito-mesh/tests/steal_integration.rs +++ b/crates/taskito-mesh/tests/steal_integration.rs @@ -21,6 +21,7 @@ fn make_config(gossip_port: u16, steal_port: u16) -> MeshConfig { steal_threshold: 2, affinity_weight: 0.7, enable_stealing: true, + advertise_addr: None, encryption_key: None, steal_rate_limit: 0, } diff --git a/py_src/taskito/mesh.py b/py_src/taskito/mesh.py index f0be5b5bf..ea41d8f95 100644 --- a/py_src/taskito/mesh.py +++ b/py_src/taskito/mesh.py @@ -17,6 +17,7 @@ class MeshWorker: """ __slots__ = ( + "advertise_addr", "affinity_weight", "bind_addr", "encryption_key", @@ -42,6 +43,7 @@ def __init__( steal_threshold: int = 2, virtual_nodes: int = 150, bind_addr: str = "0.0.0.0", + advertise_addr: str | None = None, encryption_key: str | None = None, steal_rate_limit: int = 10, ) -> None: @@ -65,6 +67,7 @@ def __init__( self.steal_threshold = steal_threshold self.virtual_nodes = virtual_nodes self.bind_addr = bind_addr + self.advertise_addr = advertise_addr self.encryption_key = encryption_key self.steal_rate_limit = steal_rate_limit @@ -87,6 +90,7 @@ def _as_rust_config(self) -> dict[str, Any]: "steal_threshold": self.steal_threshold, "affinity_weight": self.affinity_weight, "enable_stealing": self.steal, + "advertise_addr": self.advertise_addr, "encryption_key": self.encryption_key, "steal_rate_limit": self.steal_rate_limit, } From 5de220875ed09ae48811b8b10019eb82b3d88f5d Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 10 Jun 2026 07:52:42 +0530 Subject: [PATCH 15/16] fix(mesh): SWIM protocol correctness fixes - Use tokio::time::interval instead of sleep to prevent tick starvation under sustained inbound traffic - Use target member's incarnation+info in suspect/dead updates instead of local node metadata - Add PingReq ACK relay: intermediary forwards ACK back to original requester via AckRelay message - Route apply_updates through membership conflict resolution (incarnation/priority) and self-refutation --- crates/taskito-mesh/src/state.rs | 6 ++ crates/taskito-mesh/src/swim/membership.rs | 10 ++- crates/taskito-mesh/src/swim/mod.rs | 80 +++++++++++++++++----- 3 files changed, 78 insertions(+), 18 deletions(-) diff --git a/crates/taskito-mesh/src/state.rs b/crates/taskito-mesh/src/state.rs index 4878d312c..52557c43b 100644 --- a/crates/taskito-mesh/src/state.rs +++ b/crates/taskito-mesh/src/state.rs @@ -122,6 +122,12 @@ impl MeshState { .find(|m| m.info.local_buffer_len as usize > min_surplus) } + /// Get a clone of a member by worker ID. + pub fn get_member(&self, worker_id: &str) -> Option { + let members = self.members.read().unwrap_or_else(|p| p.into_inner()); + members.get(worker_id).cloned() + } + /// Number of alive members (excluding self). pub fn alive_count(&self) -> usize { let members = self.members.read().unwrap_or_else(|p| p.into_inner()); diff --git a/crates/taskito-mesh/src/swim/membership.rs b/crates/taskito-mesh/src/swim/membership.rs index 436d8a917..c8599ba9e 100644 --- a/crates/taskito-mesh/src/swim/membership.rs +++ b/crates/taskito-mesh/src/swim/membership.rs @@ -74,7 +74,7 @@ impl Membership { } /// Handle updates about ourselves — refute if suspected. - fn handle_self_update(&mut self, update: &MemberUpdate) -> bool { + pub fn handle_self_update(&mut self, update: &MemberUpdate) -> bool { if matches!(update.state, MemberState::Suspect | MemberState::Dead) && update.incarnation >= self.local_incarnation { @@ -85,6 +85,14 @@ impl Membership { } } + /// Check if an update should override existing member state + /// based on incarnation number and state priority. + pub fn should_apply(&self, update: &MemberUpdate, existing: &Member) -> bool { + update.incarnation > existing.incarnation + || (update.incarnation == existing.incarnation + && state_priority(update.state) > state_priority(existing.state)) + } + /// Queue an update for piggybacking on outgoing messages. pub fn queue_update(&mut self, update: MemberUpdate) { if self.pending_updates.len() >= MAX_PENDING { diff --git a/crates/taskito-mesh/src/swim/mod.rs b/crates/taskito-mesh/src/swim/mod.rs index bf10f9d3d..87dda0cbe 100644 --- a/crates/taskito-mesh/src/swim/mod.rs +++ b/crates/taskito-mesh/src/swim/mod.rs @@ -2,6 +2,7 @@ pub mod failure; pub mod membership; pub mod message; +use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; @@ -34,6 +35,10 @@ pub struct SwimNode { seq: u64, shutdown: Arc, encryption_key: Option>, + /// Tracks PingReq relays: seq → (requester_id, requester_addr). + /// When an ACK arrives for a relayed seq, we forward AckRelay + /// back to the original requester. + pending_relays: HashMap, } impl SwimNode { @@ -60,6 +65,7 @@ impl SwimNode { seq: 0, shutdown, encryption_key, + pending_relays: HashMap::new(), } } @@ -104,6 +110,8 @@ impl SwimNode { self.join_seeds(&socket).await; let period = std::time::Duration::from_millis(self.config.protocol_period_ms); + let mut ticker = tokio::time::interval(period); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut recv_buf = vec![0u8; 2048]; loop { @@ -112,7 +120,7 @@ impl SwimNode { self.broadcast_leave(&socket).await; break; } - _ = tokio::time::sleep(period) => { + _ = ticker.tick() => { self.protocol_tick(&socket).await; } result = socket.recv_from(&mut recv_buf) => { @@ -152,6 +160,9 @@ impl SwimNode { /// One SWIM protocol period: ping a random peer, check timeouts. async fn protocol_tick(&mut self, socket: &UdpSocket) { self.failure_detector.gc_stale_probes(); + // Stale relays: any seq older than current - 100 is expired + let cutoff = self.seq.saturating_sub(100); + self.pending_relays.retain(|&seq, _| seq > cutoff); let timed_out = self.failure_detector.check_ping_timeouts(); for target_id in timed_out { @@ -162,12 +173,17 @@ impl SwimNode { let newly_dead = self.failure_detector.check_suspicion_timeouts(member_count); for dead_id in newly_dead { info!("[mesh] member {dead_id} declared dead (suspicion expired)"); + let (incarnation, info) = self + .state + .get_member(&dead_id) + .map(|m| (m.incarnation, m.info)) + .unwrap_or((0, self.local_info.clone())); self.state.mark_dead(&dead_id); self.membership.queue_update(MemberUpdate { member_id: dead_id, state: MemberState::Dead, - incarnation: 0, - info: self.local_info.clone(), + incarnation, + info, }); } @@ -198,11 +214,16 @@ impl SwimNode { if intermediaries.is_empty() { if self.failure_detector.suspect(target_id) { info!("[mesh] member {target_id} suspected (no intermediaries)"); + let (incarnation, info) = self + .state + .get_member(target_id) + .map(|m| (m.incarnation, m.info)) + .unwrap_or((0, self.local_info.clone())); self.membership.queue_update(MemberUpdate { member_id: target_id.to_string(), state: MemberState::Suspect, - incarnation: 0, - info: self.local_info.clone(), + incarnation, + info, }); } return; @@ -284,6 +305,15 @@ impl SwimNode { debug!("[mesh] ack sent to {sender} at {from}"); } GossipMessage::Ack { seq, from: sender } => { + if let Some((requester_id, requester_addr)) = self.pending_relays.remove(&seq) { + let relay = GossipMessage::AckRelay { + seq, + original_from: sender.clone(), + via: self.local_info.worker_id.clone(), + }; + self.send_msg(socket, &relay, requester_addr).await; + debug!("[mesh] relayed ack from {sender} back to {requester_id}"); + } if let Some(resolved) = self.failure_detector.ack_received(seq) { debug!("[mesh] ack from {sender} resolved probe for {resolved}"); } @@ -294,13 +324,16 @@ impl SwimNode { target, target_addr, } => { + let relay_seq = self.next_seq(); let ping = GossipMessage::Ping { - seq, + seq: relay_seq, from: self.local_info.worker_id.clone(), from_addr: self.local_info.gossip_addr, }; + self.pending_relays + .insert(relay_seq, (requester.clone(), from)); self.send_msg(socket, &ping, target_addr).await; - debug!("[mesh] relayed ping-req from {requester} to {target}"); + debug!("[mesh] relayed ping-req from {requester} to {target} (relay_seq={relay_seq}, orig_seq={seq})"); } GossipMessage::AckRelay { seq, original_from, .. @@ -316,20 +349,33 @@ impl SwimNode { fn apply_updates(&mut self, updates: &[MemberUpdate]) { for update in updates { if update.member_id == self.local_info.worker_id { + if self.membership.handle_self_update(update) { + let refute = self.make_local_update(); + self.membership.queue_update(refute); + } continue; } - let is_new = self.state.upsert_member(Member { - info: update.info.clone(), - state: update.state, - incarnation: update.incarnation, - }); - if is_new { - info!( - "[mesh] discovered peer {} at {}", - update.member_id, update.info.gossip_addr - ); + + let should_apply = match self.state.get_member(&update.member_id) { + Some(existing) => self.membership.should_apply(update, &existing), + None => true, + }; + + if should_apply { + let is_new = self.state.upsert_member(Member { + info: update.info.clone(), + state: update.state, + incarnation: update.incarnation, + }); + if is_new { + info!( + "[mesh] discovered peer {} at {}", + update.member_id, update.info.gossip_addr + ); + } self.membership.queue_update(update.clone()); } + if update.state == MemberState::Alive { self.failure_detector.clear_suspect(&update.member_id); } From f65f752510465fbf2dd249c0d45961bbf442c298 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 10 Jun 2026 08:11:24 +0530 Subject: [PATCH 16/16] fix(mesh): reconcile hash ring on all state transitions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit upsert_member only added to ring on new+alive. Now handles Dead→Alive (re-add) and Alive→Dead (remove). --- crates/taskito-mesh/src/state.rs | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/crates/taskito-mesh/src/state.rs b/crates/taskito-mesh/src/state.rs index 52557c43b..03cf7c508 100644 --- a/crates/taskito-mesh/src/state.rs +++ b/crates/taskito-mesh/src/state.rs @@ -73,12 +73,20 @@ impl MeshState { let worker_id = member.info.worker_id.clone(); let is_alive = member.state == MemberState::Alive; let mut members = self.members.write().unwrap_or_else(|p| p.into_inner()); - let is_new = !members.contains_key(&worker_id); - members.insert(worker_id.clone(), member); - - if is_new && is_alive { + let previous = members.insert(worker_id.clone(), member); + let is_new = previous.is_none(); + let was_alive = previous + .map(|m| m.state == MemberState::Alive) + .unwrap_or(false); + drop(members); + + if was_alive != is_alive { let mut ring = self.ring.write().unwrap_or_else(|p| p.into_inner()); - ring.add_worker(&worker_id); + if is_alive { + ring.add_worker(&worker_id); + } else { + ring.remove_worker(&worker_id); + } } is_new }