use super::{
driver::Driver,
idle::Searcher,
mailbox::{Cancel, Forward, Mailbox, Message},
pool::{MAX_WORKERS, Pool, Threads},
request::{RequestOutput, RetiredResources},
sleep::{Sleep, Timers},
spinner::{Config as SpinnerConfig, Spinner},
task::{AfterPoll, Ready, RootWaker, Runnable},
tasks::Tasks,
timeout::TimeoutWheel,
waker::SUBMISSION_SEQ_MASK,
};
#[cfg(feature = "external")]
use crate::Pacer;
use crate::{
BlobLayout, BlobVersion, Clock, Error, Execution, Handle, METRICS_PREFIX, Name, SinkOf,
Spawner as _, StreamOf, Supervisor as _, child_label,
network::{
iouring::{Config as NetworkConfig, Network},
metered::Network as MeteredNetwork,
},
prefixed_name,
process::metered::Metrics as ProcessMetrics,
signal::Signal,
storage::{
iouring::{Config as StorageConfig, Storage},
metered::Storage as MeteredStorage,
},
telemetry::metrics::{
CounterFamily, Gauge, GaugeFamily, Metric, Register, Registered, Registry, add_attribute,
raw, task::Label, validate_label,
},
utils::{self, FactoryGuard, MetricHandle, Panicker, signal::Stopper, supervision::Tree},
};
use commonware_macros::select;
use commonware_parallel::{Rayon, rayon::ThreadPoolBuilder};
use commonware_utils::{
channel::oneshot,
sync::{Condvar, Mutex},
sys_rng,
};
use governor::clock::{Clock as GClock, ReasonablyRealtime};
use rand_core::{Rng, TryCryptoRng, TryRng};
use std::{
any::Any,
cell::RefCell,
convert::Infallible,
env,
future::{Future, poll_fn},
mem,
net::{IpAddr, SocketAddr, ToSocketAddrs},
num::{NonZeroU64, NonZeroUsize},
ops::RangeInclusive,
panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
path::PathBuf,
pin::{Pin, pin},
ptr,
rc::Rc,
sync::{Arc, Weak},
task::{Context as TaskContext, Poll, Waker},
time::{Duration, Instant, SystemTime},
};
const BATCH_SIZE: usize = 64;
type BoxedTask = Pin<Box<dyn Future<Output = ()> + Send>>;
#[derive(Clone, Debug)]
pub struct RingConfig {
pub size: u32,
pub timeout_wheel_tick: Duration,
}
impl Default for RingConfig {
fn default() -> Self {
Self {
size: 128,
timeout_wheel_tick: Duration::from_millis(5),
}
}
}
#[derive(Clone)]
pub struct Config {
worker_threads: usize,
global_queue_interval: u32,
ring_config: RingConfig,
idle_spinner: SpinnerConfig,
thread_stack_size: usize,
catch_panics: bool,
storage_directory: PathBuf,
storage_blob_layouts: RangeInclusive<BlobLayout>,
tcp_nodelay: Option<bool>,
zero_linger: bool,
connect_timeout: Duration,
read_write_timeout: Duration,
read_buffer_size: usize,
}
impl Config {
pub fn new() -> Self {
let ring_config = RingConfig {
size: if cfg!(test) { 128 } else { 1024 },
..RingConfig::default()
};
let suffix = sys_rng().next_u64();
Self {
worker_threads: 2,
global_queue_interval: 31,
ring_config,
idle_spinner: SpinnerConfig::default(),
thread_stack_size: utils::thread::system_thread_stack_size(),
catch_panics: false,
storage_directory: env::temp_dir().join(format!("commonware_iouring_runtime_{suffix}")),
storage_blob_layouts: BlobLayout::ALL,
tcp_nodelay: Some(true),
zero_linger: true,
connect_timeout: Duration::from_secs(10),
read_write_timeout: Duration::from_secs(60),
read_buffer_size: 64 * 1024,
}
}
pub const fn with_worker_threads(mut self, workers: usize) -> Self {
self.worker_threads = workers;
self
}
pub const fn with_global_queue_interval(mut self, interval: u32) -> Self {
self.global_queue_interval = interval;
self
}
pub const fn with_ring_config(mut self, config: RingConfig) -> Self {
self.ring_config = config;
self
}
pub const fn with_idle_spinner(mut self, config: SpinnerConfig) -> Self {
self.idle_spinner = config;
self
}
pub const fn with_thread_stack_size(mut self, size: usize) -> Self {
self.thread_stack_size = size;
self
}
pub const fn with_catch_panics(mut self, catch: bool) -> Self {
self.catch_panics = catch;
self
}
pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = timeout;
self
}
pub const fn with_read_write_timeout(mut self, timeout: Duration) -> Self {
self.read_write_timeout = timeout;
self
}
pub const fn with_tcp_nodelay(mut self, enabled: Option<bool>) -> Self {
self.tcp_nodelay = enabled;
self
}
pub const fn with_zero_linger(mut self, enabled: bool) -> Self {
self.zero_linger = enabled;
self
}
pub const fn with_read_buffer_size(mut self, size: usize) -> Self {
self.read_buffer_size = size;
self
}
pub fn with_storage_directory(mut self, directory: impl Into<PathBuf>) -> Self {
self.storage_directory = directory.into();
self
}
pub fn with_storage_blob_layouts(mut self, layouts: RangeInclusive<BlobLayout>) -> Self {
assert!(
!layouts.is_empty(),
"storage blob layouts must be non-empty"
);
self.storage_blob_layouts = layouts;
self
}
pub const fn worker_threads(&self) -> usize {
self.worker_threads
}
pub const fn global_queue_interval(&self) -> u32 {
self.global_queue_interval
}
pub const fn ring_config(&self) -> &RingConfig {
&self.ring_config
}
pub const fn idle_spinner(&self) -> &SpinnerConfig {
&self.idle_spinner
}
pub const fn thread_stack_size(&self) -> usize {
self.thread_stack_size
}
pub const fn catch_panics(&self) -> bool {
self.catch_panics
}
pub const fn connect_timeout(&self) -> Duration {
self.connect_timeout
}
pub const fn read_write_timeout(&self) -> Duration {
self.read_write_timeout
}
pub const fn tcp_nodelay(&self) -> Option<bool> {
self.tcp_nodelay
}
pub const fn zero_linger(&self) -> bool {
self.zero_linger
}
pub const fn read_buffer_size(&self) -> usize {
self.read_buffer_size
}
pub const fn storage_directory(&self) -> &PathBuf {
&self.storage_directory
}
pub const fn storage_blob_layouts(&self) -> &RangeInclusive<BlobLayout> {
&self.storage_blob_layouts
}
fn validate(&mut self) {
assert!(
(1..=MAX_WORKERS).contains(&self.worker_threads),
"worker threads ({}) must be between 1 and {MAX_WORKERS}",
self.worker_threads,
);
assert!(
self.global_queue_interval != 0,
"global queue interval must be nonzero"
);
assert!(self.ring_config.size != 0, "ring size must be nonzero");
self.ring_config.size = self
.ring_config
.size
.checked_next_power_of_two()
.expect("ring size overflow");
assert!(self.ring_config.size <= 32_768, "ring size exceeds 32768");
assert!(
self.idle_spinner.budget_us <= self.idle_spinner.max_budget_us,
"spinner budget_us ({}) must not exceed max_budget_us ({})",
self.idle_spinner.budget_us,
self.idle_spinner.max_budget_us,
);
assert!(
!self.storage_blob_layouts.is_empty(),
"storage blob layouts must be non-empty"
);
assert!(
!self.connect_timeout.is_zero(),
"connect timeout must be nonzero"
);
assert!(
!self.read_write_timeout.is_zero(),
"read/write timeout must be nonzero"
);
TimeoutWheel::validate_layout(self.max_timeout(), self.ring_config.timeout_wheel_tick)
.expect("invalid io_uring operation deadline layout");
}
fn max_timeout(&self) -> Duration {
self.connect_timeout.max(self.read_write_timeout)
}
}
impl Default for Config {
fn default() -> Self {
Self::new()
}
}
struct TaskMetrics {
tasks_spawned: CounterFamily<Label>,
tasks_running: GaugeFamily<Label>,
}
impl TaskMetrics {
fn new(registry: &mut impl Register) -> Self {
Self {
tasks_spawned: registry.register(
"tasks_spawned",
"Total number of tasks spawned",
raw::Family::default(),
),
tasks_running: registry.register(
"tasks_running",
"Number of tasks currently running",
raw::Family::default(),
),
}
}
}
#[derive(Default)]
struct Workers {
state: Mutex<WorkerCount>,
idle: Condvar,
}
#[derive(Default)]
struct WorkerCount {
closed: bool,
active: usize,
}
impl Workers {
fn reserve(self: &Arc<Self>) -> Option<ActiveWorker> {
let mut state = self.state.lock();
if state.closed {
return None;
}
state.active += 1;
Some(ActiveWorker(self.clone()))
}
fn close(&self) {
self.state.lock().closed = true;
}
fn wait(&self) {
let mut state = self.state.lock();
while state.active != 0 {
self.idle.wait(&mut state);
}
}
}
struct ActiveWorker(
Arc<Workers>,
);
impl Drop for ActiveWorker {
fn drop(&mut self) {
let mut state = self.0.state.lock();
state.active -= 1;
if state.active == 0 {
self.0.idle.notify_all();
}
drop(state);
#[cfg(test)]
tests::after_release(&self.0);
}
}
pub struct Shared {
pub cfg: Config,
pub pool: Arc<Pool>,
registry: Registry,
metrics: TaskMetrics,
pending_operations: Gauge,
shutdown: Mutex<Stopper>,
panicker: Panicker,
workers: Arc<Workers>,
tasks: Tasks,
storage: MeteredStorage<Storage>,
network: MeteredNetwork<Network>,
}
struct Launch {
task: BoxedTask,
shared: Arc<Shared>,
active: ActiveWorker,
}
impl Shared {
fn launch(self: &Arc<Self>, task: BoxedTask, active: ActiveWorker) {
let payload = Launch {
task,
shared: self.clone(),
active,
};
#[cfg(test)]
let payload = tests::before_launch(payload);
utils::thread::spawn(self.cfg.thread_stack_size, move || {
Worker::run_task(payload)
});
}
}
pub struct Context {
name: String,
attributes: Vec<(String, String)>,
shared: Arc<Shared>,
tree: Arc<Tree>,
execution: Execution,
}
impl Context {
fn metrics(&self) -> &TaskMetrics {
&self.shared.metrics
}
}
impl crate::Spawner for Context {
fn dedicated(mut self) -> Self {
self.execution = Execution::Dedicated;
self
}
fn shared(mut self, blocking: bool) -> Self {
self.execution = Execution::Shared(blocking);
self
}
fn spawn<F, Fut, T>(mut self, f: F) -> Handle<T>
where
F: FnOnce(Self) -> Fut + Send + 'static,
Fut: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let (_, metric) = spawn_metrics!(self);
let parent = self.tree.clone();
let execution = self.execution;
self.execution = Execution::default();
let (child, aborted) = Tree::child(&parent);
if aborted {
return Handle::closed(metric);
}
self.tree = child;
let shared = self.shared.clone();
let reservation = if matches!(execution, Execution::Dedicated | Execution::Shared(true)) {
let Some(reservation) = shared.workers.reserve() else {
return Handle::closed(metric);
};
Some(reservation)
} else {
if shared.tasks.is_closed() {
return Handle::closed(metric);
}
None
};
let guard = FactoryGuard::new(&parent, metric);
let future = f(self);
let (future, handle) = Handle::init(
future,
guard.disarm(),
shared.panicker.clone(),
parent.clone(),
);
if let Some(reservation) = reservation {
shared.launch(Box::pin(future), reservation);
} else if let Err(task) = shared.tasks.register(future, Arc::downgrade(&shared.pool)) {
parent.abort();
task.clear();
}
handle
}
async fn stop(self, value: i32, timeout: Option<Duration>) -> Result<(), Error> {
let resolved = self.shared.shutdown.lock().stop(value);
let timeout = timeout.map_or_else(
|| futures::future::Either::Right(futures::future::pending()),
|duration| futures::future::Either::Left(self.sleep(duration)),
);
select! {
result = resolved => result.map_err(|_| Error::Closed),
_ = timeout => Err(Error::Timeout),
}
}
fn stopped(&self) -> Signal {
self.shared.shutdown.lock().stopped()
}
}
impl crate::Strategizer for Context {
fn strategy(&self, parallelism: NonZeroUsize) -> Rayon {
let pool = ThreadPoolBuilder::new()
.num_threads(parallelism.get())
.stack_size(self.shared.cfg.thread_stack_size)
.build()
.expect("failed to create io_uring Rayon thread pool");
Rayon::with_pool(Arc::new(pool))
}
}
impl crate::Supervisor for Context {
fn child(&self, label: &'static str) -> Self {
let (tree, _) = Tree::child(&self.tree);
Self {
name: child_label(&self.name, label),
attributes: self.attributes.clone(),
shared: self.shared.clone(),
tree,
execution: Execution::default(),
}
}
fn with_attribute(mut self, key: &'static str, value: impl std::fmt::Display) -> Self {
validate_label(key);
add_attribute(&mut self.attributes, key, value);
self
}
fn name(&self) -> Name {
Name {
label: self.name.clone(),
attributes: self.attributes.clone(),
}
}
}
impl crate::Metrics for Context {
fn register<N: Into<String>, H: Into<String>, M: Metric>(
&self,
name: N,
help: H,
metric: M,
) -> Registered<M> {
self.shared.registry.register(
prefixed_name(&self.name, &name.into()),
help.into(),
self.attributes.clone(),
Arc::new(metric),
)
}
fn encode(&self) -> String {
self.shared.registry.encode()
}
}
impl Clock for Context {
fn current(&self) -> SystemTime {
SystemTime::now()
}
fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static + use<> {
Sleep::new(duration)
}
fn sleep_until(
&self,
deadline: SystemTime,
) -> impl Future<Output = ()> + Send + 'static + use<> {
Sleep::until(deadline)
}
}
#[cfg(feature = "external")]
impl Pacer for Context {
fn pace<'a, F, T>(
&'a self,
_latency: Duration,
future: F,
) -> impl Future<Output = T> + Send + 'a
where
F: Future<Output = T> + Send + 'a,
T: Send + 'a,
{
future
}
}
impl GClock for Context {
type Instant = SystemTime;
fn now(&self) -> Self::Instant {
self.current()
}
}
impl ReasonablyRealtime for Context {}
impl crate::Network for Context {
type Listener = <MeteredNetwork<Network> as crate::Network>::Listener;
async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
self.shared.network.bind(socket).await
}
async fn dial(&self, socket: SocketAddr) -> Result<(SinkOf<Self>, StreamOf<Self>), Error> {
self.shared.network.dial(socket).await
}
}
impl crate::Resolver for Context {
async fn resolve(&self, host: &str) -> Result<Vec<IpAddr>, Error> {
let host = host.to_owned();
self.child("resolver")
.shared(true)
.spawn(move |_| async move {
(host.as_str(), 0)
.to_socket_addrs()
.map(|addresses| addresses.map(|address| address.ip()).collect())
.map_err(|error| Error::ResolveFailed(error.to_string()))
})
.await
.map_err(|error| Error::ResolveFailed(error.to_string()))?
}
}
impl TryRng for Context {
type Error = Infallible;
fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
Ok(sys_rng().next_u32())
}
fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
Ok(sys_rng().next_u64())
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
sys_rng().fill_bytes(dest);
Ok(())
}
}
impl TryCryptoRng for Context {}
impl crate::Storage for Context {
type Blob = <MeteredStorage<Storage> as crate::Storage>::Blob;
async fn open_versioned(
&self,
partition: &str,
name: &[u8],
versions: RangeInclusive<BlobVersion>,
) -> Result<(Self::Blob, u64, BlobVersion), Error> {
self.shared
.storage
.open_versioned(partition, name, versions)
.await
}
async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
self.shared.storage.remove(partition, name).await
}
async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
self.shared.storage.scan(partition).await
}
}
type Completion = (oneshot::Sender<Result<(), Error>>, Result<(), Error>);
type OperationResult = (
oneshot::Sender<Result<RequestOutput, Error>>,
Result<RequestOutput, Error>,
);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Role {
Pool(u32),
Dedicated,
}
impl Role {
pub const fn is_pool(self) -> bool {
matches!(self, Self::Pool(_))
}
pub const fn owns_runner(self) -> bool {
matches!(self, Self::Pool(0))
}
}
struct Seat {
index: u32,
tasks: NonZeroU64,
ready: Ready,
}
pub struct Local {
seat: Option<Seat>,
pub driver: Option<Driver>,
pub timers: Timers,
pub closing: bool,
pub now: Instant,
pub root_ready: bool,
pub mailbox: Arc<Mailbox>,
pub deferred: Deferred,
shared: Arc<Shared>,
reported_pending: usize,
}
impl Local {
fn new(shared: Arc<Shared>, role: Role) -> std::io::Result<Self> {
#[cfg(test)]
tests::before_startup(&shared.workers)?;
let mailbox = match role {
Role::Pool(index) => shared.pool.mailbox(index).clone(),
Role::Dedicated => Arc::new(Mailbox::new()?),
};
let now = Instant::now();
let driver = Driver::new(
&shared.cfg.ring_config,
shared.cfg.max_timeout(),
mailbox.waker.clone(),
now,
)?;
let seat = match role {
Role::Pool(index) => Some(Seat {
index,
tasks: shared.tasks.id(),
ready: shared.pool.queue(index).owner(),
}),
Role::Dedicated => None,
};
Ok(Self {
seat,
driver: Some(driver),
timers: Timers::default(),
closing: false,
now,
root_ready: true,
mailbox,
deferred: Deferred::default(),
shared,
reported_pending: 0,
})
}
pub fn current() -> Option<Rc<RefCell<Self>>> {
CURRENT
.try_with(|current| current.borrow().clone())
.ok()
.flatten()
}
#[inline(always)]
pub fn owner(mailbox: &Weak<Mailbox>) -> Option<Rc<RefCell<Self>>> {
let local = Self::current()?;
let matches = ptr::eq(Arc::as_ptr(&local.borrow().mailbox), mailbox.as_ptr());
matches.then_some(local)
}
#[inline]
pub fn with_current<R>(f: impl FnOnce(Option<&mut Self>) -> R) -> R {
let mut f = Some(f);
let result = CURRENT.try_with(|current| {
let current = current.borrow();
let local = current.as_ref()?;
let f = f.take().expect("the closure must be taken only once");
Some(f(Some(&mut local.borrow_mut())))
});
if let Ok(Some(result)) = result {
return result;
}
let f = f.take().expect("the closure must be taken only once");
f(None)
}
#[inline]
pub fn serves(&self, owner: NonZeroU64) -> bool {
self.seat.as_ref().is_some_and(|seat| seat.tasks == owner) && !self.closing
}
pub const fn role(&self) -> Role {
match &self.seat {
Some(seat) => Role::Pool(seat.index),
None => Role::Dedicated,
}
}
pub fn workers(&self) -> usize {
self.shared.pool.workers()
}
pub const fn ready(&self) -> &Ready {
&self
.seat
.as_ref()
.expect("the worker must be a pool worker")
.ready
}
const fn ready_mut(&mut self) -> &mut Ready {
&mut self
.seat
.as_mut()
.expect("the worker must be a pool worker")
.ready
}
#[inline]
pub fn enqueue(&mut self, runnable: Runnable) {
let Self { seat, shared, .. } = self;
let seat = seat.as_mut().expect("the worker must be a pool worker");
shared.pool.enqueue(&mut seat.ready, runnable);
}
pub fn cancel(mailbox: &Weak<Mailbox>, cancel: Cancel) {
if let Some(local) = Self::owner(mailbox) {
local.borrow_mut().release(cancel);
return;
}
if let Some(mailbox) = mailbox.upgrade() {
let _ = mailbox.send(Message::Cancel(cancel));
}
}
fn release(&mut self, cancel: Cancel) {
match cancel {
Cancel::Waiter(id) => self.driver.as_mut().unwrap().orphan(id, &mut self.deferred),
Cancel::Timer(id) => self.timers.cancel(id, &mut self.deferred),
}
}
fn update_pending(&mut self) {
let pending = self.driver.as_ref().unwrap().len();
if pending > self.reported_pending {
self.shared
.pending_operations
.inc_by((pending - self.reported_pending) as _);
} else if pending < self.reported_pending {
self.shared
.pending_operations
.dec_by((self.reported_pending - pending) as _);
}
self.reported_pending = pending;
}
fn is_ready(&self) -> bool {
let seat = self.seat.as_ref();
seat.is_some_and(|seat| !seat.ready.is_empty())
|| self.root_ready
|| !self.deferred.is_empty()
|| (seat.is_some() && self.shared.pool.has_global())
}
fn next_deadline(&mut self) -> Option<Instant> {
[
self.driver.as_mut().unwrap().next_deadline(),
self.timers.next_deadline(),
]
.into_iter()
.flatten()
.min()
}
}
thread_local! {
static CURRENT: RefCell<Option<Rc<RefCell<Local>>>> = const { RefCell::new(None) };
}
struct Scope;
impl Scope {
fn assert_vacant() {
CURRENT.with(|current| {
assert!(
current.borrow().is_none(),
"nested io_uring runtime entry is not supported"
);
});
}
fn install(local: Rc<RefCell<Local>>) -> Self {
CURRENT.with(|current| {
let mut current = current.borrow_mut();
assert!(current.is_none(), "nested io_uring worker scope");
*current = Some(local);
});
Self
}
}
impl Drop for Scope {
fn drop(&mut self) {
CURRENT.with(|current| {
current.borrow_mut().take();
});
}
}
pub type Panic = Box<dyn Any + Send + 'static>;
#[derive(Default)]
pub struct Panics {
first: Option<Panic>,
}
impl Panics {
pub fn contain<T>(f: impl FnOnce() -> T) -> Option<T> {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
let result: std::thread::Result<T> = Ok(f());
} else {
let result = catch_unwind(AssertUnwindSafe(f));
}
}
match result {
Ok(output) => Some(output),
Err(panic) => {
if let Err(secondary) = catch_unwind(AssertUnwindSafe(|| drop(panic))) {
mem::forget(secondary);
}
None
}
}
}
fn run(&mut self, f: impl FnOnce()) {
if let Err(panic) = catch_unwind(AssertUnwindSafe(f)) {
self.retain(panic);
}
}
fn retain(&mut self, panic: Panic) {
if self.first.is_none() {
self.first = Some(panic);
} else {
mem::forget(panic);
}
}
fn take(&mut self) -> Option<Panic> {
self.first.take()
}
}
impl Drop for Panics {
fn drop(&mut self) {
mem::forget(self.first.take());
}
}
struct RetirementGuard;
impl Drop for RetirementGuard {
fn drop(&mut self) {
std::process::abort();
}
}
#[derive(Default)]
pub struct Deferred {
pub wakes: Vec<Waker>,
pub drops: Vec<Waker>,
pub outputs: Vec<RequestOutput>,
pub resources: Vec<RetiredResources>,
pub completions: Vec<Completion>,
pub results: Vec<OperationResult>,
}
impl Deferred {
pub const fn is_empty(&self) -> bool {
self.wakes.is_empty()
&& self.drops.is_empty()
&& self.outputs.is_empty()
&& self.resources.is_empty()
&& self.completions.is_empty()
&& self.results.is_empty()
}
fn run(&mut self, panics: &mut Panics) {
for waker in self.wakes.drain(..) {
panics.run(|| waker.wake());
}
for waker in self.drops.drain(..) {
panics.run(|| drop(waker));
}
for output in self.outputs.drain(..) {
panics.run(|| drop(output));
}
for resources in self.resources.drain(..) {
panics.run(|| drop(resources));
}
for (sender, output) in self.completions.drain(..) {
panics.run(|| {
let _ = sender.send(output);
});
}
for (sender, output) in self.results.drain(..) {
panics.run(|| {
let _ = sender.send(output);
});
}
}
}
pub struct Worker {
local: Rc<RefCell<Local>>,
scope: Option<Scope>,
deferred: Deferred,
panics: Panics,
inbox: Vec<Message>,
processed_seq: u32,
until_global: u32,
global_queue_interval: u32,
finished: bool,
}
impl Worker {
fn new(local: Local) -> Self {
if local.role().is_pool() {
local.shared.pool.enter();
}
let global_queue_interval = local.shared.cfg.global_queue_interval;
let local = Rc::new(RefCell::new(local));
let mut worker = Self {
local: local.clone(),
scope: None,
deferred: Deferred::default(),
panics: Panics::default(),
inbox: Vec::new(),
processed_seq: 0,
until_global: global_queue_interval,
global_queue_interval,
finished: false,
};
worker.scope = Some(Scope::install(local));
worker
}
pub fn run<F, Fut>(
shared: Arc<Shared>,
role: Role,
build: F,
service: Option<BoxedTask>,
) -> Result<(Self, Option<Fut::Output>), Panic>
where
F: FnOnce() -> Fut,
Fut: Future,
{
let local = Local::new(shared.clone(), role).map_err(|error| {
catch_unwind(AssertUnwindSafe(|| {
panic!(
"failed to create native io_uring worker (Linux 6.1 with SINGLE_ISSUER and DEFER_TASKRUN is required): {error}"
)
}))
.expect_err("a startup error must panic")
})?;
let mut worker = Self::new(local);
let mailbox = worker.local.borrow().mailbox.clone();
if let Some(service) = service
&& let Err(service) = shared.tasks.register(service, Arc::downgrade(&shared.pool))
{
worker.panics.run(|| service.clear());
}
let root_waker = RootWaker::new(Arc::downgrade(&mailbox)).into();
let result = catch_unwind(AssertUnwindSafe(|| {
let root = build();
worker.drive(pin!(root), &root_waker)
}));
let output = match result {
Ok(output) => output,
Err(panic) => {
worker.panics.retain(panic);
None
}
};
if role.owns_runner() {
shared.workers.close();
}
worker.begin_close();
Ok((worker, output))
}
fn run_task(payload: Launch) {
let Launch {
active,
shared,
task,
} = payload;
let result = catch_unwind(AssertUnwindSafe(|| {
let root = TaskRoot { task: Some(task) };
let (mut worker, output) = Self::run(shared.clone(), Role::Dedicated, || root, None)?;
worker.cleanup();
worker.result(output)
}));
match result {
Ok(Ok(())) => {}
Ok(Err(panic)) | Err(panic) => shared.panicker.notify(panic),
}
drop(shared);
drop(active);
}
fn callbacks(&mut self, searcher: Option<&mut Searcher<'_>>) {
if self.local.borrow().deferred.is_empty() {
return;
}
if let Some(searcher) = searcher {
searcher.leave_search();
}
mem::swap(&mut self.deferred, &mut self.local.borrow_mut().deferred);
self.deferred.run(&mut self.panics);
}
const fn acknowledge_batch(&mut self) {
self.processed_seq = self.processed_seq.wrapping_add(1) & SUBMISSION_SEQ_MASK;
}
pub fn take_panic(&mut self) -> Option<Panic> {
self.panics.take()
}
fn begin_close(&mut self) {
let mailbox = {
let mut local = self.local.borrow_mut();
let role = local.role();
if role.owns_runner() {
local.shared.tasks.close();
local.shared.pool.close();
}
local.closing = true;
if role.is_pool() {
return;
}
local.mailbox.clone()
};
self.close_mailbox(&mailbox);
}
fn close_mailbox(&mut self, mailbox: &Mailbox) {
let messages = mailbox.close();
if !messages.is_empty() {
self.acknowledge_batch();
}
self.inbox.extend(messages);
}
pub fn cleanup(&mut self) {
if self.finished {
return;
}
let retirement = RetirementGuard;
self.begin_close();
let (role, shared) = {
let local = self.local.borrow();
(local.role(), local.shared.clone())
};
if role.is_pool() {
shared.pool.wait_closed();
}
if role.owns_runner() {
shared.pool.stop();
}
if let Role::Pool(index) = role {
self.local.borrow_mut().ready_mut().discard();
for task in shared.tasks.drain(index as usize) {
Panics::contain(|| task.clear());
}
shared.pool.finish();
let mailbox = self.local.borrow().mailbox.clone();
self.close_mailbox(&mailbox);
}
for message in self.inbox.drain(..) {
Panics::contain(|| drop(message));
}
{
let mut local = self.local.borrow_mut();
let Local {
driver,
timers,
deferred,
..
} = &mut *local;
driver.as_mut().unwrap().close(deferred);
timers.clear(deferred);
local.update_pending();
}
self.callbacks(None);
loop {
self.service(false);
self.callbacks(None);
let mut local = self.local.borrow_mut();
if local.driver.as_ref().unwrap().is_empty() {
assert!(
local.deferred.is_empty(),
"callbacks refilled deferred work after close"
);
break;
}
let deadline = local.next_deadline();
let parked = local
.driver
.as_mut()
.unwrap()
.park(self.processed_seq, deadline);
drop(local);
parked.expect("io_uring shutdown wait failed");
}
mem::forget(retirement);
let driver = self.local.borrow_mut().driver.take();
self.panics.run(|| drop(driver));
self.finished = true;
self.scope.take();
}
fn result<T>(&mut self, output: Option<T>) -> Result<T, Panic> {
if let Some(panic) = self.panics.take() {
self.panics.run(|| drop(output));
return Err(panic);
}
Ok(output.expect("worker root ended without an output or failure"))
}
fn messages(&mut self, mailbox: &Arc<Mailbox>) {
if self.inbox.is_empty()
&& mailbox.waker.pending(self.processed_seq)
&& mailbox.take(&mut self.inbox)
{
self.acknowledge_batch();
self.inbox.reverse();
}
for _ in 0..BATCH_SIZE {
let Some(message) = self.inbox.pop() else {
break;
};
match message {
Message::WakeRoot => self.local.borrow_mut().root_ready = true,
Message::Forward(forward) => {
let mut local = self.local.borrow_mut();
let Local {
driver,
timers,
deferred,
..
} = &mut *local;
match forward {
Forward::Waiter(id, sender) => {
driver.as_mut().unwrap().forward(id, sender, deferred);
}
Forward::Timer(id, sender) => timers.forward(id, sender, deferred),
}
}
Message::Cancel(cancel) => self.local.borrow_mut().release(cancel),
}
}
}
fn service(&mut self, defer_kernel_service: bool) -> bool {
let mut local = self.local.borrow_mut();
local.now = Instant::now();
let Local {
now,
driver,
timers,
deferred,
..
} = &mut *local;
let result = driver
.as_mut()
.unwrap()
.service(*now, defer_kernel_service, deferred);
let outcome = result.expect("io_uring driver service failed");
timers.expire(*now, deferred);
local.update_pending();
outcome
}
#[inline]
#[must_use]
fn next_runnable(&mut self, pool: &Pool) -> Option<Runnable> {
self.until_global -= 1;
if self.until_global == 0 {
self.until_global = self.global_queue_interval;
} else {
if let Some(runnable) = self.local.borrow_mut().ready_mut().pop() {
return Some(runnable);
}
if !pool.has_global() {
return None;
}
}
Self::take_global(&self.local, pool)
}
#[inline(never)]
#[must_use]
fn take_global(local: &RefCell<Local>, pool: &Pool) -> Option<Runnable> {
let mut local = local.borrow_mut();
let ready = local.ready_mut();
pool.take(ready).or_else(|| ready.pop())
}
fn search_before_park(&self, searcher: &mut Searcher<'_>) -> bool {
if searcher.search(self.local.borrow_mut().ready_mut()) {
return true;
}
#[cfg(test)]
tests::at_park(
searcher.pool(),
searcher.index(),
tests::ParkPoint::BeforeIdle,
);
if searcher.publish_and_look(self.local.borrow_mut().ready_mut()) {
return true;
}
#[cfg(test)]
tests::at_park(
searcher.pool(),
searcher.index(),
tests::ParkPoint::BeforeWait,
);
false
}
fn drive<Fut: Future>(
&mut self,
mut root: Pin<&mut Fut>,
root_waker: &Waker,
) -> Option<Fut::Output> {
let (mailbox, shared, spinner_cfg, role) = {
let local = self.local.borrow();
(
local.mailbox.clone(),
local.shared.clone(),
local.shared.cfg.idle_spinner.clone(),
local.role(),
)
};
let (pool, index) = match role {
Role::Pool(index) => (Some(&*shared.pool), Some(index)),
Role::Dedicated => (None, None),
};
let owner = shared.tasks.id();
let mut searcher = index.map(|index| Searcher::new(&shared.pool, index));
let mut spinner = Spinner::new(&spinner_cfg, || {
mailbox.waker.pending(self.processed_seq) || mailbox.waker.signalled()
});
let max_spin =
Duration::from_micros(spinner_cfg.max_budget_us.try_into().unwrap_or(u64::MAX));
loop {
if self.panics.first.is_some() {
return None;
}
if let (Some(pool), Some(searcher)) = (pool, &mut searcher) {
for _ in 0..BATCH_SIZE {
let Some(runnable) = self.next_runnable(pool) else {
break;
};
searcher.leave_search();
match runnable.poll(owner) {
AfterPoll::Done => {}
AfterPoll::Requeue(runnable) => self.local.borrow_mut().enqueue(runnable),
AfterPoll::Retire(task) => drop(shared.tasks.remove(&task)),
}
}
}
let poll_root = mem::take(&mut self.local.borrow_mut().root_ready);
if poll_root {
if let Some(searcher) = &mut searcher {
searcher.leave_search();
}
let mut cx = TaskContext::from_waker(root_waker);
if let Poll::Ready(output) = root.as_mut().poll(&mut cx) {
return Some(output);
}
}
self.callbacks(searcher.as_mut());
if self.panics.first.is_some() {
return None;
}
let defer = !self.local.borrow().is_ready()
&& self.inbox.is_empty()
&& !mailbox.waker.pending(self.processed_seq);
let kernel_deferred = self.service(defer);
self.messages(&mailbox);
self.callbacks(searcher.as_mut());
if self.panics.first.is_some() {
return None;
}
#[cfg(test)]
if let (Some(pool), Some(index)) = (pool, index) {
tests::at_park(pool, index, tests::ParkPoint::AfterMessages);
}
let (ready, needs_kernel, pending_submissions, deadline) = {
let mut local = self.local.borrow_mut();
(
local.is_ready(),
!local.driver.as_ref().unwrap().is_empty(),
local.driver.as_ref().unwrap().has_pending_submissions(),
local.next_deadline(),
)
};
if ready
|| pending_submissions
|| !self.inbox.is_empty()
|| mailbox.waker.pending(self.processed_seq)
{
if let Some(searcher) = &mut searcher {
searcher.leave_search();
}
if (kernel_deferred && needs_kernel) || (defer && pending_submissions) {
self.service(false);
self.callbacks(None);
}
continue;
}
let now = Instant::now();
if deadline.is_some_and(|deadline| deadline <= now) {
if let Some(searcher) = &mut searcher {
searcher.leave_search();
}
if kernel_deferred && needs_kernel {
self.service(false);
self.callbacks(None);
}
continue;
}
if let Some(searcher) = &mut searcher
&& self.search_before_park(searcher)
{
searcher.leave_search();
if kernel_deferred && needs_kernel {
self.service(false);
self.callbacks(None);
}
continue;
}
#[cfg(test)]
tests::before_park();
if needs_kernel {
let result = self
.local
.borrow_mut()
.driver
.as_mut()
.unwrap()
.park(self.processed_seq, deadline);
if let Some(searcher) = &mut searcher {
searcher.finish_park();
}
if !result.expect("io_uring kernel wait failed") {
self.service(false);
self.callbacks(searcher.as_mut());
}
} else {
let near_deadline = deadline
.is_some_and(|deadline| deadline.saturating_duration_since(now) <= max_spin);
if !near_deadline
&& spinner.spin(|| {
mailbox.waker.pending(self.processed_seq) || mailbox.waker.signalled()
})
{
mailbox.waker.consume_signal();
} else if let Some(duration) = mailbox.waker.park_idle(self.processed_seq, deadline)
{
spinner.on_wake(duration);
}
if let Some(searcher) = &mut searcher {
searcher.finish_park();
}
}
}
}
}
impl Drop for Worker {
fn drop(&mut self) {
self.cleanup();
}
}
struct TaskRoot {
task: Option<BoxedTask>,
}
impl Future for TaskRoot {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
Panics::contain(|| self.task.as_mut().unwrap().as_mut().poll(cx)).unwrap_or(Poll::Ready(()))
}
}
impl Drop for TaskRoot {
fn drop(&mut self) {
Panics::contain(|| drop(self.task.take()));
}
}
pub struct Runner {
cfg: Config,
}
impl Runner {
pub const fn new(cfg: Config) -> Self {
Self { cfg }
}
}
impl Default for Runner {
fn default() -> Self {
Self::new(Config::default())
}
}
impl crate::Runner for Runner {
type Context = Context;
fn start<F, Fut>(mut self, f: F) -> Fut::Output
where
F: FnOnce(Context) -> Fut,
Fut: Future,
{
Scope::assert_vacant();
self.cfg.validate();
let mut registry = Registry::new();
let mut runtime_registry = registry.sub_registry(METRICS_PREFIX);
let metrics = TaskMetrics::new(&mut runtime_registry);
let pending_operations = runtime_registry.register(
"pending_operations",
"Number of active logical requests across io_uring workers",
raw::Gauge::default(),
);
let process = ProcessMetrics::init(&mut runtime_registry);
let storage = Storage::new(StorageConfig {
storage_directory: self.cfg.storage_directory.clone(),
blob_layouts: self.cfg.storage_blob_layouts.clone(),
});
crate::storage::sync(&self.cfg.storage_directory).unwrap_or_else(|error| {
panic!(
"failed to sync storage filesystem at startup ({}): {error}",
self.cfg.storage_directory.display()
);
});
let storage = MeteredStorage::new(storage, &mut runtime_registry);
let network = MeteredNetwork::new(
Network::new(NetworkConfig {
tcp_nodelay: self.cfg.tcp_nodelay,
zero_linger: self.cfg.zero_linger,
connect_timeout: self.cfg.connect_timeout,
read_write_timeout: self.cfg.read_write_timeout,
read_buffer_size: self.cfg.read_buffer_size,
}),
&mut runtime_registry,
);
let (panicker, tasks) = Panicker::new(self.cfg.catch_panics);
let mailboxes = (0..self.cfg.worker_threads)
.map(|_| Mailbox::new().map(Arc::new))
.collect::<Result<Vec<_>, _>>()
.unwrap_or_else(|error| panic!("failed to create io_uring worker mailbox: {error}"));
let workers = mailboxes.len();
let shared = Arc::new(Shared {
cfg: self.cfg,
pool: Arc::new(Pool::new(mailboxes)),
registry,
metrics,
pending_operations,
shutdown: Mutex::new(Stopper::default()),
panicker,
workers: Arc::new(Workers::default()),
tasks: Tasks::new(workers),
storage,
network,
});
let label = Label::root();
shared.metrics.tasks_spawned.get_or_create(&label).inc();
let metric = MetricHandle::new(shared.metrics.tasks_running.get_or_create(&label).clone());
let tree = Tree::root();
let context_shared = shared.clone();
let context_tree = tree.clone();
let mut threads = Threads::default();
let (pool_panicker, mut pool_failures) = Panicker::new(false);
let output = Worker::run(
shared.clone(),
Role::Pool(0),
|| {
threads.start(&shared, &pool_panicker);
let root = tasks.interrupt(f(Context {
name: label.name().to_string(),
attributes: Vec::new(),
shared: context_shared,
tree: context_tree,
execution: Execution::default(),
}));
let failures = &mut pool_failures;
async move {
let mut root = pin!(root);
poll_fn(|cx| {
if let Some(panic) = failures.try_take() {
resume_unwind(panic);
}
root.as_mut().poll(cx)
})
.await
}
},
Some(Box::pin(process.collect(Sleep::new))),
)
.and_then(|(mut worker, output)| {
#[cfg(test)]
tests::before_abort();
worker.panics.run(|| tree.abort());
worker.cleanup();
worker.panics.run(|| threads.join());
if let Some(panic) = pool_failures.try_take() {
worker.panics.retain(panic);
}
shared.workers.wait();
worker.result(output)
});
metric.finish();
match output {
Ok(output) => output,
Err(panic) => resume_unwind(panic),
}
}
}
#[cfg(test)]
#[path = "tests.rs"]
pub mod tests;