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;
pub trait WaitContext {
fn fd_wait(&mut self, fd: i32, 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);
}
#[repr(C)]
pub(crate) struct FdWait {
fd: i32,
events: u32,
index: usize,
}
impl FdWait {
pub fn new(fd: i32, events: u32) -> Self {
debug_assert!((events & (POLLIN | POLLOUT)) > 0);
Self {
fd,
events,
index: usize::MAX,
}
}
}
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.events) {
Ok(index) => self.index = index,
Err(e) => return Poll::Ready(Err(e)),
}
return Poll::Pending;
}
match ctx.fd_awaked(self.index, self.events) {
Some(_) => Poll::Ready(Ok(())),
None => Poll::Pending,
}
}
}
pub(crate) struct FdDel(i32);
impl FdDel {
pub fn new(fd: i32) -> Self {
Self(fd)
}
}
impl Future for FdDel {
type Output = ();
fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
ctx.fd_del(self.0);
Poll::Ready(())
}
}
#[repr(C)]
pub(crate) struct FdSet<'a> {
cnt: usize,
stats: &'a mut [FdState],
}
#[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 MAX_SIZE: usize = 32;
const FD_ABORT: u32 = 0x8000_0000;
pub fn new() -> Self {
Self {
stats: &mut [],
cnt: 0,
}
}
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(&mut self, worker: &mut Worker) {
for (index, stat) in (&self.stats).iter().enumerate() {
if index == self.cnt {
break;
}
if stat.fd > -1 {
let _ = unsafe { worker.sched.del_fd_event(stat.fd) };
}
}
}
pub fn fd_del(&mut self, fd: i32, worker: &mut Worker) {
for (index, stat) in self.stats.iter_mut().enumerate() {
if index == self.cnt {
break;
}
if stat.fd == fd {
let _ = unsafe { worker.sched.del_fd_event(fd) };
stat.fd = -1;
self.cnt -= 1;
self.stats[self.cnt].index = index as u32;
break;
}
}
}
fn fd_find(&self, fd: i32) -> Option<usize> {
for (index, stat) in self.stats.iter().enumerate() {
if index == self.cnt {
return None;
}
if stat.fd == fd {
return Some(index);
}
}
None
}
fn extend(&mut self, worker: &mut Worker) -> Result<()> {
let len = if self.stats.len() == 0 { 2 } else { self.stats.len() * 2 };
if len > Self::MAX_SIZE {
return Err(Error::new(hierr::ERANGE));
}
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)))
})?;
self.stats = new_stats.leak().0;
Ok(())
}
fn fd_add(&mut self, fd: i32, worker: &mut Worker) -> Result<usize> {
if let Some(index) = self.fd_find(fd) {
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(
&mut self,
fd: i32,
events: u32,
e: &Event,
worker: &mut Worker,
) -> Result<usize> {
debug_assert!((events & (POLLIN | POLLOUT)) > 0);
let index = self.fd_add(fd, 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.sched.add_fd_event(e, new_events, stat.fd, index as u16)? };
} else {
unsafe { worker.sched.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()
}
}