Lazify sparse registry support

rs/BUILD.bazel

@@ -43,6 +43,8 @@ bzl_library(
         "//rs/private:downloader",
         "//rs/private:git_repository",
         "//rs/private:lint_flags",
+        "//rs/private:registry_config_repository",
+        "//rs/private:registry_utils",
         "//rs/private:repository_utils",
         "//rs/private:resolver",
         "//rs/private:select_utils",

rs/extensions.bzl

@@ -7,9 +7,11 @@ load("//rs/private:cargo_credentials.bzl", "load_cargo_credentials")
 load("//rs/private:cfg_parser.bzl", "cfg_matches_expr_for_cfg_attrs", "triple_to_cfg_attrs")
 load("//rs/private:crate_git_repository.bzl", "crate_git_repository")
 load("//rs/private:crate_repository.bzl", "crate_repository", "local_crate_repository")
-load("//rs/private:downloader.bzl", "download_metadata_for_git_crates", "download_sparse_registry_configs", "new_downloader_state", "parse_git_url", "sharded_path", "start_crate_registry_downloads", "start_github_downloads")
+load("//rs/private:downloader.bzl", "download_metadata_for_git_crates", "new_downloader_state", "parse_git_url", "start_crate_registry_downloads", "start_github_downloads")
 load("//rs/private:git_repository.bzl", "git_repository")
 load("//rs/private:lint_flags.bzl", "cargo_toml_lint_flags")
+load("//rs/private:registry_config_repository.bzl", "registry_config_repository")
+load("//rs/private:registry_utils.bzl", "CRATES_IO_REGISTRY", "registry_config_repo_name")
 load("//rs/private:repository_utils.bzl", "render_select")
 load("//rs/private:resolver.bzl", "resolve")
 load("//rs/private:select_utils.bzl", "compute_select")
@@ -251,7 +253,6 @@ def _generate_hub_and_spokes(
         cargo_lock_path,
         workspace_cargo_toml_json,
         all_packages,
-        sparse_registry_configs,
         platform_triples,
         cargo_credentials,
         cargo_config,
@@ -270,7 +271,6 @@ def _generate_hub_and_spokes(
         cargo_lock_path (path): Cargo.lock path
         workspace_cargo_toml_json (dict): Parsed workspace Cargo.toml
         all_packages: list[package]: from cargo lock parsing
-        sparse_registry_configs: dict[source, sparse registry config]
         platform_triples (list[string]): Triples to resolve for
         cargo_credentials (dict): Mapping of registry to auth token.
         cargo_config (label): .cargo/config.toml file
@@ -779,25 +779,19 @@ crate.annotation(
 
         if source.startswith("sparse+"):
             checksum = package["checksum"]
-            url = sparse_registry_configs[source].format(**{
-                "crate": crate_name,
-                "version": version,
-                "prefix": sharded_path(crate_name),
-                "lowerprefix": sharded_path(crate_name.lower()),
-                "sha256-checksum": checksum,
-            })
 
             if dry_run:
                 continue
 
             qualifiers = {}
-            if source != "sparse+https://index.crates.io/":
+            if source != CRATES_IO_REGISTRY:
                 qualifiers["repository_url"] = source.split("+", 1)[1]
 
             crate_repository(
                 name = repo_name,
-                url = url,
-                strip_prefix = "%s-%s" % (crate_name, version),
+                crate_name = crate_name,
+                version = version,
+                registry_config = "@%s//:dl" % registry_config_repo_name(hub_name, source),
                 sbom_extra_qualifiers = qualifiers,
                 checksum = checksum,
                 # The repository will need to recompute these, but this lets us avoid serializing them.
@@ -1233,17 +1227,33 @@ def _crate_impl(mctx):
                 cargo_credentials = {}
 
             cargo_credentials_by_hub_name[cfg.name] = cargo_credentials
-            start_crate_registry_downloads(mctx, downloader_state, annotations, packages_by_hub_name[cfg.name], cargo_credentials, cfg.debug)
+            packages = packages_by_hub_name[cfg.name]
+            registry_sources = set()
+
+            for package in packages:
+                source = package.get("source")
+                if source == "registry+https://github.com/rust-lang/crates.io-index":
+                    source = CRATES_IO_REGISTRY
+                    package["source"] = source
+
+                if source and source.startswith("sparse+"):
+                    registry_sources.add(source)
+
+            start_crate_registry_downloads(mctx, downloader_state, annotations, packages, cargo_credentials, cfg.debug)
+
+            for source in sorted(registry_sources):
+                registry_config_repository(
+                    name = registry_config_repo_name(cfg.name, source),
+                    source = source,
+                    cargo_config = cfg.cargo_config,
+                    use_home_cargo_credentials = cfg.use_home_cargo_credentials,
+                )
 
     for fetch_state in downloader_state.in_flight_git_crate_fetches_by_url.values():
         fetch_state.download_token.wait()
 
     download_metadata_for_git_crates(mctx, downloader_state, annotations_by_hub_name)
 
-    # TODO(zbarsky): Unfortunate that we block on the download for crates.io even though it's well-known.
-    # Should we hardcode it?
-    sparse_registry_configs = download_sparse_registry_configs(mctx, downloader_state)
-
     facts = {}
     direct_deps = []
     direct_dev_deps = []
@@ -1263,9 +1273,9 @@ def _crate_impl(mctx):
 
             if cfg.debug:
                 for _ in range(25):
-                    _generate_hub_and_spokes(mctx, cfg.name, annotations, suggested_annotation_snippet_paths, cargo_path, cfg.cargo_lock, cargo_toml_by_hub_name[cfg.name], hub_packages, sparse_registry_configs, cfg.platform_triples, cargo_credentials, cfg.cargo_config, cfg.validate_lockfile, cfg.debug, cfg.use_legacy_rules_rust_platforms, dry_run = True)
+                    _generate_hub_and_spokes(mctx, cfg.name, annotations, suggested_annotation_snippet_paths, cargo_path, cfg.cargo_lock, cargo_toml_by_hub_name[cfg.name], hub_packages, cfg.platform_triples, cargo_credentials, cfg.cargo_config, cfg.validate_lockfile, cfg.debug, cfg.use_legacy_rules_rust_platforms, dry_run = True)
 
-            facts |= _generate_hub_and_spokes(mctx, cfg.name, annotations, suggested_annotation_snippet_paths, cargo_path, cfg.cargo_lock, cargo_toml_by_hub_name[cfg.name], hub_packages, sparse_registry_configs, cfg.platform_triples, cargo_credentials, cfg.cargo_config, cfg.validate_lockfile, cfg.debug, cfg.use_legacy_rules_rust_platforms)
+            facts |= _generate_hub_and_spokes(mctx, cfg.name, annotations, suggested_annotation_snippet_paths, cargo_path, cfg.cargo_lock, cargo_toml_by_hub_name[cfg.name], hub_packages, cfg.platform_triples, cargo_credentials, cfg.cargo_config, cfg.validate_lockfile, cfg.debug, cfg.use_legacy_rules_rust_platforms)
 
     # Lay down the git repos we will need; per-crate git_repository can clone from these.
     git_sources = set()

rs/private/BUILD.bazel

@@ -48,6 +48,7 @@ bzl_library(
     visibility = ["//rs:__subpackages__"],
     deps = [
         ":cargo_credentials",
+        ":registry_utils",
         ":repository_utils",
         ":toml2json",
         "@bazel_tools//tools/build_defs/repo:cache.bzl",
@@ -62,11 +63,25 @@ bzl_library(
     deps = [
         ":annotations",
         ":cargo_credentials",
+        ":registry_utils",
         ":toml2json",
         "@bazel_tools//tools/build_defs/repo:git_worker.bzl",
     ],
 )
 
+bzl_library(
+    name = "registry_config_repository",
+    srcs = ["registry_config_repository.bzl"],
+    visibility = ["//rs:__subpackages__"],
+    deps = [":cargo_credentials"],
+)
+
+bzl_library(
+    name = "registry_utils",
+    srcs = ["registry_utils.bzl"],
+    visibility = ["//rs:__subpackages__"],
+)
+
 bzl_library(
     name = "repository_utils",
     srcs = ["repository_utils.bzl"],

rs/private/crate_repository.bzl

@@ -2,6 +2,7 @@ load("@bazel_features//:features.bzl", "bazel_features")
 load("@bazel_tools//tools/build_defs/repo:cache.bzl", "get_default_canonical_id")
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "patch")
 load(":cargo_credentials.bzl", "load_cargo_credentials", "registry_auth_headers")
+load(":registry_utils.bzl", "sharded_path")
 load(":repository_utils.bzl", "common_attrs", "generate_build_file")
 load(":toml2json.bzl", "run_toml2json")
 
@@ -15,13 +16,27 @@ def _crate_repository_impl(rctx):
     else:
         headers = {}
 
+    crate_name = rctx.attr.crate_name
+    version = rctx.attr.version
+    sha256 = rctx.attr.checksum
+
+    dl = rctx.read(rctx.attr.registry_config)
+
+    url = dl.format(**{
+        "crate": crate_name,
+        "version": version,
+        "prefix": sharded_path(crate_name),
+        "lowerprefix": sharded_path(crate_name.lower()),
+        "sha256-checksum": sha256,
+    })
+
     rctx.download_and_extract(
-        rctx.attr.url,
+        url,
         type = "tar.gz",
-        canonical_id = get_default_canonical_id(rctx, urls = [rctx.attr.url]),
+        canonical_id = get_default_canonical_id(rctx, urls = [url]),
         headers = headers,
-        strip_prefix = rctx.attr.strip_prefix,
-        sha256 = rctx.attr.checksum,
+        strip_prefix = "%s-%s" % (crate_name, version),
+        sha256 = sha256,
     )
 
     patch(rctx)
@@ -35,11 +50,13 @@ def _crate_repository_impl(rctx):
 crate_repository = repository_rule(
     implementation = _crate_repository_impl,
     attrs = {
-        "url": attr.string(mandatory = True),
+        "crate_name": attr.string(mandatory = True),
+        "version": attr.string(mandatory = True),
         "cargo_config": attr.label(),
         "source": attr.string(),
         "use_home_cargo_credentials": attr.bool(),
         "checksum": attr.string(),
+        "registry_config": attr.label(allow_single_file = True, mandatory = True),
         "sbom_extra_qualifiers": attr.string_dict(),
     } | common_attrs,
 )

rs/private/downloader.bzl

@@ -1,10 +1,9 @@
 load("@bazel_tools//tools/build_defs/repo:git_worker.bzl", "git_repo")
-load(":cargo_credentials.bzl", "registry_auth_headers")
 load(":annotations.bzl", "annotation_for")
+load(":cargo_credentials.bzl", "registry_auth_headers")
+load(":registry_utils.bzl", "CRATES_IO_REGISTRY", "sharded_path")
 load(":toml2json.bzl", "run_toml2json")
 
-CRATES_IO_REGISTRY = "sparse+https://index.crates.io/"
-
 def parse_git_url(url):
     # Drop query params (?rev=...) and keep only before '#'
     parts = url.split("#")
@@ -33,22 +32,8 @@ def _github_source_to_raw_content_base_url(url):
 def _sanitize_path_fragment(path):
     return path.replace("/", "_").replace(":", "_")
 
-def sharded_path(crate):
-    # crates.io-index sharding rules (ASCII names)
-    n = len(crate)
-    if n == 0:
-        fail("empty crate name")
-    if n == 1:
-        return "1/" + crate
-    if n == 2:
-        return "2/" + crate
-    if n == 3:
-        return "3/%s/%s" % (crate[0], crate)
-    return "%s/%s/%s" % (crate[0:2], crate[2:4], crate)
-
 def new_downloader_state():
     return struct(
-        in_flight_sparse_registry_configs_by_source = {},
         in_flight_registry_fetches_by_crate = {},
         in_flight_git_crate_fetches_by_url = {},
         pending_git_clones_by_source = {},
@@ -108,25 +93,6 @@ def start_crate_registry_downloads(
         if source == "registry+https://github.com/rust-lang/crates.io-index":
             source = CRATES_IO_REGISTRY
             package["source"] = source
-            # We hardcode the response for crates.io to avoid a fetch in
-            # the common case when not using a custom registry.
-            # TODO(zbarsky): This could be solved more cleanly by using a repository rule
-            # to do these fetches, thus making them lazy.
-        elif source.startswith("sparse+") and source not in state.in_flight_sparse_registry_configs_by_source:
-            registry = source.removeprefix("sparse+")
-
-            state.in_flight_sparse_registry_configs_by_source[source] = mctx.download(
-                registry + "config.json",
-                _sanitize_path_fragment(source) + "config.json",
-                headers = registry_auth_headers(cargo_credentials, source),
-                block = False,
-            )
-
-    for package in packages:
-        source = package.get("source")
-        if not source:
-            continue
-
         name = package["name"]
         version = package["version"]
 
@@ -256,7 +222,8 @@ def _workspace_member_by_package_name_from_local_clone(mctx, workspace_cargo_tom
 
         member_cargo_toml_json = run_toml2json(
             mctx,
-            workspace_cargo_toml_path_str.replace("Cargo.toml", member + "/Cargo.toml"))
+            workspace_cargo_toml_path_str.replace("Cargo.toml", member + "/Cargo.toml"),
+        )
         package_name = member_cargo_toml_json.get("package", {}).get("name")
         if package_name:
             member_by_package_name[package_name] = member
@@ -389,26 +356,3 @@ def download_metadata_for_git_crates(
                 package["cargo_toml_json"] = cargo_toml_json
                 if cargo_toml_json.get("workspace"):
                     package["workspace_cargo_toml_json"] = cargo_toml_json
-
-def download_sparse_registry_configs(mctx, state):
-    # Hardcoded one to avoid the fetch...
-    sparse_registry_configs = {
-        CRATES_IO_REGISTRY: "https://static.crates.io/crates/{crate}/{version}/download",
-    }
-
-    for source, token in state.in_flight_sparse_registry_configs_by_source.items():
-        token.wait()
-        dl = json.decode(mctx.read(_sanitize_path_fragment(source) + "config.json"))["dl"]
-
-        if not (
-            "{crate}" in dl or
-            "{version}" in dl or
-            "{sha256-checksum}" in dl or
-            "{prefix}" in dl or
-            "{lowerprefix}" in dl
-        ):
-            dl += "/{crate}/{version}/download"
-
-        sparse_registry_configs[source] = dl
-
-    return sparse_registry_configs

rs/private/registry_config_repository.bzl

@@ -0,0 +1,42 @@
+load(":cargo_credentials.bzl", "load_cargo_credentials", "registry_auth_headers")
+
+def _registry_config_repository_impl(rctx):
+    # TODO(zbarsky): Is there a better way than fetching this in every crate repository?
+    if rctx.attr.use_home_cargo_credentials:
+        headers = registry_auth_headers(
+            load_cargo_credentials(rctx, rctx.attr.cargo_config),
+            rctx.attr.source,
+        )
+    else:
+        headers = {}
+
+    rctx.download(
+        rctx.attr.source.removeprefix("sparse+") + "config.json",
+        "config.json",
+        headers = headers,
+    )
+
+    dl = json.decode(rctx.read("config.json"))["dl"]
+    if not (
+        "{crate}" in dl or
+        "{version}" in dl or
+        "{sha256-checksum}" in dl or
+        "{prefix}" in dl or
+        "{lowerprefix}" in dl
+    ):
+        dl += "/{crate}/{version}/download"
+
+    rctx.file("dl", dl)
+    rctx.file("BUILD.bazel", "exports_files(['dl'])")
+
+    # Registry config can change upstream, so this repository is intentionally not reproducible.
+    return rctx.repo_metadata(reproducible = False)
+
+registry_config_repository = repository_rule(
+    implementation = _registry_config_repository_impl,
+    attrs = {
+        "source": attr.string(mandatory = True),
+        "cargo_config": attr.label(),
+        "use_home_cargo_credentials": attr.bool(),
+    },
+)

rs/private/registry_utils.bzl

@@ -0,0 +1,19 @@
+CRATES_IO_REGISTRY = "sparse+https://index.crates.io/"
+
+def registry_config_repo_name(hub_name, source):
+    return hub_name + "_" + registry_repo_name(source)
+
+def registry_repo_name(source):
+    return source.removeprefix("sparse+").replace(":", "_").replace("/", "_")
+
+def sharded_path(crate):
+    n = len(crate)
+    if n == 0:
+        fail("empty crate name")
+    if n == 1:
+        return "1/" + crate
+    if n == 2:
+        return "2/" + crate
+    if n == 3:
+        return "3/%s/%s" % (crate[0], crate)
+    return "%s/%s/%s" % (crate[0:2], crate[2:4], crate)