use std::{
io,
task::{Context, Poll, Waker},
};
use crate::sys::token::Completion;
use dambi::{LiveKey, Slab};
use super::RequestId;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct WakerKey(usize);
impl WakerKey {
fn as_usize(self) -> usize {
self.0
}
}
pub(crate) struct WakerSlots {
slots: Vec<Option<Waker>>,
}
impl WakerSlots {
pub(crate) fn with_capacity(min_cap: usize) -> Self {
Self {
slots: vec![None; min_cap],
}
}
pub(crate) fn ensure_index(&mut self, index: usize) {
if index >= self.slots.len() {
self.slots.resize_with(index + 1, || None);
}
}
pub(crate) fn put(&mut self, index: usize, waker: &Waker) -> WakerKey {
debug_assert!(
index < self.slots.len(),
"dope invariant violated: waker index out of bounds"
);
let slot = &mut self.slots[index];
if let Some(old) = slot {
if !old.will_wake(waker) {
*old = waker.clone();
}
} else {
*slot = Some(waker.clone());
}
WakerKey(index)
}
pub(crate) fn update(&mut self, key: WakerKey, waker: &Waker) {
let index = key.as_usize();
debug_assert!(
index < self.slots.len(),
"dope invariant violated: waker index out of bounds"
);
let Some(old) = self.slots[index].as_mut() else {
panic!("dope invariant violated: waiting request missing waker");
};
if !old.will_wake(waker) {
*old = waker.clone();
}
}
pub(crate) fn take(&mut self, key: WakerKey) -> Waker {
let index = key.as_usize();
debug_assert!(
index < self.slots.len(),
"dope invariant violated: waker index out of bounds"
);
self.slots[index]
.take()
.unwrap_or_else(|| panic!("dope invariant violated: waiting request missing waker"))
}
pub(crate) fn clear_existing(&mut self, index: usize) {
debug_assert!(
index < self.slots.len(),
"dope invariant violated: waker index out of bounds"
);
let _ = self.slots[index].take();
}
}
const DROPPED_INLINE_CAP: usize = 256;
type DroppedSlot = super::small_slot::SmallSlot<DROPPED_INLINE_CAP>;
pub(crate) struct DroppedData {
slots: Vec<DroppedSlot>,
}
impl DroppedData {
pub(in crate::driver) fn with_capacity(min_cap: usize) -> Self {
let mut slots = Vec::with_capacity(min_cap);
slots.resize_with(min_cap, DroppedSlot::new);
Self { slots }
}
fn ensure_index(&mut self, index: usize) {
if index < self.slots.len() {
return;
}
self.slots.resize_with(index + 1, DroppedSlot::new);
}
pub(in crate::driver) fn put<T: 'static>(&mut self, index: usize, value: T) {
self.ensure_index(index);
self.slots[index].put(value);
}
pub(in crate::driver) fn clear(&mut self, index: usize) {
if index >= self.slots.len() {
return;
}
self.slots[index].clear();
}
}
pub(in crate::driver) enum RequestEntry {
Submitted,
Waiting(WakerKey),
Dropped,
Completed(Completion),
}
impl RequestEntry {
fn submitted() -> Self {
Self::Submitted
}
fn complete(&mut self, completion: Completion) -> Result<Option<WakerKey>, ()> {
match self {
Self::Submitted => {
*self = Self::Completed(completion);
Ok(None)
}
Self::Waiting(key) => {
let key = *key;
*self = Self::Completed(completion);
Ok(Some(key))
}
Self::Dropped => Err(()),
Self::Completed(_) => panic!("dope invariant violated: request completion repeated"),
}
}
fn discard<T: 'static>(
&mut self,
dropped: &mut DroppedData,
index: usize,
data: &mut Option<T>,
) {
match self {
Self::Submitted | Self::Waiting(_) => {
if let Some(value) = data.take() {
dropped.put(index, value);
} else {
dropped.clear(index);
}
*self = Self::Dropped;
}
Self::Dropped | Self::Completed(_) => {}
}
}
}
pub(crate) struct Requests {
slots: Slab<RequestEntry>,
epochs: Vec<u32>,
wakers: WakerSlots,
dropped: DroppedData,
}
impl Requests {
pub(in crate::driver) fn with_capacity(min_cap: usize) -> Self {
Self {
slots: Slab::with_capacity(min_cap),
epochs: Vec::with_capacity(min_cap),
wakers: WakerSlots::with_capacity(min_cap),
dropped: DroppedData::with_capacity(min_cap),
}
}
fn dropped() -> Completion {
Completion {
result: Err(io::Error::other("op dropped")),
flags: 0,
}
}
fn live_key(&self, id: RequestId) -> LiveKey {
let index = id.index();
let epoch = id.epoch();
let slot_epoch = self
.epochs
.get(index)
.copied()
.unwrap_or_else(|| panic!("dope invariant violated: request epoch missing"));
assert!(
slot_epoch == epoch,
"dope invariant violated: request epoch mismatch",
);
self.slots
.live_key(index)
.unwrap_or_else(|| panic!("dope invariant violated: request slot missing"))
}
fn remove(&mut self, key: LiveKey) -> RequestEntry {
let index = key.as_usize();
self.dropped.clear(index);
self.slots.remove_live(key)
}
fn remove_completed(&mut self, id: RequestId) -> Completion {
match self.remove(self.live_key(id)) {
RequestEntry::Completed(completion) => completion,
_ => panic!("dope invariant violated: request remove lifecycle mismatch"),
}
}
fn ensure_index(&mut self, index: usize) {
self.wakers.ensure_index(index);
if self.epochs.len() <= index {
self.epochs.resize(index + 1, 0);
}
}
fn next_epoch(&mut self, index: usize) -> u32 {
self.ensure_index(index);
let epoch = self.epochs[index].wrapping_add(1) & RequestId::EPOCH_MASK;
let epoch = epoch.max(1);
self.epochs[index] = epoch;
epoch
}
fn insert_entry(&mut self, entry: RequestEntry) -> RequestId {
let index = self.slots.insert(entry);
let epoch = self.next_epoch(index);
self.dropped.clear(index);
RequestId::new(index, epoch)
}
pub(in crate::driver) fn insert(&mut self) -> RequestId {
self.insert_entry(RequestEntry::submitted())
}
pub(in crate::driver) fn rollback(&mut self, index: RequestId) {
let index = index.index();
self.dropped.clear(index);
let _ = self.slots.remove(index);
self.wakers.clear_existing(index);
}
pub(in crate::driver) fn poll(
&mut self,
index: RequestId,
cx: &mut Context<'_>,
) -> Poll<Completion> {
let key = self.live_key(index);
let slot = index.index();
match self.slots.get_live(key) {
RequestEntry::Submitted => {
let waker = self.wakers.put(slot, cx.waker());
*self.slots.get_live(key) = RequestEntry::Waiting(waker);
Poll::Pending
}
RequestEntry::Waiting(waker) => {
self.wakers.update(*waker, cx.waker());
Poll::Pending
}
RequestEntry::Completed(_) => Poll::Ready(self.remove_completed(index)),
RequestEntry::Dropped => Poll::Ready(Self::dropped()),
}
}
pub(in crate::driver) fn discard<T: 'static>(
&mut self,
index: RequestId,
data: &mut Option<T>,
) {
let key = self.live_key(index);
let slot = index.index();
match self.slots.get_live(key) {
RequestEntry::Submitted => {
let (slots, dropped) = (&mut self.slots, &mut self.dropped);
slots.get_live(key).discard(dropped, slot, data);
}
RequestEntry::Waiting(waker) => {
let waker = *waker;
let (slots, dropped) = (&mut self.slots, &mut self.dropped);
slots.get_live(key).discard(dropped, slot, data);
let _ = self.wakers.take(waker);
}
RequestEntry::Completed(_) => match self.remove(key) {
RequestEntry::Completed(_) => {}
_ => panic!("dope invariant violated: request remove lifecycle mismatch"),
},
RequestEntry::Dropped => {}
}
}
pub(in crate::driver) fn on_request(&mut self, id: RequestId, result: i32, flags: u32) {
if io_uring::cqueue::notif(flags) {
return;
}
let index = id.index();
let Some(entry) = self.slots.get_mut(index) else {
return;
};
if self.epochs.get(index).copied() != Some(id.epoch()) {
return;
}
let completion = Completion {
result: if result >= 0 {
Ok(result as u32)
} else {
Err(io::Error::from_raw_os_error(-result))
},
flags,
};
match entry.complete(completion) {
Ok(Some(waker)) => {
self.wakers.take(waker).wake();
}
Ok(None) => {}
Err(()) => {
self.wakers.clear_existing(index);
match self.remove(self.live_key(id)) {
RequestEntry::Dropped => {}
_ => panic!("dope invariant violated: request remove lifecycle mismatch"),
}
}
}
}
}