RedBear-OS/src/event.rs

use alloc::sync::Arc;
use core::sync::atomic::{AtomicUsize, Ordering};
use hashbrown::{hash_map::DefaultHashBuilder, HashMap};
use smallvec::SmallVec;
use syscall::data::GlobalSchemes;

use crate::{
    context,
    scheme::{self, SchemeExt, SchemeId},
    sync::{
        CleanLockToken, LockToken, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard,
        WaitCondition, WaitQueue, L0, L1, L2,
    },
    syscall::{
        data::Event,
        error::{Error, Result, EAGAIN, EBADF, EINVAL, EINTR},
        flag::{EVENT_READ, EVENT_WRITE, EventFlags},
        usercopy::{UserSliceRo, UserSliceWo},
    },
};

int_like!(EventQueueId, AtomicEventQueueId, usize, AtomicUsize);

pub struct EventQueue {
    id: EventQueueId,
    queue: WaitQueue<Event>,
}

const EVENTFD_COUNTER_MAX: u64 = u64::MAX - 1;
const EVENTFD_TAG_BIT: usize = 1usize << (usize::BITS - 1);

pub struct EventCounter {
    id: usize,
    counter: Mutex<L1, u64>,
    read_condition: WaitCondition,
    write_condition: WaitCondition,
    semaphore: bool,
}

impl EventQueue {
    pub fn new(id: EventQueueId) -> EventQueue {
        EventQueue {
            id,
            queue: WaitQueue::new(),
        }
    }

    pub fn is_currently_empty(&self, token: &mut CleanLockToken) -> bool {
        self.queue.is_currently_empty(token)
    }

    pub fn read(&self, buf: UserSliceWo, block: bool, token: &mut CleanLockToken) -> Result<usize> {
        self.queue
            .receive_into_user(buf, block, "EventQueue::read", token)
    }

    pub fn write(&self, events: &[Event], token: &mut CleanLockToken) -> Result<usize> {
        for event in events {
            let file = {
                let context_ref = context::current();
                let mut context = context_ref.read(token.token());
                let (context, mut token) = context.token_split();
                let files = context.files.read(token.token());
                match files.get(event.id).ok_or(Error::new(EBADF))? {
                    Some(file) => file.clone(),
                    None => return Err(Error::new(EBADF)),
                }
            };

            let (scheme, number) = {
                let description = file.description.read(token.token());
                (description.scheme, description.number)
            };

            if scheme == GlobalSchemes::Event.scheme_id() && number == self.id.into() {
                // Do not allow recursively registering the same event queue
                //TODO: should we also disallow event queues that contain this event queue?
                return Err(Error::new(EBADF));
            }

            register(
                RegKey { scheme, number },
                QueueKey {
                    queue: self.id,
                    id: event.id,
                    data: event.data,
                },
                event.flags,
                token,
            );

            let flags = sync(RegKey { scheme, number }, token)?;
            if !flags.is_empty() {
                trigger(scheme, number, flags, token);
            }
        }

        Ok(events.len())
    }

    pub fn into_drop(self, token: LockToken<'_, L1>) {
        self.queue.condition.into_drop_locked(token);
    }
}

impl EventCounter {
    pub fn new(id: usize, init: u64, semaphore: bool) -> EventCounter {
        EventCounter {
            id,
            counter: Mutex::new(init),
            read_condition: WaitCondition::new(),
            write_condition: WaitCondition::new(),
            semaphore,
        }
    }

    pub fn is_readable(&self, token: &mut CleanLockToken) -> bool {
        *self.counter.lock(token.token()) > 0
    }

    pub fn is_writable(&self, token: &mut CleanLockToken) -> bool {
        *self.counter.lock(token.token()) < EVENTFD_COUNTER_MAX
    }

