Skip to content

Commit e4e8216

Browse files
james7132james7132
authored
Avoid boxing large futures (#88)
1 parent 214bc16 commit e4e8216

File tree

3 files changed

+99
-102
lines changed

3 files changed

+99
-102
lines changed

‎src/raw.rs‎

Lines changed: 65 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,11 @@
11
use alloc::alloc::Layout as StdLayout;
2-
use core::cell::UnsafeCell;
32
use core::future::Future;
43
use core::mem::{self, ManuallyDrop};
54
use core::pin::Pin;
65
use core::ptr::NonNull;
76
use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
87

9-
#[cfg(not(feature = "portable-atomic"))]
10-
use core::sync::atomic::AtomicUsize;
118
use core::sync::atomic::Ordering;
12-
#[cfg(feature = "portable-atomic")]
13-
use portable_atomic::AtomicUsize;
149

1510
use crate::header::Header;
1611
use crate::runnable::{Schedule, ScheduleInfo};
@@ -99,7 +94,7 @@ impl<F, T, S, M> Clone for RawTask<F, T, S, M> {
9994
}
10095

10196
impl<F, T, S, M> RawTask<F, T, S, M> {
102-
const TASK_LAYOUT: TaskLayout = Self::eval_task_layout();
97+
pub(crate) const TASK_LAYOUT: TaskLayout = Self::eval_task_layout();
10398

10499
/// Computes the memory layout for a task.
105100
#[inline]
@@ -131,81 +126,81 @@ impl<F, T, S, M> RawTask<F, T, S, M> {
131126
}
132127
}
133128

129+
/// Allocates a task with the given `future` and `schedule` function.
130+
///
131+
/// It is assumed that initially only the `Runnable` and the `Task` exist.
132+
///
133+
/// Use a macro to brute force inlining to minimize stack copies of potentially
134+
/// large futures.
135+
macro_rules! allocate_task {
136+
($f:tt, $s:tt, $m:tt, $builder:ident, $schedule:ident, $raw:ident => $future:block) => {{
137+
let allocation =
138+
alloc::alloc::alloc(RawTask::<$f, <$f as Future>::Output, $s, $m>::TASK_LAYOUT.layout);
139+
// Allocate enough space for the entire task.
140+
let ptr = NonNull::new(allocation as *mut ()).unwrap_or_else(|| crate::utils::abort());
141+
142+
let $raw = RawTask::<$f, <$f as Future>::Output, $s, $m>::from_ptr(ptr.as_ptr());
143+
144+
let crate::Builder {
145+
metadata,
146+
#[cfg(feature = "std")]
147+
propagate_panic,
148+
} = $builder;
149+
150+
// Write the header as the first field of the task.
151+
($raw.header as *mut Header<$m>).write(Header {
152+
#[cfg(not(feature = "portable-atomic"))]
153+
state: core::sync::atomic::AtomicUsize::new(SCHEDULED | TASK | REFERENCE),
154+
#[cfg(feature = "portable-atomic")]
155+
state: portable_atomic::AtomicUsize::new(SCHEDULED | TASK | REFERENCE),
156+
awaiter: core::cell::UnsafeCell::new(None),
157+
vtable: &RawTask::<$f, <$f as Future>::Output, $s, $m>::TASK_VTABLE,
158+
metadata,
159+
#[cfg(feature = "std")]
160+
propagate_panic,
161+
});
162+
163+
// Write the schedule function as the third field of the task.
164+
($raw.schedule as *mut S).write($schedule);
165+
166+
// Explicitly avoid using abort_on_panic here to avoid extra stack
167+
// copies of the future on lower optimization levels.
168+
let bomb = crate::utils::Bomb;
169+
170+
// Generate the future, now that the metadata has been pinned in place.
171+
// Write the future as the fourth field of the task.
172+
$raw.future.write($future);
173+
// (&(*raw.header).metadata)
174+
175+
mem::forget(bomb);
176+
ptr
177+
}};
178+
}
179+
180+
pub(crate) use allocate_task;
181+
134182
impl<F, T, S, M> RawTask<F, T, S, M>
135183
where
136184
F: Future<Output = T>,
137185
S: Schedule<M>,
138186
{
139-
const RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
187+
pub(crate) const RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
140188
Self::clone_waker,
141189
Self::wake,
142190
Self::wake_by_ref,
143191
Self::drop_waker,
144192
);
145193

146-
/// Allocates a task with the given `future` and `schedule` function.
147-
///
148-
/// It is assumed that initially only the `Runnable` and the `Task` exist.
149-
pub(crate) fn allocate<'a, Gen: FnOnce(&'a M) -> F>(
150-
future: Gen,
151-
schedule: S,
152-
builder: crate::Builder<M>,
153-
) -> NonNull<()>
154-
where
155-
F: 'a,
156-
M: 'a,
157-
{
158-
// Compute the layout of the task for allocation. Abort if the computation fails.
159-
//
160-
// n.b. notgull: task_layout now automatically aborts instead of panicking
161-
let task_layout = Self::task_layout();
162-
163-
unsafe {
164-
// Allocate enough space for the entire task.
165-
let ptr = match NonNull::new(alloc::alloc::alloc(task_layout.layout) as *mut ()) {
166-
None => abort(),
167-
Some(p) => p,
168-
};
169-
170-
let raw = Self::from_ptr(ptr.as_ptr());
171-
172-
let crate::Builder {
173-
metadata,
174-
#[cfg(feature = "std")]
175-
propagate_panic,
176-
} = builder;
177-
178-
// Write the header as the first field of the task.
179-
(raw.header as *mut Header<M>).write(Header {
180-
state: AtomicUsize::new(SCHEDULED | TASK | REFERENCE),
181-
awaiter: UnsafeCell::new(None),
182-
vtable: &TaskVTable {
183-
schedule: Self::schedule,
184-
drop_future: Self::drop_future,
185-
get_output: Self::get_output,
186-
drop_ref: Self::drop_ref,
187-
destroy: Self::destroy,
188-
run: Self::run,
189-
clone_waker: Self::clone_waker,
190-
layout_info: &Self::TASK_LAYOUT,
191-
},
192-
metadata,
193-
#[cfg(feature = "std")]
194-
propagate_panic,
195-
});
196-
197-
// Write the schedule function as the third field of the task.
198-
(raw.schedule as *mut S).write(schedule);
199-
200-
// Generate the future, now that the metadata has been pinned in place.
201-
let future = abort_on_panic(|| future(&(*raw.header).metadata));
202-
203-
// Write the future as the fourth field of the task.
204-
raw.future.write(future);
205-
206-
ptr
207-
}
208-
}
194+
pub(crate) const TASK_VTABLE: TaskVTable = TaskVTable {
195+
schedule: Self::schedule,
196+
drop_future: Self::drop_future,
197+
get_output: Self::get_output,
198+
drop_ref: Self::drop_ref,
199+
destroy: Self::destroy,
200+
run: Self::run,
201+
clone_waker: Self::clone_waker,
202+
layout_info: &Self::TASK_LAYOUT,
203+
};
209204

