Skip to content

Commit 921e8a0

Browse files
author
Stjepan Glavina
authored
Make the crate no_std (#15)
* Remove crossbeam * Rewrite notification and registration * Make crate no_std * Add a comment
1 parent 119f82f commit 921e8a0

File tree

8 files changed

+164
-107
lines changed

8 files changed

+164
-107
lines changed

Cargo.toml

+5-1
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,13 @@ categories = ["asynchronous", "concurrency"]
1313
readme = "README.md"
1414

1515
[dependencies]
16-
crossbeam-utils = "0.7.0"
16+
libc = "0.2.66"
17+
18+
[target.'cfg(windows)'.dependencies]
19+
winapi = { version = "0.3.8", features = ["processthreadsapi"] }
1720

1821
[dev-dependencies]
1922
crossbeam = "0.7.3"
23+
crossbeam-utils = "0.7.0"
2024
futures = "0.3.1"
2125
lazy_static = "1.4.0"

src/header.rs

+87-52
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
1-
use std::alloc::Layout;
2-
use std::cell::Cell;
3-
use std::fmt;
4-
use std::sync::atomic::{AtomicUsize, Ordering};
5-
use std::task::Waker;
6-
7-
use crossbeam_utils::Backoff;
1+
use core::alloc::Layout;
2+
use core::cell::UnsafeCell;
3+
use core::fmt;
4+
use core::sync::atomic::{AtomicUsize, Ordering};
5+
use core::task::Waker;
86

97
use crate::raw::TaskVTable;
108
use crate::state::*;
@@ -22,7 +20,7 @@ pub(crate) struct Header {
2220
/// The task that is blocked on the `JoinHandle`.
2321
///
2422
/// This waker needs to be woken up once the task completes or is closed.
25-
pub(crate) awaiter: Cell<Option<Waker>>,
23+
pub(crate) awaiter: UnsafeCell<Option<Waker>>,
2624

2725
/// The virtual table.
2826
///
@@ -55,7 +53,7 @@ impl Header {
5553
Ok(_) => {
5654
// Notify the awaiter that the task has been closed.
5755
if state & AWAITER != 0 {
58-
self.notify();
56+
self.notify(None);
5957
}
6058

6159
break;
@@ -67,68 +65,105 @@ impl Header {
6765

6866
/// Notifies the awaiter blocked on this task.
6967
///
70-
/// If there is a registered waker, it will be removed from the header and woken up.
68+
/// If the awaiter is the same as the current waker, it will not be notified.
7169
#[inline]
72-
pub(crate) fn notify(&self) {
73-
if let Some(waker) = self.swap_awaiter(None) {
74-
// We need a safeguard against panics because waking can panic.
75-
abort_on_panic(|| {
76-
waker.wake();
77-
});
78-
}
79-
}
70+
pub(crate) fn notify(&self, current: Option<&Waker>) {
71+
// Mark the awaiter as being notified.
72+
let state = self.state.fetch_or(NOTIFYING, Ordering::AcqRel);
8073

81-
/// Notifies the awaiter blocked on this task, unless its waker matches `current`.
82-
///
83-
/// If there is a registered waker, it will be removed from the header in any case.
84-
#[inline]
85-
pub(crate) fn notify_unless(&self, current: &Waker) {
86-
if let Some(waker) = self.swap_awaiter(None) {
87-
if !waker.will_wake(current) {
74+
// If the awaiter was not being notified nor registered...
75+
if state & (NOTIFYING | REGISTERING) == 0 {
76+
// Take the waker out.
77+
let waker = unsafe { (*self.awaiter.get()).take() };
78+
79+
// Mark the state as not being notified anymore nor containing an awaiter.
80+
self.state
81+
.fetch_and(!NOTIFYING & !AWAITER, Ordering::Release);
82+
83+
if let Some(w) = waker {
8884
// We need a safeguard against panics because waking can panic.
89-
abort_on_panic(|| {
90-
waker.wake();
85+
abort_on_panic(|| match current {
86+
None => w.wake(),
87+
Some(c) if !w.will_wake(c) => w.wake(),
88+
Some(_) => {}
9189
});
9290
}
9391
}
9492
}
9593

96-
/// Swaps the awaiter for a new waker and returns the previous value.
94+
/// Registers a new awaiter blocked on this task.
95+
///
96+
/// This method is called when `JoinHandle` is polled and the task has not completed.
9797
#[inline]
98-
pub(crate) fn swap_awaiter(&self, new: Option<Waker>) -> Option<Waker> {
99-
let new_is_none = new.is_none();
98+
pub(crate) fn register(&self, waker: &Waker) {
99+
// Load the state and synchronize with it.
100+
let mut state = self.state.fetch_or(0, Ordering::Acquire);
100101

101-
// We're about to try acquiring the lock in a loop. If it's already being held by another
102-
// thread, we'll have to spin for a while so it's best to employ a backoff strategy.
103-
let backoff = Backoff::new();
104102
loop {
105-
// Acquire the lock. If we're storing an awaiter, then also set the awaiter flag.
106-
let state = if new_is_none {
107-
self.state.fetch_or(LOCKED, Ordering::Acquire)
108-
} else {
109-
self.state.fetch_or(LOCKED | AWAITER, Ordering::Acquire)
110-
};
103+
// There can't be two concurrent registrations because `JoinHandle` can only be polled
104+
// by a unique pinned reference.
105+
debug_assert!(state & REGISTERING == 0);
106+
107+
// If the state is being notified at this moment, just wake and return without
108+
// registering.
109+
if state & NOTIFYING != 0 {
110+
waker.wake_by_ref();
111+
return;
112+
}
111113

112-
// If the lock was acquired, break from the loop.
113-
if state & LOCKED == 0 {
114-
break;
114+
// Mark the state to let other threads know we're registering a new awaiter.
115+
match self.state.compare_exchange_weak(
116+
state,
117+
state | REGISTERING,
118+
Ordering::AcqRel,
119+
Ordering::Acquire,
120+
) {
121+
Ok(_) => {
122+
state |= REGISTERING;
123+
break;
124+
}
125+
Err(s) => state = s,
115126
}
127+
}
116128

117-
// Snooze for a little while because the lock is held by another thread.
118-
backoff.snooze();
129+
// Put the waker into the awaiter field.
130+
unsafe {
131+
abort_on_panic(|| (*self.awaiter.get()) = Some(waker.clone()));
119132
}
120133

121-
// Replace the awaiter.
122-
let old = self.awaiter.replace(new);
134+
// This variable will contain the newly registered waker if a notification comes in before
135+
// we complete registration.
136+
let mut waker = None;
137+
138+
loop {
139+
// If there was a notification, take the waker out of the awaiter field.
140+
if state & NOTIFYING != 0 {
141+
if let Some(w) = unsafe { (*self.awaiter.get()).take() } {
142+
waker = Some(w);
143+
}
144+
}
145+
146+
// The new state is not being notified nor registered, but there might or might not be
147+
// an awaiter depending on whether there was a concurrent notification.
148+
let new = if waker.is_none() {
149+
(state & !NOTIFYING & !REGISTERING) | AWAITER
150+
} else {
151+
state & !NOTIFYING & !REGISTERING & !AWAITER
152+
};
123153

124-
// Release the lock. If we've cleared the awaiter, then also unset the awaiter flag.
125-
if new_is_none {
126-
self.state.fetch_and(!LOCKED & !AWAITER, Ordering::Release);
127-
} else {
128-
self.state.fetch_and(!LOCKED, Ordering::Release);
154+
match self
155+
.state
156+
.compare_exchange_weak(state, new, Ordering::AcqRel, Ordering::Acquire)
157+
{
158+
Ok(_) => break,
159+
Err(s) => state = s,
160+
}
129161
}
130162

131-
old
163+
// If there was a notification during registration, wake the awaiter now.
164+
if let Some(w) = waker {
165+
abort_on_panic(|| w.wake());
166+
}
132167
}
133168

134169
/// Returns the offset at which the tag of type `T` is stored.

src/join_handle.rs

+11-11
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
use std::fmt;
2-
use std::future::Future;
3-
use std::marker::{PhantomData, Unpin};
4-
use std::pin::Pin;
5-
use std::ptr::NonNull;
6-
use std::sync::atomic::Ordering;
7-
use std::task::{Context, Poll};
1+
use core::fmt;
2+
use core::future::Future;
3+
use core::marker::{PhantomData, Unpin};
4+
use core::pin::Pin;
5+
use core::ptr::NonNull;
6+
use core::sync::atomic::Ordering;
7+
use core::task::{Context, Poll};
88

99
use crate::header::Header;
1010
use crate::state::*;
@@ -71,7 +71,7 @@ impl<R, T> JoinHandle<R, T> {
7171

7272
// Notify the awaiter that the task has been closed.
7373
if state & AWAITER != 0 {
74-
(*header).notify();
74+
(*header).notify(None);
7575
}
7676

7777
break;
@@ -190,7 +190,7 @@ impl<R, T> Future for JoinHandle<R, T> {
190190
if state & CLOSED != 0 {
191191
// Even though the awaiter is most likely the current task, it could also be
192192
// another task.
193-
(*header).notify_unless(cx.waker());
193+
(*header).notify(Some(cx.waker()));
194194
return Poll::Ready(None);
195195
}
196196

@@ -199,7 +199,7 @@ impl<R, T> Future for JoinHandle<R, T> {
199199
// Replace the waker with one associated with the current task. We need a
200200
// safeguard against panics because dropping the previous waker can panic.
201201
abort_on_panic(|| {
202-
(*header).swap_awaiter(Some(cx.waker().clone()));
202+
(*header).register(cx.waker());
203203
});
204204

205205
// Reload the state after registering. It is possible that the task became
@@ -230,7 +230,7 @@ impl<R, T> Future for JoinHandle<R, T> {
230230
// Notify the awaiter. Even though the awaiter is most likely the current
231231
// task, it could also be another task.
232232
if state & AWAITER != 0 {
233-
(*header).notify_unless(cx.waker());
233+
(*header).notify(Some(cx.waker()));
234234
}
235235

236236
// Take the output from the task.

src/lib.rs

+3
Original file line numberDiff line numberDiff line change
@@ -97,10 +97,13 @@
9797
//! [`Task`]: struct.Task.html
9898
//! [`JoinHandle`]: struct.JoinHandle.html
9999
100+
#![no_std]
100101
#![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]
101102
#![doc(test(attr(deny(rust_2018_idioms, warnings))))]
102103
#![doc(test(attr(allow(unused_extern_crates, unused_variables))))]
103104

105+
extern crate alloc;
106+
104107
mod header;
105108
mod join_handle;
106109
mod raw;

src/raw.rs

+18-18
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
use std::alloc::{self, Layout};
2-
use std::cell::Cell;
3-
use std::future::Future;
4-
use std::marker::PhantomData;
5-
use std::mem::{self, ManuallyDrop};
6-
use std::pin::Pin;
7-
use std::ptr::NonNull;
8-
use std::sync::atomic::{AtomicUsize, Ordering};
9-
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
1+
use alloc::alloc::Layout;
2+
use core::cell::UnsafeCell;
3+
use core::future::Future;
4+
use core::marker::PhantomData;
5+
use core::mem::{self, ManuallyDrop};
6+
use core::pin::Pin;
7+
use core::ptr::NonNull;
8+
use core::sync::atomic::{AtomicUsize, Ordering};
9+
use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
1010

1111
use crate::header::Header;
1212
use crate::state::*;
@@ -107,8 +107,8 @@ where
107107

108108
unsafe {
109109
// Allocate enough space for the entire task.
110-
let raw_task = match NonNull::new(alloc::alloc(task_layout.layout) as *mut ()) {
111-
None => std::process::abort(),
110+
let raw_task = match NonNull::new(alloc::alloc::alloc(task_layout.layout) as *mut ()) {
111+
None => libc::abort(),
112112
Some(p) => p,
113113
};
114114

@@ -117,7 +117,7 @@ where
117117
// Write the header as the first field of the task.
118118
(raw.header as *mut Header).write(Header {
119119
state: AtomicUsize::new(SCHEDULED | HANDLE | REFERENCE),
120-
awaiter: Cell::new(None),
120+
awaiter: UnsafeCell::new(None),
121121
vtable: &TaskVTable {
122122
raw_waker_vtable: RawWakerVTable::new(
123123
Self::clone_waker,
@@ -307,7 +307,7 @@ where
307307
if state & RUNNING == 0 {
308308
// If the reference count overflowed, abort.
309309
if state > isize::max_value() as usize {
310-
std::process::abort();
310+
libc::abort();
311311
}
312312

313313
// Schedule the task. There is no need to call `Self::schedule(ptr)`
@@ -339,7 +339,7 @@ where
339339

340340
// If the reference count overflowed, abort.
341341
if state > isize::max_value() as usize {
342-
std::process::abort();
342+
libc::abort();
343343
}
344344

345345
RawWaker::new(ptr, raw_waker_vtable)
@@ -449,7 +449,7 @@ where
449449
});
450450

451451
// Finally, deallocate the memory reserved by the task.
452-
alloc::dealloc(ptr as *mut u8, task_layout.layout);
452+
alloc::alloc::dealloc(ptr as *mut u8, task_layout.layout);
453453
}
454454

455455
/// Runs a task.
@@ -474,7 +474,7 @@ where
474474
if state & CLOSED != 0 {
475475
// Notify the awaiter that the task has been closed.
476476
if state & AWAITER != 0 {
477-
(*raw.header).notify();
477+
(*raw.header).notify(None);
478478
}
479479

480480
// Drop the future.
@@ -542,7 +542,7 @@ where
542542

543543
// Notify the awaiter that the task has been completed.
544544
if state & AWAITER != 0 {
545-
(*raw.header).notify();
545+
(*raw.header).notify(None);
546546
}
547547

548548
// Drop the task reference.
@@ -649,7 +649,7 @@ where
649649

650650
// Notify the awaiter that the task has been closed.
651651
if state & AWAITER != 0 {
652-
(*raw.header).notify();
652+
(*raw.header).notify(None);
653653
}
654654

655655
// Drop the task reference.

src/state.rs

+10-4
Original file line numberDiff line numberDiff line change
@@ -48,10 +48,16 @@ pub(crate) const HANDLE: usize = 1 << 4;
4848
/// check that tells us if we need to wake anyone without acquiring the lock inside the task.
4949
pub(crate) const AWAITER: usize = 1 << 5;
5050

51-
/// Set if the awaiter is locked.
51+
/// Set if an awaiter is being registered.
5252
///
53-
/// This lock is acquired before a new awaiter is registered or the existing one is woken up.
54-
pub(crate) const LOCKED: usize = 1 << 6;
53+
/// This flag is set when `JoinHandle` is polled and we are registering a new awaiter.
54+
pub(crate) const REGISTERING: usize = 1 << 6;
55+
56+
/// Set if the awaiter is being notified.
57+
///
58+
/// This flag is set when notifying the awaiter. If an awaiter is concurrently registered and
59+
/// notified, whichever side came first will take over the reposibility of resolving the race.
60+
pub(crate) const NOTIFYING: usize = 1 << 7;
5561

5662
/// A single reference.
5763
///
@@ -61,4 +67,4 @@ pub(crate) const LOCKED: usize = 1 << 6;
6167
///
6268
/// Note that the reference counter only tracks the `Task` and `Waker`s. The `JoinHandle` is
6369
/// tracked separately by the `HANDLE` flag.
64-
pub(crate) const REFERENCE: usize = 1 << 7;
70+
pub(crate) const REFERENCE: usize = 1 << 8;

0 commit comments

Comments
 (0)