Skip to main content

wowlab_tidy/languages/rust/rules/performance/
async_loop_no_yield.rs

1use ra_ap_syntax::{AstNode, ast};
2
3use super::support;
4use crate::{AstCtx, Example, Violation};
5
6#[rustfmt::skip]
7const EXAMPLES: &[Example] = &[
8    Example {
9        label: "async for loop with call but no await",
10        code: "async fn f(xs: &[u32]) { for x in xs { process(*x); } }",
11        pass: false,
12    },
13    Example {
14        label: "loop in async block without await",
15        code: "fn g() { let _fut = async { loop { step(); } }; }",
16        pass: false,
17    },
18    Example {
19        label: "three statements without calls or await",
20        code: "async fn f(xs: &[u32]) { for x in xs { let a = *x; let b = a + 1; let _c = b - a; } }",
21        pass: false,
22    },
23    Example {
24        label: "loop yields via yield_now().await",
25        code: "async fn f(xs: &[u32]) { for x in xs { process(*x); tokio::task::yield_now().await; } }",
26        pass: true,
27    },
28    Example {
29        label: "sync fn loop is fine",
30        code: "fn f(xs: &[u32]) { for x in xs { process(*x); } }",
31        pass: true,
32    },
33    Example {
34        label: "tiny accumulation loop without calls",
35        code: "async fn f(xs: &[u32]) -> u32 { let mut total = 0; for x in xs { total += x; } total }",
36        pass: true,
37    },
38    Example {
39        label: "loop inside sync closure in async fn",
40        code: "async fn f(xs: &[u32]) { let sum = || { for x in xs { process(*x); } }; sum(); }",
41        pass: true,
42    },
43    Example {
44        label: "await nested deeper in loop body",
45        code: "async fn f(n: u32) { for _ in 0..n { if n > 1 { step().await; } } }",
46        pass: true,
47    },
48    Example {
49        label: "async loop in test module",
50        code: "#[cfg(test)]\nmod tests {\n    async fn t(xs: &[u32]) { for x in xs { process(*x); } }\n}",
51        pass: true,
52    },
53];
54
55crate::ast_rule!(
56    async_loop_no_yield,
57    "Flag loops in async contexts whose bodies never `.await` (CPU-bound work without yield points).",
58    "CPU-bound async loops must cooperatively yield (yield_now().await) so they do not starve the runtime.",
59);
60
61fn check_async_loop_no_yield(ctx: &AstCtx<'_>) -> Vec<Violation> {
62    let mut violations = Vec::new();
63
64    for node in ctx.nodes::<ast::ForExpr>() {
65        check_loop(ctx, &node, &mut violations);
66    }
67
68    for node in ctx.nodes::<ast::WhileExpr>() {
69        check_loop(ctx, &node, &mut violations);
70    }
71
72    for node in ctx.nodes::<ast::LoopExpr>() {
73        check_loop(ctx, &node, &mut violations);
74    }
75
76    violations
77}
78
79const MIN_LOOP_BODY_STMTS: usize = 3;
80
81const MSG: &str = "loop in async context has no .await — CPU-bound iterations starve the runtime; add yield_now().await between chunks";
82
83fn check_loop<N>(ctx: &AstCtx<'_>, node: &N, violations: &mut Vec<Violation>)
84where
85    N: AstNode,
86{
87    if ctx.is_in_test(node) || !support::is_in_async_context(node) {
88        return;
89    }
90
91    let Some(body) = support::loop_body(node) else {
92        return;
93    };
94    let has_await = body
95        .syntax()
96        .descendants()
97        .any(|syntax| ast::AwaitExpr::can_cast(syntax.kind()));
98
99    if has_await {
100        return;
101    }
102
103    let has_call = body.syntax().descendants().any(|syntax| {
104        ast::CallExpr::can_cast(syntax.kind()) || ast::MethodCallExpr::can_cast(syntax.kind())
105    });
106    let statements = body
107        .stmt_list()
108        .map_or(0, |statements| statements.statements().count());
109
110    if has_call || statements >= MIN_LOOP_BODY_STMTS {
111        violations.push(ctx.violation(node, MSG));
112    }
113}
114
115crate::tidy_ast_test!(check_async_loop_no_yield, {
116    crate::example_tests!(EXAMPLES, check_async_loop_no_yield);
117});