    pub fn read(&self, buf: UserSliceWo, block: bool, token: &mut CleanLockToken) -> Result<usize> {
        if buf.len() < core::mem::size_of::<u64>() {
            return Err(Error::new(EINVAL));
        }

        loop {
            let counter = self.counter.lock(token.token());
            let (mut counter, mut token) = counter.into_split();

            if *counter > 0 {
                let value = if self.semaphore {
                    *counter -= 1;
                    1
                } else {
                    let value = *counter;
                    *counter = 0;
                    value
                };

                buf.limit(core::mem::size_of::<u64>())
                    .ok_or(Error::new(EINVAL))?
                    .copy_from_slice(&value.to_ne_bytes())?;

                trigger_locked(
                    GlobalSchemes::Event.scheme_id(),
                    self.id,
                    EVENT_WRITE,
                    token.token(),
                );
                self.write_condition.notify_locked(token.token());

                return Ok(core::mem::size_of::<u64>());
            }

            if !block {
                return Err(Error::new(EAGAIN));
            }

            if !self
                .read_condition
                .wait(counter, "EventCounter::read", &mut token)
            {
                return Err(Error::new(EINTR));
            }
        }
    }

    pub fn write(&self, buf: UserSliceRo, block: bool, token: &mut CleanLockToken) -> Result<usize> {
        if buf.len() != core::mem::size_of::<u64>() {
            return Err(Error::new(EINVAL));
        }

        let value = unsafe { buf.read_exact::<u64>()? };
        if value == u64::MAX {
            return Err(Error::new(EINVAL));
        }

        loop {
            let counter = self.counter.lock(token.token());
            let (mut counter, mut token) = counter.into_split();

            if EVENTFD_COUNTER_MAX - *counter >= value {
                let was_zero = *counter == 0;
                *counter += value;

                if was_zero && value != 0 {
                    trigger_locked(
                        GlobalSchemes::Event.scheme_id(),
                        self.id,
                        EVENT_READ,
                        token.token(),
                    );
                    self.read_condition.notify_locked(token.token());
                }

                return Ok(core::mem::size_of::<u64>());
            }

            if !block {
                return Err(Error::new(EAGAIN));
            }

            if !self
                .write_condition
                .wait(counter, "EventCounter::write", &mut token)
            {
                return Err(Error::new(EINTR));
            }
        }
    }

    pub fn into_drop(self, _token: LockToken<'_, L1>) {
        drop(self);
    }
}

pub type EventQueueList = HashMap<EventQueueId, Arc<EventQueue>>;
pub type EventCounterList = HashMap<usize, Arc<EventCounter>>;

// Next queue id
static NEXT_QUEUE_ID: AtomicUsize = AtomicUsize::new(0);
static NEXT_COUNTER_ID: AtomicUsize = AtomicUsize::new(0);

/// Get next queue id
pub fn next_queue_id() -> EventQueueId {
    EventQueueId::from(NEXT_QUEUE_ID.fetch_add(1, Ordering::SeqCst))
}

pub fn next_counter_id() -> usize {
    EVENTFD_TAG_BIT | NEXT_COUNTER_ID.fetch_add(1, Ordering::SeqCst)
}

pub fn is_counter_id(id: usize) -> bool {
    id & EVENTFD_TAG_BIT != 0
}

// Current event queues
static QUEUES: RwLock<L2, EventQueueList> =
    RwLock::new(EventQueueList::with_hasher(DefaultHashBuilder::new()));
static COUNTERS: RwLock<L2, EventCounterList> =
    RwLock::new(EventCounterList::with_hasher(DefaultHashBuilder::new()));

