use std::cell::Cell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread::{self, Thread, ThreadId};
use std::time::Instant;
use crossbeam_utils::Backoff;
use select::Selected;
#[derive(Clone)]
pub struct Context {
inner: Arc<Inner>,
}
struct Inner {
select: AtomicUsize,
packet: AtomicUsize,
thread: Thread,
thread_id: ThreadId,
}
impl Context {
#[inline]
pub fn with<F, R>(f: F) -> R
where
F: FnOnce(&Context) -> R,
{
thread_local! {
static CONTEXT: Cell<Option<Context>> = Cell::new(Some(Context::new()));
}
let mut f = Some(f);
let mut f = move |cx: &Context| -> R {
let f = f.take().unwrap();
f(cx)
};
CONTEXT
.try_with(|cell| match cell.take() {
None => f(&Context::new()),
Some(cx) => {
cx.reset();
let res = f(&cx);
cell.set(Some(cx));
res
}
}).unwrap_or_else(|_| f(&Context::new()))
}
#[cold]
fn new() -> Context {
Context {
inner: Arc::new(Inner {
select: AtomicUsize::new(Selected::Waiting.into()),
packet: AtomicUsize::new(0),
thread: thread::current(),
thread_id: thread::current().id(),
}),
}
}
#[inline]
fn reset(&self) {
self.inner
.select
.store(Selected::Waiting.into(), Ordering::Release);
self.inner.packet.store(0, Ordering::Release);
}
#[inline]
pub fn try_select(&self, select: Selected) -> Result<(), Selected> {
self.inner
.select
.compare_exchange(
Selected::Waiting.into(),
select.into(),
Ordering::AcqRel,
Ordering::Acquire,
)
.map(|_| ())
.map_err(|e| e.into())
}
#[inline]
pub fn selected(&self) -> Selected {
Selected::from(self.inner.select.load(Ordering::Acquire))
}
#[inline]
pub fn store_packet(&self, packet: usize) {
if packet != 0 {
self.inner.packet.store(packet, Ordering::Release);
}
}
#[inline]
pub fn wait_packet(&self) -> usize {
let backoff = Backoff::new();
loop {
let packet = self.inner.packet.load(Ordering::Acquire);
if packet != 0 {
return packet;
}
backoff.snooze();
}
}
#[inline]
pub fn wait_until(&self, deadline: Option<Instant>) -> Selected {
let backoff = Backoff::new();
loop {
let sel = Selected::from(self.inner.select.load(Ordering::Acquire));
if sel != Selected::Waiting {
return sel;
}
if backoff.is_complete() {
break;
} else {
backoff.snooze();
}
}
loop {
let sel = Selected::from(self.inner.select.load(Ordering::Acquire));
if sel != Selected::Waiting {
return sel;
}
if let Some(end) = deadline {
let now = Instant::now();
if now < end {
thread::park_timeout(end - now);
} else {
return match self.try_select(Selected::Aborted) {
Ok(()) => Selected::Aborted,
Err(s) => s,
};
}
} else {
thread::park();
}
}
}
#[inline]
pub fn unpark(&self) {
self.inner.thread.unpark();
}
#[inline]
pub fn thread_id(&self) -> ThreadId {
self.inner.thread_id
}
}