jay-compositor 1.15.0

use crate::async_engine::AsyncEngine;
use crate::async_engine::SpawnedFuture;
use crate::cpu_worker::AsyncCpuWork;
use crate::cpu_worker::CompletedWork;
use crate::cpu_worker::CpuWork;
use crate::cpu_worker::WorkCompletion;
use crate::io_uring::IoUring;
use crate::io_uring::IoUringError;
use crate::io_uring::IoUringTaskId;
use std::any::Any;
use std::ptr;
use std::rc::Rc;
use std::slice;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::Relaxed;
use thiserror::Error;
use uapi::Fd;
use uapi::c;

#[derive(Debug, Error)]
pub enum ReadWriteJobError {
    #[error("An io_uring error occurred")]
    IoUring(#[source] IoUringError),
    #[error("The job was cancelled")]
    Cancelled,
    #[error("Tried to operate outside the bounds of the file descriptor")]
    OutOfBounds,
}

pub struct ReadWriteWork {
    cancel: Arc<CancelState>,
    config: Option<Box<ReadWriteWorkConfig>>,
}

unsafe impl Send for ReadWriteWork {}

impl ReadWriteWork {
    pub unsafe fn new() -> Self {
        let cancel = Arc::new(CancelState::default());
        ReadWriteWork {
            cancel: cancel.clone(),
            config: Some(Box::new(ReadWriteWorkConfig {
                fd: -1,
                offset: 0,
                ptr: ptr::null_mut(),
                len: 0,
                write: false,
                cancel,
                result: None,
            })),
        }
    }

    pub fn config(&mut self) -> &mut ReadWriteWorkConfig {
        self.config.as_mut().unwrap()
    }
}

pub struct ReadWriteWorkConfig {
    pub fd: c::c_int,
    pub offset: usize,
    pub ptr: *mut u8,
    pub len: usize,
    pub write: bool,
    pub result: Option<Result<(), ReadWriteJobError>>,
    cancel: Arc<CancelState>,
}

#[derive(Default)]
struct CancelState {
    cancelled: AtomicBool,
    cancel_id: AtomicU64,
}

impl CpuWork for ReadWriteWork {
    fn run(&mut self) -> Option<Box<dyn AsyncCpuWork>> {
        self.cancel.cancelled.store(false, Relaxed);
        self.cancel.cancel_id.store(0, Relaxed);
        self.config.take().map(|b| b as _)
    }

    fn cancel_async(&mut self, ring: &Rc<IoUring>) {
        self.cancel.cancelled.store(true, Relaxed);
        let id = self.cancel.cancel_id.load(Relaxed);
        if id != 0 {
            ring.cancel(IoUringTaskId::from_raw(id));
        }
    }

    fn async_work_done(&mut self, work: Box<dyn AsyncCpuWork>) {
        let work = (work as Box<dyn Any>).downcast().unwrap();
        self.config = Some(work);
    }
}

impl AsyncCpuWork for ReadWriteWorkConfig {
    fn run(
        mut self: Box<Self>,
        eng: &Rc<AsyncEngine>,
        ring: &Rc<IoUring>,
        completion: WorkCompletion,
    ) -> SpawnedFuture<CompletedWork> {
        let ring = ring.clone();
        eng.spawn("shm read/write", async move {
            let res = loop {
                if self.cancel.cancelled.load(Relaxed) {
                    break Err(ReadWriteJobError::Cancelled);
                }
                if self.len == 0 {
                    break Ok(());
                }
                let res = if self.write {
                    ring.write_no_cancel(
                        Fd::new(self.fd),
                        self.offset,
                        unsafe { slice::from_raw_parts(self.ptr, self.len) },
                        None,
                        |id| self.cancel.cancel_id.store(id.raw(), Relaxed),
                    )
                    .await
                } else {
                    ring.read_no_cancel(
                        Fd::new(self.fd),
                        self.offset,
                        unsafe { slice::from_raw_parts_mut(self.ptr, self.len) },
                        |id| self.cancel.cancel_id.store(id.raw(), Relaxed),
                    )
                    .await
                };
                match res {
                    Ok(0) => break Err(ReadWriteJobError::OutOfBounds),
                    Ok(n) => {
                        self.len -= n;
                        self.offset += n;
                        unsafe {
                            self.ptr = self.ptr.add(n);
                        }
                    }
                    Err(e) => break Err(ReadWriteJobError::IoUring(e)),
                }
            };
            self.result = Some(res);
            completion.complete(self)
        })
    }
}