lancer / kernel / src / ipc / notification.rs

use crate::cap::pool::{ObjectPool, POOL};
use crate::cap::table::CapRef;
use crate::error::KernelError;
use crate::ipc::{IpcOutcome, IpcResult};
use crate::proc::{BlockedReason, ProcessManager};
use crate::types::{Generation, MAX_PIDS, ObjPhys, Pid};
use lancer_core::header::NONE_SENTINEL;
use lancer_core::object_layout::{EndpointObject, NotificationObject};

pub const BOUND_NOTIFICATION_BADGE: u64 = (-128i64) as u64;

pub type SignalResult = Result<bool, KernelError>;
pub type PollResult = Result<u64, KernelError>;

const MAX_NOTIFICATION_WAITERS: usize = 4;

fn has_waiters(notif: &NotificationObject) -> bool {
    notif.waiter_count > 0
}

pub fn signal_inner(notif: &mut NotificationObject, bits: u64) -> bool {
    debug_assert!(
        !x86_64::instructions::interrupts::are_enabled(),
        "signal_inner requires interrupts disabled"
    );
    notif.word |= bits;
    has_waiters(notif)
}

pub fn drain_and_wake(notif: &mut NotificationObject, ptable: &mut ProcessManager) {
    let woken_word = notif.word;
    let mut failed_count = 0u8;
    let mut failed_pids = [NONE_SENTINEL; MAX_NOTIFICATION_WAITERS];

    (0..notif.waiter_count as usize).for_each(|i| {
        let raw = notif.waiters[i];
        match Pid::try_new(raw) {
            Some(pid) => {
                let proof = ptable[pid].blocked_proof();
                match ptable.unblock_and_enqueue(pid, proof) {
                    Ok(()) => {
                        let exec = ptable.exec_mut(pid).unwrap();
                        exec.saved_context.rax = 0;
                        exec.saved_context.rdx = woken_word;
                        exec.seal_context();
                    }
                    Err(e) => {
                        crate::kprintln!(
                            "[notify] BUG: failed to unblock waiter pid {}: {:?}",
                            pid.raw(),
                            e
                        );
                        failed_pids[failed_count as usize] = raw;
                        failed_count += 1;
                    }
                }
            }
            None => {}
        }
    });

    notif.waiter_count = failed_count;
    (0..MAX_NOTIFICATION_WAITERS).for_each(|i| {
        notif.waiters[i] = match i < failed_count as usize {
            true => failed_pids[i],
            false => NONE_SENTINEL,
        };
    });
    if failed_count == 0 {
        notif.word = 0;
    }
}

pub fn do_signal(cap: &CapRef, bits: u64, ptable: &mut ProcessManager) -> SignalResult {
    let notif_id = cap.phys();
    let notif_gen = cap.generation();

    let mut woke_any = false;

    {
        let mut pool = POOL.lock();
        let notif = pool.write_as::<NotificationObject>(notif_id, notif_gen)?;

        if signal_inner(notif, bits) {
            woke_any = true;
            drain_and_wake(notif, ptable);
        }
    }

    let mut pool = POOL.lock();
    if wake_bound_receivers_with_pool(notif_id, notif_gen, bits, ptable, &mut pool) {
        woke_any = true;
    }

    Ok(woke_any)
}

pub fn do_wait(cap: &CapRef, pid: Pid, ptable: &mut ProcessManager) -> IpcResult<u64> {
    let mut pool = POOL.lock();
    let notif = pool.write_as::<NotificationObject>(cap.phys(), cap.generation())?;

    if notif.word != 0 {
        let val = notif.word;
        notif.word = 0;
        Ok(IpcOutcome::Done(val))
    } else {
        let phys = cap.phys();
        let proof = ptable[pid].block_on(BlockedReason::WaitingNotification(
            phys,
            cap.generation(),
        ))?;
        match (notif.waiter_count as usize) < MAX_NOTIFICATION_WAITERS {
            true => {
                notif.waiters[notif.waiter_count as usize] = pid.raw();
                notif.waiter_count += 1;
                Ok(IpcOutcome::Blocked)
            }
            false => {
                ptable.unblock_and_enqueue(pid, proof)?;
                Err(KernelError::ResourceExhausted)
            }
        }
    }
}

pub fn do_poll(cap: &CapRef) -> PollResult {
    let mut pool = POOL.lock();
    let notif = pool.write_as::<NotificationObject>(cap.phys(), cap.generation())?;

    let val = notif.word;
    notif.word = 0;
    Ok(val)
}

pub fn wake_bound_receivers_with_pool(
    notif_id: ObjPhys,
    notif_gen: Generation,
    woken_word: u64,
    ptable: &mut ProcessManager,
    pool: &mut ObjectPool,
) -> bool {
    (0..MAX_PIDS as u32).fold(false, |woke, idx| {
        let pid = Pid::new(idx);
        let Some(sched) = ptable.get(pid) else {
            return woke;
        };
        let matches_bound = sched
            .bound_notification()
            .is_some_and(|(id, g)| id == notif_id && g == notif_gen);
        if !matches_bound {
            return woke;
        }
        let is_receiving = matches!(sched.blocked_reason(), Some(BlockedReason::Receiving(_, _)));
        if !is_receiving {
            return woke;
        }

        let (ep_id, ep_gen) = match sched.blocked_reason() {
            Some(BlockedReason::Receiving(id, generation)) => (id, generation),
            _ => return woke,
        };

        let proof = ptable[pid].blocked_proof();
        match ptable.unblock_and_enqueue(pid, proof) {
            Ok(()) => {
                let exec = ptable.exec_mut(pid).unwrap();
                exec.saved_context.rax = BOUND_NOTIFICATION_BADGE;
                exec.saved_context.rdx = woken_word;
                exec.seal_context();

                if let Ok(ep) = pool.write_as::<EndpointObject>(ep_id, ep_gen) {
                    super::endpoint::remove_from_recv(ep, pid, ptable);
                }
                true
            }
            Err(e) => {
                crate::kprintln!(
                    "[notify] BUG: failed to unblock bound receiver pid {}: {:?}",
                    pid.raw(),
                    e
                );
                woke
            }
        }
    })
}