pgo: write the profile before exiting on Windows

Cargo.toml

@@ -538,6 +538,7 @@ renamed_and_removed_lints = { level = "allow", priority = -1 }
 # https://doc.rust-lang.org/nightly/rustc/check-cfg/cargo-specifics.html
 unexpected_cfgs = { level = "warn", check-cfg = [
   'cfg(fuzzing)',
+  'cfg(pgo_training)',
   'cfg(target_os, values("cygwin"))',
   'cfg(wasi_runner)',
 ] }

src/bin/coreutils.rs

@@ -4,12 +4,11 @@
 // file that was distributed with this source code.
 
 use clap::Command;
-use coreutils::validation;
+use coreutils::validation::{self, exit};
 use itertools::Itertools as _;
 use std::cmp;
 use std::ffi::OsString;
 use std::io::{self, Write};
-use std::process;
 use uucore::{Args, error::strip_errno};
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -45,7 +44,7 @@ Currently defined functions:
         && e.kind() != io::ErrorKind::BrokenPipe
     {
         let _ = writeln!(io::stderr(), "coreutils: {}", strip_errno(&e));
-        process::exit(1);
+        exit(1);
     }
 }
 
@@ -57,7 +56,7 @@ fn main() {
     let binary = validation::binary_path(&mut args);
     let binary_as_util = validation::name(&binary).unwrap_or_else(|| {
         usage(&utils, "<unknown binary name>");
-        process::exit(0);
+        exit(0);
     });
 
     // binary name ends with util name?
@@ -88,7 +87,7 @@ fn main() {
                 // we should fail with additional args https://github.com/uutils/coreutils/issues/11383#issuecomment-4082564058
                 if args.next().is_some() {
                     let _ = writeln!(io::stderr(), "coreutils: invalid argument");
-                    process::exit(1);
+                    exit(1);
                 }
                 let mut out = io::stdout().lock();
                 for util in utils.keys() {
@@ -96,19 +95,19 @@ fn main() {
                         && e.kind() != io::ErrorKind::BrokenPipe
                     {
                         let _ = writeln!(io::stderr(), "coreutils: {}", strip_errno(&e));
-                        process::exit(1);
+                        exit(1);
                     }
                 }
-                process::exit(0);
+                exit(0);
             }
             "--version" | "-V" => {
                 if let Err(e) = writeln!(io::stdout(), "coreutils {VERSION} (multi-call binary)")
                     && e.kind() != io::ErrorKind::BrokenPipe
                 {
                     let _ = writeln!(io::stderr(), "coreutils: {}", strip_errno(&e));
-                    process::exit(1);
+                    exit(1);
                 }
-                process::exit(0);
+                exit(0);
             }
             // Not a special command: fallthrough to calling a util
             _ => {}
@@ -121,12 +120,12 @@ fn main() {
             // Could be something like:
             // #[cfg(not(feature = "only_english"))]
             validation::setup_localization_or_exit(util);
-            process::exit(uumain(vec![util_os].into_iter().chain(args)));
+            exit(uumain(vec![util_os].into_iter().chain(args)));
         }
         // GNU coreutils --help string shows help for coreutils
         if util == "--help" || util == "-h" {
             usage(&utils, binary_as_util);
-            process::exit(0);
+            exit(0);
         } else if util.starts_with('-') {
             // Argument looks like an option but wasn't recognized
             validation::unrecognized_option(binary_as_util, &util_os);
@@ -140,5 +139,5 @@ fn main() {
     } else {
         let _ = writeln!(io::stderr(), "coreutils: missing argument");
     }
-    process::exit(1);
+    exit(1);
 }

src/common/validation.rs

@@ -3,7 +3,7 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
-// spell-checker:ignore memfd_create prefixcat rsplit testcat
+// spell-checker:ignore memfd_create prefixcat rsplit testcat atexit profraw Cprofile
 
 use std::ffi::{OsStr, OsString};
 use std::io::{Write, stderr};
@@ -14,6 +14,31 @@ use uucore::Args;
 use uucore::display::Quotable;
 use uucore::locale;
 
+// The instrumented binary built by `util/build-pgo.sh` (`--cfg pgo_training`)
+// has to flush its own profile counters on Windows. Every exit below goes
+// through `std::process::exit`, which is `libc::exit` on Unix but
+// `ExitProcess` on Windows; the latter skips the `atexit` handler the LLVM
+// profiling runtime writes the counters from, so the training runs would
+// leave nothing but empty `.profraw` files behind.
+#[cfg(all(pgo_training, windows))]
+unsafe extern "C" {
+    fn __llvm_profile_write_file() -> i32;
+}
+
+/// Terminates the process with `code`, flushing the PGO counters when the
+/// binary was built for profile training on Windows.
+pub fn exit(code: i32) -> ! {
+    #[cfg(all(pgo_training, windows))]
+    // SAFETY: `__llvm_profile_write_file` is provided by the LLVM profiling
+    // runtime, which is linked in whenever this cfg is set (the same build
+    // passes `-Cprofile-generate`). It takes no arguments and only writes the
+    // counter file named by `LLVM_PROFILE_FILE`.
+    unsafe {
+        __llvm_profile_write_file();
+    }
+    process::exit(code)
+}
+
 /// Gets all available utilities including "coreutils"
 #[allow(clippy::type_complexity)]
 pub fn get_all_utilities<T: Args>(
@@ -31,7 +56,7 @@ pub fn not_found(util: &OsStr) -> ! {
         "coreutils: unknown program '{}'",
         util.maybe_quote()
     );
-    process::exit(1);
+    exit(1);
 }
 
 /// Prints an "unrecognized option" error and exits
@@ -41,7 +66,7 @@ pub fn unrecognized_option(binary_name: &str, option: &OsStr) -> ! {
         "{binary_name}: unrecognized option '{}'",
         option.to_string_lossy()
     );
-    process::exit(1);
+    exit(1);
 }
 
 /// Sets up localization for a utility with proper error handling
@@ -55,7 +80,7 @@ pub fn setup_localization_or_exit(util_name: &str) {
             } => eprintln!("Localization parse error at {snippet}: {err_msg}"),
             other => eprintln!("Could not init the localization system: {other}"),
         }
-        process::exit(99)
+        exit(99)
     });
 }
 

util/build-pgo.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-# spell-checker:ignore (jargon) profdata profraw sysroot rustlib nullglob aeiou nocheck CGU Cprofile cygpath
+# spell-checker:ignore (jargon) profdata profraw sysroot rustlib nullglob aeiou nocheck CGU Cprofile cygpath atexit
 #
 # Build uutils coreutils with Profile-Guided Optimization.
 #
@@ -108,7 +108,12 @@ mkdir -p "$PROFILE_DIR"
 # final one. Inlining happens before instrumentation, so building it with LTO
 # off yields counters for functions that no longer exist once the final build
 # runs whole-program codegen, and the profile is then largely wasted on it.
-cargo_build "$INSTR_DIR" "-Cprofile-generate=${PROFILE_DIR}"
+#
+# `--cfg pgo_training` makes the binary flush its own counters before exiting:
+# it always leaves through `std::process::exit`, which on Windows is
+# `ExitProcess` and skips the `atexit` handler the profiling runtime would
+# otherwise write the profile from.
+cargo_build "$INSTR_DIR" "-Cprofile-generate=${PROFILE_DIR} --cfg pgo_training"
 
 BIN="${INSTR_DIR}/${TARGET}/release/coreutils${EXE}"
 [ -x "$BIN" ] || { echo "instrumented binary not found: ${BIN}" >&2; exit 1; }