Skip to content

Commit

Permalink
Make BiLock strict-provenance compatible (#2716)
Browse files Browse the repository at this point in the history
  • Loading branch information
taiki-e committed Mar 11, 2023
1 parent 4266221 commit 0ffaaa0
Showing 1 changed file with 30 additions and 16 deletions.
46 changes: 30 additions & 16 deletions futures-util/src/lock/bilock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
use alloc::boxed::Box;
use alloc::sync::Arc;
use core::cell::UnsafeCell;
use core::fmt;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::AtomicPtr;
use core::sync::atomic::Ordering::SeqCst;
use core::{fmt, ptr};
#[cfg(feature = "bilock")]
use futures_core::future::Future;
use futures_core::task::{Context, Poll, Waker};
Expand Down Expand Up @@ -41,7 +41,7 @@ pub struct BiLock<T> {

#[derive(Debug)]
struct Inner<T> {
state: AtomicUsize,
state: AtomicPtr<Waker>,
value: Option<UnsafeCell<T>>,
}

Expand All @@ -61,7 +61,10 @@ impl<T> BiLock<T> {
/// Similarly, reuniting the lock and extracting the inner value is only
/// possible when `T` is `Unpin`.
pub fn new(t: T) -> (Self, Self) {
let arc = Arc::new(Inner { state: AtomicUsize::new(0), value: Some(UnsafeCell::new(t)) });
let arc = Arc::new(Inner {
state: AtomicPtr::new(ptr::null_mut()),
value: Some(UnsafeCell::new(t)),
});

(Self { arc: arc.clone() }, Self { arc })
}
Expand All @@ -87,7 +90,8 @@ impl<T> BiLock<T> {
pub fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<BiLockGuard<'_, T>> {
let mut waker = None;
loop {
match self.arc.state.swap(1, SeqCst) {
let n = self.arc.state.swap(invalid_ptr(1), SeqCst);
match n as usize {
// Woohoo, we grabbed the lock!
0 => return Poll::Ready(BiLockGuard { bilock: self }),

Expand All @@ -96,27 +100,27 @@ impl<T> BiLock<T> {

// A task was previously blocked on this lock, likely our task,
// so we need to update that task.
n => unsafe {
let mut prev = Box::from_raw(n as *mut Waker);
_ => unsafe {
let mut prev = Box::from_raw(n);
*prev = cx.waker().clone();
waker = Some(prev);
},
}

// type ascription for safety's sake!
let me: Box<Waker> = waker.take().unwrap_or_else(|| Box::new(cx.waker().clone()));
let me = Box::into_raw(me) as usize;
let me = Box::into_raw(me);

match self.arc.state.compare_exchange(1, me, SeqCst, SeqCst) {
match self.arc.state.compare_exchange(invalid_ptr(1), me, SeqCst, SeqCst) {
// The lock is still locked, but we've now parked ourselves, so
// just report that we're scheduled to receive a notification.
Ok(_) => return Poll::Pending,

// Oops, looks like the lock was unlocked after our swap above
// and before the compare_exchange. Deallocate what we just
// allocated and go through the loop again.
Err(0) => unsafe {
waker = Some(Box::from_raw(me as *mut Waker));
Err(n) if n.is_null() => unsafe {
waker = Some(Box::from_raw(me));
},

// The top of this loop set the previous state to 1, so if we
Expand All @@ -125,7 +129,7 @@ impl<T> BiLock<T> {
// but we're trying to acquire the lock and there's only one
// other reference of the lock, so it should be impossible for
// that task to ever block itself.
Err(n) => panic!("invalid state: {}", n),
Err(n) => panic!("invalid state: {}", n as usize),
}
}
}
Expand Down Expand Up @@ -164,7 +168,8 @@ impl<T> BiLock<T> {
}

fn unlock(&self) {
match self.arc.state.swap(0, SeqCst) {
let n = self.arc.state.swap(ptr::null_mut(), SeqCst);
match n as usize {
// we've locked the lock, shouldn't be possible for us to see an
// unlocked lock.
0 => panic!("invalid unlocked state"),
Expand All @@ -174,8 +179,8 @@ impl<T> BiLock<T> {

// Another task has parked themselves on this lock, let's wake them
// up as its now their turn.
n => unsafe {
Box::from_raw(n as *mut Waker).wake();
_ => unsafe {
Box::from_raw(n).wake();
},
}
}
Expand All @@ -189,7 +194,7 @@ impl<T: Unpin> Inner<T> {

impl<T> Drop for Inner<T> {
fn drop(&mut self) {
assert_eq!(self.state.load(SeqCst), 0);
assert!(self.state.load(SeqCst).is_null());
}
}

Expand Down Expand Up @@ -277,3 +282,12 @@ impl<'a, T> Future for BiLockAcquire<'a, T> {
self.bilock.poll_lock(cx)
}
}

// Based on core::ptr::invalid_mut. Equivalent to `addr as *mut T`, but is strict-provenance compatible.
#[allow(clippy::useless_transmute)]
#[inline]
fn invalid_ptr<T>(addr: usize) -> *mut T {
// SAFETY: every valid integer is also a valid pointer (as long as you don't dereference that
// pointer).
unsafe { core::mem::transmute(addr) }
}

0 comments on commit 0ffaaa0

Please sign in to comment.