Search code examples
rustasync-awaittraitsrust-futures

Make returned Future Send if parameters are Send


Can I propagate the Send trait of function parameters to its return type, so that the return type is impl Send if and only if the parameters are?

Details:

An async function has a nice feature. Its returned Future is automatically Send if it can be. In the following example, the async function will create a Future that is Send, if the inputs to the function are Send.

struct MyStruct;

impl MyStruct {
    // This async fn returns an `impl Future<Output=T> + Send` if `T` is Send.
    // Otherwise, it returns an `impl Future<Output=T>` without `Send`.
    async fn func<T>(&self, t: T) -> T {
        t
    }
}

fn assert_is_send(_v: impl Send) {}

fn main() {
    // This works
    assert_is_send(MyStruct.func(4u64));
    // And the following correctly fails
    assert_is_send(MyStruct.func(std::rc::Rc::new(4u64)));
}

playground

Now, I want to move such a function into a trait, which requires using async-trait (which is some codegen that effectively writes my async fn as a function returning Pin<Box<dyn Future>>) or doing something similar manually. Is there a way to write this in a way to retain this auto-Send behavior where the returned Future is made Send if T is Send? The following example implements it as two separate functions:

use std::pin::Pin;
use std::future::Future;

struct MyStruct;
impl MyStruct {
    fn func_send<T: 'static + Send>(&self, t: T) -> Pin<Box<dyn Future<Output = T> + Send>> {
        Box::pin(async{t})
    }
    
    fn func_not_send<T: 'static>(&self, t: T) -> Pin<Box<dyn Future<Output = T>>> {
        Box::pin(async{t})
    }
}

fn assert_is_send(_v: impl Send) {}

fn main() {
    // This works
    assert_is_send(MyStruct.func_send(4u64));
    // And the following correctly fails
    // assert_is_send(MyStruct.func(std::rc::Rc::new(4u64)));
}

playground

But actually, I don't want them to be separate. I want them to be one function similar to how async fn does it automatically. Something along the lines of

use std::pin::Pin;
use std::future::Future;

struct MyStruct;
impl MyStruct {
    fn func<T: 'static + ?Send>(&self, t: T) -> Pin<Box<dyn Future<Output = T> + ?Send>> {
        Box::pin(async{t})
    }
}

fn assert_is_send(_v: impl Send) {}

fn main() {
    // This should
    assert_is_send(MyStruct.func(4u64));
    // And this should fail
    assert_is_send(MyStruct.func(std::rc::Rc::new(4u64)));
}

Is something like this possible in Rust? I'm ok with writing the async-trait magic manually and modifying it instead of using the async-trait crate if that is a way to make it work.

Some ideas I had but they haven't really borne fruit yet:

  • Use min-specialization to specialize on Send? But doesn't seem like that feature is going to be stabilized anytime soon so maybe not the best option.
  • Return a custom MyFuture type instead of just impl Future and somehow impl Send for MyFuture where T: Send? Would probably be difficult though since I would have to be able to name that Future and async code usually produces impl Future types that cannot be named.
  • Writing a procedural macro that adds + Send to the return type if it recognizes that the input type is Send. Actually, can procedural macros detect if a certain type implements Send? My guess would be it's not possible since they just work on token streams.

Solution

  • (2) is the only way that could work.

    There are two ways to make it work:

    1. Write the future manually, without the help of async and .await. But that means writing the future manually:
    enum ConditionalSendFut<T> {
        Start { t: T },
        Done,
    }
    
    impl<T> Unpin for ConditionalSendFut<T> {}
    
    impl<T> Future for ConditionalSendFut<T> {
        type Output = T;
    
        fn poll(mut self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<Self::Output> {
            match &mut *self {
                Self::Start { .. } => {
                    let t = match std::mem::replace(&mut *self, Self::Done) {
                        Self::Start { t } => t,
                        _ => unreachable!(),
                    };
                    Poll::Ready(t)
                }
                Self::Done => Poll::Pending,
            }
        }
    }
    
    struct MyStruct;
    impl MyStruct {
        fn func<T: 'static>(&self, t: T) -> ConditionalSendFut<T> {
            ConditionalSendFut::Start { t }
        }
    }
    

    Playground.

    1. Store a Pin<Box<dyn Future<Output = T>>> and conditionally impl Send on the future. But this requires unsafe code and manually ensuring that you don't hold other non-Send types across .await points:
    struct ConditionalSendFut<T>(Pin<Box<dyn Future<Output = T>>>);
    
    // SAFETY: The only non-`Send` type we're holding across an `.await`
    // point is `T`.
    unsafe impl<T: Send> Send for ConditionalSendFut<T> {}
    
    impl<T> Future for ConditionalSendFut<T> {
        type Output = T;
    
        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
            self.0.as_mut().poll(context)
        }
    }
    
    struct MyStruct;
    impl MyStruct {
        fn func<T: 'static>(&self, t: T) -> ConditionalSendFut<T> {
            ConditionalSendFut(Box::pin(async { t }))
        }
    }
    

    Playground.

    (1) cannot work with traits, as each impl will have a different future. This leaves us with (2) only. I would not recommend it, but it is possible.

    It is very likely that when async fns in traits will be stable there will be a mechanism to that (what is talked about currently is to impl them conditionally and use bounds on use sites to require them) but currently there is no such thing, even on the nightly implementation of async fns in traits.