430 lines
13 KiB
Rust
430 lines
13 KiB
Rust
pub mod jobs;
|
|
#[cfg(test)]
|
|
mod tests;
|
|
|
|
use {
|
|
crate::{
|
|
async_engine::{AsyncEngine, SpawnedFuture},
|
|
io_uring::IoUring,
|
|
utils::{
|
|
buf::TypedBuf, copyhashmap::CopyHashMap, errorfmt::ErrorFmt, oserror::OsError,
|
|
ptr_ext::MutPtrExt, queue::AsyncQueue, stack::Stack,
|
|
},
|
|
},
|
|
parking_lot::Mutex,
|
|
std::{
|
|
any::Any,
|
|
cell::{Cell, RefCell},
|
|
collections::VecDeque,
|
|
mem,
|
|
ptr::NonNull,
|
|
rc::Rc,
|
|
sync::Arc,
|
|
thread,
|
|
},
|
|
thiserror::Error,
|
|
uapi::{c, OwnedFd},
|
|
};
|
|
|
|
pub trait CpuJob {
|
|
fn work(&mut self) -> &mut dyn CpuWork;
|
|
fn completed(self: Box<Self>);
|
|
}
|
|
|
|
pub trait CpuWork: Send {
|
|
fn run(&mut self) -> Option<Box<dyn AsyncCpuWork>>;
|
|
|
|
fn cancel_async(&mut self, ring: &Rc<IoUring>) {
|
|
let _ = ring;
|
|
unreachable!();
|
|
}
|
|
|
|
fn async_work_done(&mut self, work: Box<dyn AsyncCpuWork>) {
|
|
let _ = work;
|
|
unreachable!();
|
|
}
|
|
}
|
|
|
|
pub trait AsyncCpuWork {
|
|
fn run(
|
|
self: Box<Self>,
|
|
eng: &Rc<AsyncEngine>,
|
|
ring: &Rc<IoUring>,
|
|
completion: WorkCompletion,
|
|
) -> SpawnedFuture<CompletedWork>;
|
|
|
|
fn into_any(self: Box<Self>) -> Box<dyn Any>;
|
|
}
|
|
|
|
pub struct WorkCompletion {
|
|
worker: Rc<Worker>,
|
|
id: CpuJobId,
|
|
}
|
|
|
|
pub struct CompletedWork(());
|
|
|
|
impl WorkCompletion {
|
|
pub fn complete(self, work: Box<dyn AsyncCpuWork>) -> CompletedWork {
|
|
let job = self.worker.async_jobs.remove(&self.id).unwrap();
|
|
unsafe {
|
|
job.work.deref_mut().async_work_done(work);
|
|
}
|
|
self.worker.send_completion(self.id);
|
|
CompletedWork(())
|
|
}
|
|
}
|
|
|
|
pub struct CpuWorker {
|
|
data: Rc<CpuWorkerData>,
|
|
_completions_listener: SpawnedFuture<()>,
|
|
_job_enqueuer: SpawnedFuture<()>,
|
|
}
|
|
|
|
#[must_use]
|
|
pub struct PendingJob {
|
|
id: CpuJobId,
|
|
thread_data: Rc<CpuWorkerData>,
|
|
job_data: Rc<PendingJobData>,
|
|
}
|
|
|
|
#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
|
|
enum PendingJobState {
|
|
#[default]
|
|
Waiting,
|
|
Abandoned,
|
|
Completed,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct PendingJobData {
|
|
job: Cell<Option<NonNull<dyn CpuJob>>>,
|
|
state: Cell<PendingJobState>,
|
|
}
|
|
|
|
enum Job {
|
|
New {
|
|
id: CpuJobId,
|
|
work: *mut dyn CpuWork,
|
|
},
|
|
Cancel {
|
|
id: CpuJobId,
|
|
},
|
|
}
|
|
|
|
unsafe impl Send for Job {}
|
|
|
|
struct CpuWorkerData {
|
|
next: CpuJobIds,
|
|
jobs_to_enqueue: AsyncQueue<Job>,
|
|
new_jobs: Arc<Mutex<VecDeque<Job>>>,
|
|
have_new_jobs: Rc<OwnedFd>,
|
|
completed_jobs_remote: Arc<Mutex<VecDeque<CpuJobId>>>,
|
|
completed_jobs_local: RefCell<VecDeque<CpuJobId>>,
|
|
have_completed_jobs: Rc<OwnedFd>,
|
|
pending_jobs: CopyHashMap<CpuJobId, Rc<PendingJobData>>,
|
|
ring: Rc<IoUring>,
|
|
_stop: OwnedFd,
|
|
pending_job_data_cache: Stack<Rc<PendingJobData>>,
|
|
}
|
|
|
|
linear_ids!(CpuJobIds, CpuJobId, u64);
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum CpuWorkerError {
|
|
#[error("Could not create a pipe")]
|
|
Pipe(#[source] OsError),
|
|
#[error("Could not create an eventfd")]
|
|
EventFd(#[source] OsError),
|
|
#[error("Could not dup an eventfd")]
|
|
Dup(#[source] OsError),
|
|
}
|
|
|
|
impl PendingJob {
|
|
pub fn detach(self) {
|
|
match self.job_data.state.get() {
|
|
PendingJobState::Waiting => {
|
|
self.job_data.state.set(PendingJobState::Abandoned);
|
|
}
|
|
PendingJobState::Abandoned => {
|
|
unreachable!();
|
|
}
|
|
PendingJobState::Completed => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for CpuWorker {
|
|
fn drop(&mut self) {
|
|
self.data.do_equeue_jobs();
|
|
if self.data.pending_jobs.is_not_empty() {
|
|
log::warn!("CpuWorker dropped with pending jobs. Completed jobs will not be triggered.")
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for PendingJob {
|
|
fn drop(&mut self) {
|
|
match self.job_data.state.get() {
|
|
PendingJobState::Waiting => {
|
|
log::warn!("PendingJob dropped before completion. Blocking.");
|
|
let data = &self.thread_data;
|
|
let id = self.id;
|
|
self.job_data.state.set(PendingJobState::Abandoned);
|
|
data.jobs_to_enqueue.push(Job::Cancel { id });
|
|
data.do_equeue_jobs();
|
|
let mut buf = 0u64;
|
|
while data.pending_jobs.contains(&id) {
|
|
if let Err(e) = uapi::read(data.have_completed_jobs.raw(), &mut buf) {
|
|
panic!("Could not wait for job completions: {}", ErrorFmt(e));
|
|
}
|
|
data.dispatch_completions();
|
|
}
|
|
}
|
|
PendingJobState::Abandoned => {}
|
|
PendingJobState::Completed => {
|
|
self.thread_data
|
|
.pending_job_data_cache
|
|
.push(self.job_data.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl CpuWorkerData {
|
|
async fn wait_for_completions(self: Rc<Self>) {
|
|
let mut buf = TypedBuf::<u64>::new();
|
|
loop {
|
|
if let Err(e) = self.ring.read(&self.have_completed_jobs, buf.buf()).await {
|
|
log::error!("Could not wait for job completions: {}", ErrorFmt(e));
|
|
return;
|
|
}
|
|
self.dispatch_completions();
|
|
}
|
|
}
|
|
|
|
fn dispatch_completions(&self) {
|
|
let completions = &mut *self.completed_jobs_local.borrow_mut();
|
|
mem::swap(completions, &mut *self.completed_jobs_remote.lock());
|
|
while let Some(id) = completions.pop_front() {
|
|
let job_data = self.pending_jobs.remove(&id).unwrap();
|
|
let job = job_data.job.take().unwrap();
|
|
let job = unsafe { Box::from_raw(job.as_ptr()) };
|
|
match job_data.state.get() {
|
|
PendingJobState::Waiting => {
|
|
job_data.state.set(PendingJobState::Completed);
|
|
job.completed();
|
|
}
|
|
PendingJobState::Abandoned => {
|
|
self.pending_job_data_cache.push(job_data);
|
|
}
|
|
PendingJobState::Completed => {
|
|
unreachable!();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn equeue_jobs(self: Rc<Self>) {
|
|
loop {
|
|
self.jobs_to_enqueue.non_empty().await;
|
|
self.do_equeue_jobs();
|
|
}
|
|
}
|
|
|
|
fn do_equeue_jobs(&self) {
|
|
self.jobs_to_enqueue.move_to(&mut self.new_jobs.lock());
|
|
if let Err(e) = uapi::eventfd_write(self.have_new_jobs.raw(), 1) {
|
|
panic!("Could not signal eventfd: {}", ErrorFmt(e));
|
|
}
|
|
}
|
|
}
|
|
|
|
impl CpuWorker {
|
|
pub fn new(ring: &Rc<IoUring>, eng: &Rc<AsyncEngine>) -> Result<Self, CpuWorkerError> {
|
|
let new_jobs: Arc<Mutex<VecDeque<Job>>> = Default::default();
|
|
let completed_jobs: Arc<Mutex<VecDeque<CpuJobId>>> = Default::default();
|
|
let (stop_read, stop_write) =
|
|
uapi::pipe2(c::O_CLOEXEC).map_err(|e| CpuWorkerError::Pipe(e.into()))?;
|
|
let have_new_jobs =
|
|
uapi::eventfd(0, c::EFD_CLOEXEC).map_err(|e| CpuWorkerError::EventFd(e.into()))?;
|
|
let have_completed_jobs =
|
|
uapi::eventfd(0, c::EFD_CLOEXEC).map_err(|e| CpuWorkerError::EventFd(e.into()))?;
|
|
thread::Builder::new()
|
|
.name("cpu worker".to_string())
|
|
.spawn({
|
|
let new_jobs = new_jobs.clone();
|
|
let completed_jobs = completed_jobs.clone();
|
|
let have_new_jobs = uapi::fcntl_dupfd_cloexec(have_new_jobs.raw(), 0)
|
|
.map_err(|e| CpuWorkerError::Dup(e.into()))?;
|
|
let have_completed_jobs = uapi::fcntl_dupfd_cloexec(have_completed_jobs.raw(), 0)
|
|
.map_err(|e| CpuWorkerError::Dup(e.into()))?;
|
|
move || {
|
|
work(
|
|
new_jobs,
|
|
completed_jobs,
|
|
stop_write,
|
|
have_new_jobs,
|
|
have_completed_jobs,
|
|
)
|
|
}
|
|
})
|
|
.unwrap();
|
|
let data = Rc::new(CpuWorkerData {
|
|
next: Default::default(),
|
|
jobs_to_enqueue: Default::default(),
|
|
new_jobs,
|
|
have_new_jobs: Rc::new(have_new_jobs),
|
|
completed_jobs_remote: completed_jobs,
|
|
completed_jobs_local: Default::default(),
|
|
have_completed_jobs: Rc::new(have_completed_jobs),
|
|
pending_jobs: Default::default(),
|
|
ring: ring.clone(),
|
|
_stop: stop_read,
|
|
pending_job_data_cache: Default::default(),
|
|
});
|
|
Ok(Self {
|
|
_completions_listener: eng.spawn(data.clone().wait_for_completions()),
|
|
_job_enqueuer: eng.spawn(data.clone().equeue_jobs()),
|
|
data,
|
|
})
|
|
}
|
|
|
|
pub fn submit(&self, job: Box<dyn CpuJob>) -> PendingJob {
|
|
let mut job = NonNull::from(Box::leak(job));
|
|
let id = self.data.next.next();
|
|
self.data.jobs_to_enqueue.push(Job::New {
|
|
id,
|
|
work: unsafe { job.as_mut().work() },
|
|
});
|
|
let job_data = self.data.pending_job_data_cache.pop().unwrap_or_default();
|
|
job_data.job.set(Some(job));
|
|
job_data.state.set(PendingJobState::Waiting);
|
|
self.data.pending_jobs.set(id, job_data.clone());
|
|
PendingJob {
|
|
id,
|
|
thread_data: self.data.clone(),
|
|
job_data,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn work(
|
|
new_jobs: Arc<Mutex<VecDeque<Job>>>,
|
|
completed_jobs: Arc<Mutex<VecDeque<CpuJobId>>>,
|
|
stop: OwnedFd,
|
|
have_new_jobs: OwnedFd,
|
|
have_completed_jobs: OwnedFd,
|
|
) {
|
|
let eng = AsyncEngine::new();
|
|
let ring = IoUring::new(&eng, 32).unwrap();
|
|
let worker = Rc::new(Worker {
|
|
eng,
|
|
ring,
|
|
completed_jobs,
|
|
have_completed_jobs,
|
|
async_jobs: Default::default(),
|
|
stopped: Cell::new(false),
|
|
});
|
|
let _stop_listener = worker.eng.spawn(worker.clone().handle_stop(stop));
|
|
let _new_job_listener = worker
|
|
.eng
|
|
.spawn(worker.clone().handle_new_jobs(new_jobs, have_new_jobs));
|
|
if let Err(e) = worker.ring.run() {
|
|
panic!("io_uring failed: {}", ErrorFmt(e));
|
|
}
|
|
}
|
|
|
|
struct Worker {
|
|
eng: Rc<AsyncEngine>,
|
|
ring: Rc<IoUring>,
|
|
completed_jobs: Arc<Mutex<VecDeque<CpuJobId>>>,
|
|
have_completed_jobs: OwnedFd,
|
|
async_jobs: CopyHashMap<CpuJobId, AsyncJob>,
|
|
stopped: Cell<bool>,
|
|
}
|
|
|
|
struct AsyncJob {
|
|
_future: SpawnedFuture<CompletedWork>,
|
|
work: *mut dyn CpuWork,
|
|
}
|
|
|
|
impl Worker {
|
|
async fn handle_stop(self: Rc<Self>, stop: OwnedFd) {
|
|
let stop = Rc::new(stop);
|
|
if let Err(e) = self.ring.poll(&stop, 0).await {
|
|
log::error!(
|
|
"Could not wait for stop fd to become readable: {}",
|
|
ErrorFmt(e)
|
|
);
|
|
} else {
|
|
assert!(self.async_jobs.is_empty());
|
|
self.stopped.set(true);
|
|
self.ring.stop();
|
|
}
|
|
}
|
|
|
|
async fn handle_new_jobs(
|
|
self: Rc<Self>,
|
|
jobs_remote: Arc<Mutex<VecDeque<Job>>>,
|
|
new_jobs: OwnedFd,
|
|
) {
|
|
let mut buf = TypedBuf::<u64>::new();
|
|
let new_jobs = Rc::new(new_jobs);
|
|
let mut jobs = VecDeque::new();
|
|
loop {
|
|
if let Err(e) = self.ring.read(&new_jobs, buf.buf()).await {
|
|
if self.stopped.get() {
|
|
return;
|
|
}
|
|
panic!(
|
|
"Could not wait for new jobs fd to be signaled: {}",
|
|
ErrorFmt(e),
|
|
);
|
|
}
|
|
mem::swap(&mut jobs, &mut *jobs_remote.lock());
|
|
while let Some(job) = jobs.pop_front() {
|
|
self.handle_new_job(job);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn handle_new_job(self: &Rc<Self>, job: Job) {
|
|
match job {
|
|
Job::Cancel { id } => {
|
|
let mut jobs = self.async_jobs.lock();
|
|
if let Some(job) = jobs.get_mut(&id) {
|
|
unsafe {
|
|
job.work.deref_mut().cancel_async(&self.ring);
|
|
}
|
|
}
|
|
}
|
|
Job::New { id, work } => match unsafe { work.deref_mut() }.run() {
|
|
None => {
|
|
self.send_completion(id);
|
|
return;
|
|
}
|
|
Some(w) => {
|
|
let completion = WorkCompletion {
|
|
worker: self.clone(),
|
|
id,
|
|
};
|
|
let future = w.run(&self.eng, &self.ring, completion);
|
|
self.async_jobs.set(
|
|
id,
|
|
AsyncJob {
|
|
_future: future,
|
|
work,
|
|
},
|
|
);
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
fn send_completion(&self, id: CpuJobId) {
|
|
self.completed_jobs.lock().push_back(id);
|
|
if let Err(e) = uapi::eventfd_write(self.have_completed_jobs.raw(), 1) {
|
|
panic!("Could not signal job completion: {}", ErrorFmt(e));
|
|
}
|
|
}
|
|
}
|