Skip to content

Commit 15d8eb8

Browse files
cjhopmanfacebook-github-bot
authored andcommitted
Move ExecutionContext and CancellationNotificationFuture+state into shared_state.rs
Summary: This now puts all the shared state objects for managing cancellation state into one place. I'll refactor these a bit to make the interactions clearer. Reviewed By: JakobDegen Differential Revision: D65362033 fbshipit-source-id: 2aa7cd12f2216a720438ac0c6b6cc02ad621c7ec
1 parent 1010400 commit 15d8eb8

File tree

4 files changed

+263
-268
lines changed

4 files changed

+263
-268
lines changed

app/buck2_futures/src/cancellation.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@ use dupe::Dupe;
1919
use futures::FutureExt;
2020
use once_cell::sync::Lazy;
2121

22-
use crate::details::cancellable_future::context::ExecutionContextInner;
23-
use crate::details::cancellable_future::CancellationNotificationData;
24-
use crate::details::cancellable_future::CancellationNotificationFuture;
2522
use crate::details::cancellation_context::CancellationContextInner;
2623
use crate::details::cancellation_context::ExplicitCancellationContext;
2724
use crate::details::cancellation_context::ExplicitCriticalSectionGuard;
25+
use crate::details::shared_state::CancellationNotificationData;
26+
use crate::details::shared_state::CancellationNotificationFuture;
27+
use crate::details::shared_state::ExecutionContextInner;
2828
use crate::details::shared_state::SharedState;
2929

3030
static NEVER_CANCELLED: Lazy<CancellationContext> =

app/buck2_futures/src/details/cancellable_future.rs

Lines changed: 1 addition & 262 deletions
Original file line numberDiff line numberDiff line change
@@ -15,21 +15,15 @@
1515
1616
use std::future::Future;
1717
use std::pin::Pin;
18-
use std::sync::atomic::AtomicU8;
19-
use std::sync::atomic::Ordering;
20-
use std::sync::Arc;
2118
use std::task::Context;
2219
use std::task::Poll;
2320

24-
use dupe::Dupe;
2521
use futures::future::BoxFuture;
26-
use futures::task::AtomicWaker;
27-
use parking_lot::Mutex;
2822
use pin_project::pin_project;
29-
use slab::Slab;
3023

3124
use crate::cancellation::CancellationContext;
3225
use crate::cancellation::CancellationHandle;
26+
use crate::details::shared_state::ExecutionContextOuter;
3327
use crate::details::shared_state::SharedState;
3428
use crate::drop_on_ready::DropOnReadyFuture;
3529
use crate::owning_future::OwningFuture;
@@ -161,261 +155,6 @@ impl<T> Future for ExplicitlyCancellableFutureInner<T> {
161155
}
162156
}
163157

