Author: NikolaRHristov

Source/Queue/StealingQueue.rs

@@ -1,123 +1,147 @@
+//! A generic, priority-aware, work-stealing queue implementation.
+//!
+//! This module is self-contained and can be used by any scheduler or
+//! application to manage and distribute tasks of any type that can be
+//! prioritized.
+
 #![allow(non_snake_case, non_camel_case_types)]
 
 use std::sync::Arc;
 
 use crossbeam_deque::{Injector, Stealer, Worker};
 use rand::seq::SliceRandom;
 
+/// Defines a contract for types that can be prioritized by the queue.
 pub trait Prioritized {
-	type P: PartialEq + Eq + Copy;
-
-	fn GetPriority(&self) -> Self::P;
+	/// The type of the priority value used by the implementor.
+	type Kind: PartialEq + Eq + Copy;
+	/// A method to retrieve the priority of the item.
+	fn Rank(&self) -> Self::Kind;
 }
 
+/// Defines the internal priority levels used by the generic queue.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Priority {
 	High,
-
 	Normal,
-
 	Low,
 }
 
+/// Holds the queue components that are safe to share across all threads.
+///
+/// This includes global injectors for submitting new tasks and stealers for
+/// taking tasks from other workers, organized by priority level.
 struct Share<T> {
+	/// Global, multi-producer queues for each priority.
 	Injector:(Injector<T>, Injector<T>, Injector<T>),
-
+	/// Share handles for stealing tasks from each worker's queue.
 	Stealer:(Vec<Stealer<T>>, Vec<Stealer<T>>, Vec<Stealer<T>>),
 }
 