/// Get the event queues list, const
pub fn queues(token: LockToken<'_, L0>) -> RwLockReadGuard<'_, L2, EventQueueList> {
    QUEUES.read(token)
}

/// Get the event queues list, mutable
pub fn queues_mut(token: LockToken<'_, L0>) -> RwLockWriteGuard<'_, L2, EventQueueList> {
    QUEUES.write(token)
}

pub fn counters(token: LockToken<'_, L0>) -> RwLockReadGuard<'_, L2, EventCounterList> {
    COUNTERS.read(token)
}

pub fn counters_mut(token: LockToken<'_, L0>) -> RwLockWriteGuard<'_, L2, EventCounterList> {
    COUNTERS.write(token)
}

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct RegKey {
    pub scheme: SchemeId,
    pub number: usize,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct QueueKey {
    pub queue: EventQueueId,
    pub id: usize,
    pub data: usize,
}

type Registry = HashMap<RegKey, HashMap<QueueKey, EventFlags>>;

static REGISTRY: RwLock<L2, Registry> =
    RwLock::new(HashMap::with_hasher(DefaultHashBuilder::new()));

pub fn register(
    reg_key: RegKey,
    queue_key: QueueKey,
    flags: EventFlags,
    token: &mut CleanLockToken,
) {
    let mut registry = REGISTRY.write(token.token());

    let entry = registry.entry(reg_key).or_default();

    if flags.is_empty() {
        entry.remove(&queue_key);
    } else {
        entry.insert(queue_key, flags);
    }
}

pub fn sync(reg_key: RegKey, token: &mut CleanLockToken) -> Result<EventFlags> {
    let mut flags = EventFlags::empty();

    {
        let registry = REGISTRY.read(token.token());
        if let Some(queue_list) = registry.get(&reg_key) {
            for (_queue_key, &queue_flags) in queue_list.iter() {
                flags |= queue_flags;
            }
        }
    }

    let scheme = scheme::get_scheme(token.token(), reg_key.scheme)?;

    scheme.fevent(reg_key.number, flags, token)
}

pub fn unregister_file(scheme: SchemeId, number: usize, token: &mut CleanLockToken) {
    let mut registry = REGISTRY.write(token.token());
    registry.remove(&RegKey { scheme, number });
}

//TODO: Implement unregister_queue
// pub fn unregister_queue(scheme: SchemeId, number: usize) {
//
// }

const MAX_EVENT: usize = 8;

#[must_use]
fn trigger_inner(
    scheme: SchemeId,
    number: usize,
    flags: EventFlags,
    todo: &mut SmallVec<[EventQueueId; MAX_EVENT]>,
    offset: &mut usize,
    mut token: LockToken<'_, L1>,
) -> bool {
    let mut matching_keys: SmallVec<[(QueueKey, EventFlags); MAX_EVENT]> = SmallVec::new();
    let mut full = false;

    {
        let registry = REGISTRY.read(token.token());
        if let Some(queue_list) = registry.get(&RegKey { scheme, number }) {
            for (queue_key, &queue_flags) in queue_list.iter().skip(*offset) {
                let common_flags = flags & queue_flags;
                if !common_flags.is_empty() {
                    if matching_keys.len() == matching_keys.inline_size() {
                        full = true;
                        break;
                    }
                    matching_keys.push((queue_key.clone(), common_flags));
                }
                *offset += 1;
            }
        }
    }

    while let Some((queue_key, common_flags)) = matching_keys.pop() {
        let Some(queue) = QUEUES.read(token.token()).get(&queue_key.queue).cloned() else {
            continue;
        };

        let event = Event {
            id: queue_key.id,
            flags: common_flags,
            data: queue_key.data,
        };

        todo.push(queue_key.queue);
        queue.queue.send_locked(event, token.token());
        if let Some(queue) = Arc::into_inner(queue) {
            queue.into_drop(token.token());
        }
    }

    full
}

pub fn trigger(scheme: SchemeId, number: usize, flags: EventFlags, token: &mut CleanLockToken) {
    trigger_locked(scheme, number, flags, token.token().downgrade());
}

pub fn trigger_locked(
    scheme: SchemeId,
    number: usize,
    flags: EventFlags,
    mut token: LockToken<'_, L1>,
) {
    let mut todo = SmallVec::<[EventQueueId; MAX_EVENT]>::new();
    let mut done = SmallVec::<[EventQueueId; MAX_EVENT]>::new();

    // First trigger with the original file
    let mut offset = 0;
    while trigger_inner(scheme, number, flags, &mut todo, &mut offset, token.token()) {}

    // Handle triggers on queues
    while let Some(queue_id) = todo.pop() {
        if let Err(insert_idx) = done.binary_search(&queue_id) {
            done.insert(insert_idx, queue_id);
            let mut offset = 0;
            while trigger_inner(
                GlobalSchemes::Event.scheme_id(),
                queue_id.into(),
                EventFlags::EVENT_READ,
                &mut todo,
                &mut offset,
                token.token(),
            ) {}
        }
    }
}