use super::super::Pool;
use super::super::core::{CodecAck, PushError, WriteDispatch};
use super::super::slot::{ConnFlags, CoreSlot};
use super::Server;
use super::state::SlotEgress;
use crate::CodecLayer;
use crate::SlotId;
use crate::WriteBufStorage;
use crate::protocol::server::ServerProtocol;
use dope_core::IoVec;
use dope_core::driver::{DriverOps, PoolDriver};
use dope_core::profile::Profile;
use dope_transport::Transport;
enum WriteEventOutcome {
Idle,
More,
Close,
}
impl<H: ServerProtocol, T: Transport, F: Profile, B: PoolDriver> super::super::PoolWrite<B>
for Pool<Server<H, T>, F, B>
{
fn on_write_event(
&mut self,
driver: &mut B::Driver,
token: B::Token,
result: i32,
notif: bool,
) {
let (slot_idx, result) = match self.core.classify_write(token, result, notif) {
WriteDispatch::Skipped => return,
WriteDispatch::Ready { slot_idx, result } => (slot_idx, result),
};
if <H::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
if !self.core.slots[slot_idx]
.conn
.flags
.contains(ConnFlags::LIVE)
{
return;
}
if result < 0 {
self.core.free_slot(driver, slot_idx);
return;
}
let n = result as u32;
match self.advance_write_state(slot_idx, n) {
WriteEventOutcome::Close => self.core.free_slot(driver, slot_idx),
WriteEventOutcome::More => self.submit_slot_write(driver, slot_idx),
WriteEventOutcome::Idle => {}
}
} else {
self.handle_write_event_codec(driver, slot_idx, result);
}
}
}
impl<H: ServerProtocol, T: Transport, F: Profile, B: PoolDriver> Pool<Server<H, T>, F, B> {
fn handle_write_event_codec(&mut self, driver: &mut B::Driver, slot_idx: usize, result: i32) {
match self.core.handle_codec_inflight(slot_idx, result) {
CodecAck::NoOp => {}
CodecAck::Err => self.core.free_slot(driver, slot_idx),
CodecAck::Progressed => {
self.submit_slot_write(driver, slot_idx);
let close = self.core.slots[slot_idx]
.conn
.flags
.contains(ConnFlags::CLOSE_AFTER)
|| self.core.codec_close_pending(slot_idx);
if self.core.codec_idle(slot_idx) && close {
self.core.free_slot(driver, slot_idx);
}
}
}
}
#[inline(always)]
fn advance_write_state(&mut self, slot_idx: usize, n: u32) -> WriteEventOutcome {
let CoreSlot { conn, user, .. } = &mut self.core.slots[slot_idx];
let user = user.get_mut();
if user.egress.is_stream() {
if conn.write_tail < conn.write_inflight_end {
let new_tail = (conn.write_tail as u32).saturating_add(n);
conn.write_tail = new_tail.min(u16::MAX as u32) as u16;
if conn.write_tail >= conn.write_inflight_end {
conn.write_head = 0;
conn.write_tail = 0;
conn.write_inflight_end = 0;
return WriteEventOutcome::More;
}
return WriteEventOutcome::Idle;
}
let SlotEgress::Stream(stream_state) = &mut user.egress else {
return WriteEventOutcome::Close;
};
stream_state.advance(n);
if stream_state.eof && stream_state.cur_chunk.is_none() {
user.egress = SlotEgress::Idle;
if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
return WriteEventOutcome::Close;
}
return WriteEventOutcome::Idle;
}
return WriteEventOutcome::More;
}
if !conn.pending_body_ptr.is_null() {
let new_tail = conn.large_tail.saturating_add(n);
conn.large_tail = new_tail;
if conn.large_tail >= conn.large_inflight {
conn.write_head = 0;
conn.write_tail = 0;
conn.write_inflight_end = 0;
conn.pending_body_ptr = std::ptr::null();
conn.pending_body_len = 0;
conn.large_inflight = 0;
conn.large_tail = 0;
if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
return WriteEventOutcome::Close;
}
}
return WriteEventOutcome::Idle;
}
if conn.large_inflight > 0 {
let new_tail = conn.large_tail.saturating_add(n);
conn.large_tail = new_tail;
if conn.large_tail >= conn.large_inflight {
conn.large_inflight = 0;
conn.large_tail = 0;
if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
return WriteEventOutcome::Close;
}
if user.egress.gather_iov_count().is_some() {
return WriteEventOutcome::More;
}
}
return WriteEventOutcome::Idle;
}
let new_tail = (conn.write_tail as u32).saturating_add(n);
conn.write_tail = new_tail.min(u16::MAX as u32) as u16;
if conn.write_tail >= conn.write_inflight_end {
conn.write_head = 0;
conn.write_tail = 0;
conn.write_inflight_end = 0;
if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
return WriteEventOutcome::Close;
}
}
WriteEventOutcome::Idle
}
#[inline(always)]
fn try_submit_gather(&mut self, driver: &mut B::Driver, slot_idx: usize) -> bool {
let Some(iov_count) = self.core.slots[slot_idx]
.user
.get()
.egress
.gather_iov_count()
else {
return false;
};
if self.core.slots[slot_idx].io.write.has_in_flight {
return true;
}
let entry = &mut self.core.slots[slot_idx];
entry.buf.pending_msghdr.msg_iovlen = iov_count as _;
let msg_ptr: *const libc::msghdr = &entry.buf.pending_msghdr;
let submitted = unsafe { self.driver_submit_send_msg(driver, slot_idx, msg_ptr) };
if submitted {
self.core.slots[slot_idx].user.get_mut().egress = SlotEgress::Idle;
}
true
}
pub fn push_response(
&mut self,
driver: &mut B::Driver,
conn_id: SlotId,
bytes: &[u8],
close_after: bool,
) -> Result<(), PushError> {
debug_assert!(
H::SUPPORTS_PARK,
"push_response requires ServerProtocol::SUPPORTS_PARK = true",
);
let slot_idx = self.core.resolve_live_slot(conn_id)?;
let entry = &mut self.core.slots[slot_idx];
let head = entry.conn.write_head as usize;
let room = entry.buf.write_buf.as_slice().len().saturating_sub(head);
if bytes.len() > room {
return Err(PushError::BufferFull);
}
entry.buf.write_buf.as_mut_slice()[head..head + bytes.len()].copy_from_slice(bytes);
entry.conn.write_head = (head + bytes.len()) as u16;
if close_after {
entry.conn.flags.set(ConnFlags::CLOSE_AFTER, true);
}
entry.conn.flags.set(ConnFlags::PARKED, false);
self.submit_slot_write(driver, slot_idx);
Ok(())
}
#[inline(always)]
pub(super) fn submit_slot_write(&mut self, driver: &mut B::Driver, slot_idx: usize) {
if !<H::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
return self.submit_slot_write_codec(driver, slot_idx);
}
if self.try_submit_gather(driver, slot_idx) {
return;
}
self.submit_slot_write_plain(driver, slot_idx);
}
fn submit_slot_write_plain(&mut self, driver: &mut B::Driver, slot_idx: usize) {
if self.core.slots[slot_idx].user.get().egress.is_stream()
&& self.core.slots[slot_idx].conn.pending_body_ptr.is_null()
{
self.encode_stream_chunks_to_buf(slot_idx);
}
let stream_in_progress = self.core.slots[slot_idx].user.get().egress.is_stream();
let CoreSlot { conn, buf, io, .. } = &mut self.core.slots[slot_idx];
let head = conn.write_head as usize;
let inflight_end = conn.write_inflight_end as usize;
if head <= inflight_end {
if stream_in_progress {
return self.submit_stream_write(driver, slot_idx);
}
return;
}
if io.write.has_in_flight {
return;
}
if !conn.pending_body_ptr.is_null() {
let hdr_ptr = buf.write_buf.as_mut_slice()[inflight_end..head].as_ptr();
let hdr_len = head - inflight_end;
let body_ptr = conn.pending_body_ptr;
let body_len = conn.pending_body_len as usize;
buf.pending_iovs[0] = IoVec::raw(hdr_ptr, hdr_len);
buf.pending_iovs[1] = IoVec::raw(body_ptr, body_len);
buf.pending_msghdr.msg_iovlen = 2;
let msg_ptr: *const libc::msghdr = &buf.pending_msghdr;
conn.large_inflight = (hdr_len as u32).saturating_add(body_len as u32);
conn.large_tail = 0;
conn.write_inflight_end = head as u16;
let _ = unsafe { self.driver_submit_send_msg(driver, slot_idx, msg_ptr) };
} else {
let bytes_ptr = buf.write_buf.as_mut_slice()[inflight_end..head].as_ptr();
let bytes_len = (head - inflight_end) as u32;
let _ = self.driver_submit_send(driver, slot_idx, bytes_ptr, bytes_len);
self.core.slots[slot_idx].conn.write_inflight_end = head as u16;
}
}
#[inline(always)]
pub(super) fn driver_submit_send(
&mut self,
driver: &mut B::Driver,
slot_idx: usize,
ptr: *const u8,
len: u32,
) -> bool {
let external = self.core.slot_id_base + slot_idx;
let conn_fd = self.core.slots[slot_idx].conn.fd;
let write_slot = &mut self.core.slots[slot_idx].io.write;
let (token, epoch) = write_slot.next_token::<B::Token>(external);
if driver.submit_send_tagged(token, conn_fd, ptr, len) {
self.core.slots[slot_idx].io.write.mark_inflight(epoch);
true
} else {
false
}
}
#[inline(always)]
pub(super) unsafe fn driver_submit_send_msg(
&mut self,
driver: &mut B::Driver,
slot_idx: usize,
msg: *const libc::msghdr,
) -> bool {
let external = self.core.slot_id_base + slot_idx;
let conn_fd = self.core.slots[slot_idx].conn.fd;
let write_slot = &mut self.core.slots[slot_idx].io.write;
let (token, epoch) = write_slot.next_token::<B::Token>(external);
if unsafe { driver.submit_send_msg_tagged(token, conn_fd, msg) } {
self.core.slots[slot_idx].io.write.mark_inflight(epoch);
true
} else {
false
}
}
fn submit_slot_write_codec(&mut self, driver: &mut B::Driver, slot_idx: usize) {
let stream_in_progress = self.core.slots[slot_idx].user.get().egress.is_stream();
if stream_in_progress && self.core.slots[slot_idx].conn.pending_body_ptr.is_null() {
self.encode_stream_chunks_to_buf(slot_idx);
}
self.core.encrypt_write_buf(slot_idx);
if !self.core.slots[slot_idx].conn.pending_body_ptr.is_null() {
let body_ptr = self.core.slots[slot_idx].conn.pending_body_ptr;
let body_len = self.core.slots[slot_idx].conn.pending_body_len as usize;
let body = unsafe { std::slice::from_raw_parts(body_ptr, body_len) };
<H::CodecLayer as CodecLayer>::process_outbound(
self.core.slots[slot_idx].codec_state.get_mut(),
body,
);
self.core.slots[slot_idx].conn.pending_body_ptr = std::ptr::null();
self.core.slots[slot_idx].conn.pending_body_len = 0;
self.core.slots[slot_idx].conn.large_inflight = 0;
self.core.slots[slot_idx].conn.large_tail = 0;
}
self.core.try_submit_codec(driver, slot_idx);
}
}