-pub struct StealingQueue<T:Prioritized<P = Priority>> {
+/// A generic, priority-aware, work-stealing queue.
+///
+/// This is the public-facing entry point for submitting tasks. It is generic
+/// over any task type `T` that implements the `Prioritized` trait.
+pub struct StealingQueue<T:Prioritized<Kind = Priority>> {
+	/// A shared, thread-safe pointer to the queue's shared components.
 	Share:Arc<Share<T>>,
 }
 
+/// Contains all necessary components for a single worker thread to operate.
+///
+/// This includes the thread-local `Worker` deques, which are not safe to share,
+/// making this context object the sole owner of a worker's private queues.
 pub struct Context<T> {
+	/// A unique identifier for the worker, used to avoid self-stealing.
 	pub Identifier:usize,
-
+	/// Thread-local work queues for each priority level.
 	Local:(Worker<T>, Worker<T>, Worker<T>),
-
+	/// A reference to the shared components of the entire queue system.
 	Share:Arc<Share<T>>,
 }
 
-impl<T:Prioritized<P = Priority>> StealingQueue<T> {
-	pub fn New(Count:usize) -> (Self, Vec<Context<T>>) {
+impl<T:Prioritized<Kind = Priority>> StealingQueue<T> {
+	/// Creates a complete work-stealing queue system.
+	///
+	/// This function initializes all the necessary queues, both shared and
+	/// thread-local, for a given number of workers.
+	///
+	/// Returns a tuple containing:
+	/// 1. The public-facing `StealingQueue` for submitting new tasks.
+	/// 2. A `Vec` of `Context` objects, one for each worker thread to own.
+	pub fn Create(Count:usize) -> (Self, Vec<Context<T>>) {
 		let mut High:Vec<Worker<T>> = Vec::with_capacity(Count);
-
 		let mut Normal:Vec<Worker<T>> = Vec::with_capacity(Count);
-
 		let mut Low:Vec<Worker<T>> = Vec::with_capacity(Count);
 
+		// For each priority level, create a thread-local worker queue and its
+		// corresponding shared stealer.
 		let StealerHigh:Vec<Stealer<T>> = (0..Count)
 			.map(|_| {
 				let Worker = Worker::new_fifo();
-
 				let Stealer = Worker.stealer();
-
 				High.push(Worker);
-
 				Stealer
 			})
 			.collect();
 
 		let StealerNormal:Vec<Stealer<T>> = (0..Count)
 			.map(|_| {
 				let Worker = Worker::new_fifo();
-
 				let Stealer = Worker.stealer();
-
 				Normal.push(Worker);
-
 				Stealer
 			})
 			.collect();
 
 		let StealerLow:Vec<Stealer<T>> = (0..Count)
 			.map(|_| {
 				let Worker = Worker::new_fifo();
-
 				let Stealer = Worker.stealer();
-
 				Low.push(Worker);
-
 				Stealer
 			})
 			.collect();
 
-		let Shared = Arc::new(Share {
+		// Bundle all shared components into an Arc for safe sharing.
+		let Share = Arc::new(Share {
 			Injector:(Injector::new(), Injector::new(), Injector::new()),
-
 			Stealer:(StealerHigh, StealerNormal, StealerLow),
 		});
 
-		let mut Context = Vec::with_capacity(Count);
-
-		for Id in 0..Count {
-			Context.push(Context {
-				Identifier:Id,
-
+		// Create a unique context for each worker, giving it ownership of its
+		// local queues and a reference to the shared components.
+		let mut Contexts = Vec::with_capacity(Count);
+		for Identifier in 0..Count {
+			Contexts.push(Context {
+				Identifier,
 				Local:(High.remove(0), Normal.remove(0), Low.remove(0)),
-
-				Share:Shared.clone(),
+				Share:Share.clone(),
 			});
 		}
 
-		let Queue = Self { Share:Shared };
-
-		(Queue, Context)
+		let Queue = Self { Share };
+		(Queue, Contexts)
 	}
 
+	/// Submits a new task to the appropriate global queue based on its
+	/// priority. This method is thread-safe and can be called from any
+	/// context.
 	pub fn Submit(&self, Task:T) {
-		match Task.GetPriority() {
+		match Task.Rank() {
 			Priority::High => self.Share.Injector.0.push(Task),
-
 			Priority::Normal => self.Share.Injector.1.push(Task),
-
 			Priority::Low => self.Share.Injector.2.push(Task),
 		}
 	}
 }
 
 impl<T> Context<T> {
-	pub fn NextTask(&self) -> Option<T> {
+	/// Finds the next available task for the worker to execute.
+	// This method implements the complete work-finding logic:
+	/// 1. Check local queue (high to low priority).
+	/// 2. Steal from the system (high to low priority).
+	pub fn Next(&self) -> Option<T> {
 		self.Local
 			.0
 			.pop()
@@ -128,25 +152,27 @@ impl<T> Context<T> {
 			.or_else(|| self.Steal(&self.Share.Injector.2, &self.Share.Stealer.2, &self.Local.2))
 	}
 
-	fn Steal<'a>(&self, Injector:&'a Injector<T>, Stealer:&'a [Stealer<T>], Local:&'a Worker<T>) -> Option<T> {
+	/// Attempts to steal a task from a specific priority set.
+	///
+	/// It first tries to steal a batch from the global injector queue. If that
+	/// fails, it attempts to steal from a randomly chosen peer worker to ensure
+	/// fair distribution and avoid contention hotspots.
+	fn Steal<'a>(&self, Injector:&'a Injector<T>, Stealers:&'a [Stealer<T>], Local:&'a Worker<T>) -> Option<T> {
 		if Injector.steal_batch_and_pop(Local).is_success() {
 			return Local.pop();
 		}
 
-		let mut Index:Vec<usize> = (0..Stealer.len()).collect();
+		let mut Indices:Vec<usize> = (0..Stealers.len()).collect();
+		Indices.shuffle(&mut rand::rng());
 
-		Index.shuffle(&mut rand::rng());
-
-		for i in Index {
-			if i == self.Identifier {
+		for Index in Indices {
+			if Index == self.Identifier {
 				continue;
 			}
-
-			if Stealer[i].steal_batch_and_pop(Local).is_success() {
+			if Stealers[Index].steal_batch_and_pop(Local).is_success() {
 				return Local.pop();
 			}
 		}
-
 		None
 	}
 }

Source/Queue/mod.rs

@@ -1,3 +1,10 @@
+//! Declares the `StealingQueue` module.
+//!
+//! This module encapsulates all logic for the generic, priority-aware,
+//! work-stealing queue system, which is the foundational component of the
+//! scheduler.
+
 #![allow(non_snake_case, non_camel_case_types)]
 
+/// Provides the generic, priority-aware, work-stealing queue implementation.
 pub mod StealingQueue;

Source/Scheduler/Scheduler.rs

@@ -1,3 +1,7 @@
+//! Manages the pool of workers and the task queue system.
+
+#![allow(non_snake_case, non_camel_case_types)]
+
 use std::{
 	collections::HashMap,
 	future::Future,
@@ -10,74 +14,86 @@ use std::{
 use log::{error, info, warn};
 use tokio::task::JoinHandle;
 
-use super::Worker::Worker;
+use super::{SchedulerBuilder::Concurrency, Worker::Worker};
 use crate::{
-	Queue::StealingQueue::StealingQueue as StealingQueueStruct,
-	Scheduler::SchedulerBuilder::Concurrency as ConcurrencyEnum,
-	Task::{Priority::Enum as PriorityEnum, Task::Struct as TaskStruct},
+	Queue::StealingQueue::StealingQueue,
+	Task::{Priority::Priority, Task::Task},
 };
 
+/// Manages a pool of worker threads and a work-stealing queue to execute
+/// tasks efficiently. This is the primary public-facing struct of the library.
 pub struct Scheduler {
-	Queue:StealingQueueStruct<TaskStruct>,
-
+	/// The underlying work-stealing queue system used for task submission.
+	Queue:StealingQueue<Task>,
+	/// Handles to the spawned worker threads, used for graceful shutdown.
 	Handle:Vec<JoinHandle<()>>,
-
+	/// An atomic flag to signal all workers to shut down.
 	Running:Arc<AtomicBool>,
 }
 
 impl Scheduler {
-	pub fn Start(number_of_workers:usize, _queue_configs:HashMap<String, ConcurrencyEnum>) -> Self {
-		info!("[Scheduler] Starting scheduler with {} worker threads.", number_of_workers);
-
+	/// Creates and starts a new scheduler with a given configuration.
+	///
+	/// This is a crate-private function, intended to be called only by the
+	/// `SchedulerBuilder`'s `Build` method.
+	pub(crate) fn Create(Count:usize, _Configuration:HashMap<String, Concurrency>) -> Self {
+		info!("[Scheduler] Create with {} workers.", Count);
 		let Running = Arc::new(AtomicBool::new(true));
 
-		let (Queue, WorkerContexts) = StealingQueueStruct::<TaskStruct>::New(number_of_workers);
-
-		let mut Handle = Vec::with_capacity(number_of_workers);
+		// Create the entire queue system and retrieve the contexts for each worker.
+		let (Queue, Contexts) = StealingQueue::<Task>::Create(Count);
 
-		for Context in WorkerContexts.into_iter() {
+		let mut Handle = Vec::with_capacity(Count);
+		// Spawn an asynchronous task for each worker.
+		for Context in Contexts.into_iter() {
 			let Running = Running.clone();
-
 			Handle.push(tokio::spawn(async move {
-				let WorkerInstance = Worker::New(Context, Running);
-
-				WorkerInstance.Run().await;
+				// Each task creates and runs a worker, consuming its context.
+				Worker::Create(Context, Running).Run().await;
 			}));
 		}
 
 		Self { Queue, Handle, Running }
 	}
 
-	pub fn Submit<F>(&self, future_instance:F, task_priority:PriorityEnum)
+	/// Submits a new task to the scheduler's global queue.
+	///
+	/// The task will be picked up by the next available worker according to its
+	/// priority and the work-stealing logic.
+	pub fn Submit<F>(&self, Operation:F, Priority:Priority)
 	where
 		F: Future<Output = ()> + Send + 'static, {
-		self.Queue.Submit(TaskStruct::New(future_instance, task_priority));
+		self.Queue.Submit(Task::Create(Operation, Priority));
 	}
 
-	pub async fn ShutDown(&mut self) {
+	/// Asynchronously shuts down the scheduler.
+	///
+	/// This method signals all worker threads to stop their loops and then
+	/// waits for each one to complete its current task and exit gracefully.
+	pub async fn Stop(&mut self) {
 		if !self.Running.swap(false, Ordering::Relaxed) {
-			info!("[Scheduler] ShutDown already initiated.");
-
+			info!("[Scheduler] Stop already initiated.");
 			return;
 		}
 
-		info!("[Scheduler] Shutting down worker threads...");
-
-		for handle in self.Handle.drain(..) {
-			if let Err(e) = handle.await {
-				error!("[Scheduler] Error joining worker task during shutdown: {}", e);
+		info!("[Scheduler] Stopping worker threads...");
+		for Handle in self.Handle.drain(..) {
+			if let Err(Error) = Handle.await {
+				error!("[Scheduler] Error joining worker: {}", Error);
 			}
 		}
-
-		info!("[Scheduler] All workers shut down successfully.");
+		info!("[Scheduler] All workers stopped.");
 	}
 }
 
 impl Drop for Scheduler {
+	/// Ensures workers are signaled to stop if the `Scheduler` is dropped.
+	///
+	/// This prevents orphaned worker threads if the user forgets to call the
+	/// explicit `Stop` method.
 	fn drop(&mut self) {
 		if self.Running.load(Ordering::Relaxed) {
-			warn!("[Scheduler] Scheduler dropped without explicit shutdown. Signaling workers to stop.");
-
+			warn!("[Scheduler] Dropped without explicit stop. Signaling workers.");
 			self.Running.store(false, Ordering::Relaxed);
 		}
 	}