xtask/tasks/fmt/
rustfmt.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use super::FmtPass;
5use crate::fs_helpers::git_diffed;
6use crate::shell::XtaskShell;
7use crate::tasks::fmt::FmtCtx;
8use anyhow::Context;
9use std::collections::BTreeMap;
10use std::collections::BTreeSet;
11use std::path::Path;
12use std::path::PathBuf;
13use std::sync::atomic::AtomicBool;
14use std::sync::atomic::AtomicUsize;
15use std::sync::atomic::Ordering;
16
17/// Windows caps a process's command line at 32767 characters. Batches are
18/// normally kept far below this to spread work across threads; this is only a
19/// backstop for unusually long paths.
20const MAX_FILE_ARG_BYTES: usize = 24 * 1024;
21
22/// Spawning `rustfmt` costs more than formatting a handful of files, so batches
23/// below this size are not worth the extra parallelism. Measured on this repo,
24/// one process per file is ~3x slower than batching.
25const MIN_BATCH_FILES: usize = 8;
26
27/// A set of files to format together under a common edition.
28struct Group {
29    edition: Option<String>,
30    files: Vec<PathBuf>,
31}
32
33/// Metadata about one workspace package.
34struct Package {
35    /// The directory containing the package's `Cargo.toml`.
36    dir: PathBuf,
37    edition: String,
38    /// The crate root of each of the package's targets.
39    roots: Vec<PathBuf>,
40}
41
42pub struct Rustfmt;
43
44impl FmtPass for Rustfmt {
45    fn run(self, ctx: FmtCtx) -> anyhow::Result<()> {
46        let FmtCtx {
47            ctx,
48            fix,
49            only_diffed,
50        } = ctx;
51        let sh = XtaskShell::new()?;
52        let rust_toolchain = sh.var("RUST_TOOLCHAIN").map(|s| format!("+{s}")).ok();
53        let fmt_check = (!fix).then_some("--check");
54
55        let packages = workspace_packages(&sh)?;
56
57        let groups = if only_diffed {
58            let mut files = git_diffed(ctx.in_git_hook)?;
59            files.retain(|f| f.extension().unwrap_or_default() == "rs");
60            group_diffed(&packages, files)
61        } else {
62            // Deliberately avoid `cargo fmt`: it spawns a single `rustfmt` with
63            // every crate root in the workspace on the command line, which
64            // overflows the Windows command line length limit and leaves all
65            // but one core idle. Replicate what it does, but in batches.
66            group_roots(&packages)
67        };
68
69        if run_rustfmt(rust_toolchain.as_deref(), fmt_check, &groups)? {
70            anyhow::bail!("found formatting issues");
71        }
72
73        Ok(())
74    }
75}
76
77/// Group every package's crate roots by edition. `rustfmt` walks `mod`
78/// declarations itself, so these roots transitively cover every file that
79/// `cargo fmt` would format.
80fn group_roots(packages: &[Package]) -> Vec<Group> {
81    let mut by_edition: BTreeMap<&str, BTreeSet<&Path>> = BTreeMap::new();
82    for package in packages {
83        by_edition
84            .entry(&package.edition)
85            .or_default()
86            .extend(package.roots.iter().map(PathBuf::as_path));
87    }
88
89    by_edition
90        .into_iter()
91        .map(|(edition, files)| Group {
92            edition: Some(edition.to_owned()),
93            files: files.into_iter().map(Path::to_path_buf).collect(),
94        })
95        .collect()
96}
97
98/// Group diffed files by the edition of the innermost package containing them,
99/// so that they are parsed the same way as in a full run.
100fn group_diffed(packages: &[Package], files: Vec<PathBuf>) -> Vec<Group> {
101    let mut by_edition: BTreeMap<Option<&str>, Vec<PathBuf>> = BTreeMap::new();
102    for file in files {
103        let edition = packages
104            .iter()
105            .filter(|p| file.starts_with(&p.dir))
106            .max_by_key(|p| p.dir.components().count())
107            .map(|p| p.edition.as_str());
108
109        by_edition.entry(edition).or_default().push(file);
110    }
111
112    by_edition
113        .into_iter()
114        .map(|(edition, files)| Group {
115            edition: edition.map(str::to_owned),
116            files,
117        })
118        .collect()
119}
120
121/// Run `rustfmt` over each group, splitting the groups into batches that run in
122/// parallel.
123///
124/// Returns whether any invocation reported formatting issues.
125fn run_rustfmt(
126    rust_toolchain: Option<&str>,
127    fmt_check: Option<&str>,
128    groups: &[Group],
129) -> anyhow::Result<bool> {
130    let threads = std::thread::available_parallelism().map_or(1, |n| n.get());
131    let total_files: usize = groups.iter().map(|g| g.files.len()).sum();
132    let max_files = batch_size(total_files, threads);
133
134    let batches = groups
135        .iter()
136        .flat_map(|g| {
137            batch_files(&g.files, MAX_FILE_ARG_BYTES, max_files)
138                .into_iter()
139                .map(move |files| (g.edition.as_deref(), files))
140        })
141        .collect::<Vec<_>>();
142
143    let next = AtomicUsize::new(0);
144    let failed = AtomicBool::new(false);
145
146    std::thread::scope(|scope| -> anyhow::Result<()> {
147        let handles = (0..threads.min(batches.len()))
148            .map(|_| {
149                scope.spawn(|| -> anyhow::Result<()> {
150                    // `xshell::Shell` isn't `Sync`, so give each thread its own.
151                    let sh = XtaskShell::new()?;
152                    while let Some((edition, files)) =
153                        batches.get(next.fetch_add(1, Ordering::Relaxed))
154                    {
155                        let res = sh
156                            .cmd("rustfmt")
157                            .args(rust_toolchain)
158                            .args(fmt_check)
159                            .args(edition.map(|e| format!("--edition={e}")))
160                            .args(*files)
161                            .quiet()
162                            .run();
163
164                        if res.is_err() {
165                            failed.store(true, Ordering::Relaxed);
166                        }
167                    }
168                    Ok(())
169                })
170            })
171            .collect::<Vec<_>>();
172
173        for handle in handles {
174            handle.join().unwrap()?;
175        }
176        Ok(())
177    })?;
178
179    Ok(failed.load(Ordering::Relaxed))
180}
181
182/// Choose how many files to put in each batch: enough batches to keep every
183/// thread busy through the end of the run, but not so few files per batch that
184/// process startup dominates.
185fn batch_size(total_files: usize, threads: usize) -> usize {
186    total_files.div_ceil(threads * 4).max(MIN_BATCH_FILES)
187}
188
189/// Split `files` into batches that fit in a command line and are small enough
190/// to spread across the available threads.
191fn batch_files(files: &[PathBuf], max_bytes: usize, max_files: usize) -> Vec<&[PathBuf]> {
192    let mut batches = Vec::new();
193    let mut start = 0;
194    let mut len = 0;
195
196    for (i, file) in files.iter().enumerate() {
197        // +1 for the argument separator
198        let file_len = file.as_os_str().len() + 1;
199        if i > start && (len + file_len > max_bytes || i - start >= max_files) {
200            batches.push(&files[start..i]);
201            start = i;
202            len = 0;
203        }
204        len += file_len;
205    }
206
207    if start < files.len() {
208        batches.push(&files[start..]);
209    }
210
211    batches
212}
213
214/// Collect the edition, directory, and crate roots of every workspace package.
215fn workspace_packages(sh: &XtaskShell) -> anyhow::Result<Vec<Package>> {
216    #[derive(serde::Deserialize)]
217    struct Metadata {
218        packages: Vec<MetadataPackage>,
219    }
220
221    #[derive(serde::Deserialize)]
222    struct MetadataPackage {
223        edition: String,
224        manifest_path: PathBuf,
225        targets: Vec<Target>,
226    }
227
228    #[derive(serde::Deserialize)]
229    struct Target {
230        src_path: PathBuf,
231    }
232
233    let output = sh
234        .cmd("cargo")
235        .args(["metadata", "--no-deps", "--format-version", "1"])
236        .quiet()
237        .output()
238        .context("failed to run cargo metadata")?;
239
240    let metadata: Metadata =
241        serde_json::from_slice(&output.stdout).context("failed to parse cargo metadata")?;
242
243    let cwd = std::env::current_dir()?;
244    // Shorten the paths as much as possible, since the command line length is
245    // the constraint being worked around here.
246    let shorten = |path: &Path| path.strip_prefix(&cwd).unwrap_or(path).to_path_buf();
247
248    let mut packages = Vec::new();
249    for package in metadata.packages {
250        let dir = package
251            .manifest_path
252            .parent()
253            .context("manifest path has no parent")?;
254
255        let mut roots: Vec<PathBuf> = package
256            .targets
257            .iter()
258            .map(|t| shorten(&t.src_path))
259            .collect();
260        roots.sort();
261        roots.dedup();
262
263        packages.push(Package {
264            dir: shorten(dir),
265            edition: package.edition,
266            roots,
267        });
268    }
269
270    Ok(packages)
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    const LIMIT: usize = 24 * 1024;
278
279    fn test_files() -> Vec<PathBuf> {
280        (0..10_000)
281            .map(|i| PathBuf::from(format!("some/moderately/long/path/to/crate{i}/src/lib.rs")))
282            .collect()
283    }
284
285    #[test]
286    fn batches_stay_under_limit() {
287        let files = test_files();
288        let batches = batch_files(&files, LIMIT, usize::MAX);
289        assert!(batches.len() > 1);
290        assert_eq!(batches.iter().map(|b| b.len()).sum::<usize>(), files.len());
291
292        for batch in batches {
293            let len: usize = batch.iter().map(|f| f.as_os_str().len() + 1).sum();
294            assert!(len <= LIMIT, "batch too long: {len}");
295        }
296    }
297
298    #[test]
299    fn batches_stay_under_file_count() {
300        let files = test_files();
301        let batches = batch_files(&files, usize::MAX, 64);
302        assert_eq!(batches.len(), files.len().div_ceil(64));
303        assert!(batches.iter().all(|b| b.len() <= 64));
304    }
305
306    #[test]
307    fn single_oversized_file_still_batched() {
308        let files = vec![PathBuf::from("a".repeat(LIMIT * 2))];
309        assert_eq!(batch_files(&files, LIMIT, usize::MAX).len(), 1);
310    }
311
312    #[test]
313    fn unlimited_produces_one_batch() {
314        let files = test_files();
315        assert_eq!(batch_files(&files, usize::MAX, usize::MAX).len(), 1);
316    }
317
318    #[test]
319    fn empty_input_produces_no_batches() {
320        assert!(batch_files(&[], LIMIT, usize::MAX).is_empty());
321    }
322
323    #[test]
324    fn small_runs_use_a_single_batch() {
325        assert_eq!(batch_size(8, 8), MIN_BATCH_FILES);
326        assert_eq!(batch_size(1, 8), MIN_BATCH_FILES);
327    }
328
329    #[test]
330    fn large_runs_spread_across_threads() {
331        assert_eq!(batch_size(483, 8), 16);
332        assert!(483_usize.div_ceil(batch_size(483, 8)) >= 8);
333    }
334
335    #[test]
336    fn diffed_files_use_innermost_package_edition() {
337        let packages = vec![
338            Package {
339                dir: PathBuf::from(""),
340                edition: "2021".to_string(),
341                roots: Vec::new(),
342            },
343            Package {
344                dir: PathBuf::from("vm/devices/net"),
345                edition: "2024".to_string(),
346                roots: Vec::new(),
347            },
348        ];
349
350        let groups = group_diffed(
351            &packages,
352            vec![
353                PathBuf::from("xtask/src/main.rs"),
354                PathBuf::from("vm/devices/net/netvsp/src/lib.rs"),
355            ],
356        );
357
358        let editions: Vec<_> = groups
359            .iter()
360            .map(|g| (g.edition.as_deref(), g.files.clone()))
361            .collect();
362        assert_eq!(
363            editions,
364            vec![
365                (Some("2021"), vec![PathBuf::from("xtask/src/main.rs")]),
366                (
367                    Some("2024"),
368                    vec![PathBuf::from("vm/devices/net/netvsp/src/lib.rs")]
369                ),
370            ]
371        );
372    }
373}
xtask/tasks/fmt/rustfmt.rs

xtask/tasks/fmt/
rustfmt.rs
