use core::{
sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering},
task::Waker,
time::Duration,
};
use ax_errno::AxResult;
use ax_task::future::{block_on, poll_io, timeout};
use axpoll::{IoEvents, Pollable};
use crate::{
get_service,
options::{Configurable, GetSocketOption, SetSocketOption},
};
pub(crate) struct GeneralOptions {
nonblock: AtomicBool,
reuse_address: AtomicBool,
send_timeout_nanos: AtomicU64,
recv_timeout_nanos: AtomicU64,
device_mask: AtomicU32,
}
impl Default for GeneralOptions {
fn default() -> Self {
Self::new()
}
}
impl GeneralOptions {
pub fn new() -> Self {
Self {
nonblock: AtomicBool::new(false),
reuse_address: AtomicBool::new(false),
send_timeout_nanos: AtomicU64::new(0),
recv_timeout_nanos: AtomicU64::new(0),
device_mask: AtomicU32::new(0),
}
}
pub fn nonblocking(&self) -> bool {
self.nonblock.load(Ordering::Relaxed)
}
pub fn reuse_address(&self) -> bool {
self.reuse_address.load(Ordering::Relaxed)
}
pub fn send_timeout(&self) -> Option<Duration> {
let nanos = self.send_timeout_nanos.load(Ordering::Relaxed);
(nanos > 0).then(|| Duration::from_nanos(nanos))
}
pub fn recv_timeout(&self) -> Option<Duration> {
let nanos = self.recv_timeout_nanos.load(Ordering::Relaxed);
(nanos > 0).then(|| Duration::from_nanos(nanos))
}
pub fn set_device_mask(&self, mask: u32) {
self.device_mask.store(mask, Ordering::Release);
}
pub fn device_mask(&self) -> u32 {
self.device_mask.load(Ordering::Acquire)
}
pub fn register_waker(&self, waker: &Waker) {
get_service().register_waker(self.device_mask(), waker);
}
pub fn send_poller<P: Pollable, F: FnMut() -> AxResult<T>, T>(
&self,
pollable: &P,
f: F,
) -> AxResult<T> {
block_on(timeout(
self.send_timeout(),
poll_io(pollable, IoEvents::OUT, self.nonblocking(), f),
))?
}
pub fn recv_poller<P: Pollable, F: FnMut() -> AxResult<T>, T>(
&self,
pollable: &P,
f: F,
) -> AxResult<T> {
block_on(timeout(
self.recv_timeout(),
poll_io(pollable, IoEvents::IN, self.nonblocking(), f),
))?
}
}
impl Configurable for GeneralOptions {
fn get_option_inner(&self, option: &mut GetSocketOption) -> AxResult<bool> {
use GetSocketOption as O;
match option {
O::Error(error) => {
**error = 0;
}
O::NonBlocking(nonblock) => {
**nonblock = self.nonblocking();
}
O::ReuseAddress(reuse) => {
**reuse = self.reuse_address();
}
O::SendTimeout(timeout) => {
**timeout = Duration::from_nanos(self.send_timeout_nanos.load(Ordering::Relaxed));
}
O::ReceiveTimeout(timeout) => {
**timeout = Duration::from_nanos(self.recv_timeout_nanos.load(Ordering::Relaxed));
}
_ => return Ok(false),
}
Ok(true)
}
fn set_option_inner(&self, option: SetSocketOption) -> AxResult<bool> {
use SetSocketOption as O;
match option {
O::NonBlocking(nonblock) => {
self.nonblock.store(*nonblock, Ordering::Relaxed);
}
O::ReuseAddress(reuse) => {
self.reuse_address.store(*reuse, Ordering::Relaxed);
}
O::SendTimeout(timeout) => {
self.send_timeout_nanos
.store(timeout.as_nanos() as u64, Ordering::Relaxed);
}
O::ReceiveTimeout(timeout) => {
self.recv_timeout_nanos
.store(timeout.as_nanos() as u64, Ordering::Relaxed);
}
O::SendBuffer(_) | O::ReceiveBuffer(_) => {
}
_ => return Ok(false),
}
Ok(true)
}
}