210205
/// Creates a `RawTask` from a raw task pointer.
211206
#[inline]

‎src/runnable.rs‎

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,8 @@ use core::ptr::NonNull;
66
use core::sync::atomic::Ordering;
77
use core::task::Waker;
88

9-
use alloc::boxed::Box;
10-
119
use crate::header::Header;
12-
use crate::raw::RawTask;
10+
use crate::raw::{allocate_task, RawTask};
1311
use crate::state::*;
1412
use crate::Task;
1513

@@ -277,6 +275,23 @@ impl Builder<()> {
277275
}
278276
}
279277

278+
// Use a macro to brute force inlining to minimize stack copies of potentially
279+
// large futures.
280+
macro_rules! spawn_unchecked {
281+
($f:tt, $s:tt, $m:tt, $builder:ident, $schedule:ident, $raw:ident => $future:block) => {{
282+
let ptr = allocate_task!($f, $s, $m, $builder, $schedule, $raw => $future);
283+
284+
#[allow(unused_unsafe)]
285+
// SAFTETY: The task was just allocated above.
286+
let runnable = unsafe { Runnable::from_raw(ptr) };
287+
let task = Task {
288+
ptr,
289+
_marker: PhantomData,
290+
};
291+
(runnable, task)
292+
}};
293+
}
294+
280295
impl<M> Builder<M> {
281296
/// Propagates panics that occur in the task.
282297
///
@@ -366,7 +381,9 @@ impl<M> Builder<M> {
366381
Fut::Output: Send + 'static,
367382
S: Schedule<M> + Send + Sync + 'static,
368383
{
369-
unsafe { self.spawn_unchecked(future, schedule) }
384+
unsafe {
385+
spawn_unchecked!(Fut, S, M, self, schedule, raw => { future(&(*raw.header).metadata) })
386+
}
370387
}
371388

372389
/// Creates a new thread-local task.
@@ -512,24 +529,7 @@ impl<M> Builder<M> {
512529
S: Schedule<M>,
513530
M: 'a,
514531
{
515-
// Allocate large futures on the heap.
516-
let ptr = if mem::size_of::<Fut>() >= 2048 {
517-
let future = |meta| {
518-
let future = future(meta);
519-
Box::pin(future)
520-
};
521-
522-
RawTask::<_, Fut::Output, S, M>::allocate(future, schedule, self)
523-
} else {
524-
RawTask::<Fut, Fut::Output, S, M>::allocate(future, schedule, self)
525-
};
526-
527-
let runnable = Runnable::from_raw(ptr);
528-
let task = Task {
529-
ptr,
530-
_marker: PhantomData,
531-
};
532-
(runnable, task)
532+
spawn_unchecked!(Fut, S, M, self, schedule, raw => { future(&(*raw.header).metadata) })
533533
}
534534
}
535535

@@ -570,7 +570,8 @@ where
570570
F::Output: Send + 'static,
571571
S: Schedule + Send + Sync + 'static,
572572
{
573-
unsafe { spawn_unchecked(future, schedule) }
573+
let builder = Builder::new();
574+
unsafe { spawn_unchecked!(F, S, (), builder, schedule, raw => { future }) }
574575
}
575576

576577
/// Creates a new thread-local task.
@@ -650,7 +651,8 @@ where
650651
F: Future,
651652
S: Schedule,
652653
{
653-
Builder::new().spawn_unchecked(move |()| future, schedule)
654+
let builder = Builder::new();
655+
spawn_unchecked!(F, S, (), builder, schedule, raw => { future })
654656
}
655657

656658
/// A handle to a runnable task.

‎src/utils.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,19 @@ pub(crate) fn abort() -> ! {
1717
panic!("aborting the process");
1818
}
1919

20+
pub(crate) struct Bomb;
21+
22+
impl Drop for Bomb {
23+
fn drop(&mut self) {
24+
abort();
25+
}
26+
}
27+
2028
/// Calls a function and aborts if it panics.
2129
///
2230
/// This is useful in unsafe code where we can't recover from panics.
2331
#[inline]
2432
pub(crate) fn abort_on_panic<T>(f: impl FnOnce() -> T) -> T {
25-
struct Bomb;
26-
27-
impl Drop for Bomb {
28-
fn drop(&mut self) {
29-
abort();
30-
}
31-
}
32-
3333
let bomb = Bomb;
3434
let t = f();
3535
mem::forget(bomb);

0 commit comments

Comments
 (0)