Change Future into Async Function

Change Future into Async Function

(Jin Qing’s Column, Dec., 2024)

Source code: connection.rs

Change struct StreamUpgrade into async function to simplify the code.

StreamUpgrade is used as a Future, which can be replaced by an async function.

impl<UserData, TOk, TErr> Future for StreamUpgrade<UserData, TOk, TErr> {
    type Output = (UserData, Result<TOk, StreamUpgradeError<TErr>>);
    ...
}

Change the constructor StreamUpgrade::new_inbound() into async function

Original constructor and poll()

impl<UserData, TOk, TErr> StreamUpgrade<UserData, TOk, TErr> {
    fn new_inbound<Upgrade>(
        substream: SubstreamBox,
        protocol: SubstreamProtocol<Upgrade, UserData>,
        counter: ActiveStreamCounter,
    ) -> Self
    where
        Upgrade: InboundUpgradeSend<Output = TOk, Error = TErr>,
    {
        let timeout = *protocol.timeout();
        let (upgrade, open_info) = protocol.into_upgrade();
        let protocols = upgrade.protocol_info();

        Self {
            user_data: Some(open_info),
            timeout: Delay::new(timeout),
            upgrade: Box::pin(async move {
                let (info, stream) =
                    multistream_select::listener_select_proto(substream, protocols)
                        .await
                        .map_err(to_stream_upgrade_error)?;

                let output = upgrade
                    .upgrade_inbound(Stream::new(stream, counter), info)
                    .await
                    .map_err(StreamUpgradeError::Apply)?;

                Ok(output)
            }),
        }
    }
}

impl<UserData, TOk, TErr> Future for StreamUpgrade<UserData, TOk, TErr> {
    type Output = (UserData, Result<TOk, StreamUpgradeError<TErr>>);

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match self.timeout.poll_unpin(cx) {
            Poll::Ready(()) => {
                return Poll::Ready((
                    self.user_data
                        .take()
                        .expect("Future not to be polled again once ready."),
                    Err(StreamUpgradeError::Timeout),
                ))
            }

            Poll::Pending => {}
        }

        let result = futures::ready!(self.upgrade.poll_unpin(cx));
        let user_data = self
            .user_data
            .take()
            .expect("Future not to be polled again once ready.");

        Poll::Ready((user_data, result))
    }
}

New async function:

async fn new_inbound_upgrade_stream<UserData, TOk, TErr, Upgrade>(
    substream: SubstreamBox,
    protocol: SubstreamProtocol<Upgrade, UserData>,
    counter: ActiveStreamCounter,
) -> (UserData, Result<TOk, StreamUpgradeError<TErr>>)
where
    Upgrade: InboundUpgradeSend<Output = TOk, Error = TErr>,
{
    let timeout = *protocol.timeout();
    let (upgrade, open_info) = protocol.into_upgrade();
    let fut = new_inbound_upgrade_stream_inner(substream, upgrade, counter);
    let res = tokio::time::timeout(timeout, fut).await;
    if let Ok(res) = res {
        return (open_info, res);
    }
    (open_info, Err(StreamUpgradeError::Timeout))
}

async fn new_inbound_upgrade_stream_inner<TOk, TErr, Upgrade>(
    substream: SubstreamBox,
    upgrade: Upgrade,
    counter: ActiveStreamCounter,
) -> Result<TOk, StreamUpgradeError<TErr>>
where
    Upgrade: InboundUpgradeSend<Output = TOk, Error = TErr>,
{
    let protocols = upgrade.protocol_info();
    let (info, stream) = multistream_select::listener_select_proto(substream, protocols)
        .await
        .map_err(to_stream_upgrade_error)?;

    let output = upgrade
        .upgrade_inbound(Stream::new(stream, counter), info)
        .await
        .map_err(StreamUpgradeError::Apply)?;

    Ok(output)
}

Change the usage of StreamUpgrade

Change

/// Futures that upgrade incoming substreams.
    negotiating_in: FuturesUnordered<
        StreamUpgrade<
            THandler::InboundOpenInfo,
            <THandler::InboundProtocol as InboundUpgradeSend>::Output,
            <THandler::InboundProtocol as InboundUpgradeSend>::Error,
        >,
    >,

into

negotiating_in: FuturesUnordered<Pin<Box<dyn
        Future<Output = (
            THandler::InboundOpenInfo,
            Result<
                <THandler::InboundProtocol as InboundUpgradeSend>::Output,
                StreamUpgradeError<<THandler::InboundProtocol as InboundUpgradeSend>::Error>,
            >,
        )> + Send,
    >>>,

And change

negotiating_in.push(StreamUpgrade::new_inbound(
        substream,
        protocol,
        stream_counter.clone(),
    ));

into

negotiating_in.push(new_inbound_upgrade_stream(
        substream,
        protocol,
        stream_counter.clone(),
    ).boxed());