#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))]
use crate::loom::cell::UnsafeCell;
use crate::loom::future::AtomicWaker;
use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize};
use crate::loom::thread;
use std::cmp;
use std::fmt;
use std::ptr::{self, NonNull};
use std::sync::atomic::Ordering::{self, AcqRel, Acquire, Relaxed, Release};
use std::task::Poll::{Pending, Ready};
use std::task::{Context, Poll};
use std::usize;
pub(crate) struct Semaphore {
state: AtomicUsize,
head: UnsafeCell<NonNull<Waiter>>,
rx_lock: AtomicUsize,
stub: Box<Waiter>,
}
#[derive(Debug)]
pub(crate) struct Permit {
waiter: Option<Box<Waiter>>,
state: PermitState,
}
#[derive(Debug)]
pub(crate) struct AcquireError(());
#[derive(Debug)]
pub(crate) enum TryAcquireError {
Closed,
NoPermits,
}
#[derive(Debug)]
struct Waiter {
state: AtomicUsize,
waker: AtomicWaker,
next: AtomicPtr<Waiter>,
}
#[derive(Copy, Clone)]
struct SemState(usize);
#[derive(Debug, Copy, Clone)]
enum PermitState {
Waiting(u16),
Acquired(u16),
}
#[derive(Debug, Copy, Clone)]
struct WaiterState(usize);
const QUEUED: usize = 0b001;
const CLOSED: usize = 0b10;
const DROPPED: usize = 0b100;
const PERMIT_ONE: usize = 0b1000;
const PERMIT_MASK: usize = usize::MAX - (PERMIT_ONE - 1);
const PERMIT_SHIFT: u32 = PERMIT_ONE.trailing_zeros();
const NUM_FLAG: usize = 0b01;
const CLOSED_FLAG: usize = 0b10;
const MAX_PERMITS: usize = usize::MAX >> NUM_SHIFT;
const NUM_SHIFT: usize = 2;
impl Semaphore {
pub(crate) fn new(permits: usize) -> Semaphore {
let stub = Box::new(Waiter::new());
let ptr = NonNull::from(&*stub);
debug_assert!(ptr.as_ptr() as usize & NUM_FLAG == 0);
let state = SemState::new(permits, &stub);
Semaphore {
state: AtomicUsize::new(state.to_usize()),
head: UnsafeCell::new(ptr),
rx_lock: AtomicUsize::new(0),
stub,
}
}
pub(crate) fn available_permits(&self) -> usize {
let curr = SemState(self.state.load(Acquire));
curr.available_permits()
}
fn poll_acquire(
&self,
cx: &mut Context<'_>,
num_permits: u16,
permit: &mut Permit,
) -> Poll<Result<(), AcquireError>> {
self.poll_acquire2(num_permits, || {
let waiter = permit.waiter.get_or_insert_with(|| Box::new(Waiter::new()));
waiter.waker.register_by_ref(cx.waker());
Some(NonNull::from(&**waiter))
})
}
fn try_acquire(&self, num_permits: u16) -> Result<(), TryAcquireError> {
match self.poll_acquire2(num_permits, || None) {
Poll::Ready(res) => res.map_err(to_try_acquire),
Poll::Pending => Err(TryAcquireError::NoPermits),
}
}
fn poll_acquire2<F>(
&self,
num_permits: u16,
mut get_waiter: F,
) -> Poll<Result<(), AcquireError>>
where
F: FnMut() -> Option<NonNull<Waiter>>,
{
let num_permits = num_permits as usize;
let mut curr = SemState(self.state.load(Acquire));
let mut maybe_waiter: Option<NonNull<Waiter>> = None;
macro_rules! revert_to_idle {
() => {
if let Some(waiter) = maybe_waiter {
unsafe { waiter.as_ref() }.revert_to_idle();
}
};
}
loop {
let mut next = curr;
if curr.is_closed() {
revert_to_idle!();
return Ready(Err(AcquireError::closed()));
}
let acquired = next.acquire_permits(num_permits, &self.stub);
if !acquired {
debug_assert!(curr.waiter().is_some() || curr.available_permits() < num_permits);
if let Some(waiter) = maybe_waiter.as_ref() {
let w = unsafe { waiter.as_ref() };
w.set_permits_to_acquire(num_permits - curr.available_permits());
} else {
if let Some(waiter) = get_waiter() {
let w = unsafe { waiter.as_ref() };
if !w.to_queued(num_permits - curr.available_permits()) {
return Pending;
}
maybe_waiter = Some(waiter);
} else {
return Pending;
}
}
next.set_waiter(maybe_waiter.unwrap());
}
debug_assert_ne!(curr.0, 0);
debug_assert_ne!(next.0, 0);
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => {
if acquired {
revert_to_idle!();
return Ready(Ok(()));
} else {
let prev_waiter =
curr.waiter().unwrap_or_else(|| NonNull::from(&*self.stub));
let waiter = maybe_waiter.unwrap();
unsafe {
prev_waiter.as_ref().store_next(waiter);
}
return Pending;
}
}
Err(actual) => {
curr = SemState(actual);
}
}
}
}
pub(crate) fn close(&self) {
let prev = self.rx_lock.fetch_or(1, AcqRel);
if prev != 0 {
return;
}
self.add_permits_locked(0, true);
}
pub(crate) fn add_permits(&self, n: usize) {
if n == 0 {
return;
}
let prev = self.rx_lock.fetch_add(n << 1, AcqRel);
if prev != 0 {
return;
}
self.add_permits_locked(n, false);
}
fn add_permits_locked(&self, mut rem: usize, mut closed: bool) {
while rem > 0 || 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);
closed = false;
actual
} else {
let actual = self.rx_lock.fetch_sub(n, AcqRel);
closed = actual & 1 == 1;
actual
};
rem = (actual >> 1) - rem;
}
}
fn add_permits_locked2(&self, mut n: usize, closed: bool) {
if closed {
n = usize::MAX;
}
'outer: while n > 0 {
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 {
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) {
thread::yield_now();
continue 'outer;
}
if closed {
debug_assert!(curr.is_closed(), "state = {:?}", curr);
return;
}
let mut next = curr;
next.release_permits(n, &self.stub);
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => return,
Err(actual) => {
curr = SemState(actual);
}
}
}
}
};
self.head.with_mut(|head| *head = next);
head = next;
next_ptr = next.as_ref().next.load(Acquire);
}
if !head.as_ref().assign_permits(&mut n, closed) {
assert_eq!(n, 0);
return;
}
if let Some(next) = NonNull::new(next_ptr) {
self.head.with_mut(|head| *head = next);
self.remove_queued(head, closed);
continue 'outer;
}
let state = SemState::load(&self.state, Acquire);
let tail = state.waiter().unwrap();
if tail != head {
thread::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);
self.remove_queued(head, closed);
continue 'outer;
}
thread::yield_now();
}
}
}
fn remove_queued(&self, waiter: NonNull<Waiter>, closed: bool) {
let mut curr = WaiterState(unsafe { waiter.as_ref() }.state.load(Acquire));
loop {
if curr.is_dropped() {
let _ = unsafe { Box::from_raw(waiter.as_ptr()) };
return;
}
if curr.permits_to_acquire() > 0 {
unsafe {
self.push_waiter(waiter, closed);
}
return;
}
let mut next = curr;
next.unset_queued();
let w = unsafe { waiter.as_ref() };
match w.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => return,
Err(actual) => {
curr = WaiterState(actual);
}
}
}
}
unsafe fn push_stub(&self, closed: bool) {
self.push_waiter(self.stub(), closed);
}
unsafe fn push_waiter(&self, waiter: NonNull<Waiter>, closed: bool) {
waiter.as_ref().next.store(ptr::null_mut(), Relaxed);
let next = SemState::new_ptr(waiter, closed);
let prev = SemState(self.state.swap(next.0, AcqRel));
debug_assert_eq!(closed, prev.is_closed());
let prev = prev.waiter().unwrap();
debug_assert_ne!(prev, waiter);
prev.as_ref().next.store(waiter.as_ptr(), Release);
}
fn stub(&self) -> NonNull<Waiter> {
unsafe { NonNull::new_unchecked(&*self.stub as *const _ as *mut _) }
}
}
impl Drop for Semaphore {
fn drop(&mut self) {
self.close();
}
}
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(crate) fn new() -> Permit {
use PermitState::Acquired;
Permit {
waiter: None,
state: Acquired(0),
}
}
#[allow(dead_code)]
pub(crate) fn is_acquired(&self) -> bool {
match self.state {
PermitState::Acquired(num) if num > 0 => true,
_ => false,
}
}
pub(crate) fn poll_acquire(
&mut self,
cx: &mut Context<'_>,
num_permits: u16,
semaphore: &Semaphore,
) -> Poll<Result<(), AcquireError>> {
use std::cmp::Ordering::*;
use PermitState::*;
match self.state {
Waiting(requested) => {
let waiter = self.waiter.as_ref().unwrap();
match requested.cmp(&num_permits) {
Less => {
let delta = num_permits - requested;
if !waiter.try_inc_permits_to_acquire(delta as usize) {
let waiter = NonNull::from(&**waiter);
let _ = semaphore.poll_acquire2(delta, || Some(waiter))?;
}
self.state = Waiting(num_permits);
}
Greater => {
let delta = requested - num_permits;
let to_release = waiter.try_dec_permits_to_acquire(delta as usize);
semaphore.add_permits(to_release);
self.state = Waiting(num_permits);
}
Equal => {}
}
if waiter.permits_to_acquire()? == 0 {
self.state = Acquired(requested);
return Ready(Ok(()));
}
waiter.waker.register_by_ref(cx.waker());
if waiter.permits_to_acquire()? == 0 {
self.state = Acquired(requested);
return Ready(Ok(()));
}
Pending
}
Acquired(acquired) => {
if acquired >= num_permits {
Ready(Ok(()))
} else {
match semaphore.poll_acquire(cx, num_permits - acquired, self)? {
Ready(()) => {
self.state = Acquired(num_permits);
Ready(Ok(()))
}
Pending => {
self.state = Waiting(num_permits);
Pending
}
}
}
}
}
}
pub(crate) fn try_acquire(
&mut self,
num_permits: u16,
semaphore: &Semaphore,
) -> Result<(), TryAcquireError> {
use PermitState::*;
match self.state {
Waiting(requested) => {
let waiter = self.waiter.as_ref().unwrap();
if requested > num_permits {
let delta = requested - num_permits;
let to_release = waiter.try_dec_permits_to_acquire(delta as usize);
semaphore.add_permits(to_release);
self.state = Waiting(num_permits);
}
let res = waiter.permits_to_acquire().map_err(to_try_acquire)?;
if res == 0 {
if requested < num_permits {
semaphore.try_acquire(num_permits - requested)?;
}
self.state = Acquired(num_permits);
Ok(())
} else {
Err(TryAcquireError::NoPermits)
}
}
Acquired(acquired) => {
if acquired < num_permits {
semaphore.try_acquire(num_permits - acquired)?;
self.state = Acquired(num_permits);
}
Ok(())
}
}
}
pub(crate) fn release(&mut self, n: u16, semaphore: &Semaphore) {
let n = self.forget(n);
semaphore.add_permits(n as usize);
}
pub(crate) fn forget(&mut self, n: u16) -> u16 {
use PermitState::*;
match self.state {
Waiting(requested) => {
let n = cmp::min(n, requested);
let acquired = self
.waiter
.as_ref()
.unwrap()
.try_dec_permits_to_acquire(n as usize) as u16;
if n == requested {
self.state = Acquired(0);
} else if acquired == requested - n {
self.state = Waiting(acquired);
} else {
self.state = Waiting(requested - n);
}
acquired
}
Acquired(acquired) => {
let n = cmp::min(n, acquired);
self.state = Acquired(acquired - n);
n
}
}
}
}
impl Default for Permit {
fn default() -> Self {
Self::new()
}
}
impl Drop for Permit {
fn drop(&mut self) {
if let Some(waiter) = self.waiter.take() {
let state = WaiterState(waiter.state.fetch_or(DROPPED, AcqRel));
if state.is_queued() {
std::mem::forget(waiter);
}
}
}
}
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 {
write!(fmt, "semaphore closed")
}
}
impl std::error::Error for AcquireError {}
impl TryAcquireError {
pub(crate) fn is_closed(&self) -> bool {
match self {
TryAcquireError::Closed => true,
_ => false,
}
}
pub(crate) fn is_no_permits(&self) -> bool {
match self {
TryAcquireError::NoPermits => true,
_ => false,
}
}
}
impl fmt::Display for TryAcquireError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TryAcquireError::Closed => write!(fmt, "semaphore closed"),
TryAcquireError::NoPermits => write!(fmt, "no permits available"),
}
}
}
impl std::error::Error for TryAcquireError {}
impl Waiter {
fn new() -> Waiter {
Waiter {
state: AtomicUsize::new(0),
waker: AtomicWaker::new(),
next: AtomicPtr::new(ptr::null_mut()),
}
}
fn permits_to_acquire(&self) -> Result<usize, AcquireError> {
let state = WaiterState(self.state.load(Acquire));
if state.is_closed() {
Err(AcquireError(()))
} else {
Ok(state.permits_to_acquire())
}
}
fn try_inc_permits_to_acquire(&self, n: usize) -> bool {
let mut curr = WaiterState(self.state.load(Acquire));
loop {
if !curr.is_queued() {
assert_eq!(0, curr.permits_to_acquire());
return false;
}
let mut next = curr;
next.set_permits_to_acquire(n + curr.permits_to_acquire());
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => return true,
Err(actual) => curr = WaiterState(actual),
}
}
}
fn try_dec_permits_to_acquire(&self, n: usize) -> usize {
let mut curr = WaiterState(self.state.load(Acquire));
loop {
if curr.is_closed() {
return 0;
}
if !curr.is_queued() {
assert_eq!(0, curr.permits_to_acquire());
}
let delta = cmp::min(n, curr.permits_to_acquire());
let rem = curr.permits_to_acquire() - delta;
let mut next = curr;
next.set_permits_to_acquire(rem);
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => return n - delta,
Err(actual) => curr = WaiterState(actual),
}
}
}
fn to_queued(&self, num_permits: usize) -> bool {
let mut curr = WaiterState(self.state.load(Acquire));
debug_assert_eq!(curr.permits_to_acquire(), 0);
loop {
let mut next = curr;
next.set_permits_to_acquire(num_permits);
next.set_queued();
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => {
if curr.is_queued() {
return false;
} else {
self.next.store(ptr::null_mut(), Relaxed);
return true;
}
}
Err(actual) => curr = WaiterState(actual),
}
}
}
fn set_permits_to_acquire(&self, num_permits: usize) {
debug_assert!(WaiterState(self.state.load(Acquire)).is_queued());
let mut state = WaiterState(QUEUED);
state.set_permits_to_acquire(num_permits);
self.state.store(state.0, Release);
}
fn assign_permits(&self, n: &mut usize, closed: bool) -> bool {
let mut curr = WaiterState(self.state.load(Acquire));
loop {
let mut next = curr;
let assign = cmp::min(curr.permits_to_acquire(), *n);
next.set_permits_to_acquire(curr.permits_to_acquire() - assign);
if closed {
next.set_closed();
}
match self.state.compare_exchange(curr.0, next.0, AcqRel, Acquire) {
Ok(_) => {
*n -= assign;
if next.permits_to_acquire() == 0 {
if curr.permits_to_acquire() > 0 {
self.waker.wake();
}
return true;
} else {
return false;
}
}
Err(actual) => curr = WaiterState(actual),
}
}
}
fn revert_to_idle(&self) {
self.state.store(0, Relaxed);
}
fn store_next(&self, next: NonNull<Waiter>) {
self.next.store(next.as_ptr(), Release);
}
}
impl SemState {
fn new(permits: usize, stub: &Waiter) -> 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<Waiter>, 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: &Waiter) -> bool {
!self.has_available_permits() && !self.is_stub(stub)
}
fn acquire_permits(&mut self, num: usize, stub: &Waiter) -> bool {
debug_assert!(num > 0);
if self.available_permits() < num {
return false;
}
debug_assert!(self.waiter().is_none());
self.0 -= num << NUM_SHIFT;
if self.0 == NUM_FLAG {
self.0 = stub as *const _ as usize;
}
true
}
fn release_permits(&mut self, permits: usize, stub: &Waiter) {
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<Waiter>> {
if self.is_waiter() {
let waiter = NonNull::new(self.as_ptr()).expect("null pointer stored");
Some(waiter)
} else {
None
}
}
fn as_ptr(self) -> *mut Waiter {
(self.0 & !CLOSED_FLAG) as *mut Waiter
}
fn set_waiter(&mut self, waiter: NonNull<Waiter>) {
let waiter = waiter.as_ptr() as usize;
debug_assert!(!self.is_closed());
self.0 = waiter;
}
fn is_stub(self, stub: &Waiter) -> bool {
self.as_ptr() as usize == stub as *const _ as usize
}
fn load(cell: &AtomicUsize, ordering: Ordering) -> SemState {
let value = cell.load(ordering);
SemState(value)
}
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 WaiterState {
fn permits_to_acquire(self) -> usize {
self.0 >> PERMIT_SHIFT
}
fn set_permits_to_acquire(&mut self, val: usize) {
self.0 = (val << PERMIT_SHIFT) | (self.0 & !PERMIT_MASK)
}
fn is_queued(self) -> bool {
self.0 & QUEUED == QUEUED
}
fn set_queued(&mut self) {
self.0 |= QUEUED;
}
fn is_closed(self) -> bool {
self.0 & CLOSED == CLOSED
}
fn set_closed(&mut self) {
self.0 |= CLOSED;
}
fn unset_queued(&mut self) {
assert!(self.is_queued());
self.0 -= QUEUED;
}
fn is_dropped(self) -> bool {
self.0 & DROPPED == DROPPED
}
}