use once_cell::unsync::Lazy;
use std::alloc::Layout;
use spin::Mutex;
use std::collections::VecDeque;
use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::ptr;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::cell::RefCell;
const REGION_SIZE: usize = 256 * 1024;
const REGION_TRIES: usize = 2;
thread_local! {
static POOL: Lazy<RefCell<Pool>> = Lazy::new(|| RefCell::new(Pool::new()));
}
fn align_up(mut v: usize, a: usize) -> usize {
while v % a != 0 {
v = v.next_power_of_two();
}
v
}
struct Pool {
queue: VecDeque<Arc<Mutex<Region>>>,
}
impl Pool {
fn new() -> Self {
let mut queue = VecDeque::new();
for _ in 0..REGION_TRIES {
queue.push_front(Arc::new(Mutex::new(Region::new())));
}
Self { queue }
}
fn allocate<T, F: Future<Output = T> + Send + 'static>(
&mut self,
f: F,
) -> (
&'static mut (dyn Future<Output = T> + Send),
Arc<Mutex<Region>>,
) {
let mut f = Some(f);
for _ in 0..REGION_TRIES {
let region = self.queue.pop_front().unwrap();
let mut r = region.lock();
match r.allocate(f.take().unwrap()) {
Ok(ptr) => {
drop(r);
self.queue.push_front(region.clone());
return (unsafe { mem::transmute(ptr) }, region);
}
Err(v) => {
drop(r);
self.queue.push_back(region);
f = Some(v);
}
}
}
self.queue.push_front(Arc::new(Mutex::new(Region::new())));
self.allocate(f.unwrap())
}
}
struct Region {
start: usize,
next: usize,
refs: usize,
}
impl Region {
fn new() -> Self {
let start: Box<[u8; REGION_SIZE]> = Box::new([0; REGION_SIZE]);
let start = Box::into_raw(start) as usize;
Self {
start,
next: start,
refs: 0,
}
}
fn deallocate(&mut self) {
self.refs -= 1;
if self.refs == 0 {
self.next = self.start;
}
}
fn allocate<T, F: Future<Output = T> + Send + 'static>(
&mut self,
f: F,
) -> Result<*mut (dyn Future<Output = T> + Send), F> {
let layout = Layout::for_value(&f);
let aligned_next = align_up(self.next, layout.align());
let potential_end = aligned_next + layout.size();
if potential_end > self.start + REGION_SIZE {
return Err(f);
}
unsafe {
self.refs += 1;
self.next = aligned_next + layout.size();
ptr::write(aligned_next as *mut F, f);
let faked_box = Box::from_raw(aligned_next as *mut F);
let faked_box = faked_box as Box<dyn Future<Output = T> + Send>;
Ok(Box::into_raw(faked_box))
}
}
}
pub struct DynFuture<T: 'static> {
inner: Pin<&'static mut (dyn Future<Output = T> + Send)>,
region: Arc<Mutex<Region>>,
}
impl<T> Drop for DynFuture<T> {
fn drop(&mut self) {
self.region.lock().deallocate();
}
}
impl<T> DynFuture<T> {
pub fn new(f: impl Future<Output = T> + Send + 'static) -> Self {
let (inner, region) = POOL.with(|pool| pool.borrow_mut().allocate(f));
unsafe {
Self {
inner: Pin::new_unchecked(inner),
region,
}
}
}
}
impl<T> Future for DynFuture<T> {
type Output = T;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
self.inner.as_mut().poll(cx)
}
}