125 lines
5.1 KiB
Rust
125 lines
5.1 KiB
Rust
#![deny(warnings)]
|
|
#![allow(dead_code)]
|
|
//! [TSM.ID].[11031972] -- Platform X Ecosystem
|
|
//! xcu-thread-weaver -- Work-Stealing Thread Pool & NPU Scheduler
|
|
use std::collections::VecDeque;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
#[derive(Debug)]
|
|
pub enum WeaverError { QueueFull(String), WorkerPanicked(String), InvalidConfig(String) }
|
|
impl std::fmt::Display for WeaverError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self { Self::QueueFull(e) => write!(f, "Full: {e}"), Self::WorkerPanicked(e) => write!(f, "Panic: {e}"), Self::InvalidConfig(e) => write!(f, "Config: {e}") }
|
|
}
|
|
}
|
|
impl std::error::Error for WeaverError {}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
|
pub enum TaskPriority { Critical = 0, High = 1, Normal = 2, Low = 3, Background = 4 }
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct WeaverTask { pub id: String, pub priority: TaskPriority, pub data: Vec<u8>, pub created_at: u64 }
|
|
|
|
struct WorkerQueue { deque: VecDeque<WeaverTask>, processed: u64, stolen_from: u64 }
|
|
impl WorkerQueue {
|
|
fn new() -> Self { Self { deque: VecDeque::new(), processed: 0, stolen_from: 0 } }
|
|
fn push(&mut self, task: WeaverTask) {
|
|
// Insert sorted by priority
|
|
let pos = self.deque.iter().position(|t| t.priority > task.priority).unwrap_or(self.deque.len());
|
|
self.deque.insert(pos, task);
|
|
}
|
|
fn pop(&mut self) -> Option<WeaverTask> { let t = self.deque.pop_front(); if t.is_some() { self.processed += 1; } t }
|
|
fn steal(&mut self) -> Option<WeaverTask> { let t = self.deque.pop_back(); if t.is_some() { self.stolen_from += 1; } t }
|
|
fn len(&self) -> usize { self.deque.len() }
|
|
}
|
|
|
|
pub struct ThreadWeaver {
|
|
queues: Vec<Arc<Mutex<WorkerQueue>>>,
|
|
worker_count: usize,
|
|
next_worker: Mutex<usize>,
|
|
}
|
|
|
|
impl ThreadWeaver {
|
|
pub fn new(workers: usize) -> Result<Self, WeaverError> {
|
|
if workers == 0 { return Err(WeaverError::InvalidConfig("Workers must be > 0".into())); }
|
|
let queues: Vec<_> = (0..workers).map(|_| Arc::new(Mutex::new(WorkerQueue::new()))).collect();
|
|
Ok(Self { queues, worker_count: workers, next_worker: Mutex::new(0) })
|
|
}
|
|
|
|
/// Submit task (round-robin assignment)
|
|
pub fn submit(&self, task: WeaverTask) -> Result<usize, WeaverError> {
|
|
let mut next = self.next_worker.lock().map_err(|_| WeaverError::QueueFull("Lock".into()))?;
|
|
let idx = *next % self.worker_count;
|
|
*next = (*next + 1) % self.worker_count;
|
|
drop(next);
|
|
if let Ok(mut q) = self.queues[idx].lock() { q.push(task); Ok(idx) }
|
|
else { Err(WeaverError::QueueFull("Worker lock".into())) }
|
|
}
|
|
|
|
/// Worker picks next task (own queue first, then steal)
|
|
pub fn pick_task(&self, worker_id: usize) -> Option<WeaverTask> {
|
|
// Try own queue first
|
|
if let Ok(mut q) = self.queues[worker_id].lock() {
|
|
if let Some(task) = q.pop() { return Some(task); }
|
|
}
|
|
// Work stealing: try other queues
|
|
for i in 0..self.worker_count {
|
|
if i == worker_id { continue; }
|
|
if let Ok(mut q) = self.queues[i].lock() {
|
|
if q.len() > 1 { // Only steal if other has > 1 task
|
|
if let Some(task) = q.steal() { return Some(task); }
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
/// Get queue stats
|
|
pub fn stats(&self) -> Vec<(usize, u64, u64)> {
|
|
self.queues.iter().enumerate().map(|(_i, q)| {
|
|
if let Ok(q) = q.lock() { (q.len(), q.processed, q.stolen_from) }
|
|
else { (0, 0, 0) }
|
|
}).collect()
|
|
}
|
|
|
|
pub fn total_pending(&self) -> usize {
|
|
self.queues.iter().map(|q| q.lock().map(|q| q.len()).unwrap_or(0)).sum()
|
|
}
|
|
|
|
pub fn total_processed(&self) -> u64 {
|
|
self.queues.iter().map(|q| q.lock().map(|q| q.processed).unwrap_or(0)).sum()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
#[test]
|
|
fn test_submit_pick() {
|
|
let w = ThreadWeaver::new(4).unwrap();
|
|
w.submit(WeaverTask { id: "t1".into(), priority: TaskPriority::Normal, data: vec![], created_at: 0 }).unwrap();
|
|
let task = w.pick_task(0).unwrap();
|
|
assert_eq!(task.id, "t1");
|
|
}
|
|
#[test]
|
|
fn test_priority_ordering() {
|
|
let w = ThreadWeaver::new(1).unwrap();
|
|
w.submit(WeaverTask { id: "low".into(), priority: TaskPriority::Low, data: vec![], created_at: 0 }).unwrap();
|
|
w.submit(WeaverTask { id: "crit".into(), priority: TaskPriority::Critical, data: vec![], created_at: 0 }).unwrap();
|
|
w.submit(WeaverTask { id: "high".into(), priority: TaskPriority::Critical, data: vec![], created_at: 0 }).unwrap();
|
|
let first = w.pick_task(0).unwrap();
|
|
assert_eq!(first.id, "crit"); // Critical first
|
|
}
|
|
#[test]
|
|
fn test_work_stealing() {
|
|
let w = ThreadWeaver::new(2).unwrap();
|
|
// Submit 3 tasks to worker 0
|
|
for i in 0..3 {
|
|
let _idx = w.submit(WeaverTask { id: format!("t{i}"), priority: TaskPriority::Normal, data: vec![], created_at: 0 }).unwrap();
|
|
}
|
|
// Worker 1 has nothing, should steal
|
|
let stolen = w.pick_task(1);
|
|
assert!(stolen.is_some());
|
|
}
|
|
}
|