use crate::WriteBufStorage;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::os::fd::AsRawFd;
use dambi::MmapSlab;
use dope_core::driver::{
BackendToken, DriverLifecycle, DriverOps, FixedBuffers, PoolDriver, ProvidedRingOps,
};
use dope_core::{FdHandle, FixedFd};
use super::slot::{ConnFlags, CoreSlot};
pub(super) use crate::CodecLayer;
pub(crate) use crate::SlotId;
use dope_transport::Transport;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PushError {
ConnClosed,
GenerationMismatch,
BufferFull,
}
pub(crate) enum RecvDispatch {
Skipped,
EofClose {
slot_idx: usize,
},
Errno {
slot_idx: usize,
errno: i32,
more: bool,
},
Chunk {
slot_idx: usize,
ptr: *const u8,
len: usize,
bid: u16,
more: bool,
},
DrainOnly {
slot_idx: usize,
more: bool,
},
}
pub(crate) enum WriteDispatch {
Skipped,
Ready { slot_idx: usize, result: i32 },
}
pub(crate) enum CodecAck {
NoOp,
Err,
Progressed,
}
#[derive(Default, Clone, Copy)]
pub(super) struct RecvSlotState {
pub(super) epoch_seq: u32,
pub(super) epoch: u32,
pub(super) armed: bool,
pub(super) canceling: bool,
}
#[derive(Default, Clone, Copy)]
pub(super) struct WriteSlotState {
pub(super) epoch_seq: u32,
pub(super) in_flight_epoch: u32,
pub(super) has_in_flight: bool,
}
impl WriteSlotState {
#[inline(always)]
pub(super) fn next_token<T: BackendToken>(&mut self, external: usize) -> (T, u32) {
let token = T::next_for_slot(external, &mut self.epoch_seq);
let (_, epoch) = token.parts();
(token, epoch)
}
#[inline(always)]
pub(super) fn mark_inflight(&mut self, epoch: u32) {
self.has_in_flight = true;
self.in_flight_epoch = epoch;
}
}
#[derive(Default, Clone, Copy)]
pub(crate) struct ConnIoSlot {
pub(super) recv: RecvSlotState,
pub(super) write: WriteSlotState,
}
pub struct PoolCore<
T: Transport,
L: CodecLayer,
W: WriteBufStorage,
U: Default + 'static,
B: PoolDriver,
> {
pub(crate) slots: MmapSlab<CoreSlot<L, W, U>>,
pub(crate) free_conns: Vec<u16>,
pub(crate) recv_rearm_pending: VecDeque<u16>,
pub(crate) slot_id_base: usize,
pub(super) _markers: PhantomData<(T, L, W, B)>,
}
impl<T: Transport, L: CodecLayer, W: WriteBufStorage, U: Default + 'static, B: PoolDriver>
PoolCore<T, L, W, U, B>
{
pub(crate) fn new(max_conn: usize, slot_id_base: usize) -> std::io::Result<Self> {
let mut slots = MmapSlab::<CoreSlot<L, W, U>>::new_zeroed(max_conn)?;
slots.prewarm();
for slot in slots.iter_mut() {
slot.codec_state.write(L::ConnState::default());
slot.user.write(U::default());
}
let mut free_conns: Vec<u16> = (0..max_conn as u16).rev().collect();
free_conns.shrink_to_fit();
Ok(Self {
slots,
free_conns,
recv_rearm_pending: VecDeque::with_capacity(max_conn),
slot_id_base,
_markers: PhantomData,
})
}
#[inline(always)]
pub(crate) fn internal_slot(&self, external: usize) -> Option<usize> {
external
.checked_sub(self.slot_id_base)
.filter(|&i| i < self.slots.len())
}
#[cold]
#[inline(never)]
pub(crate) fn allocate_slot(
&mut self,
driver: &mut B::Driver,
fd_handle: FdHandle,
) -> Option<(usize, SlotId)> {
let fixed_idx = match driver.register_fixed(fd_handle.as_fd()) {
Ok(Some(idx)) => idx,
_ => return None,
};
let Some(slot) = self.free_conns.pop() else {
driver.unregister_fixed(fixed_idx);
return None;
};
let slot_idx = slot as usize;
let fixed_fd = fd_handle.fixed_io(fixed_idx);
std::mem::forget(fd_handle.into_owned());
let next_gen = {
let entry = &mut self.slots[slot_idx];
entry.conn.reset_for_reuse();
entry.conn.fd = fixed_fd;
entry.conn.flags.set(ConnFlags::LIVE, true);
entry.io.write.has_in_flight = false;
entry.buf.pending_msghdr = unsafe { std::mem::zeroed() };
entry.buf.pending_msghdr.msg_iov = entry.buf.pending_iovs.as_mut_ptr();
entry.conn.generation
};
*self.slots[slot_idx].codec_state.get_mut() = L::ConnState::default();
*self.slots[slot_idx].user.get_mut() = U::default();
if !self.arm_recv(driver, slot_idx) {
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(fixed_fd.as_fd().as_raw_fd()) };
self.slots[slot_idx].conn.flags = ConnFlags::empty();
self.free_conns.push(slot);
return None;
}
Some((slot_idx, SlotId::from_parts(slot as u32, next_gen)))
}
pub(crate) fn handle_codec_inflight(&mut self, slot_idx: usize, result: i32) -> CodecAck {
if !self.slots[slot_idx].conn.flags.contains(ConnFlags::LIVE) {
let _ = self.consume_codec_inflight(slot_idx, usize::MAX);
return CodecAck::NoOp;
}
if result < 0 {
let _ = self.consume_codec_inflight(slot_idx, usize::MAX);
return CodecAck::Err;
}
let n = result as usize;
let _more = self.consume_codec_inflight(slot_idx, n);
CodecAck::Progressed
}
pub(crate) fn resolve_live_slot(&self, conn_id: SlotId) -> Result<usize, PushError> {
let slot_idx = conn_id.slot() as usize;
if slot_idx >= self.slots.len() {
return Err(PushError::ConnClosed);
}
let entry = &self.slots[slot_idx];
if entry.conn.generation != conn_id.generation() {
return Err(PushError::GenerationMismatch);
}
if !entry.conn.flags.contains(ConnFlags::LIVE) {
return Err(PushError::ConnClosed);
}
Ok(slot_idx)
}
pub(crate) fn classify_write(
&mut self,
token: B::Token,
result: i32,
notif: bool,
) -> WriteDispatch {
let (external, epoch) = token.parts();
let Some(slot_idx) = self.internal_slot(external) else {
return WriteDispatch::Skipped;
};
let write_slot = &mut self.slots[slot_idx].io.write;
if !write_slot.has_in_flight || write_slot.in_flight_epoch != epoch {
return WriteDispatch::Skipped;
}
if notif {
return WriteDispatch::Skipped;
}
write_slot.has_in_flight = false;
WriteDispatch::Ready { slot_idx, result }
}
pub(crate) fn classify_recv(
&mut self,
driver: &mut B::Driver,
token: B::Token,
result: i32,
more: bool,
bid: Option<u16>,
) -> RecvDispatch {
let (external, epoch) = token.parts();
let Some(slot_idx) = self.internal_slot(external) else {
if let Some(b) = bid {
driver.provided_ring().defer(b);
}
return RecvDispatch::Skipped;
};
if self.slots[slot_idx].io.recv.epoch != epoch {
if let Some(b) = bid {
driver.provided_ring().defer(b);
}
return RecvDispatch::Skipped;
}
let live = self.slots[slot_idx].conn.flags.contains(ConnFlags::LIVE);
let canceling = self.slots[slot_idx].io.recv.canceling;
if !live || canceling {
if let Some(b) = bid {
driver.provided_ring().defer(b);
}
if !more {
self.slots[slot_idx].io.recv.armed = false;
}
return RecvDispatch::Skipped;
}
if result < 0 {
if let Some(b) = bid {
driver.provided_ring().defer(b);
}
return RecvDispatch::Errno {
slot_idx,
errno: -result,
more,
};
}
if result == 0 {
if let Some(b) = bid {
driver.provided_ring().defer(b);
}
self.slots[slot_idx].io.recv.armed = false;
return RecvDispatch::EofClose { slot_idx };
}
let Some(b) = bid else {
return RecvDispatch::DrainOnly { slot_idx, more };
};
let len = result as usize;
let (ptr, _) = driver.provided_ring().ptr_len(b);
RecvDispatch::Chunk {
slot_idx,
ptr,
len,
bid: b,
more,
}
}
pub(crate) fn arm_recv(&mut self, driver: &mut B::Driver, slot_idx: usize) -> bool {
let buf_group = driver.buf_group();
let external = self.slot_id_base + slot_idx;
let entry = &mut self.slots[slot_idx];
let token =
<B::Token as BackendToken>::next_for_slot(external, &mut entry.io.recv.epoch_seq);
let (_, epoch) = token.parts();
let fd = entry.conn.fd;
let armed = driver.arm_with_retry(|d| d.arm_recv_multi(token, fd, buf_group));
let recv = &mut self.slots[slot_idx].io.recv;
recv.epoch = epoch;
recv.canceling = false;
if armed {
recv.armed = true;
return true;
}
recv.armed = false;
self.recv_rearm_pending.push_back(slot_idx as u16);
false
}
#[inline(always)]
pub(crate) fn flush_recv_rearm(&mut self, driver: &mut B::Driver) {
if self.recv_rearm_pending.is_empty() {
return;
}
self.flush_recv_rearm_slow(driver);
}
#[cold]
#[inline(never)]
fn flush_recv_rearm_slow(&mut self, driver: &mut B::Driver) {
let n = self.recv_rearm_pending.len();
for _ in 0..n {
let Some(slot_idx) = self.recv_rearm_pending.pop_front() else {
return;
};
let slot_idx = slot_idx as usize;
let entry = &self.slots[slot_idx];
if !entry.conn.flags.contains(ConnFlags::LIVE)
|| entry.io.recv.armed
|| entry.io.recv.canceling
{
continue;
}
let _ = self.arm_recv(driver, slot_idx);
}
}
pub(crate) fn reserve_slot(&mut self) -> Option<(usize, SlotId)> {
let slot = self.free_conns.pop()?;
let slot_idx = slot as usize;
let entry = &mut self.slots[slot_idx];
entry.conn.reset_for_reuse();
entry.io.write.has_in_flight = false;
let next_gen = entry.conn.generation;
*entry.codec_state.get_mut() = L::ConnState::default();
*entry.user.get_mut() = U::default();
Some((slot_idx, SlotId::from_parts(slot as u32, next_gen)))
}
pub(crate) fn activate_slot(
&mut self,
driver: &mut B::Driver,
slot_idx: usize,
fixed_fd: FixedFd,
) -> bool {
{
let entry = &mut self.slots[slot_idx];
entry.conn.fd = fixed_fd;
entry.conn.flags.set(ConnFlags::LIVE, true);
entry.buf.pending_msghdr = unsafe { std::mem::zeroed() };
entry.buf.pending_msghdr.msg_iov = entry.buf.pending_iovs.as_mut_ptr();
}
if !self.arm_recv(driver, slot_idx) {
self.slots[slot_idx].conn.flags = ConnFlags::empty();
return false;
}
true
}
pub(crate) fn release_reserved(&mut self, slot_idx: usize) {
*self.slots[slot_idx].user.get_mut() = U::default();
let entry = &mut self.slots[slot_idx];
entry.conn.flags = ConnFlags::empty();
self.free_conns.push(slot_idx as u16);
}
pub(crate) fn free_slot(&mut self, driver: &mut B::Driver, slot_idx: usize) {
*self.slots[slot_idx].user.get_mut() = U::default();
self.close_slot_core(driver, slot_idx);
}
fn close_slot_core(&mut self, driver: &mut B::Driver, slot_idx: usize) {
let entry = &mut self.slots[slot_idx];
if !entry.conn.flags.contains(ConnFlags::LIVE) {
return;
}
let fixed_idx = entry.conn.fd.fixed_index();
let fd_raw = std::os::fd::AsRawFd::as_raw_fd(&entry.conn.fd.as_fd());
if entry.io.recv.armed {
let token = <B::Token as BackendToken>::from_index_raw(
self.slot_id_base + slot_idx,
entry.io.recv.epoch,
);
entry.io.recv.canceling = true;
driver.cancel_recv(token);
} else {
entry.io.recv.canceling = false;
}
entry.io.recv.armed = false;
entry.io.write.has_in_flight = false;
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(fd_raw) };
entry.conn.flags = ConnFlags::empty();
self.free_conns.push(slot_idx as u16);
}
pub(crate) fn try_submit_codec(&mut self, driver: &mut B::Driver, slot_idx: usize) {
if self.slots[slot_idx].io.write.has_in_flight {
return;
}
let Some((ptr, len)) = L::take_inflight(self.slots[slot_idx].codec_state.get_mut()) else {
return;
};
let entry = &mut self.slots[slot_idx];
let external = self.slot_id_base + slot_idx;
let (token, epoch) = entry.io.write.next_token::<B::Token>(external);
let conn_fd = entry.conn.fd;
if driver.submit_send_tagged(token, conn_fd, ptr, len) {
entry.io.write.mark_inflight(epoch);
}
}
pub(crate) fn process_codec_inbound(
&mut self,
driver: &mut B::Driver,
slot_idx: usize,
src: &[u8],
) -> Option<Vec<u8>> {
L::process_inbound(self.slots[slot_idx].codec_state.get_mut(), src);
self.try_submit_codec(driver, slot_idx);
if L::close_pending(self.slots[slot_idx].codec_state.get()) {
self.slots[slot_idx]
.conn
.flags
.set(ConnFlags::CLOSE_AFTER, true);
}
if !L::is_active(self.slots[slot_idx].codec_state.get()) {
return None;
}
let plaintext = L::take_plaintext(self.slots[slot_idx].codec_state.get_mut());
if plaintext.is_empty() {
L::install_plaintext(self.slots[slot_idx].codec_state.get_mut(), plaintext);
return None;
}
Some(plaintext)
}
pub(crate) fn install_plaintext(&mut self, slot_idx: usize, plaintext: Vec<u8>) {
L::install_plaintext(self.slots[slot_idx].codec_state.get_mut(), plaintext);
}
pub(crate) fn encrypt_write_buf(&mut self, slot_idx: usize) {
let head = self.slots[slot_idx].conn.write_head as usize;
let inflight_end = self.slots[slot_idx].conn.write_inflight_end as usize;
if head <= inflight_end {
return;
}
let entry = &mut self.slots[slot_idx];
let plaintext = &entry.buf.write_buf.as_mut_slice()[inflight_end..head];
L::process_outbound(entry.codec_state.get_mut(), plaintext);
entry.conn.write_head = 0;
entry.conn.write_tail = 0;
entry.conn.write_inflight_end = 0;
}
pub(crate) fn consume_codec_inflight(&mut self, slot_idx: usize, n: usize) -> bool {
L::consume_inflight(self.slots[slot_idx].codec_state.get_mut(), n)
}
pub(crate) fn codec_idle(&self, slot_idx: usize) -> bool {
let state = self.slots[slot_idx].codec_state.get();
L::wire_slice(state).is_empty() && !L::has_inflight(state)
}
pub(crate) fn codec_close_pending(&self, slot_idx: usize) -> bool {
L::close_pending(self.slots[slot_idx].codec_state.get())
}
}
impl<T: Transport, L: CodecLayer, W: WriteBufStorage, U: Default + 'static, B: PoolDriver> Drop
for PoolCore<T, L, W, U, B>
{
fn drop(&mut self) {
for slot in self.slots.iter_mut() {
unsafe {
slot.codec_state.drop_in_place();
slot.user.drop_in_place();
}
}
}
}