Custom SelectAndCreate Future where the condition is async

What I'm trying to achieve:

I'm trying to select futures based on a condition, where the condition is a future itself:

let condition = async move { true };
let left = || async move { 42 };
let right = || async move { 1337 };

let fut = SelectAndCreate::new(condition, left, right);

let x = fut.await;
assert_eq!(x, 42);

For more context, here is what I want to use this for.

What I have tried:

Use BoxFuture to make things easier

I know that I can use futures::future::BoxFuture to just box a dyn Future, but this would cause allocations on the heap, which I'm trying to avoid for performance reasons. (Also I think it's a good challenge to try to optimize it as much as possible). So I want to avoid doing this if possible.

Use a custom future

Here is my custom implementation: Link

My problem with this implementation is that I'm using unsafe code and I'm not 100% sure if I use it safely:

#[pin_project::pin_project]
struct SelectAndCreate<Cond, Left, Right, LFut, RFut, T>
where
    Cond: Future<Output = bool>,

    Left: FnOnce() -> LFut,
    Right: FnOnce() -> RFut,

    LFut: Future<Output = T>,
    RFut: Future<Output = T>,
{
    #[pin]
    condition: Cond,

    #[pin]
    future: Option<Either<LFut, RFut>>,

    // INV: This is Some(...) when future is None
    left: Option<Left>,

    // INV: This is Some(...) when future is None
    right: Option<Right>,

    _marker: PhantomData<(LFut, RFut, T)>,
}

Troubleling unsafe block:

unsafe {
	let future_ref = this.future.as_mut();
	// Is this safe?
	let future = Pin::get_unchecked_mut(future_ref);
	*future = Some(fut);
}

Is it safe here to use Pin::get_unchecked_mut? If it is, what is the reasoning and if it isn't why not?

Is there another approach that I can try?

I think you have to create a type that implements Future from SelectAndCreate, not implement Future directly on SelectAndCreate. At least I have no idea how to implement Future::poll where we have to poll the future returned by an FnOnce() -> Future. Calling the function to get the future to poll it will move it in the process, which we can't do when we have a Pin<&mut Self>. The easiest way to construct such a future would be through an async method, something like this, maybe?

use std::future::Future;

struct SelectAndCreate<C, L, R> {
    condition: C,
    left: L,
    right: R,
}

impl<C, L, LFut, R, RFut, T> SelectAndCreate<C, L, R>
where
    C: Future<Output = bool>,
    L: FnOnce() -> LFut,
    LFut: Future<Output = T>,
    R: FnOnce() -> RFut,
    RFut: Future<Output = T>,
{
    fn new(condition: C, left: L, right: R) -> Self {
        Self {
            condition,
            left,
            right,
        }
    }

    async fn wait(self) -> T {
        if self.condition.await {
            (self.left)().await
        } else {
            (self.right)().await
        }
    }
}

#[tokio::main]
async fn main() {
    let condition = async move { true };
    let left = || async move { 42 };
    let right = || async move { 1337 };

    let fut = SelectAndCreate::new(condition, left, right);

    let x = fut.wait().await;
    assert_eq!(x, 42);
}

Playground.

I think it is safe because your Option is None and only the Some variant needs to be pinned. Moreover once you create the Some variant it actually remains pinned.

A different approach would be to not take a FnOnce that creates the Future, but instead directly the Future. The FnOnce is kinda useless because async blocks are already lazy, so you don't actually start executing the branches when you create them (which would be the normal reason for a sync version of this struct). If you do this the implementation can be made fully safe. Rust Playground (I used pin_project_lite just because pin_project is not available on the playground, but they are equivalent)

Use Pin::set()

2 Likes

Thank you!
I don't know how I missed that but it does exactly what I want :slight_smile:

Thanks, but the whole reason for writing this class was to return a custom future inside a trait implementation of another library, where the return type has to be a future.
But this is actually a good idea in general and I'll keep it mind!

I managed to solve my issue using Pin::set like @Hyeonu suggested.

Thanks, sadly I can't get rid of the FnOnce, as in the use-case both of the futures left and right depend on a variable where they take the ownership of it, so I can't construct both futures at the same time.

But I managed to solve my issue using Pin::set which does what I want and isn't unsafe.

Don't you have the same problem with two closures? You cannot move the variable in both of them. And whatever solution you have with the closures should work with two async blocks.

The thing with closures is, my custom future struct can take the value (which the closures will take ownership of), wait for the condition and the execute only the right closure, meaning the value is only moved once.
Like this is what I actually need it for. I have two tower::Service, which have a function "call": fn call(&mut self, req: Request) -> Self::Future. As they both need to have the ownership of req: Request I can't custruct the futures at the same time.

Does your actual closures take that value as a parameter? Because I can see that making a difference with a future.

If instead the closure is parameterless you'll have to move the request into both of the closures, which you also cannot do.