| 1 | #!/usr/bin/env python3 |
| 2 | """Ratchet for blocking calls that could land on Tokio workers (#6149). |
| 3 | |
| 4 | Codewhale's convention: async code must not run blocking operations inline. |
| 5 | `std::fs`/`thread::sleep` (and friends) are fine inside `spawn_blocking`, |
| 6 | on dedicated `std::thread`s, and in synchronous entry points — but every |
| 7 | unprotected call site is one careless caller away from parking a runtime |
| 8 | worker. This check counts the sites that are NOT already inside a blocking |
| 9 | scope (`spawn_blocking`, `spawn_blocking_supervised`, `std::thread::spawn`, |
| 10 | `thread::Builder`) or test code, and fails if any file exceeds its recorded |
| 11 | budget in `check-blocking-calls-budget.json`. |
| 12 | |
| 13 | Fix the call site — wrap the work in `spawn_blocking` (the established |
| 14 | pattern, ~80 sites) or switch to `tokio::fs`/`tokio::time` — or, if the site |
| 15 | is genuinely only reachable from synchronous code, acknowledge the debt by |
| 16 | raising the file's budget. |
| 17 | |
| 18 | Run `python3 scripts/check-blocking-calls-budget.py --update` to regenerate |
| 19 | the budget after removing sites or after an intentional addition. |
| 20 | """ |
| 21 | |
| 22 | from __future__ import annotations |
| 23 | |
| 24 | import json |
| 25 | import re |
| 26 | import sys |
| 27 | from pathlib import Path |
| 28 | |
| 29 | ROOT = Path(__file__).resolve().parents[1] |
| 30 | CRATES = ROOT / "crates" |
| 31 | BUDGET_PATH = Path(__file__).with_suffix(".json") |
| 32 | |
| 33 | PATTERNS = { |
| 34 | "thread_sleep": re.compile(r"\bthread::sleep\s*\("), |
| 35 | "std_fs": re.compile( |
| 36 | r"\bstd::fs::(?:read|read_to_string|write|create_dir|create_dir_all|" |
| 37 | r"remove_file|remove_dir|remove_dir_all|copy|rename|metadata|" |
| 38 | r"symlink_metadata|read_dir|canonicalize|exists|set_permissions|" |
| 39 | r"hard_link|soft_link|symlink|File|OpenOptions|DirBuilder)\b" |
| 40 | ), |
| 41 | } |
| 42 | |
| 43 | ATTR_RE = re.compile(r"#\s*\[([^\]]*)\]") |
| 44 | FN_RE = re.compile( |
| 45 | r"\b(?:pub(?:\([^)]*\))?\s+)?(?:unsafe\s+)?(?:extern\s+\"[^\"]*\"\s+)?" |
| 46 | r"(async\s+)?fn\s+([A-Za-z_][\w]*)" |
| 47 | ) |
| 48 | MOD_RE = re.compile(r"\bmod\s+([A-Za-z_][\w]*)") |
| 49 | |
| 50 | TOKEN_RE = re.compile( |
| 51 | r"#\s*\[[^\]]*\]" |
| 52 | r"|\bmod\s+\w+" |
| 53 | r"|\bimpl\b" |
| 54 | r"|\basync\s+move\s*\{" |
| 55 | r"|\basync\s*\{" |
| 56 | r"|\b(?:pub(?:\([^)]*\))?\s+)?(?:unsafe\s+)?(?:async\s+)?fn\s+\w+" |
| 57 | r"|spawn_blocking(?:_supervised)?" |
| 58 | r"|thread::spawn" |
| 59 | r"|thread::Builder::new" |
| 60 | r"|[{}]" |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def strip_comments_and_strings(text: str) -> str: |
| 65 | """Blank out comments and string/char literal contents, keeping newlines.""" |
| 66 | out = list(text) |
| 67 | i, n = 0, len(text) |
| 68 | line_comment = block_comment = in_str = in_char = in_raw = False |
| 69 | block_depth = 0 |
| 70 | raw_hashes = 0 |
| 71 | while i < n: |
| 72 | c = text[i] |
| 73 | if line_comment: |
| 74 | if c == "\n": |
| 75 | line_comment = False |
| 76 | else: |
| 77 | out[i] = " " |
| 78 | i += 1 |
| 79 | continue |
| 80 | if block_comment: |
| 81 | if text[i : i + 2] == "/*": |
| 82 | block_depth += 1 |
| 83 | out[i] = out[i + 1] = " " |
| 84 | i += 2 |
| 85 | continue |
| 86 | if text[i : i + 2] == "*/": |
| 87 | block_depth -= 1 |
| 88 | out[i] = out[i + 1] = " " |
| 89 | i += 2 |
| 90 | if block_depth == 0: |
| 91 | block_comment = False |
| 92 | continue |
| 93 | if c != "\n": |
| 94 | out[i] = " " |
| 95 | i += 1 |
| 96 | continue |
| 97 | if in_str: |
| 98 | if c == "\\": |
| 99 | out[i] = out[i + 1] = " " |
| 100 | i += 2 |
| 101 | continue |
| 102 | if c == '"': |
| 103 | in_str = False |
| 104 | elif c != "\n": |
| 105 | out[i] = " " |
| 106 | i += 1 |
| 107 | continue |
| 108 | if in_char: |
| 109 | if c == "\\": |
| 110 | out[i] = out[i + 1] = " " |
| 111 | i += 2 |
| 112 | continue |
| 113 | if c == "'": |
| 114 | in_char = False |
| 115 | elif c != "\n": |
| 116 | out[i] = " " |
| 117 | i += 1 |
| 118 | continue |
| 119 | if in_raw: |
| 120 | if c == '"' and text[i + 1 : i + 1 + raw_hashes] == "#" * raw_hashes: |
| 121 | for j in range(1 + raw_hashes): |
| 122 | out[i + j] = " " |
| 123 | i += 1 + raw_hashes |
| 124 | in_raw = False |
| 125 | continue |
| 126 | if c != "\n": |
| 127 | out[i] = " " |
| 128 | i += 1 |
| 129 | continue |
| 130 | if text[i : i + 2] == "//": |
| 131 | line_comment = True |
| 132 | out[i] = out[i + 1] = " " |
| 133 | i += 2 |
| 134 | continue |
| 135 | if text[i : i + 2] == "/*": |
| 136 | block_comment = True |
| 137 | block_depth = 1 |
| 138 | out[i] = out[i + 1] = " " |
| 139 | i += 2 |
| 140 | continue |
| 141 | if c == "r": |
| 142 | m = re.match(r'r(#+)"', text[i:]) |
| 143 | if m: |
| 144 | raw_hashes = len(m.group(1)) |
| 145 | in_raw = True |
| 146 | for j in range(2 + raw_hashes): |
| 147 | out[i + j] = " " |
| 148 | i += 2 + raw_hashes |
| 149 | continue |
| 150 | if c == '"': |
| 151 | in_str = True |
| 152 | out[i] = " " |
| 153 | i += 1 |
| 154 | continue |
| 155 | if c == "'" and re.match(r"'(?:\\.|[^'\\])'", text[i:]): |
| 156 | in_char = True |
| 157 | out[i] = " " |
| 158 | i += 1 |
| 159 | continue |
| 160 | i += 1 |
| 161 | return "".join(out) |
| 162 | |
| 163 | |
| 164 | def file_counts(path: Path) -> dict[str, int]: |
| 165 | """Count unprotected blocking-call sites in one Rust source file.""" |
| 166 | text = path.read_text(encoding="utf-8", errors="replace") |
| 167 | code = strip_comments_and_strings(text) |
| 168 | counts = {name: 0 for name in PATTERNS} |
| 169 | # Scope stack: entries are dicts {kind, open_depth} where kind is |
| 170 | # 'test', 'blocking', 'fn', 'mod', or 'impl'. A hit counts only when the |
| 171 | # innermost enclosing scope is neither test code nor a blocking pool / |
| 172 | # dedicated-thread closure. |
| 173 | stack: list[dict] = [] |
| 174 | pending_attr_test = False |
| 175 | pending_blocking = False |
| 176 | depth = 0 |
| 177 | for line in code.split("\n"): |
| 178 | # Interleave pattern hits and scope tokens in column order so a |
| 179 | # one-liner like `fn f() { thread::sleep(..) }` sees the fn scope. |
| 180 | events: list[tuple[int, str, object]] = [] |
| 181 | for name, pat in PATTERNS.items(): |
| 182 | for m in pat.finditer(line): |
| 183 | events.append((m.start(), "hit", name)) |
| 184 | for m in TOKEN_RE.finditer(line): |
| 185 | events.append((m.start(), "tok", m.group(0))) |
| 186 | events.sort(key=lambda e: e[0]) |
| 187 | for _col, kind, payload in events: |
| 188 | if kind == "hit": |
| 189 | if not any(s["kind"] in ("test", "blocking") for s in stack): |
| 190 | counts[payload] += 1 # type: ignore[index] |
| 191 | continue |
| 192 | tok = payload # type: ignore[assignment] |
| 193 | if tok.startswith("#"): |
| 194 | inner = tok[tok.index("[") + 1 : -1] |
| 195 | if "test" in inner: |
| 196 | pending_attr_test = True |
| 197 | continue |
| 198 | if tok == "{": |
| 199 | depth += 1 |
| 200 | if pending_blocking: |
| 201 | stack.append({"kind": "blocking", "open": depth}) |
| 202 | elif stack and stack[-1]["open"] is None: |
| 203 | stack[-1]["open"] = depth |
| 204 | pending_blocking = False |
| 205 | continue |
| 206 | if tok == "}": |
| 207 | while stack and stack[-1]["open"] == depth: |
| 208 | stack.pop() |
| 209 | depth -= 1 |
| 210 | continue |
| 211 | if "spawn_blocking" in tok or tok in ("thread::spawn", "thread::Builder::new"): |
| 212 | pending_blocking = True |
| 213 | continue |
| 214 | if tok.startswith("async") and tok.endswith("{"): |
| 215 | stack.append({"kind": "fn", "open": depth + 1}) |
| 216 | depth += 1 |
| 217 | pending_attr_test = False |
| 218 | pending_blocking = False |
| 219 | continue |
| 220 | fm = FN_RE.match(tok) |
| 221 | if fm: |
| 222 | kind = "test" if pending_attr_test else "fn" |
| 223 | stack.append({"kind": kind, "open": None}) |
| 224 | pending_attr_test = False |
| 225 | pending_blocking = False |
| 226 | continue |
| 227 | mm = MOD_RE.match(tok) |
| 228 | if mm: |
| 229 | name = mm.group(1) |
| 230 | kind = "test" if (pending_attr_test or name.startswith("test")) else "mod" |
| 231 | stack.append({"kind": kind, "open": None}) |
| 232 | pending_attr_test = False |
| 233 | pending_blocking = False |
| 234 | continue |
| 235 | if tok == "impl": |
| 236 | stack.append({"kind": "impl", "open": None}) |
| 237 | pending_attr_test = False |
| 238 | pending_blocking = False |
| 239 | continue |
| 240 | return {k: v for k, v in counts.items() if v} |
| 241 | |
| 242 | |
| 243 | CFG_TEST_MOD = re.compile( |
| 244 | r"#\[cfg\(test\)\]\s*(?:pub(?:\([^)]*\))?\s+)?mod\s+([A-Za-z_][A-Za-z0-9_]*)\s*;" |
| 245 | ) |
| 246 | |
| 247 | |
| 248 | def cfg_test_module_files() -> set[Path]: |
| 249 | """Files that are entire `#[cfg(test)]` modules declared by another file. |
| 250 | |
| 251 | The per-file scanner recognises test scope it can see *inside* a file — |
| 252 | `#[test]`, `mod tests`, `fn test_*`. It cannot see that a whole file is |
| 253 | test-only, because that fact lives in the parent's `#[cfg(test)] mod |
| 254 | foo;` declaration. Extracting a test suite into its own file therefore |
| 255 | made an untouched call site look new (#6209 follow-up: PR #6096's |
| 256 | `session_export_*_tests.rs`). Excluding these keeps the ratchet's stated |
| 257 | contract — "not ... test code" — instead of taxing the extraction. |
| 258 | """ |
| 259 | excluded: set[Path] = set() |
| 260 | for path in CRATES.rglob("*.rs"): |
| 261 | try: |
| 262 | text = path.read_text(encoding="utf-8", errors="ignore") |
| 263 | except OSError: |
| 264 | continue |
| 265 | for name in CFG_TEST_MOD.findall(text): |
| 266 | for candidate in (path.parent / f"{name}.rs", path.parent / name / "mod.rs"): |
| 267 | if candidate.is_file(): |
| 268 | excluded.add(candidate.resolve()) |
| 269 | return excluded |
| 270 | |
| 271 | |
| 272 | def collect_current() -> dict[str, dict[str, int]]: |
| 273 | budget: dict[str, dict[str, int]] = {} |
| 274 | test_only = cfg_test_module_files() |
| 275 | for path in sorted(CRATES.rglob("*.rs")): |
| 276 | if path.resolve() in test_only: |
| 277 | continue |
| 278 | try: |
| 279 | counts = file_counts(path) |
| 280 | except OSError: |
| 281 | continue |
| 282 | if counts: |
| 283 | rel = str(path.relative_to(ROOT)) |
| 284 | budget[rel] = counts |
| 285 | return budget |
| 286 | |
| 287 | |
| 288 | def main() -> int: |
| 289 | update = "--update" in sys.argv |
| 290 | current = collect_current() |
| 291 | if update: |
| 292 | BUDGET_PATH.write_text( |
| 293 | json.dumps(current, indent=2, sort_keys=True) + "\n", |
| 294 | encoding="utf-8", |
| 295 | ) |
| 296 | total = sum(sum(v.values()) for v in current.values()) |
| 297 | print(f"wrote {BUDGET_PATH.name}: {total} sites across {len(current)} files") |
| 298 | return 0 |
| 299 | |
| 300 | if not BUDGET_PATH.exists(): |
| 301 | print(f"missing {BUDGET_PATH.name}; run with --update to create it", file=sys.stderr) |
| 302 | return 2 |
| 303 | budget = json.loads(BUDGET_PATH.read_text(encoding="utf-8")) |
| 304 | |
| 305 | failures: list[str] = [] |
| 306 | savings: list[str] = [] |
| 307 | for path, counts in sorted(current.items()): |
| 308 | allowed = budget.get(path, {}) |
| 309 | for name, count in counts.items(): |
| 310 | limit = allowed.get(name, 0) |
| 311 | if count > limit: |
| 312 | failures.append( |
| 313 | f"{path}: {name} sites {count} > budget {limit}" |
| 314 | ) |
| 315 | elif count < limit: |
| 316 | savings.append( |
| 317 | f"{path}: {name} sites {count} < budget {limit} — tighten with --update" |
| 318 | ) |
| 319 | for path, counts in sorted(budget.items()): |
| 320 | if path not in current: |
| 321 | savings.append(f"{path}: file clean — tighten with --update") |
| 322 | else: |
| 323 | for name in counts: |
| 324 | if name not in current[path]: |
| 325 | savings.append( |
| 326 | f"{path}: {name} sites 0 < budget {counts[name]} — tighten with --update" |
| 327 | ) |
| 328 | |
| 329 | for line in savings: |
| 330 | print(line) |
| 331 | if failures: |
| 332 | print( |
| 333 | "\nBlocking-call budget exceeded — new `thread::sleep`/`std::fs` call " |
| 334 | "sites appeared outside spawn_blocking/dedicated-thread/test scopes:", |
| 335 | file=sys.stderr, |
| 336 | ) |
| 337 | for line in failures: |
| 338 | print(f" {line}", file=sys.stderr) |
| 339 | print( |
| 340 | "Move the work into `tokio::task::spawn_blocking` (or use tokio::fs " |
| 341 | "/ tokio::time), or raise the budget with --update if the site can " |
| 342 | "only run on synchronous code. See #6149.", |
| 343 | file=sys.stderr, |
| 344 | ) |
| 345 | return 1 |
| 346 | total = sum(sum(v.values()) for v in current.values()) |
| 347 | print(f"blocking-call budget: {total} sites across {len(current)} files, within budget") |
| 348 | return 0 |
| 349 | |
| 350 | |
| 351 | if __name__ == "__main__": |
| 352 | raise SystemExit(main()) |
| 353 |