pgo_tools: add auto-update script for llvm

This CL adds an auto-update script for LLVM PGO profiles. It's intended
to run regularly on Chrotomation.

BUG=b:334876457
TEST=Ran the script with --dry-run

Change-Id: If778e0b352bffc888e793ccd0061390469ce2f3f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5501549
Tested-by: George Burgess <gbiv@chromium.org>
Reviewed-by: Bob Haarman <inglorion@chromium.org>
Commit-Queue: George Burgess <gbiv@chromium.org>

5 files changed, 438 insertions, 4 deletions

diff --git a/cros_utils/git_utils.py b/cros_utils/git_utils.py
index 99a4964f..33d52140 100644
--- a/cros_utils/git_utils.py
+++ b/cros_utils/git_utils.py
@@ -144,9 +144,23 @@ def try_set_autosubmit_labels(cwd: Path, cl_id: int) -> None:
 
 
 @contextlib.contextmanager
-def create_worktree(git_directory: Path) -> Generator[Path, None, None]:
-    """Creates a temp worktree of `git_directory`, yielding the result."""
-    with tempfile.TemporaryDirectory(prefix="git_utils_worktree_") as t:
+def create_worktree(
+    git_directory: Path, in_dir: Optional[Path] = None
+) -> Generator[Path, None, None]:
+    """Creates a temp worktree of `git_directory`, yielding the result.
+
+    Args:
+        git_directory: The directory to create a worktree of.
+        in_dir: The directory to make the worktree in. If None, uses the same
+            default as tempfile.TemporaryDirectory.
+
+    Yields:
+        A worktree to work in. This is cleaned up once the contextmanager is
+        exited.
+    """
+    with tempfile.TemporaryDirectory(
+        prefix="git_utils_worktree_", dir=in_dir
+    ) as t:
         tempdir = Path(t)
         logging.info(
             "Establishing worktree of %s in %s", git_directory, tempdir
diff --git a/pgo_tools/auto_update_llvm_pgo_profile.py b/pgo_tools/auto_update_llvm_pgo_profile.py
new file mode 100644
index 00000000..16f64323
--- /dev/null
+++ b/pgo_tools/auto_update_llvm_pgo_profile.py
@@ -0,0 +1,322 @@
+# Copyright 2024 The ChromiumOS Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Automatically manages LLVM's PGO profiles.
+
+Specifically, this script:
+    - generates & uploads new PGO profiles for llvm-next, if necessary
+    - ensures that the revisions for said llvm-next profiles are in the
+      associated manifest file
+    - cleans old llvm profiles from the ebuild
+
+Run this outside of the chroot.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+import re
+import shlex
+import subprocess
+import textwrap
+from typing import Iterable, List, Optional
+
+from cros_utils import git_utils
+from llvm_tools import chroot
+from llvm_tools import get_llvm_hash
+from llvm_tools import llvm_next
+from pgo_tools import pgo_utils
+
+
+PROFDATA_REV_PREFIX = "gs://chromeos-localmirror/distfiles/llvm-profdata-r"
+PROFDATA_REV_SUFFIX = ".xz"
+
+# Path to LLVM's 9999 ebuild from chromiumos-overlay
+LLVM_EBUILD_SUBPATH = Path("sys-devel", "llvm", "llvm-9999.ebuild")
+
+
+class GsProfileCache:
+    """Caches which LLVM revisions we have profile information for (in gs://).
+
+    To use this:
+        1. Create an instance of this class using `GsProfileCache.fetch()`.
+        2. Check if we have profile information for a revision: `123 in cache`.
+        3. Inform the cache that we have information for a revision:
+           `cache.insert_rev(123)`.
+    """
+
+    def __init__(self, profiles: Iterable[int]):
+        self.profile_revs = set(profiles)
+
+    def __contains__(self, rev: int) -> bool:
+        return rev in self.profile_revs
+
+    def __len__(self) -> int:
+        return len(self.profile_revs)
+
+    def insert_rev(self, rev: int) -> None:
+        self.profile_revs.add(rev)
+
+    @classmethod
+    def fetch(cls) -> "GsProfileCache":
+        stdout = subprocess.run(
+            [
+                "gsutil",
+                "ls",
+                f"{PROFDATA_REV_PREFIX}*{PROFDATA_REV_SUFFIX}",
+            ],
+            check=True,
+            stdin=subprocess.DEVNULL,
+            stdout=subprocess.PIPE,
+            encoding="utf-8",
+        ).stdout.strip()
+
+        prof_re = re.compile(
+            re.escape(PROFDATA_REV_PREFIX)
+            + r"(\d+)"
+            + re.escape(PROFDATA_REV_SUFFIX)
+        )
+        profiles = set()
+        for line in stdout.splitlines():
+            m = prof_re.fullmatch(line)
+            if not m:
+                if not line.strip():
+                    continue
+                raise ValueError(f"Unparseable line from gs://: {line!r}")
+            profiles.add(int(m.group(1)))
+        return cls(profiles)
+
+
+def parse_args(my_dir: Path, argv: List[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--chromiumos-tree",
+        type=Path,
+        help="""
+        Path to the root of the ChromeOS tree to edit. Autodetected if not
+        specified.
+        """,
+    )
+    parser.add_argument(
+        "--clobber-llvm",
+        action="store_true",
+        help="""
+        If a profile needs to be generated and there are uncommitted changes in
+        the LLVM source directory, clobber the changes.
+        """,
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Don't actually upload CLs, or generate new benchmark profiles.",
+    )
+    opts = parser.parse_args(argv)
+
+    if not opts.chromiumos_tree:
+        opts.chromiumos_tree = chroot.FindChromeOSRootAbove(my_dir)
+    return opts
+
+
+def maybe_upload_new_llvm_next_profile(
+    *,
+    profile_cache: GsProfileCache,
+    dry_run: bool,
+    toolchain_utils: Path,
+    clobber_llvm: bool,
+) -> None:
+    llvm_next_rev = llvm_next.LLVM_NEXT_REV
+    if llvm_next_rev in profile_cache:
+        logging.info(
+            "llvm-next profile already exists in gs://; no need to make a "
+            "new one"
+        )
+        return
+
+    create_script = (
+        toolchain_utils
+        / "py"
+        / "bin"
+        / "pgo_tools"
+        / "create_chroot_and_generate_pgo_profile.py"
+    )
+    logging.info("Generating a PGO profile for LLVM r%d", llvm_next_rev)
+    cmd: pgo_utils.Command = [
+        create_script,
+        f"--rev={llvm_next_rev}",
+        "--upload",
+    ]
+    if clobber_llvm:
+        cmd.append("--clobber-llvm")
+
+    if dry_run:
+        logging.info(
+            "Skipping PGO profile generation for llvm r%d due to --dry-run. "
+            "Would run: %s",
+            llvm_next_rev,
+            shlex.join(str(x) for x in cmd),
+        )
+        profile_cache.insert_rev(llvm_next_rev)
+        return
+
+    pgo_utils.run(cmd)
+    profile_cache.insert_rev(llvm_next_rev)
+
+
+def overwrite_llvm_pgo_listing(
+    chromiumos_overlay: Path, profile_revs: List[int]
+) -> bool:
+    ebuild = chromiumos_overlay / LLVM_EBUILD_SUBPATH
+    contents = ebuild.read_text(encoding="utf-8")
+    new_pgo_listing = "\t" + "\n\t".join(str(x) for x in profile_revs)
+
+    array_start = "\nLLVM_PGO_PROFILE_REVS=(\n"
+    array_start_index = contents.index(array_start)
+    array_end_index = contents.index("\n)", array_start_index)
+
+    new_contents = (
+        contents[: array_start_index + len(array_start)]
+        + new_pgo_listing
+        + contents[array_end_index:]
+    )
+    if new_contents == contents:
+        return False
+    ebuild.write_text(new_contents, encoding="utf-8")
+    return True
+
+
+def update_llvm_ebuild_manifest(
+    chromeos_tree: Path, chromiumos_overlay: Path
+) -> None:
+    overlay_relpath = chromiumos_overlay.relative_to(chromeos_tree)
+    overlay_chroot_path = Path("/mnt") / "host" / "source" / overlay_relpath
+    llvm_9999 = overlay_chroot_path / LLVM_EBUILD_SUBPATH
+    ebuild_manifest_cmd = shlex.join(["ebuild", str(llvm_9999), "manifest"])
+    logging.info("Running `%s` in the chroot...", ebuild_manifest_cmd)
+    subprocess.run(
+        ["cros_sdk", "--", "bash", "-c", ebuild_manifest_cmd],
+        check=True,
+        cwd=chromeos_tree,
+    )
+
+
+def create_llvm_pgo_ebuild_update(
+    chromeos_tree: Path,
+    chromiumos_overlay: Path,
+    profile_cache: GsProfileCache,
+    dry_run: bool,
+) -> Optional[str]:
+    llvm_dir = get_llvm_hash.GetAndUpdateLLVMProjectInLLVMTools()
+    llvm_hash = get_llvm_hash.LLVMHash()
+    current_llvm_sha = llvm_hash.GetCrOSCurrentLLVMHash(chromeos_tree)
+    current_llvm_rev = get_llvm_hash.GetVersionFrom(llvm_dir, current_llvm_sha)
+    logging.info("Current LLVM revision is %d", current_llvm_rev)
+    want_revisions = [current_llvm_rev]
+
+    llvm_next_rev = llvm_next.LLVM_NEXT_REV
+    if current_llvm_rev != llvm_next_rev:
+        logging.info("llvm-next rev is r%d", llvm_next_rev)
+        if llvm_next_rev in profile_cache:
+            want_revisions.append(llvm_next_rev)
+        else:
+            logging.info(
+                "No PGO profile exists for r%d; skip adding to profile list",
+                llvm_next_rev,
+            )
+    logging.info(
+        "Expected LLVM PGO profile version(s) in ebuild: %s", want_revisions
+    )
+
+    made_change = overwrite_llvm_pgo_listing(chromiumos_overlay, want_revisions)
+    if not made_change:
+        logging.info("No LLVM ebuild changes made")
+        return None
+
+    # Skip the manifest update in this case, since the profile cache we're
+    # using might have had a entry inserted by the profile generation stage of
+    # this script.
+    if dry_run:
+        logging.info("Skipping manifest update; --dry-run was passed")
+    else:
+        update_llvm_ebuild_manifest(chromeos_tree, chromiumos_overlay)
+    return git_utils.commit_all_changes(
+        chromiumos_overlay,
+        textwrap.dedent(
+            """\
+            llvm: update PGO profile listing
+
+            This CL was generated by toolchain-utils'
+            pgo_tools/auto_update_llvm_pgo_profile.py script.
+
+            BUG=b:337284701
+            TEST=CQ
+            """
+        ),
+    )
+
+
+def main(argv: List[str]) -> None:
+    my_dir = Path(__file__).resolve().parent
+
+    pgo_utils.exit_if_in_chroot()
+
+    logging.basicConfig(
+        format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: "
+        "%(message)s",
+        level=logging.INFO,
+    )
+    opts = parse_args(my_dir, argv)
+
+    chromeos_root = opts.chromiumos_tree
+    chromiumos_overlay = (
+        chromeos_root / "src" / "third_party" / "chromiumos-overlay"
+    )
+    dry_run = opts.dry_run
+
+    logging.info("Populating gs:// profile cache...")
+    profile_cache = GsProfileCache.fetch()
+    logging.info("Found %d LLVM PGO profiles in gs://.", len(profile_cache))
+
+    maybe_upload_new_llvm_next_profile(
+        profile_cache=profile_cache,
+        dry_run=dry_run,
+        toolchain_utils=my_dir.parent,
+        clobber_llvm=opts.clobber_llvm,
+    )
+
+    # NOTE: `in_dir=chromeos_root` here is critical, since this function needs
+    # to enter the chroot to run `ebuild manifest`. Hence, the worktree must be
+    # trivially reachable from within the chroot.
+    with git_utils.create_worktree(
+        chromiumos_overlay, in_dir=chromeos_root
+    ) as worktree:
+        maybe_sha = create_llvm_pgo_ebuild_update(
+            chromeos_root,
+            worktree,
+            profile_cache,
+            dry_run,
+        )
+
+    if not maybe_sha:
+        logging.info("No changes made to LLVM ebuild; quit.")
+        return
+
+    if dry_run:
+        logging.info(
+            "LLVM ebuild changes committed as %s. --dry-run specified; quit.",
+            maybe_sha,
+        )
+        return
+
+    cls = git_utils.upload_to_gerrit(
+        chromiumos_overlay,
+        remote=git_utils.CROS_EXTERNAL_REMOTE,
+        branch=git_utils.CROS_MAIN_BRANCH,
+        ref=maybe_sha,
+    )
+    for cl in cls:
+        git_utils.try_set_autosubmit_labels(chromiumos_overlay, cl)
+    logging.info("%d CL(s) uploaded.", len(cls))
diff --git a/pgo_tools/auto_update_llvm_pgo_profile_test.py b/pgo_tools/auto_update_llvm_pgo_profile_test.py
new file mode 100644
index 00000000..25aafd5c
--- /dev/null
+++ b/pgo_tools/auto_update_llvm_pgo_profile_test.py
@@ -0,0 +1,94 @@
+# Copyright 2024 The ChromiumOS Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for auto_update_llvm_pgo_profile."""
+
+from pathlib import Path
+import subprocess
+import textwrap
+from unittest import mock
+
+from llvm_tools import test_helpers
+from pgo_tools import auto_update_llvm_pgo_profile
+
+
+EXAMPLE_LLVM_EBUILD_SNIPPET = """
+# foo
+# bar
+
+import baz
+
+# comments
+LLVM_PGO_PROFILE_REVS=(
+\t516547
+\t516548
+)
+# some more stuff
+"""
+
+
+class Test(test_helpers.TempDirTestCase):
+    """Tests for auto_update_llvm_pgo_profile."""
+
+    def make_tempdir_with_example_llvm_ebuild(self) -> Path:
+        cros_overlay = self.make_tempdir()
+        llvm_9999 = (
+            cros_overlay / auto_update_llvm_pgo_profile.LLVM_EBUILD_SUBPATH
+        )
+        llvm_9999.parent.mkdir(parents=True)
+        llvm_9999.write_text(EXAMPLE_LLVM_EBUILD_SNIPPET, encoding="utf-8")
+        return cros_overlay
+
+    def test_ebuild_updating_is_nop_when_revs_dont_change(self):
+        cros_overlay = self.make_tempdir_with_example_llvm_ebuild()
+        updated = auto_update_llvm_pgo_profile.overwrite_llvm_pgo_listing(
+            cros_overlay, [516547, 516548]
+        )
+        new_contents = (
+            cros_overlay / auto_update_llvm_pgo_profile.LLVM_EBUILD_SUBPATH
+        ).read_text(encoding="utf-8")
+        self.assertEqual(EXAMPLE_LLVM_EBUILD_SNIPPET, new_contents)
+        self.assertFalse(updated)
+
+    def test_ebuild_updating_works_when_rev_is_removed(self):
+        cros_overlay = self.make_tempdir_with_example_llvm_ebuild()
+        self.assertTrue(
+            auto_update_llvm_pgo_profile.overwrite_llvm_pgo_listing(
+                cros_overlay, [516547]
+            )
+        )
+        new_contents = (
+            cros_overlay / auto_update_llvm_pgo_profile.LLVM_EBUILD_SUBPATH
+        ).read_text(encoding="utf-8")
+        self.assertIn("\n\t516547\n", new_contents)
+        self.assertNotIn("\n\t516548\n", new_contents)
+
+    def test_ebuild_updating_works_when_rev_is_added(self):
+        cros_overlay = self.make_tempdir_with_example_llvm_ebuild()
+        self.assertTrue(
+            auto_update_llvm_pgo_profile.overwrite_llvm_pgo_listing(
+                cros_overlay, [516547, 516548, 516549]
+            )
+        )
+        new_contents = (
+            cros_overlay / auto_update_llvm_pgo_profile.LLVM_EBUILD_SUBPATH
+        ).read_text(encoding="utf-8")
+        self.assertIn("\n\t516547\n", new_contents)
+        self.assertIn("\n\t516548\n", new_contents)
+        self.assertIn("\n\t516549\n", new_contents)
+
+    @mock.patch.object(subprocess, "run")
+    def test_gs_parsing_works(self, mock_run):
+        run_return = mock.MagicMock()
+        run_return.stdout = textwrap.dedent(
+            """\
+            gs://chromeos-localmirror/distfiles/llvm-profdata-r1234.xz
+            gs://chromeos-localmirror/distfiles/llvm-profdata-r5678.xz
+            """
+        )
+        mock_run.return_value = run_return
+        cache = auto_update_llvm_pgo_profile.GsProfileCache.fetch()
+        self.assertEqual(len(cache), 2)
+        self.assertIn(1234, cache)
+        self.assertIn(5678, cache)
diff --git a/pgo_tools/create_chroot_and_generate_pgo_profile.py b/pgo_tools/create_chroot_and_generate_pgo_profile.py
index 6699dfe2..db1a330e 100644
--- a/pgo_tools/create_chroot_and_generate_pgo_profile.py
+++ b/pgo_tools/create_chroot_and_generate_pgo_profile.py
@@ -254,10 +254,13 @@ def main(argv: List[str]):
 
         compressed_profile_path = compress_pgo_profile(profile_path)
         upload_command = determine_upload_command(compressed_profile_path, rev)
+        friendly_upload_command = shlex.join(str(x) for x in upload_command)
         if opts.upload:
+            logging.info(
+                "Running `%s` to upload the profile...", friendly_upload_command
+            )
             pgo_utils.run(upload_command)
         else:
-            friendly_upload_command = shlex.join(str(x) for x in upload_command)
             logging.info(
                 "To upload the profile, run %r in %r",
                 friendly_upload_command,
diff --git a/py/bin/pgo_tools/auto_update_llvm_pgo_profile.py b/py/bin/pgo_tools/auto_update_llvm_pgo_profile.py
new file mode 120000
index 00000000..0f1ca492
--- /dev/null
+++ b/py/bin/pgo_tools/auto_update_llvm_pgo_profile.py
@@ -0,0 +1 @@
+../../../python_wrapper.py
\ No newline at end of file