use crate::event::{Event, Scheduler, POLLIN, POLLONESHOT, POLLOUT};
use crate::runtime::{TaskContext, Worker};
use crate::{Error, Result};
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use hipool::{Boxed, GenericAlloc};
pub trait WaitContext {
fn fd_wait(&mut self, fd: i32, index: usize, events: u32) -> Result<usize>;
fn fd_awaked(&mut self, index: usize, events: u32) -> Option<u32>;
fn fd_abort(&mut self, index: usize);
fn fd_del(&mut self, fd: i32, index: usize);
}
#[repr(C)]
pub(crate) struct FdWait<'a> {
fd: i32,
events: u32,
index: usize,
last: &'a mut usize,
}
impl<'a> FdWait<'a> {
pub fn new(fd: i32, last: &'a mut usize, events: u32) -> Self {
debug_assert!((events & (POLLIN | POLLOUT)) > 0);
Self {
fd,
events,
index: usize::MAX,
last,
}
}
}
impl Future for FdWait<'_> {
type Output = Result<()>;
fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
if ctx.aborted() {
if self.index < usize::MAX {
ctx.fd_abort(self.index);
}
return Poll::Pending;
}
if self.index == usize::MAX {
match ctx.fd_wait(self.fd, *self.last, self.events) {
Ok(index) => self.index = index,
Err(e) => return Poll::Ready(Err(e)),
}
if *self.last != self.index {
*self.last = self.index;
}
return Poll::Pending;
}
match ctx.fd_awaked(self.index, self.events) {
Some(_) => Poll::Ready(Ok(())),
None => Poll::Pending,
}
}
}
pub(crate) struct FdDel {
fd: i32,
index: usize,
}
impl FdDel {
pub fn new(fd: i32, index: usize) -> Self {
Self { fd, index }
}
}
impl Future for FdDel {
type Output = ();
fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
ctx.fd_del(self.fd, self.index);
Poll::Ready(())
}
}
#[repr(C)]
pub(crate) struct FdSet<'a> {
stats: &'a mut [FdState],
cnt: usize,
capacity: usize,
}
#[repr(C)]
#[derive(Clone)]
pub(crate) struct FdState {
fd: i32,
recv: u32,
events: u32,
index: u32,
}
impl FdState {
fn new(index: usize) -> Self {
Self {
fd: -1,
recv: 0,
events: 0,
index: index as u32,
}
}
fn set_fd(&mut self, fd: i32) {
self.fd = fd;
self.recv = 0;
self.events = 0;
}
}
impl<'a> FdSet<'a> {
const FD_ABORT: u32 = 0x8000_0000;
pub fn new(capacity: usize) -> Self {
Self {
stats: &mut [],
cnt: 0,
capacity,
}
}
pub fn handle_events(&mut self, events: u32, flags: u16) -> bool {
let index = flags as usize;
let stat = &mut self.stats[index];
stat.recv |= events;
(stat.recv & Self::FD_ABORT) == 0
}
pub fn exit<T: Scheduler>(&mut self, worker: &mut T) {
let mut cnt = 0;
for stat in (&self.stats).iter() {
if stat.fd > -1 {
let _ = unsafe { worker.del_fd_event(stat.fd) };
cnt += 1;
if cnt == self.cnt {
return;
}
}
}
}
pub fn fd_del<T: Scheduler>(&mut self, fd: i32, index: usize, worker: &mut T) {
if let Some(index) = self.fd_find(fd, index) {
let stat = &mut self.stats[index];
let _ = unsafe { worker.del_fd_event(fd) };
stat.fd = -1;
self.cnt -= 1;
self.stats[self.cnt].index = index as u32;
}
}
fn fd_find(&self, fd: i32, index: usize) -> Option<usize> {
if let Some(stat) = self.stats.get(index) {
if stat.fd == fd {
return Some(index);
}
}
let mut cnt = 0;
for (index, stat) in self.stats.iter().enumerate() {
if stat.fd == fd {
return Some(index);
}
if stat.fd > -1 {
cnt += 1;
if cnt == self.cnt {
return None;
}
}
}
None
}
fn extend<T: Worker>(&mut self, worker: &T) -> Result<()> {
if self.stats.len() >= self.capacity {
return Err(Error::new(hierr::ERANGE));
}
let mut len = if self.stats.len() == 0 {
2
} else {
self.stats.len() * 2
};
if len > self.capacity {
len = self.capacity;
}
let new_stats = Boxed::new_slice_then_in(worker.pool(), len, |n| {
Ok(self
.stats
.get(n)
.map(|v| v.clone())
.unwrap_or(FdState::new(n)))
})?;
unsafe { worker.pool().release_slice(self.stats.into()) };
self.stats = new_stats.leak().0;
Ok(())
}
fn fd_add<T: Worker>(&mut self, fd: i32, index: usize, worker: &T) -> Result<usize> {
if let Some(index) = self.fd_find(fd, index) {
return Ok(index);
}
if self.cnt == self.stats.len() {
self.extend(worker)?;
}
let index = self.stats[self.cnt].index as usize;
let stat = &mut self.stats[index];
stat.set_fd(fd);
self.cnt += 1;
Ok(index)
}
pub fn fd_wait<T: Worker>(
&mut self,
fd: i32,
index: usize,
events: u32,
e: &Event,
worker: &mut T,
) -> Result<usize> {
debug_assert!((events & (POLLIN | POLLOUT)) > 0);
let index = self.fd_add(fd, index, worker)?;
let stat = &mut self.stats[index];
let new_events = stat.events | events;
if new_events != stat.events || (new_events & POLLONESHOT) > 0 {
if stat.events == 0 {
unsafe { worker.add_fd_event(e, new_events, stat.fd, index as u16)? };
} else {
unsafe { worker.mod_fd_event(e, new_events, stat.fd, index as u16)? };
}
stat.events = new_events;
}
stat.recv &= !(events | Self::FD_ABORT);
Ok(index)
}
pub fn fd_awaked(&mut self, index: usize, events: u32) -> Option<u32> {
let stat = &self.stats[index];
let events = stat.recv as u32 & events;
if events > 0 {
Some(events)
} else {
None
}
}
pub fn fd_abort(&mut self, index: usize) {
let stat = &mut self.stats[index];
stat.recv |= Self::FD_ABORT;
}
pub fn fd_capacity(&self) -> usize {
self.stats.len()
}
}