Add `--build-branch`, `--bolt`, and `--musl` flags to update_prebuilts.py

Use `--build-branch aosp-rust-1.73.0` to fetch prebuilts from a
different branch other than `aosp-rust-toolchain`, which remains the
default.

Use `--no-bolt` to fetch the non-BOLT-optimized binaries.
Use `--no-musl` to skip downloading the MUSL binaries.

Bug: http://b/328074804
Test: ./update_prebuilts.py --no-bolt --no-musl --build-branch aosp-rust-1.73.0 --bid 11686893 -i 328074804 1.73.0b
Change-Id: I76e3d1f15a83bd21c3a1173f9b464ba5eb85e744

1 files changed, 46 insertions, 13 deletions

diff --git a/tools/update_prebuilts.py b/tools/update_prebuilts.py
index 8ea3e24..d9554d7 100755
--- a/tools/update_prebuilts.py
+++ b/tools/update_prebuilts.py
@@ -58,7 +58,7 @@ ANDROID_BP: str = "Android.bp"
 
 BRANCH_NAME_TEMPLATE: str = "rust-update-prebuilts-%s"
 
-BUILD_SERVER_BRANCH: str = "aosp-rust-toolchain"
+BUILD_SERVER_BRANCH_DEFAULT: str = "aosp-rust-toolchain"
 BUILD_SERVER_ARCHIVE_FORMAT_PATTERN: str = "rust-%s.tar.xz"
 
 HOST_ARCHIVE_PATTERN: str = "rust-%s-%s.tar.xz"
@@ -117,8 +117,10 @@ def ensure_gcert_valid() -> None:
 
 
 def fetch_build_server_artifact_strict(
-        target: str, build_id: int, build_server_pattern: str, host_name: str) -> Path:
-    result = fetch_build_server_artifact(target, build_id, build_server_pattern, host_name, None)
+        target: str, build_branch: str, build_id: int, build_server_pattern: str,
+        host_name: str) -> Path:
+    result = fetch_build_server_artifact(
+        target, build_branch, build_id, build_server_pattern, host_name, None)
 
     if result is None:
         sys.exit(1)
