prometeu-runtime/crates/console/prometeu-system/src/services/async_work.rs

use std::collections::VecDeque;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, RecvTimeoutError};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AsyncWorkJobId(u64);

impl AsyncWorkJobId {
    pub const ZERO: Self = Self(0);

    pub fn new(value: u64) -> Self {
        Self(value)
    }

    pub fn get(self) -> u64 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AsyncWorkJobKind {
    Asset,
    Memcard,
    Fs,
    Test,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AsyncWorkPriority {
    MemcardCommitWrite,
    FsWriteConfig,
    AssetLoad,
    NonCriticalReadList,
}

impl AsyncWorkPriority {
    const fn rank(self) -> u8 {
        match self {
            Self::MemcardCommitWrite => 0,
            Self::FsWriteConfig => 1,
            Self::AssetLoad => 2,
            Self::NonCriticalReadList => 3,
        }
    }
}

impl AsyncWorkJobKind {
    pub const fn default_priority(self) -> AsyncWorkPriority {
        match self {
            Self::Memcard => AsyncWorkPriority::MemcardCommitWrite,
            Self::Fs => AsyncWorkPriority::FsWriteConfig,
            Self::Asset => AsyncWorkPriority::AssetLoad,
            Self::Test => AsyncWorkPriority::NonCriticalReadList,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AsyncWorkJobOutcome {
    Completed,
    Canceled,
    Failed,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AsyncWorkLaneError {
    Shutdown,
    WorkerPanic,
    ShutdownTimeout,
    InternalFailure,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AsyncWorkLaneConfig {
    pub shutdown_timeout: Duration,
}

impl Default for AsyncWorkLaneConfig {
    fn default() -> Self {
        Self { shutdown_timeout: Duration::from_millis(250) }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AsyncWorkActiveJob {
    pub id: AsyncWorkJobId,
    pub kind: AsyncWorkJobKind,
    pub priority: AsyncWorkPriority,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct AsyncWorkLaneTelemetry {
    pub submitted_jobs: u64,
    pub started_jobs: u64,
    pub completed_jobs: u64,
    pub canceled_jobs: u64,
    pub failed_jobs: u64,
    pub shutdown_discards: u64,
    pub pending_depth: usize,
    pub max_pending_depth: usize,
    pub active_job_id: Option<AsyncWorkJobId>,
    pub last_submitted_job_id: Option<AsyncWorkJobId>,
    pub last_started_job_id: Option<AsyncWorkJobId>,
    pub last_closed_job_id: Option<AsyncWorkJobId>,
}

impl AsyncWorkLaneTelemetry {
    fn record_submitted(&mut self, job_id: AsyncWorkJobId, pending_depth: usize) {
        self.submitted_jobs += 1;
        self.pending_depth = pending_depth;
        self.max_pending_depth = self.max_pending_depth.max(pending_depth);
        self.last_submitted_job_id = Some(job_id);
    }

    fn record_started(&mut self, job_id: AsyncWorkJobId, pending_depth: usize) {
        self.started_jobs += 1;
        self.pending_depth = pending_depth;
        self.active_job_id = Some(job_id);
        self.last_started_job_id = Some(job_id);
    }

    fn record_closed(&mut self, job_id: AsyncWorkJobId, outcome: AsyncWorkJobOutcome) {
        match outcome {
            AsyncWorkJobOutcome::Completed => self.completed_jobs += 1,
            AsyncWorkJobOutcome::Canceled => self.canceled_jobs += 1,
            AsyncWorkJobOutcome::Failed => self.failed_jobs += 1,
        }
        self.active_job_id = None;
        self.last_closed_job_id = Some(job_id);
    }

    fn record_shutdown_discards(&mut self, count: usize) {
        self.shutdown_discards += count as u64;
        self.pending_depth = 0;
    }
}

#[derive(Clone, Debug)]
pub struct AsyncWorkCancelToken {
    canceled: Arc<AtomicBool>,
}

impl AsyncWorkCancelToken {
    fn new() -> Self {
        Self { canceled: Arc::new(AtomicBool::new(false)) }
    }

    pub fn cancel(&self) {
        self.canceled.store(true, Ordering::Relaxed);
    }

    pub fn is_canceled(&self) -> bool {
        self.canceled.load(Ordering::Relaxed)
    }
}

#[derive(Clone, Debug)]
pub struct AsyncWorkJobContext {
    pub job_id: AsyncWorkJobId,
    pub kind: AsyncWorkJobKind,
    pub priority: AsyncWorkPriority,
    cancel_token: AsyncWorkCancelToken,
}

impl AsyncWorkJobContext {
    pub fn is_canceled(&self) -> bool {
        self.cancel_token.is_canceled()
    }
}

type AsyncWorkFn = Box<dyn FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static>;

struct PendingAsyncWorkJob {
    id: AsyncWorkJobId,
    kind: AsyncWorkJobKind,
    priority: AsyncWorkPriority,
    cancel_token: AsyncWorkCancelToken,
    work: AsyncWorkFn,
}

struct RunningAsyncWorkJob {
    id: AsyncWorkJobId,
    kind: AsyncWorkJobKind,
    priority: AsyncWorkPriority,
    cancel_token: AsyncWorkCancelToken,
    work: AsyncWorkFn,
}

#[derive(Default)]
struct AsyncWorkLaneState {
    next_job_id: u64,
    pending: VecDeque<PendingAsyncWorkJob>,
    active: Option<AsyncWorkActiveJob>,
    active_cancel_token: Option<AsyncWorkCancelToken>,
    shutdown_requested: bool,
    telemetry: AsyncWorkLaneTelemetry,
}

#[derive(Default)]
pub struct AsyncWorkLane {
    state: Mutex<AsyncWorkLaneState>,
    ready: Condvar,
}

impl AsyncWorkLane {
    pub fn submit(
        &self,
        kind: AsyncWorkJobKind,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        self.submit_with_priority(kind, kind.default_priority(), work)
    }

    pub fn submit_with_priority(
        &self,
        kind: AsyncWorkJobKind,
        priority: AsyncWorkPriority,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        let mut state = self.state.lock().map_err(|_| AsyncWorkLaneError::InternalFailure)?;
        if state.shutdown_requested {
            return Err(AsyncWorkLaneError::Shutdown);
        }

        state.next_job_id += 1;
        let id = AsyncWorkJobId::new(state.next_job_id);
        let cancel_token = AsyncWorkCancelToken::new();
        let job = PendingAsyncWorkJob { id, kind, priority, cancel_token, work: Box::new(work) };
        let insert_at = state
            .pending
            .iter()
            .position(|pending| priority.rank() < pending.priority.rank())
            .unwrap_or(state.pending.len());
        state.pending.insert(insert_at, job);
        let pending_depth = state.pending.len();
        state.telemetry.record_submitted(id, pending_depth);
        self.ready.notify_one();
        Ok(id)
    }

    pub fn submit_memcard_commit_write(
        &self,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        self.submit_with_priority(
            AsyncWorkJobKind::Memcard,
            AsyncWorkPriority::MemcardCommitWrite,
            work,
        )
    }

    pub fn submit_fs_write_config(
        &self,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        self.submit_with_priority(AsyncWorkJobKind::Fs, AsyncWorkPriority::FsWriteConfig, work)
    }

    pub fn submit_asset_load(
        &self,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        self.submit_with_priority(AsyncWorkJobKind::Asset, AsyncWorkPriority::AssetLoad, work)
    }

    pub fn submit_non_critical_read_list(
        &self,
        kind: AsyncWorkJobKind,
        work: impl FnOnce(AsyncWorkJobContext) -> AsyncWorkJobOutcome + Send + 'static,
    ) -> Result<AsyncWorkJobId, AsyncWorkLaneError> {
        self.submit_with_priority(kind, AsyncWorkPriority::NonCriticalReadList, work)
    }

    pub fn request_cancel(&self, job_id: AsyncWorkJobId) -> bool {
        let state = self.state.lock().unwrap();
        if let Some(token) = state
            .active
            .filter(|active| active.id == job_id)
            .and(state.active_cancel_token.as_ref())
        {
            token.cancel();
            return true;
        }

        for pending in &state.pending {
            if pending.id == job_id {
                pending.cancel_token.cancel();
                return true;
            }
        }

        false
    }

    pub fn active_job(&self) -> Option<AsyncWorkActiveJob> {
        self.state.lock().unwrap().active
    }

    pub fn pending_depth(&self) -> usize {
        self.state.lock().unwrap().pending.len()
    }

    pub fn telemetry(&self) -> AsyncWorkLaneTelemetry {
        self.state.lock().unwrap().telemetry
    }

    pub fn request_shutdown(&self) {
        let mut state = self.state.lock().unwrap();
        state.shutdown_requested = true;
        if let Some(token) = &state.active_cancel_token {
            token.cancel();
        }
        self.ready.notify_all();
    }

    fn wait_take(&self) -> Option<RunningAsyncWorkJob> {
        self.wait_take_with_hook(|| {})
    }

    fn wait_take_with_hook(&self, mut before_wait: impl FnMut()) -> Option<RunningAsyncWorkJob> {
        let mut state = self.state.lock().unwrap();
        loop {
            if state.shutdown_requested {
                let discarded = state.pending.len();
                state.pending.clear();
                state.telemetry.record_shutdown_discards(discarded);
                return None;
            }

            if let Some(pending) = state.pending.pop_front() {
                let active = AsyncWorkActiveJob {
                    id: pending.id,
                    kind: pending.kind,
                    priority: pending.priority,
                };
                state.active = Some(active);
                state.active_cancel_token = Some(pending.cancel_token.clone());
                let pending_depth = state.pending.len();
                state.telemetry.record_started(pending.id, pending_depth);
                return Some(RunningAsyncWorkJob {
                    id: pending.id,
                    kind: pending.kind,
                    priority: pending.priority,
                    cancel_token: pending.cancel_token,
                    work: pending.work,
                });
            }

            before_wait();
            state = self.ready.wait(state).unwrap();
        }
    }

    fn record_finished(&self, job_id: AsyncWorkJobId, outcome: AsyncWorkJobOutcome) {
        let mut state = self.state.lock().unwrap();
        if state.active.is_some_and(|active| active.id == job_id) {
            state.active = None;
            state.active_cancel_token = None;
        }
        state.telemetry.record_closed(job_id, outcome);
    }
}

pub struct AsyncWorkLaneController {
    config: AsyncWorkLaneConfig,
    lane: Arc<AsyncWorkLane>,
    handle: Option<JoinHandle<()>>,
    done_rx: Receiver<Result<(), AsyncWorkLaneError>>,
}

impl AsyncWorkLaneController {
    pub fn start(config: AsyncWorkLaneConfig, lane: Arc<AsyncWorkLane>) -> Self {
        let worker_lane = Arc::clone(&lane);
        let (done_tx, done_rx) = mpsc::channel();
        let handle = thread::spawn(move || {
            let result = catch_unwind(AssertUnwindSafe(|| run_async_work_loop(worker_lane)))
                .map_err(|_| AsyncWorkLaneError::WorkerPanic);
            let _ = done_tx.send(result);
        });

        Self { config, lane, handle: Some(handle), done_rx }
    }

    pub fn lane(&self) -> Arc<AsyncWorkLane> {
        Arc::clone(&self.lane)
    }

    pub fn stop(&mut self) -> Result<(), AsyncWorkLaneError> {
        self.lane.request_shutdown();
        match self.done_rx.recv_timeout(self.config.shutdown_timeout) {
            Ok(result) => {
                if let Some(handle) = self.handle.take()
                    && handle.join().is_err()
                {
                    return Err(AsyncWorkLaneError::WorkerPanic);
                }
                result
            }
            Err(RecvTimeoutError::Timeout) => Err(AsyncWorkLaneError::ShutdownTimeout),
            Err(RecvTimeoutError::Disconnected) => Err(AsyncWorkLaneError::InternalFailure),
        }
    }
}

impl Drop for AsyncWorkLaneController {
    fn drop(&mut self) {
        if self.handle.is_some() {
            let _ = self.stop();
        }
    }
}

fn run_async_work_loop(lane: Arc<AsyncWorkLane>) {
    while let Some(job) = lane.wait_take() {
        let context = AsyncWorkJobContext {
            job_id: job.id,
            kind: job.kind,
            priority: job.priority,
            cancel_token: job.cancel_token,
        };
        let result = catch_unwind(AssertUnwindSafe(|| (job.work)(context)))
            .unwrap_or(AsyncWorkJobOutcome::Failed);
        lane.record_finished(job.id, result);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc;

    #[test]
    fn async_work_lane_runs_jobs_in_submission_order() {
        let lane = Arc::new(AsyncWorkLane::default());
        let mut controller =
            AsyncWorkLaneController::start(AsyncWorkLaneConfig::default(), Arc::clone(&lane));
        let (tx, rx) = mpsc::channel();

        lane.submit(AsyncWorkJobKind::Test, {
            let tx = tx.clone();
            move |_| {
                tx.send(1).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();
        lane.submit(AsyncWorkJobKind::Test, {
            let tx = tx.clone();
            move |_| {
                tx.send(2).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();

        assert_eq!(rx.recv().unwrap(), 1);
        assert_eq!(rx.recv().unwrap(), 2);
        controller.stop().unwrap();

        let telemetry = lane.telemetry();
        assert_eq!(telemetry.submitted_jobs, 2);
        assert_eq!(telemetry.started_jobs, 2);
        assert_eq!(telemetry.completed_jobs, 2);
    }

    #[test]
    fn async_work_lane_runs_higher_priority_io_before_assets() {
        let lane = Arc::new(AsyncWorkLane::default());
        let (tx, rx) = mpsc::channel();

        lane.submit_asset_load({
            let tx = tx.clone();
            move |ctx| {
                tx.send((ctx.kind, ctx.priority)).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();
        lane.submit_non_critical_read_list(AsyncWorkJobKind::Fs, {
            let tx = tx.clone();
            move |ctx| {
                tx.send((ctx.kind, ctx.priority)).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();
        lane.submit_memcard_commit_write({
            let tx = tx.clone();
            move |ctx| {
                tx.send((ctx.kind, ctx.priority)).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();
        lane.submit_fs_write_config({
            let tx = tx.clone();
            move |ctx| {
                tx.send((ctx.kind, ctx.priority)).unwrap();
                AsyncWorkJobOutcome::Completed
            }
        })
        .unwrap();

        let mut controller =
            AsyncWorkLaneController::start(AsyncWorkLaneConfig::default(), Arc::clone(&lane));

        assert_eq!(
            rx.recv().unwrap(),
            (AsyncWorkJobKind::Memcard, AsyncWorkPriority::MemcardCommitWrite)
        );
        assert_eq!(rx.recv().unwrap(), (AsyncWorkJobKind::Fs, AsyncWorkPriority::FsWriteConfig));
        assert_eq!(rx.recv().unwrap(), (AsyncWorkJobKind::Asset, AsyncWorkPriority::AssetLoad));
        assert_eq!(
            rx.recv().unwrap(),
            (AsyncWorkJobKind::Fs, AsyncWorkPriority::NonCriticalReadList)
        );
        controller.stop().unwrap();
    }

    #[test]
    fn async_work_lane_has_one_active_job() {
        let lane = Arc::new(AsyncWorkLane::default());
        let mut controller =
            AsyncWorkLaneController::start(AsyncWorkLaneConfig::default(), Arc::clone(&lane));
        let (first_started_tx, first_started_rx) = mpsc::channel();
        let (release_first_tx, release_first_rx) = mpsc::channel();
        let (second_started_tx, second_started_rx) = mpsc::channel();

        let first = lane
            .submit(AsyncWorkJobKind::Test, move |_| {
                first_started_tx.send(()).unwrap();
                release_first_rx.recv().unwrap();
                AsyncWorkJobOutcome::Completed
            })
            .unwrap();
        let second = lane
            .submit(AsyncWorkJobKind::Test, move |_| {
                second_started_tx.send(()).unwrap();
                AsyncWorkJobOutcome::Completed
            })
            .unwrap();

        first_started_rx.recv().unwrap();
        assert_eq!(lane.active_job().unwrap().id, first);
        assert_eq!(lane.pending_depth(), 1);
        assert!(second_started_rx.try_recv().is_err());

        release_first_tx.send(()).unwrap();
        second_started_rx.recv().unwrap();
        assert_eq!(lane.telemetry().last_started_job_id, Some(second));
        controller.stop().unwrap();
    }

    #[test]
    fn async_work_lane_supports_cooperative_cancellation() {
        let lane = Arc::new(AsyncWorkLane::default());
        let mut controller =
            AsyncWorkLaneController::start(AsyncWorkLaneConfig::default(), Arc::clone(&lane));
        let (started_tx, started_rx) = mpsc::channel();

        let job_id = lane
            .submit(AsyncWorkJobKind::Test, move |ctx| {
                started_tx.send(()).unwrap();
                while !ctx.is_canceled() {
                    std::thread::yield_now();
                }
                AsyncWorkJobOutcome::Canceled
            })
            .unwrap();

        started_rx.recv().unwrap();
        assert!(lane.request_cancel(job_id));
        controller.stop().unwrap();
        assert_eq!(lane.telemetry().canceled_jobs, 1);
    }

    #[test]
    fn async_work_lane_rejects_submit_after_shutdown() {
        let lane = Arc::new(AsyncWorkLane::default());
        lane.request_shutdown();

        let result = lane.submit(AsyncWorkJobKind::Test, |_| AsyncWorkJobOutcome::Completed);
        assert_eq!(result, Err(AsyncWorkLaneError::Shutdown));
    }

    #[test]
    fn async_work_lane_shutdown_discards_pending_jobs() {
        let lane = Arc::new(AsyncWorkLane::default());
        lane.submit(AsyncWorkJobKind::Test, |_| AsyncWorkJobOutcome::Completed).unwrap();
        assert_eq!(lane.pending_depth(), 1);

        lane.request_shutdown();
        assert!(lane.wait_take().is_none());
        assert_eq!(lane.pending_depth(), 0);
        assert_eq!(lane.telemetry().shutdown_discards, 1);
    }
}