use std::future::Future;
use std::pin::Pin;
pub(crate) type BoxFuture = Pin<Box<dyn Future<Output = ()> + 'static>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TaskId(pub(crate) u32);
enum TaskSlot {
Empty,
Parked(BoxFuture),
Ready(BoxFuture),
}
pub(crate) struct TaskSlab {
tasks: Vec<TaskSlot>,
}
impl TaskSlab {
pub(crate) fn new(max_connections: u32) -> Self {
let mut tasks = Vec::with_capacity(max_connections as usize);
for _ in 0..max_connections {
tasks.push(TaskSlot::Empty);
}
TaskSlab { tasks }
}
pub(crate) fn spawn(&mut self, conn_index: u32, future: BoxFuture) {
let idx = conn_index as usize;
debug_assert!(idx < self.tasks.len(), "conn_index out of range");
debug_assert!(
matches!(self.tasks[idx], TaskSlot::Empty),
"task already exists for conn_index {conn_index}"
);
self.tasks[idx] = TaskSlot::Ready(future);
}
pub(crate) fn take_ready(&mut self, conn_index: u32) -> Option<BoxFuture> {
let idx = conn_index as usize;
if idx >= self.tasks.len() {
return None;
}
match std::mem::replace(&mut self.tasks[idx], TaskSlot::Empty) {
TaskSlot::Ready(fut) => Some(fut),
other => {
self.tasks[idx] = other;
None
}
}
}
pub(crate) fn park(&mut self, conn_index: u32, future: BoxFuture) {
let idx = conn_index as usize;
debug_assert!(idx < self.tasks.len());
self.tasks[idx] = TaskSlot::Parked(future);
}
pub(crate) fn wake(&mut self, conn_index: u32) -> bool {
let idx = conn_index as usize;
if idx >= self.tasks.len() {
return false;
}
match std::mem::replace(&mut self.tasks[idx], TaskSlot::Empty) {
TaskSlot::Parked(fut) => {
self.tasks[idx] = TaskSlot::Ready(fut);
true
}
TaskSlot::Ready(fut) => {
self.tasks[idx] = TaskSlot::Ready(fut);
false }
TaskSlot::Empty => false,
}
}
pub(crate) fn remove(&mut self, conn_index: u32) {
let idx = conn_index as usize;
if idx < self.tasks.len() {
self.tasks[idx] = TaskSlot::Empty;
}
}
#[allow(dead_code)]
pub(crate) fn has_task(&self, conn_index: u32) -> bool {
let idx = conn_index as usize;
idx < self.tasks.len() && !matches!(self.tasks[idx], TaskSlot::Empty)
}
}
pub(crate) struct StandaloneTaskSlab {
tasks: Vec<TaskSlot>,
free_list: Vec<u32>,
}
impl StandaloneTaskSlab {
pub(crate) fn new(capacity: u32) -> Self {
let mut tasks = Vec::with_capacity(capacity as usize);
let mut free_list = Vec::with_capacity(capacity as usize);
for i in 0..capacity {
tasks.push(TaskSlot::Empty);
free_list.push(i);
}
StandaloneTaskSlab { tasks, free_list }
}
pub(crate) fn spawn(&mut self, future: BoxFuture) -> Option<u32> {
let idx = self.free_list.pop()?;
self.tasks[idx as usize] = TaskSlot::Ready(future);
Some(idx)
}
pub(crate) fn take_ready(&mut self, task_idx: u32) -> Option<BoxFuture> {
let idx = task_idx as usize;
if idx >= self.tasks.len() {
return None;
}
match std::mem::replace(&mut self.tasks[idx], TaskSlot::Empty) {
TaskSlot::Ready(fut) => Some(fut),
other => {
self.tasks[idx] = other;
None
}
}
}
pub(crate) fn park(&mut self, task_idx: u32, future: BoxFuture) {
let idx = task_idx as usize;
debug_assert!(idx < self.tasks.len());
self.tasks[idx] = TaskSlot::Parked(future);
}
pub(crate) fn wake(&mut self, task_idx: u32) -> bool {
let idx = task_idx as usize;
if idx >= self.tasks.len() {
return false;
}
match std::mem::replace(&mut self.tasks[idx], TaskSlot::Empty) {
TaskSlot::Parked(fut) => {
self.tasks[idx] = TaskSlot::Ready(fut);
true
}
TaskSlot::Ready(fut) => {
self.tasks[idx] = TaskSlot::Ready(fut);
false
}
TaskSlot::Empty => false,
}
}
pub(crate) fn remove(&mut self, task_idx: u32) {
let idx = task_idx as usize;
if idx < self.tasks.len() {
self.tasks[idx] = TaskSlot::Empty;
self.free_list.push(task_idx);
}
}
#[allow(dead_code)]
pub(crate) fn has_task(&self, task_idx: u32) -> bool {
let idx = task_idx as usize;
idx < self.tasks.len() && !matches!(self.tasks[idx], TaskSlot::Empty)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::task::{Context, Poll};
struct CountdownFuture(u32);
impl Future for CountdownFuture {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.0 == 0 {
Poll::Ready(())
} else {
self.0 -= 1;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
#[test]
fn spawn_and_take_ready() {
let mut slab = TaskSlab::new(4);
assert!(!slab.has_task(0));
slab.spawn(0, Box::pin(CountdownFuture(2)));
assert!(slab.has_task(0));
let fut = slab.take_ready(0);
assert!(fut.is_some());
assert!(!slab.has_task(0));
}
#[test]
fn park_and_wake() {
let mut slab = TaskSlab::new(4);
slab.spawn(1, Box::pin(CountdownFuture(1)));
let fut = slab.take_ready(1).unwrap();
slab.park(1, fut);
assert!(slab.has_task(1));
assert!(slab.take_ready(1).is_none());
assert!(slab.wake(1));
assert!(slab.take_ready(1).is_some());
}
#[test]
fn remove_task() {
let mut slab = TaskSlab::new(4);
slab.spawn(2, Box::pin(CountdownFuture(0)));
assert!(slab.has_task(2));
slab.remove(2);
assert!(!slab.has_task(2));
}
#[test]
fn wake_empty_slot() {
let mut slab = TaskSlab::new(4);
assert!(!slab.wake(3));
}
#[test]
fn wake_already_ready() {
let mut slab = TaskSlab::new(4);
slab.spawn(0, Box::pin(CountdownFuture(0)));
assert!(!slab.wake(0));
}
#[test]
fn standalone_spawn_and_take() {
let mut slab = StandaloneTaskSlab::new(4);
let idx = slab.spawn(Box::pin(CountdownFuture(2))).unwrap();
assert!(slab.has_task(idx));
let fut = slab.take_ready(idx);
assert!(fut.is_some());
assert!(!slab.has_task(idx));
}
#[test]
fn standalone_park_and_wake() {
let mut slab = StandaloneTaskSlab::new(4);
let idx = slab.spawn(Box::pin(CountdownFuture(1))).unwrap();
let fut = slab.take_ready(idx).unwrap();
slab.park(idx, fut);
assert!(slab.take_ready(idx).is_none());
assert!(slab.wake(idx));
assert!(slab.take_ready(idx).is_some());
}
#[test]
fn standalone_remove_returns_to_free_list() {
let mut slab = StandaloneTaskSlab::new(2);
let a = slab.spawn(Box::pin(CountdownFuture(0))).unwrap();
let b = slab.spawn(Box::pin(CountdownFuture(0))).unwrap();
assert!(slab.spawn(Box::pin(CountdownFuture(0))).is_none());
slab.remove(a);
assert!(slab.spawn(Box::pin(CountdownFuture(0))).is_some());
slab.remove(b);
}
#[test]
fn standalone_full_slab() {
let mut slab = StandaloneTaskSlab::new(1);
assert!(slab.spawn(Box::pin(CountdownFuture(0))).is_some());
assert!(slab.spawn(Box::pin(CountdownFuture(0))).is_none());
}
}