Code coverage misses `.await` lines

[Will-Low]

I tried running the LLVM code coverage report against this code using cargo llvm-cov:

async fn my_async_fn() -> bool {
    another_async_fn()
	.await
}

async fn another_async_fn() -> bool {
    true
}

#[cfg(test)]
mod tests {
    use crate::my_async_fn;

    #[tokio::test]
    async fn function_returns_true() {
        assert!(my_async_fn().await);
    }
}

I expected this to result in 100% code coverage.

Instead, the .await on line 3 is considered not covered:

Meta

rustc 1.61.0 (fe5b13d68 2022-05-18)
binary: rustc
commit-hash: fe5b13d681f25ee6474be29d748c65adcd91f69e
commit-date: 2022-05-18
host: aarch64-apple-darwin
release: 1.61.0
LLVM version: 14.0.0

[bossmc]

Playing around, it seems that the await is only treated as covered if the future yields underneath it, the following forces a yield:

#[allow(dead_code)]
#[derive(Debug, Default)]
struct Yield {
    init: bool,
}

impl std::future::Future for Yield {
    type Output = ();

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        if !self.init {
            self.init = true;
            cx.waker().wake_by_ref();
            return std::task::Poll::Pending;
        } else {
            return std::task::Poll::Ready(());
        }
    }
}

pub async fn do_things() {
    println!("Hello world");
    Yield::default().await;        // Causes a yield then completes on next `poll`
}

#[tokio::main]
pub async fn main() {
    do_things()
        .await;                    // This is the line/span we care about coverage for
}

With the Yield::default().await line in place then the await line in main is marked as hit, if it's commented out then the await is unhit.

My guess is that only the "bail out to the calling frame" logic has a BCB counter in it (or maybe there simply is no code in the other branch to count) so unless the future yields the counter is not incremented, and the span is deemed unhit.

[cameronelliott]

Does anyone have a suggestion on how to re-write the example so it could be used as a helper-type or function to actually get code coverage on await statements.

I am trying to do something like this:

async_func().await; would become Yield::new(async_func()).await

But, my rust futures abilities aren't so strong.

I'll keep at it, and post something if I figure it out.

[cameronelliott]

I did write a workaround that works for my needs.

https://github.com/cameronelliott/await-coverage-workaround.git

I created a helper type, so now I do read(...).fix_cov().await;

This will fix the coverage where the await is invoked.

The helper type doesn't get 100% coverage, and it may be possible to
write a helper that gets 100% coverage, but it is not required for my needs at the moment,

[cameronelliott]

Since I sometimes clean my repos, I will add the workaround right here!

use std::future::Future;

#[allow(dead_code)]
#[derive(Debug, Default)]
struct Yield {
    init: bool,
}

impl std::future::Future for Yield {
    type Output = ();

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        if !self.init {
            self.init = true;
            cx.waker().wake_by_ref();
            return std::task::Poll::Pending;
        } else {
            return std::task::Poll::Ready(());
        }
    }
}

trait FixCoverage {
    async fn fix_cov(self) -> <Self as Future>::Output
    where
        Self: Sized,
        Self: Future,
    {
        // this will NOT show as covered
        // but for my usage I just keep it outside of my coverage checked code
        let r = self.await;
        Yield::default().await;
        r
    }
}

impl<F, T> FixCoverage for F where F: Future<Output = T> {}

pub async fn do_things() {
    println!("Hello world");
}

#[tokio::main]
pub async fn main() {
    do_things().await; // will NOT show as covered

    do_things().fix_cov().await; // will show as covered
}

[bossmc]

Note that Yeild is available in tokio as tokio::task::yeild_now() (https://docs.rs/tokio/latest/tokio/task/fn.yield_now.html). An implementation slightly closer to what you originally pitched (where you wrap a future in an adaptor) might look like:

use std::pin::Pin;
use std::future::Future;
use std::task::{Context, Poll};
use futures::future::FutureExt;

pin_project_lite::pin_project! {
    struct AlwaysYeildAtLeastOnce<F> {
      yeilded: bool,
      #[pin]
      inner: F,
    }
}

impl<F> AlwaysYeildAtLeastOnce<F> {
    fn new(f: F) -> Self {
        Self {
            yeilded: false,
            inner: f,
        }
    }
}

impl<F> Future for AlwaysYeildAtLeastOnce<F> where F: Future {
    type Output = F::Output;
    
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let zelf = self.project();
        if *zelf.yeilded {
            zelf.inner.poll(cx)
        } else {
            *zelf.yeilded = true;
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

async fn dont_yeild() {
    println!("Not yeilding");
}

#[tokio::main]
async fn main() {
    AlwaysYeildAtLeastOnce::new(dont_yeild()).await;
    assert!(AlwaysYeildAtLeastOnce::new(dont_yeild()).now_or_never().is_none())
}