use super::{event_list_new, Event, Poll, Scheduler, POLLIN, POLLONESHOT, POLLOUT};
use crate::Result;
use core::borrow::Borrow;
use core::cell::UnsafeCell;
use core::ptr;
use hicollections::{rbtree, List, RbTree, RbTreeNode};
use hipool::{Allocator, Boxed};
pub(crate) struct FdEvent {
rnode: RbTreeNode,
next: UnsafeCell<*mut FdEvent>,
fd: i32,
events: u32,
recv: u32,
cnt: i32,
elist: List<Event>,
}
impl Borrow<i32> for FdEvent {
fn borrow(&self) -> &i32 {
&self.fd
}
}
impl FdEvent {
pub fn new() -> Self {
Self {
fd: -1,
events: 0,
recv: 0,
cnt: 0,
elist: event_list_new(),
rnode: RbTreeNode::new(),
next: UnsafeCell::new(ptr::null_mut()),
}
}
pub fn on_event(&mut self, events: u32, sched: &mut Scheduler) {
self.recv |= events;
for e in self.elist.iter() {
let events = e.events.get() & self.recv;
if events > 0 {
e.handle(events, sched);
}
}
}
fn add_event<T: Poll>(&mut self, event: &Event, events: u32, poll: &mut T) -> Result<()> {
match self.mod_event(event, events, poll) {
Ok(_) => {
self.cnt += 1;
Ok(())
}
Err(e) => Err(e),
}
}
fn mod_event<T: Poll>(&mut self, event: &Event, events: u32, poll: &mut T) -> Result<()> {
let new_events = events | self.events;
if new_events != self.events || (new_events & POLLONESHOT) == POLLONESHOT {
if self.events == 0 {
poll.add_event(self.fd, new_events, self as *mut _ as u64)?;
} else {
poll.mod_event(self.fd, new_events, self as *mut _ as u64)?;
}
self.events = new_events;
}
self.recv &= !events;
event.events.set(events & (POLLIN | POLLOUT));
unsafe { self.elist.add_tail(event) };
Ok(())
}
fn del_event<T: Poll>(&mut self, event: &Event, poll: &mut T) -> bool {
unsafe { self.elist.del(event) };
if self.cnt == 1 {
let _ = poll.del_event(self.fd, self as *mut _ as u64);
assert!(self.elist.empty());
true
} else {
self.cnt -= 1;
false
}
}
fn next(&self) -> *mut *mut Self {
self.next.get()
}
fn init(&mut self, fd: i32) {
self.fd = fd;
}
}
pub(crate) struct FdSet<'a, A: Allocator> {
fdset: RbTree<FdEvent>,
cached: *mut FdEvent,
alloc: &'a A,
}
impl<'a, A: Allocator> FdSet<'a, A> {
pub fn new(alloc: &'a A) -> Self {
Self {
fdset: rbtree!(FdEvent, rnode),
cached: ptr::null_mut(),
alloc,
}
}
pub fn add_fd_event<T: Poll>(
&mut self,
e: &Event,
events: u32,
fd: i32,
poll: &mut T,
) -> Result<()> {
let fdevent = unsafe { &mut *self.get_fdevent(fd)? };
fdevent.add_event(e, events, poll)
}
pub fn mod_fd_event<T: Poll>(
&mut self,
e: &Event,
events: u32,
fd: i32,
poll: &mut T,
) -> Result<()> {
let fdevent = unsafe { &mut *self.get_fdevent(fd)? };
fdevent.mod_event(e, events, poll)
}
pub fn del_fd_event<T: Poll>(&mut self, e: &Event, fd: i32, poll: &mut T) -> Result<()> {
if let Some(fdevent) = self
.fdset
.find(&fd)
.map(|f| f as *const FdEvent as *mut FdEvent)
{
let ret = unsafe { &mut *fdevent }.del_event(e, poll);
if ret {
self.del_fdevent(fd);
}
}
Ok(())
}
fn new_fdevent(&mut self, fd: i32) -> Result<*mut FdEvent> {
let fdevent = if !self.cached.is_null() {
let fdevent = unsafe { &mut *self.cached };
self.cached = unsafe { *fdevent.next() };
fdevent
} else {
Boxed::new_in(self.alloc, FdEvent::new())?.leak().0
};
fdevent.init(fd);
let ret = fdevent as *mut FdEvent;
unsafe { self.fdset.insert_borrow::<i32>(fdevent, false) };
Ok(ret)
}
fn del_fdevent(&mut self, fd: i32) {
if let Some(fdevent) = self.fdset.remove(&fd) {
unsafe { *fdevent.next() = self.cached };
self.cached = (fdevent as *const FdEvent).cast_mut();
}
}
fn get_fdevent(&mut self, fd: i32) -> Result<*mut FdEvent> {
if let Some(fdevent) = self.fdset.find(&fd) {
Ok((fdevent as *const FdEvent).cast_mut())
} else {
self.new_fdevent(fd)
}
}
}