use loom::{
futures::AtomicTask,
sync::{
atomic::{AtomicPtr, AtomicUsize},
CausalCell,
},
yield_now,
};
use futures::Poll;
use std::fmt;
use std::ptr::{self, NonNull};
use std::sync::atomic::Ordering::{self, AcqRel, Acquire, Relaxed, Release};
use std::sync::Arc;
use std::usize;
pub struct Semaphore {
state: AtomicUsize,
head: CausalCell<NonNull<WaiterNode>>,
rx_lock: AtomicUsize,
stub: Box<WaiterNode>,
}
#[derive(Debug)]
pub struct Permit {
waiter: Option<Arc<WaiterNode>>,
state: PermitState,
}
#[derive(Debug)]
pub struct AcquireError(());
#[derive(Debug)]
pub struct TryAcquireError {
kind: ErrorKind,
}
#[derive(Debug)]
enum ErrorKind {
Closed,
NoPermits,
}
#[derive(Debug)]
struct WaiterNode {
state: AtomicUsize,
task: AtomicTask,
next: AtomicPtr<WaiterNode>,
}
#[derive(Copy, Clone)]
struct SemState(usize);
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum PermitState {
Idle,
Waiting,
Acquired,
}
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
#[repr(usize)]
enum NodeState {
Idle = 0,
Queued = 1,
QueuedWaiting = 2,
Assigned = 3,
Closed = 4,
}
impl Semaphore {
pub fn new(permits: usize) -> Semaphore {
let stub = Box::new(WaiterNode::new());
let ptr = NonNull::new(&*stub as *const _ as *mut _).unwrap();
debug_assert!(ptr.as_ptr() as usize & NUM_FLAG == 0);
let state = SemState::new(permits, &stub);
Semaphore {
state: AtomicUsize::new(state.to_usize()),
head: CausalCell::new(ptr),
rx_lock: AtomicUsize::new(0),
stub,
}
}
pub fn available_permits(&self) -> usize {
let curr = SemState::load(&self.state, Acquire);
curr.available_permits()
}
fn poll_permit(&self, mut permit: Option<&mut Permit>) -> Poll<(), AcquireError> {
use futures::Async::*;
let mut curr = SemState::load(&self.state, Acquire);
debug!(" + poll_permit; sem-state = {:?}", curr);
let mut maybe_strong: Option<NonNull<WaiterNode>> = None;
macro_rules! undo_strong {
() => {
if let Some(waiter) = maybe_strong {
let waiter = unsafe { Arc::from_raw(waiter.as_ptr()) };
waiter.revert_to_idle();
}
};
}
loop {
let mut next = curr;
if curr.is_closed() {
undo_strong!();
return Err(AcquireError::closed());
}
if !next.acquire_permit(&self.stub) {
debug!(" + poll_permit -- no permits");
debug_assert!(curr.waiter().is_some());
if maybe_strong.is_none() {
if let Some(ref mut permit) = permit {
let waiter = permit
.waiter
.get_or_insert_with(|| Arc::new(WaiterNode::new()));
waiter.register();
debug!(" + poll_permit -- to_queued_waiting");
if !waiter.to_queued_waiting() {
debug!(" + poll_permit; waiter already queued");
return Ok(NotReady);
}
maybe_strong = Some(WaiterNode::into_non_null(waiter.clone()));
} else {
return Ok(NotReady);
}
}
next.set_waiter(maybe_strong.unwrap());
}
debug!(" + poll_permit -- pre-CAS; next = {:?}", next);
debug_assert_ne!(curr.0, 0);
debug_assert_ne!(next.0, 0);
match next.compare_exchange(&self.state, curr, AcqRel, Acquire) {
Ok(_) => {
debug!(" + poll_permit -- CAS ok");
match curr.waiter() {
Some(prev_waiter) => {
let waiter = maybe_strong.unwrap();
unsafe {
prev_waiter.as_ref().next.store(waiter.as_ptr(), Release);
}
debug!(" + poll_permit -- waiter pushed");
return Ok(NotReady);
}
None => {
debug!(" + poll_permit -- permit acquired");
undo_strong!();
return Ok(Ready(()));
}
}
}
Err(actual) => {
curr = actual;
}
}
}
}
pub fn close(&self) {
debug!("+ Semaphore::close");
let prev = self.rx_lock.fetch_or(1, AcqRel);
debug!(" + close -- rx_lock.fetch_add(1)");
if prev != 0 {
debug!("+ close -- locked; prev = {}", prev);
return;
}
self.add_permits_locked(0, true);
}
pub fn add_permits(&self, n: usize) {
debug!(" + add_permits; n = {}", n);
if n == 0 {
return;
}
let prev = self.rx_lock.fetch_add(n << 1, AcqRel);
debug!(" + add_permits; rx_lock.fetch_add(n << 1); n = {}", n);
if prev != 0 {
debug!(" + add_permits -- locked; prev = {}", prev);
return;
}
self.add_permits_locked(n, false);
}
fn add_permits_locked(&self, mut rem: usize, mut closed: bool) {
while rem > 0 || closed {
debug!(
" + add_permits_locked -- iter; rem = {}; closed = {:?}",
rem, closed
);
if closed {
SemState::fetch_set_closed(&self.state, AcqRel);
}
self.add_permits_locked2(rem, closed);
let n = rem << 1;
let actual = if closed {
let actual = self.rx_lock.fetch_sub(n | 1, AcqRel);
debug!(
" + add_permits_locked; rx_lock.fetch_sub(n | 1); n = {}; actual={}",
n, actual
);
closed = false;
actual
} else {
let actual = self.rx_lock.fetch_sub(n, AcqRel);
debug!(
" + add_permits_locked; rx_lock.fetch_sub(n); n = {}; actual={}",
n, actual
);
closed = actual & 1 == 1;
actual
};
rem = (actual >> 1) - rem;
}
debug!(" + add_permits; done");
}
fn add_permits_locked2(&self, mut n: usize, closed: bool) {
while n > 0 || closed {
let waiter = match self.pop(n, closed) {
Some(waiter) => waiter,
None => {
return;
}
};
debug!(" + release_n -- notify");
if waiter.notify(closed) {
n = n.saturating_sub(1);
debug!(" + release_n -- dec");
}
}
}
fn pop(&self, rem: usize, closed: bool) -> Option<Arc<WaiterNode>> {
debug!(" + pop; rem = {}", rem);
'outer: loop {
unsafe {
let mut head = self.head.with(|head| *head);
let mut next_ptr = head.as_ref().next.load(Acquire);
let stub = self.stub();
if head == stub {
debug!(" + pop; head == stub");
let next = match NonNull::new(next_ptr) {
Some(next) => next,
None => {
let mut curr = SemState::load(&self.state, Acquire);
loop {
if curr.has_waiter(&self.stub) {
debug!(" + pop; inconsistent 1");
yield_now();
continue 'outer;
}
if rem == 0 {
debug_assert!(curr.is_closed(), "state = {:?}", curr);
return None;
}
let mut next = curr;
next.release_permits(rem, &self.stub);
match next.compare_exchange(&self.state, curr, AcqRel, Acquire) {
Ok(_) => return None,
Err(actual) => {
curr = actual;
}
}
}
}
};
debug!(" + pop; got next waiter");
self.head.with_mut(|head| *head = next);
head = next;
next_ptr = next.as_ref().next.load(Acquire);
}
if let Some(next) = NonNull::new(next_ptr) {
self.head.with_mut(|head| *head = next);
return Some(Arc::from_raw(head.as_ptr()));
}
let state = SemState::load(&self.state, Acquire);
let tail = state.waiter().unwrap();
if tail != head {
debug!(" + pop; inconsistent 2");
yield_now();
continue 'outer;
}
self.push_stub(closed);
next_ptr = head.as_ref().next.load(Acquire);
if let Some(next) = NonNull::new(next_ptr) {
self.head.with_mut(|head| *head = next);
return Some(Arc::from_raw(head.as_ptr()));
}
debug!(" + pop; inconsistent 3");
yield_now();
}
}
}
unsafe fn push_stub(&self, closed: bool) {
let stub = self.stub();
stub.as_ref().next.store(ptr::null_mut(), Relaxed);
let prev = SemState::new_ptr(stub, closed).swap(&self.state, AcqRel);
debug_assert_eq!(closed, prev.is_closed());
let prev = prev.waiter().unwrap();
debug_assert_ne!(prev, stub);
prev.as_ref().next.store(stub.as_ptr(), Release);
}
fn stub(&self) -> NonNull<WaiterNode> {
unsafe { NonNull::new_unchecked(&*self.stub as *const _ as *mut _) }
}
}
impl fmt::Debug for Semaphore {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("Semaphore")
.field("state", &SemState::load(&self.state, Relaxed))
.field("head", &self.head.with(|ptr| ptr))
.field("rx_lock", &self.rx_lock.load(Relaxed))
.field("stub", &self.stub)
.finish()
}
}
unsafe impl Send for Semaphore {}
unsafe impl Sync for Semaphore {}
impl Permit {
pub fn new() -> Permit {
Permit {
waiter: None,
state: PermitState::Idle,
}
}
pub fn is_acquired(&self) -> bool {
self.state == PermitState::Acquired
}
pub fn poll_acquire(&mut self, semaphore: &Semaphore) -> Poll<(), AcquireError> {
use futures::Async::*;
match self.state {
PermitState::Idle => {}
PermitState::Waiting => {
let waiter = self.waiter.as_ref().unwrap();
if waiter.acquire()? {
self.state = PermitState::Acquired;
return Ok(Ready(()));
} else {
return Ok(NotReady);
}
}
PermitState::Acquired => {
return Ok(Ready(()));
}
}
match semaphore.poll_permit(Some(self))? {
Ready(v) => {
self.state = PermitState::Acquired;
Ok(Ready(v))
}
NotReady => {
self.state = PermitState::Waiting;
Ok(NotReady)
}
}
}
pub fn try_acquire(&mut self, semaphore: &Semaphore) -> Result<(), TryAcquireError> {
use futures::Async::*;
match self.state {
PermitState::Idle => {}
PermitState::Waiting => {
let waiter = self.waiter.as_ref().unwrap();
if waiter.acquire2().map_err(to_try_acquire)? {
self.state = PermitState::Acquired;
return Ok(());
} else {
return Err(TryAcquireError::no_permits());
}
}
PermitState::Acquired => {
return Ok(());
}
}
match semaphore.poll_permit(None).map_err(to_try_acquire)? {
Ready(()) => {
self.state = PermitState::Acquired;
Ok(())
}
NotReady => Err(TryAcquireError::no_permits()),
}
}
pub fn release(&mut self, semaphore: &Semaphore) {
if self.forget2() {
semaphore.add_permits(1);
}
}
pub fn forget(&mut self) {
self.forget2();
}
fn forget2(&mut self) -> bool {
match self.state {
PermitState::Idle => false,
PermitState::Waiting => {
let ret = self.waiter.as_ref().unwrap().cancel_interest();
self.state = PermitState::Idle;
ret
}
PermitState::Acquired => {
self.state = PermitState::Idle;
true
}
}
}
}
impl AcquireError {
fn closed() -> AcquireError {
AcquireError(())
}
}
fn to_try_acquire(_: AcquireError) -> TryAcquireError {
TryAcquireError::closed()
}
impl fmt::Display for AcquireError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl ::std::error::Error for AcquireError {
fn description(&self) -> &str {
"semaphore closed"
}
}
impl TryAcquireError {
fn closed() -> TryAcquireError {
TryAcquireError {
kind: ErrorKind::Closed,
}
}
fn no_permits() -> TryAcquireError {
TryAcquireError {
kind: ErrorKind::NoPermits,
}
}
pub fn is_closed(&self) -> bool {
match self.kind {
ErrorKind::Closed => true,
_ => false,
}
}
pub fn is_no_permits(&self) -> bool {
match self.kind {
ErrorKind::NoPermits => true,
_ => false,
}
}
}
impl fmt::Display for TryAcquireError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl ::std::error::Error for TryAcquireError {
fn description(&self) -> &str {
match self.kind {
ErrorKind::Closed => "semaphore closed",
ErrorKind::NoPermits => "no permits available",
}
}
}
impl WaiterNode {
fn new() -> WaiterNode {
WaiterNode {
state: AtomicUsize::new(NodeState::new().to_usize()),
task: AtomicTask::new(),
next: AtomicPtr::new(ptr::null_mut()),
}
}
fn acquire(&self) -> Result<bool, AcquireError> {
if self.acquire2()? {
return Ok(true);
}
self.task.register();
self.acquire2()
}
fn acquire2(&self) -> Result<bool, AcquireError> {
use self::NodeState::*;
match Idle.compare_exchange(&self.state, Assigned, AcqRel, Acquire) {
Ok(_) => Ok(true),
Err(Closed) => Err(AcquireError::closed()),
Err(_) => Ok(false),
}
}
fn register(&self) {
self.task.register()
}
fn cancel_interest(&self) -> bool {
use self::NodeState::*;
match Queued.compare_exchange(&self.state, QueuedWaiting, AcqRel, Acquire) {
Ok(_) => false,
Err(Closed) => false,
Err(Assigned) => {
match self.acquire2() {
Ok(true) => true,
Ok(false) => panic!(),
Err(_) => false,
}
}
Err(state) => panic!("unexpected state = {:?}", state),
}
}
fn to_queued_waiting(&self) -> bool {
use self::NodeState::*;
let mut curr = NodeState::load(&self.state, Acquire);
loop {
debug_assert!(curr == Idle || curr == Queued, "actual = {:?}", curr);
let next = QueuedWaiting;
match next.compare_exchange(&self.state, curr, AcqRel, Acquire) {
Ok(_) => {
if curr.is_queued() {
return false;
} else {
self.next.store(ptr::null_mut(), Relaxed);
return true;
}
}
Err(actual) => {
curr = actual;
}
}
}
}
fn notify(&self, closed: bool) -> bool {
use self::NodeState::*;
let mut curr = QueuedWaiting;
loop {
let next = match curr {
Queued => Idle,
QueuedWaiting => {
if closed {
Closed
} else {
Assigned
}
}
actual => panic!("actual = {:?}", actual),
};
match next.compare_exchange(&self.state, curr, AcqRel, Acquire) {
Ok(_) => match curr {
QueuedWaiting => {
debug!(" + notify -- task notified");
self.task.notify();
return true;
}
other => {
debug!(" + notify -- not notified; state = {:?}", other);
return false;
}
},
Err(actual) => curr = actual,
}
}
}
fn revert_to_idle(&self) {
use self::NodeState::Idle;
NodeState::store(&self.state, Idle, Relaxed);
}
fn into_non_null(arc: Arc<WaiterNode>) -> NonNull<WaiterNode> {
let ptr = Arc::into_raw(arc);
unsafe { NonNull::new_unchecked(ptr as *mut _) }
}
}
const NUM_FLAG: usize = 0b01;
const CLOSED_FLAG: usize = 0b10;
const MAX_PERMITS: usize = usize::MAX >> NUM_SHIFT;
const NUM_SHIFT: usize = 2;
impl SemState {
fn new(permits: usize, stub: &WaiterNode) -> SemState {
assert!(permits <= MAX_PERMITS);
if permits > 0 {
SemState((permits << NUM_SHIFT) | NUM_FLAG)
} else {
SemState(stub as *const _ as usize)
}
}
fn new_ptr(tail: NonNull<WaiterNode>, closed: bool) -> SemState {
let mut val = tail.as_ptr() as usize;
if closed {
val |= CLOSED_FLAG;
}
SemState(val)
}
fn available_permits(&self) -> usize {
if !self.has_available_permits() {
return 0;
}
self.0 >> NUM_SHIFT
}
fn has_available_permits(&self) -> bool {
self.0 & NUM_FLAG == NUM_FLAG
}
fn has_waiter(&self, stub: &WaiterNode) -> bool {
!self.has_available_permits() && !self.is_stub(stub)
}
fn acquire_permit(&mut self, stub: &WaiterNode) -> bool {
if !self.has_available_permits() {
return false;
}
debug_assert!(self.waiter().is_none());
self.0 -= 1 << NUM_SHIFT;
if self.0 == NUM_FLAG {
self.0 = stub as *const _ as usize;
}
true
}
fn release_permits(&mut self, permits: usize, stub: &WaiterNode) {
debug_assert!(permits > 0);
if self.is_stub(stub) {
self.0 = (permits << NUM_SHIFT) | NUM_FLAG | (self.0 & CLOSED_FLAG);
return;
}
debug_assert!(self.has_available_permits());
self.0 += permits << NUM_SHIFT;
}
fn is_waiter(&self) -> bool {
self.0 & NUM_FLAG == 0
}
fn waiter(&self) -> Option<NonNull<WaiterNode>> {
if self.is_waiter() {
let waiter = NonNull::new(self.as_ptr()).expect("null pointer stored");
Some(waiter)
} else {
None
}
}
fn as_ptr(&self) -> *mut WaiterNode {
(self.0 & !CLOSED_FLAG) as *mut WaiterNode
}
fn set_waiter(&mut self, waiter: NonNull<WaiterNode>) {
let waiter = waiter.as_ptr() as usize;
debug_assert!(waiter & NUM_FLAG == 0);
debug_assert!(!self.is_closed());
self.0 = waiter;
}
fn is_stub(&self, stub: &WaiterNode) -> bool {
self.as_ptr() as usize == stub as *const _ as usize
}
fn load(cell: &AtomicUsize, ordering: Ordering) -> SemState {
let value = cell.load(ordering);
debug!(" + SemState::load; value = {}", value);
SemState(value)
}
fn swap(&self, cell: &AtomicUsize, ordering: Ordering) -> SemState {
let prev = SemState(cell.swap(self.to_usize(), ordering));
debug_assert_eq!(prev.is_closed(), self.is_closed());
prev
}
fn compare_exchange(
&self,
cell: &AtomicUsize,
prev: SemState,
success: Ordering,
failure: Ordering,
) -> Result<SemState, SemState> {
debug_assert_eq!(prev.is_closed(), self.is_closed());
let res = cell.compare_exchange(prev.to_usize(), self.to_usize(), success, failure);
debug!(
" + SemState::compare_exchange; prev = {}; next = {}; result = {:?}",
prev.to_usize(),
self.to_usize(),
res
);
res.map(SemState).map_err(SemState)
}
fn fetch_set_closed(cell: &AtomicUsize, ordering: Ordering) -> SemState {
let value = cell.fetch_or(CLOSED_FLAG, ordering);
SemState(value)
}
fn is_closed(&self) -> bool {
self.0 & CLOSED_FLAG == CLOSED_FLAG
}
fn to_usize(&self) -> usize {
self.0
}
}
impl fmt::Debug for SemState {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let mut fmt = fmt.debug_struct("SemState");
if self.is_waiter() {
fmt.field("state", &"<waiter>");
} else {
fmt.field("permits", &self.available_permits());
}
fmt.finish()
}
}
impl NodeState {
fn new() -> NodeState {
NodeState::Idle
}
fn from_usize(value: usize) -> NodeState {
use self::NodeState::*;
match value {
0 => Idle,
1 => Queued,
2 => QueuedWaiting,
3 => Assigned,
4 => Closed,
_ => panic!(),
}
}
fn load(cell: &AtomicUsize, ordering: Ordering) -> NodeState {
NodeState::from_usize(cell.load(ordering))
}
fn store(cell: &AtomicUsize, value: NodeState, ordering: Ordering) {
cell.store(value.to_usize(), ordering);
}
fn compare_exchange(
&self,
cell: &AtomicUsize,
prev: NodeState,
success: Ordering,
failure: Ordering,
) -> Result<NodeState, NodeState> {
cell.compare_exchange(prev.to_usize(), self.to_usize(), success, failure)
.map(NodeState::from_usize)
.map_err(NodeState::from_usize)
}
fn is_queued(&self) -> bool {
use self::NodeState::*;
match *self {
Queued | QueuedWaiting => true,
_ => false,
}
}
fn to_usize(&self) -> usize {
*self as usize
}
}