@@ -128,6 +130,7 @@ def fetch_build_server_artifact_strict(
 
 def fetch_build_server_artifact(
         target: str,
+        build_branch: str,
         build_id: int,
         build_server_pattern: str,
         host_name: str,
@@ -147,7 +150,7 @@ def fetch_build_server_artifact(
         build_flag = f"--bid={build_id}" if build_id else "--latest"
         result = subprocess.run([
             str(FETCH_ARTIFACT_PATH),
-            f"--branch={BUILD_SERVER_BRANCH}",
+            f"--branch={build_branch}",
             f"--target={target}",
             build_flag,
             build_server_pattern,
@@ -169,7 +172,7 @@ def get_targets() -> list[str]:
     result = subprocess.run([
         str(ANDROID_BUILD_CLI_PATH),
         "targets",
-        f"--branch={BUILD_SERVER_BRANCH}",
+        f"--branch={BUILD_SERVER_BRANCH_DEFAULT}",
     ],
                             stdout=subprocess.PIPE,
                             stderr=subprocess.DEVNULL)
@@ -187,7 +190,7 @@ def get_lkgb() -> int:
     result = subprocess.run([
         str(ANDROID_BUILD_CLI_PATH),
         "lkgb",
-        f"--branch={BUILD_SERVER_BRANCH}",
+        f"--branch={BUILD_SERVER_BRANCH_DEFAULT}",
         "--raw",
         "--custom_raw_format={o[buildId]}"
     ] + [f"--target={t}" for t in targets],
@@ -256,6 +259,16 @@ def parse_args() -> argparse.Namespace:
         type=int,
         help="Build ID to use when fetching artifacts from the build servers")
     parser.add_argument(
+        "--bolt",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="Fetch the BOLTed prebuilts (default true)")
+    parser.add_argument(
+        "--build-branch",
+        type=str,
+        default=BUILD_SERVER_BRANCH_DEFAULT,
+        help="Branch to fetch artifacts from build servers")
+    parser.add_argument(
         "--download-only", "-d", action="store_true", help="Stop after downloading the artifacts")
     parser.add_argument(
         "--no-pgo", action="store_true", help="Fetch the non-PGOed version of rustc")
@@ -273,6 +286,11 @@ def parse_args() -> argparse.Namespace:
         type=int,
         help="Issue number to include in commit message")
     parser.add_argument(
+        "--musl",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="Fetch the musl prebuilts as well (default true)")
+    parser.add_argument(
         "--overwrite",
         "-o",
         dest="overwrite",
@@ -282,7 +300,8 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def fetch_prebuilt_artifacts(bid: int, chained: bool) -> tuple[dict[str, Path], Path, list[Path]]:
+def fetch_prebuilt_artifacts(build_branch: str, bid: int, bolt: bool, musl: bool,
+                             chained: bool) -> tuple[dict[str, Path], Path, list[Path]]:
     """
     Returns a dictionary that maps target names to prebuilt artifact paths, the
     manifest used by the build server, and a list of other build server
@@ -296,7 +315,10 @@ def fetch_prebuilt_artifacts(bid: int, chained: bool) -> tuple[dict[str, Path],
 
     if chained:
         bs_target_linux_gnu = "rustc-linux_chained"
-        bs_archive_linux_gnu = "rust-linux-bolt-%s.tar.xz"
+        if bolt:
+            bs_archive_linux_gnu = "rust-linux-bolt-%s.tar.xz"
+        else:
+            bs_archive_linux_gnu = "rust-linux-%s.tar.xz"
     else:
         bs_target_linux_gnu = "rustc-linux_gnu"
         bs_archive_linux_gnu = "rust-linux-glibc-%s.tar.xz"
@@ -307,24 +329,30 @@ def fetch_prebuilt_artifacts(bid: int, chained: bool) -> tuple[dict[str, Path],
         "linux-musl-x86": ("rustc-linux_musl", "rust-linux-musl-%s.tar.xz"),
         "windows-x86": ("rustc-windows_gnu_native", "rust-windows-%s.tar.xz"),
     }
+    if not musl:
+        del build_server_target_map["linux-musl-x86"]
 
     # Fetch the host-specific prebuilt archives and build commands
     for host_target, (bs_target, bs_archive_pattern) in build_server_target_map.items():
         bs_archive_name = bs_archive_pattern % bid
         host_archive_name = HOST_ARCHIVE_PATTERN % (host_target, bid)
         prebuilt_path_map[host_target] = fetch_build_server_artifact_strict(
-            bs_target, bid, bs_archive_name, host_archive_name)
+            bs_target, build_branch, bid, bs_archive_name, host_archive_name)
 
         host_build_command_record_name = add_extension_prefix(
             BUILD_COMMAND_RECORD_NAME, f"{host_target}.{bid}")
         other_artifacts.append(
             fetch_build_server_artifact_strict(
-                bs_target, bid, BUILD_COMMAND_RECORD_NAME, host_build_command_record_name))
+                bs_target,
+                build_branch,
+                bid,
+                BUILD_COMMAND_RECORD_NAME,
+                host_build_command_record_name))
 
     # Fetch the manifest
     manifest_name: str = f"manifest_{bid}.xml"
     manifest_path: Path = fetch_build_server_artifact_strict(
-        bs_target_linux_gnu, bid, manifest_name, manifest_name)
+        bs_target_linux_gnu, build_branch, bid, manifest_name, manifest_name)
     other_artifacts.append(manifest_path)
 
     # Fetch the profiles
@@ -332,6 +360,7 @@ def fetch_prebuilt_artifacts(bid: int, chained: bool) -> tuple[dict[str, Path],
         for profile_name in PROFILE_NAMES:
             profile_path = fetch_build_server_artifact(
                 bs_target_linux_gnu,
+                build_branch,
                 bid,
                 profile_name,
                 add_extension_prefix(profile_name, str(bid)))
@@ -340,7 +369,11 @@ def fetch_prebuilt_artifacts(bid: int, chained: bool) -> tuple[dict[str, Path],
                 other_artifacts.append(profile_path)
 
         bolt_profile_path = fetch_build_server_artifact(
-            bs_target_linux_gnu, bid, "bolt-profiles.tar.xz", f"bolt-profiles-{bid}.tar.xz")
+            bs_target_linux_gnu,
+            build_branch,
+            bid,
+            "bolt-profiles.tar.xz",
+            f"bolt-profiles-{bid}.tar.xz")
 
         if bolt_profile_path is not None:
             other_artifacts.append(bolt_profile_path)
@@ -490,7 +523,7 @@ def main() -> int:
 
     print()
 
-    prebuilt_path_map, manifest_path, other_artifacts = fetch_prebuilt_artifacts(bid, not args.no_pgo)
+    prebuilt_path_map, manifest_path, other_artifacts = fetch_prebuilt_artifacts(args.build_branch, bid, args.bolt, args.musl, not args.no_pgo)
     if not args.download_only:
         update_prebuilts(
             branch_name,