fxfs/
future_with_guard.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use pin_project::pin_project;
6use std::future::Future;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10/// A wrapper around a future and a guard object where the correctness of the future requires the
11/// guard object to be held while the future is alive.
12///
13/// This is equivalent to the below code but produces a future that is almost half the size of the
14/// future that rust generates: https://github.com/rust-lang/rust/issues/108906.
15/// ```rust
16/// let guard = acquire_guard();
17/// executor.spawn(async move {
18///   let _guard = guard;
19///   task.await;
20/// });
21/// ```
22#[pin_project]
23pub struct FutureWithGuard<T, F, R>
24where
25    F: Future<Output = R> + Send + 'static,
26    T: Send + 'static,
27{
28    #[pin]
29    future: F,
30    _object: T,
31}
32
33impl<T, F, R> FutureWithGuard<T, F, R>
34where
35    F: Future<Output = R> + Send + 'static,
36    T: Send + 'static,
37{
38    pub fn new(object: T, future: F) -> Self {
39        Self { future, _object: object }
40    }
41}
42
43impl<T, F, R> Future for FutureWithGuard<T, F, R>
44where
45    F: Future<Output = R> + Send + 'static,
46    T: Send + 'static,
47{
48    type Output = R;
49
50    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
51        self.project().future.poll(cx)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[fuchsia::test]
60    fn test_custom_future_is_smaller() {
61        async fn large_future() -> [u8; 128] {
62            let held_across_await = [0; 128];
63            std::future::pending::<()>().await;
64            held_across_await
65        }
66
67        let object = 10u64;
68        let task = large_future();
69        let custom_future = FutureWithGuard::new(object, task);
70        let custom_future_size = std::mem::size_of_val(&custom_future);
71
72        let object = 10u64;
73        let task = large_future();
74        let rust_future = async move {
75            let _object = object;
76            task.await;
77        };
78        let rust_future_size = std::mem::size_of_val(&rust_future);
79
80        // The large_future is 129 bytes:
81        //   - 128 bytes for the array
82        //   - 1 byte for the discriminant.
83        //
84        // The custom_future is 144 bytes:
85        //   - 129 bytes for the large_future
86        //   - 8 bytes for the u64 object
87        //   - 7 bytes of padding
88        //
89        // The rust_future is 272 bytes:
90        //   - 129 bytes to capture the large_future
91        //   - 129 bytes to use the large_future
92        //   - 8 bytes to capture the u64 object
93        //   - 1 byte for the discriminant
94        //   - 5 bytes of padding
95        //
96        // This assert only makes sure that the custom future is not bigger than the rust generated
97        // future.
98        //
99        // If this test starts failing while updating rust, the test can safely be disabled.
100        assert!(
101            custom_future_size <= rust_future_size,
102            "custom_future_size={custom_future_size} rust_future_size={rust_future_size}"
103        );
104    }
105}