pgo_tools: clean up after PGO generationupstream-main

Failing to clean up properly here is causing issues on Chrotomation.

BUG=b:341507487
TEST=Ran the script

Change-Id: I6137a80251cbd83021e3fdf92598c9a6fd8fb8cd
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5550817
Reviewed-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Commit-Queue: George Burgess <gbiv@chromium.org>
Tested-by: George Burgess <gbiv@chromium.org>

2 files changed, 54 insertions, 1 deletions

diff --git a/cros_utils/git_utils.py b/cros_utils/git_utils.py
index f02bc907..83b8612c 100644
--- a/cros_utils/git_utils.py
+++ b/cros_utils/git_utils.py
@@ -276,3 +276,40 @@ def fetch_and_checkout(git_dir: Path, remote: str, branch: str) -> str:
         cwd=git_dir,
         stdin=subprocess.DEVNULL,
     )
+
+
+def has_discardable_changes(git_dir: Path) -> bool:
+    """Returns whether discard_changes_and_checkout will discard changes."""
+    stdout = subprocess.run(
+        ["git", "status", "--porcelain"],
+        check=True,
+        cwd=git_dir,
+        stdin=subprocess.DEVNULL,
+        stdout=subprocess.PIPE,
+    ).stdout
+    return bool(stdout.strip())
+
+
+def discard_changes_and_checkout(git_dir: Path, ref: str):
+    """Discards local changes, and checks `ref` out."""
+    subprocess.run(
+        ["git", "clean", "-fd"],
+        check=True,
+        cwd=git_dir,
+        stdin=subprocess.DEVNULL,
+    )
+    # `git reset --hard HEAD` independently of the checkout, since we may be on
+    # a branch. The goal isn't to update the potential branch to point to
+    # `ref`.
+    subprocess.run(
+        ["git", "reset", "--hard", "HEAD"],
+        check=True,
+        cwd=git_dir,
+        stdin=subprocess.DEVNULL,
+    )
+    subprocess.run(
+        ["git", "checkout", ref],
+        check=True,
+        cwd=git_dir,
+        stdin=subprocess.DEVNULL,
+    )
diff --git a/pgo_tools/auto_update_llvm_pgo_profile.py b/pgo_tools/auto_update_llvm_pgo_profile.py
index c3693e71..963c2d68 100644
--- a/pgo_tools/auto_update_llvm_pgo_profile.py
+++ b/pgo_tools/auto_update_llvm_pgo_profile.py
@@ -24,6 +24,7 @@ from typing import Iterable, List, Optional
 
 from cros_utils import git_utils
 from llvm_tools import chroot
+from llvm_tools import cros_llvm_repo
 from llvm_tools import get_llvm_hash
 from llvm_tools import llvm_next
 from pgo_tools import pgo_utils
@@ -187,7 +188,22 @@ def maybe_upload_new_llvm_next_profile(
         profile_cache.insert_rev(llvm_next_rev)
         return
 
-    pgo_utils.run(cmd)
+    llvm_project = chromiumos_tree / cros_llvm_repo.REPO_PATH
+    if not clean_llvm and git_utils.has_discardable_changes(llvm_project):
+        raise ValueError(
+            f"Uncommitted changes exist in {llvm_project}. Please get rid of "
+            "them before running this script (e.g., with "
+            "`git clean -fd && git reset --hard HEAD`)"
+        )
+
+    initial_head = git_utils.resolve_ref(git_dir=llvm_project, ref="HEAD")
+    try:
+        pgo_utils.run(cmd)
+    finally:
+        logging.info("Restoring llvm-project to its original state...")
+        git_utils.discard_changes_and_checkout(
+            git_dir=llvm_project, ref=initial_head
+        )
 
     if upload_profile:
         profile_cache.insert_rev(llvm_next_rev)