use crate::plugin::Plugin;
use crate::resource::Resources;
use crate::task_pool::TaskPool;
#[derive(Default)]
pub struct AsyncRuntimePlugin {
pub num_threads: Option<usize>,
}
impl AsyncRuntimePlugin {
pub fn new() -> Self {
Self::default()
}
pub fn with_threads(mut self, num_threads: usize) -> Self {
self.num_threads = Some(num_threads);
self
}
}
impl Plugin for AsyncRuntimePlugin {
type Dependencies = ();
fn name(&self) -> &'static str {
"AsyncRuntimePlugin"
}
fn build(&self, resources: &mut Resources) {
let pool = match self.num_threads {
Some(n) => TaskPool::new(n),
None => TaskPool::default_threads(),
};
tracing::debug!(
"AsyncRuntimePlugin: Created TaskPool with {} threads",
pool.thread_count()
);
resources.insert(pool);
}
fn cleanup(&self, resources: &mut Resources) {
if let Some(pool) = resources.remove::<TaskPool>() {
tracing::debug!("AsyncRuntimePlugin: Shutting down TaskPool");
pool.shutdown();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::EngineBuilder;
#[test]
fn test_async_runtime_plugin_registers_task_pool() {
let engine = EngineBuilder::new()
.add_plugin(AsyncRuntimePlugin::default())
.build();
assert!(engine.get::<TaskPool>().is_some());
}
#[test]
fn test_async_runtime_plugin_with_custom_threads() {
let engine = EngineBuilder::new()
.add_plugin(AsyncRuntimePlugin::new().with_threads(2))
.build();
let pool = engine.get::<TaskPool>().unwrap();
assert_eq!(pool.thread_count(), 2);
}
#[test]
fn test_task_pool_spawn() {
let engine = EngineBuilder::new()
.add_plugin(AsyncRuntimePlugin::default())
.build();
let pool = engine.get::<TaskPool>().unwrap();
let task = pool.spawn(async { 42 });
let result = pollster::block_on(task);
assert_eq!(result, 42);
}
}