164-
pub(crate) mod context {
165-
use super::*;
166-
use crate::cancellation::CriticalSectionGuard;
167-
168-
struct ExecutionContextData {
169-
cancellation_notification: CancellationNotificationData,
170-
171-
/// How many observers are preventing immediate cancellation.
172-
prevent_cancellation: usize,
173-
}
174-
175-
impl ExecutionContextData {
176-
/// Does this future not currently prevent its cancellation?
177-
fn can_exit(&self) -> bool {
178-
self.prevent_cancellation == 0
179-
}
180-
181-
fn enter_structured_cancellation(&mut self) -> CancellationNotificationData {
182-
self.prevent_cancellation += 1;
183-
184-
self.cancellation_notification.dupe()
185-
}
186-
187-
fn notify_cancelled(&mut self) {
188-
let updated = self.cancellation_notification.inner.notified.fetch_update(
189-
Ordering::SeqCst,
190-
Ordering::SeqCst,
191-
|old| match CancellationNotificationStatus::from(old) {
192-
CancellationNotificationStatus::Pending => {
193-
Some(CancellationNotificationStatus::Notified.into())
194-
}
195-
CancellationNotificationStatus::Notified => None,
196-
CancellationNotificationStatus::Disabled => None,
197-
},
198-
);
199-
if updated.is_ok() {
200-
if let Some(mut wakers) = self.cancellation_notification.inner.wakers.lock().take()
201-
{
202-
wakers.drain().for_each(|waker| waker.wake());
203-
}
204-
}
205-
}
206-
207-
fn exit_prevent_cancellation(&mut self) -> bool {
208-
self.prevent_cancellation -= 1;
209-
210-
self.prevent_cancellation == 0
211-
}
212-
213-
fn try_to_disable_cancellation(&mut self) -> bool {
214-
let maybe_updated = self.cancellation_notification.inner.notified.fetch_update(
215-
Ordering::SeqCst,
216-
Ordering::SeqCst,
217-
|old| match CancellationNotificationStatus::from(old) {
218-
CancellationNotificationStatus::Pending => {
219-
Some(CancellationNotificationStatus::Disabled.into())
220-
}
221-
CancellationNotificationStatus::Notified => None,
222-
CancellationNotificationStatus::Disabled => None,
223-
},
224-
);
225-
226-
match maybe_updated {
227-
Ok(_) => true,
228-
Err(old) => {
229-
let old = CancellationNotificationStatus::from(old);
230-
matches!(old, CancellationNotificationStatus::Disabled)
231-
}
232-
}
233-
}
234-
}
235-
236-
/// Context relating to execution of the `poll` of the future. This will contain the information
237-
/// required for the `CancellationContext` that the future holds to enter critical sections and
238-
/// structured cancellations.
239-
pub(crate) struct ExecutionContextOuter {
240-
shared: Arc<Mutex<ExecutionContextData>>,
241-
}
242-
243-
pub(crate) struct ExecutionContextInner {
244-
shared: Arc<Mutex<ExecutionContextData>>,
245-
}
246-
247-
impl ExecutionContextOuter {
248-
pub(crate) fn new() -> (ExecutionContextOuter, ExecutionContextInner) {
249-
let shared = Arc::new(Mutex::new(ExecutionContextData {
250-
cancellation_notification: {
251-
CancellationNotificationData {
252-
inner: Arc::new(CancellationNotificationDataInner {
253-
notified: Default::default(),
254-
wakers: Mutex::new(Some(Default::default())),
255-
}),
256-
}
257-
},
258-
prevent_cancellation: 0,
259-
}));
260-
(
261-
ExecutionContextOuter {
262-
shared: shared.dupe(),
263-
},
264-
ExecutionContextInner { shared },
265-
)
266-
}
267-
268-
pub(crate) fn notify_cancelled(&self) -> bool {
269-
let mut lock = self.shared.lock();
270-
if lock.can_exit() {
271-
true
272-
} else {
273-
lock.notify_cancelled();
274-
false
275-
}
276-
}
277-
278-
pub(crate) fn can_exit(&self) -> bool {
279-
self.shared.lock().can_exit()
280-
}
281-
}
282-
283-
impl ExecutionContextInner {
284-
pub(crate) fn enter_structured_cancellation(&self) -> CriticalSectionGuard {
285-
let mut shared = self.shared.lock();
286-
287-
let notification = shared.enter_structured_cancellation();
288-
289-
CriticalSectionGuard::new_explicit(self, notification)
290-
}
291-
292-
pub(crate) fn try_to_disable_cancellation(&self) -> bool {
293-
let mut shared = self.shared.lock();
294-
if shared.try_to_disable_cancellation() {
295-
true
296-
} else {
297-
// couldn't prevent cancellation, so release our hold onto the counter
298-
shared.exit_prevent_cancellation();
299-
false
300-
}
301-
}
302-
303-
pub(crate) fn exit_prevent_cancellation(&self) -> bool {
304-
let mut shared = self.shared.lock();
305-
shared.exit_prevent_cancellation()
306-
}
307-
}
308-
}
309-
use context::ExecutionContextOuter;
310-
311-
enum CancellationNotificationStatus {
312-
/// no notifications yet. maps to '0'
313-
Pending,
314-
/// notified, maps to '1'
315-
Notified,
316-
/// disabled notifications, maps to '2'
317-
Disabled,
318-
}
319-
320-
impl From<u8> for CancellationNotificationStatus {
321-
fn from(value: u8) -> Self {
322-
match value {
323-
0 => CancellationNotificationStatus::Pending,
324-
1 => CancellationNotificationStatus::Notified,
325-
2 => CancellationNotificationStatus::Disabled,
326-
_ => panic!("invalid status"),
327-
}
328-
}
329-
}
330-
331-
impl From<CancellationNotificationStatus> for u8 {
332-
fn from(value: CancellationNotificationStatus) -> Self {
333-
match value {
334-
CancellationNotificationStatus::Pending => 0,
335-
CancellationNotificationStatus::Notified => 1,
336-
CancellationNotificationStatus::Disabled => 2,
337-
}
338-
}
339-
}
340-
341-
#[derive(Clone, Dupe)]
342-
pub(crate) struct CancellationNotificationData {
343-
inner: Arc<CancellationNotificationDataInner>,
344-
}
345-
346-
struct CancellationNotificationDataInner {
347-
/// notification status per enum 'CancellationNotificationStatus'
348-
notified: AtomicU8,
349-
wakers: Mutex<Option<Slab<Arc<AtomicWaker>>>>,
350-
}
351-
352-
pub(crate) struct CancellationNotificationFuture {
353-
data: CancellationNotificationData,
354-
// index into the waker for this future held by the Slab in 'CancellationNotificationData'
355-
id: Option<usize>,
356-
// duplicate of the waker held for us to update the waker on poll without acquiring lock
357-
waker: Arc<AtomicWaker>,
358-
}
359-
360-
impl CancellationNotificationFuture {
361-
pub(crate) fn new(data: CancellationNotificationData) -> Self {
362-
let waker = Arc::new(AtomicWaker::new());
363-
let id = data
364-
.inner
365-
.wakers
366-
.lock()
367-
.as_mut()
368-
.map(|wakers| wakers.insert(waker.dupe()));
369-
CancellationNotificationFuture { data, id, waker }
370-
}
371-
372-
fn remove_waker(&mut self, id: Option<usize>) {
373-
if let Some(id) = id {
374-
self.data
375-
.inner
376-
.wakers
377-
.lock()
378-
.as_mut()
379-
.map(|wakers| wakers.remove(id));
380-
}
381-
}
382-
}
383-
384-
impl Clone for CancellationNotificationFuture {
385-
fn clone(&self) -> Self {
386-
CancellationNotificationFuture::new(self.data.dupe())
387-
}
388-
}
389-
390-
impl Dupe for CancellationNotificationFuture {}
391-
392-
impl Drop for CancellationNotificationFuture {
393-
fn drop(&mut self) {
394-
self.remove_waker(self.id);
395-
}
396-
}
397-
398-
impl Future for CancellationNotificationFuture {
399-
type Output = ();
400-
401-
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
402-
match CancellationNotificationStatus::from(self.data.inner.notified.load(Ordering::SeqCst))
403-
{
404-
CancellationNotificationStatus::Notified => {
405-
// take the id so that we don't need to lock the wakers when this future is dropped
406-
// after completion
407-
let id = self.id.take();
408-
self.remove_waker(id);
409-
Poll::Ready(())
410-
}
411-
_ => {
412-
self.waker.register(cx.waker());
413-
Poll::Pending
414-
}
415-
}
416-
}
417-
}
418-
419158
#[cfg(test)]
420159
mod tests {
421160
use std::future::Future;

app/buck2_futures/src/details/cancellation_context.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ use crate::cancellation::CancellationObserver;
1919
use crate::cancellation::CancellationObserverInner;
2020
use crate::cancellation::CriticalSectionGuard;
2121
use crate::cancellation::DisableCancellationGuard;
22-
use crate::details::cancellable_future::context::ExecutionContextInner;
23-
use crate::details::cancellable_future::CancellationNotificationData;
24-
use crate::details::cancellable_future::CancellationNotificationFuture;
22+
use crate::details::shared_state::CancellationNotificationData;
23+
use crate::details::shared_state::CancellationNotificationFuture;
24+
use crate::details::shared_state::ExecutionContextInner;
2525

2626
pub struct ExplicitCriticalSectionGuard<'a> {
2727
pub(crate) context: Option<&'a ExecutionContextInner>,

0 commit comments

Comments
 (0)