pub mod mqtt;
pub mod queue;
pub mod scheduler;
pub mod shutdown;
pub mod signal;
pub mod spawn;
pub mod websocket;
pub mod worker;
pub use mqtt::{MqttRuntime, MqttRuntimeConfig};
pub use queue::{QueueConsumer, QueueRuntime, QueueRuntimeConfig};
pub use scheduler::SchedulerRuntime;
pub use shutdown::GracefulShutdown;
pub use signal::shutdown_signal;
pub use spawn::spawn_with_token;
pub use websocket::{WebSocketRuntime, WebSocketRuntimeConfig};
pub use worker::WorkerConfig;
use std::future::Future;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
pub struct SzRuntime {
runtime: tokio::runtime::Runtime,
worker_threads: usize,
shutdown_token: CancellationToken,
}
impl SzRuntime {
pub fn new() -> Self {
Self::with_worker_threads(num_cpus::get())
}
pub fn with_worker_threads(worker_threads: usize) -> Self {
let n = worker_threads.max(1);
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(n)
.enable_all()
.thread_name("sz-rust-worker")
.build()
.expect("Failed to create tokio runtime");
Self {
runtime,
worker_threads: n,
shutdown_token: CancellationToken::new(),
}
}
pub fn worker_threads(&self) -> usize {
self.worker_threads
}
pub fn shutdown_token(&self) -> CancellationToken {
self.shutdown_token.clone()
}
pub fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.runtime.spawn(future)
}
pub fn block_on<F>(&self, future: F) -> F::Output
where
F: Future,
{
self.runtime.block_on(future)
}
pub fn shutdown_timeout(self, timeout: Duration) -> bool {
self.shutdown_token.cancel();
self.runtime.block_on(async {
let _ = tokio::time::timeout(timeout, async {
tokio::time::sleep(Duration::from_millis(10)).await;
})
.await;
});
drop(self.runtime);
true
}
pub fn handle(&self) -> tokio::runtime::Handle {
self.runtime.handle().clone()
}
}
impl Default for SzRuntime {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_default_worker_threads() {
let rt = SzRuntime::new();
assert_eq!(rt.worker_threads(), num_cpus::get());
}
#[test]
fn test_with_worker_threads_custom() {
let rt = SzRuntime::with_worker_threads(2);
assert_eq!(rt.worker_threads(), 2);
}
#[test]
fn test_with_worker_threads_zero_falls_back_to_one() {
let rt = SzRuntime::with_worker_threads(0);
assert_eq!(rt.worker_threads(), 1);
}
#[test]
fn test_spawn_and_block_on() {
let rt = SzRuntime::with_worker_threads(1);
let handle = rt.spawn(async { 42 });
let result = rt.block_on(handle).unwrap();
assert_eq!(result, 42);
}
#[test]
fn test_block_on_directly() {
let rt = SzRuntime::with_worker_threads(1);
let result = rt.block_on(async { 100 });
assert_eq!(result, 100);
}
#[test]
fn test_shutdown_token_cancellation() {
let rt = SzRuntime::with_worker_threads(1);
let token = rt.shutdown_token();
assert!(!token.is_cancelled());
assert!(rt.shutdown_timeout(Duration::from_millis(50)));
assert!(token.is_cancelled());
}
#[test]
fn test_spawn_with_token_cancellation() {
let rt = SzRuntime::with_worker_threads(1);
let token = rt.shutdown_token();
let handle = rt.spawn(async move {
token.cancelled().await;
99
});
let token2 = rt.shutdown_token();
token2.cancel();
let result = rt.block_on(handle).unwrap();
assert_eq!(result, 99);
}
#[test]
fn test_handle_can_spawn() {
let rt = SzRuntime::with_worker_threads(1);
let handle = rt.handle();
let task = handle.spawn(async { 7 });
let result = rt.block_on(task).unwrap();
assert_eq!(result, 7);
}
#[test]
fn test_default_impl_equals_new() {
let rt1 = SzRuntime::default();
let rt2 = SzRuntime::new();
assert_eq!(rt1.worker_threads(), rt2.worker_threads());
}
#[test]
fn test_multiple_runtime_instances() {
let rt1 = SzRuntime::with_worker_threads(1);
let rt2 = SzRuntime::with_worker_threads(1);
let h1 = rt1.spawn(async { 1 });
let h2 = rt2.spawn(async { 2 });
assert_eq!(rt1.block_on(h1).unwrap(), 1);
assert_eq!(rt2.block_on(h2).unwrap(), 2);
}
}