futures_ext/future/
mod.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under both the MIT license found in the
5 * LICENSE-MIT file in the root directory of this source tree and the Apache
6 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
7 * of this source tree.
8 */
9
10//! Module extending functionality of [`futures::future`] module
11
12mod abort_handle_ref;
13mod conservative_receiver;
14mod on_cancel;
15mod on_cancel_with_data;
16mod try_shared;
17
18use std::time::Duration;
19
20use anyhow::Error;
21use futures::future::Future;
22use futures::future::FutureExt;
23use futures::future::TryFuture;
24pub use shared_error::anyhow::SharedError;
25use tokio::time::Timeout;
26
27pub use self::abort_handle_ref::spawn_controlled;
28pub use self::abort_handle_ref::ControlledHandle;
29pub use self::conservative_receiver::ConservativeReceiver;
30pub use self::on_cancel::OnCancel;
31pub use self::on_cancel_with_data::CancelData;
32pub use self::on_cancel_with_data::OnCancelWithData;
33pub use self::try_shared::TryShared;
34
35/// A trait implemented by default for all Futures which extends the standard
36/// functionality.
37pub trait FbFutureExt: Future {
38    /// Construct a new [tokio::time::Timeout].
39    fn timeout(self, timeout: Duration) -> Timeout<Self>
40    where
41        Self: Sized,
42    {
43        tokio::time::timeout(timeout, self)
44    }
45
46    /// Call the `on_cancel` callback if this future is canceled (dropped
47    /// without completion).
48    fn on_cancel<F: FnOnce()>(self, on_cancel: F) -> OnCancel<Self, F>
49    where
50        Self: Sized,
51    {
52        OnCancel::new(self, on_cancel)
53    }
54
55    /// Call the `on_cancel` callback if this future is canceled (dropped
56    /// without completion).  Pass additional data extracted from the
57    /// inner future via the CancelData trait.
58    fn on_cancel_with_data<F>(self, on_cancel: F) -> OnCancelWithData<Self, F>
59    where
60        Self: Sized + CancelData,
61        F: FnOnce(Self::Data),
62    {
63        OnCancelWithData::new(self, on_cancel)
64    }
65}
66
67impl<T> FbFutureExt for T where T: Future + ?Sized {}
68
69/// A trait implemented by default for all Futures which extends the standard
70/// functionality.
71pub trait FbTryFutureExt: Future {
72    /// Create a cloneable handle to this future where all handles will resolve
73    /// to the same result.
74    ///
75    /// Similar to [futures::future::Shared], but instead works on Futures
76    /// returning Result where Err is [anyhow::Error].
77    /// This is achieved by storing [anyhow::Error] in [std::sync::Arc].
78    fn try_shared(self) -> TryShared<Self>
79    where
80        Self: TryFuture<Error = Error> + Sized,
81        <Self as TryFuture>::Ok: Clone,
82    {
83        self::try_shared::try_shared(self)
84    }
85
86    /// Convert a `Future` of `Result<Result<I, E1>, E2>` into a `Future` of
87    /// `Result<I, E1>`, assuming `E2` can convert into `E1`.
88    #[allow(clippy::type_complexity)]
89    fn flatten_err<I, E1, E2>(
90        self,
91    ) -> futures::future::Map<Self, fn(Result<Result<I, E1>, E2>) -> Result<I, E1>>
92    where
93        Self: Sized,
94        Self: Future<Output = Result<Result<I, E1>, E2>>,
95        E1: From<E2>,
96    {
97        fn flatten_err<I, E1, E2>(e: Result<Result<I, E1>, E2>) -> Result<I, E1>
98        where
99            E1: From<E2>,
100        {
101            match e {
102                Ok(Ok(i)) => Ok(i),
103                Ok(Err(e1)) => Err(e1),
104                Err(e2) => Err(E1::from(e2)),
105            }
106        }
107
108        self.map(flatten_err)
109    }
110}
111
112impl<T> FbTryFutureExt for T where T: TryFuture + ?Sized {}