Extract Bazel registry creation into a SkyFunction

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -47,6 +47,7 @@
 import com.google.devtools.build.lib.bazel.bzlmod.NonRegistryOverride;
 import com.google.devtools.build.lib.bazel.bzlmod.RegistryFactory;
 import com.google.devtools.build.lib.bazel.bzlmod.RegistryFactoryImpl;
+import com.google.devtools.build.lib.bazel.bzlmod.RegistryFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.RepoSpecFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionEvalFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionUsagesFunction;
@@ -243,9 +244,6 @@ public void workspaceInit(
             clientEnvironmentSupplier,
             directories,
             BazelSkyframeExecutorConstants.EXTERNAL_PACKAGE_HELPER);
-    RegistryFactory registryFactory =
-        new RegistryFactoryImpl(
-            directories.getWorkspace(), downloadManager, clientEnvironmentSupplier);
     singleExtensionEvalFunction =
         new SingleExtensionEvalFunction(directories, clientEnvironmentSupplier, downloadManager);
 
@@ -259,7 +257,6 @@ public void workspaceInit(
             SkyFunctions.MODULE_FILE,
             new ModuleFileFunction(
                 runtime.getRuleClassProvider().getBazelStarlarkEnvironment(),
-                registryFactory,
                 directories.getWorkspace(),
                 builtinModules))
         .addSkyFunction(SkyFunctions.BAZEL_DEP_GRAPH, new BazelDepGraphFunction())
@@ -271,8 +268,13 @@ SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction())
         .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction)
         .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
-        .addSkyFunction(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
-        .addSkyFunction(SkyFunctions.YANKED_VERSIONS, new YankedVersionsFunction(registryFactory))
+        .addSkyFunction(
+            SkyFunctions.REGISTRY,
+            new RegistryFunction(
+                new RegistryFactoryImpl(
+                    directories.getWorkspace(), downloadManager, clientEnvironmentSupplier)))
+        .addSkyFunction(SkyFunctions.REPO_SPEC, new RepoSpecFunction())
+        .addSkyFunction(SkyFunctions.YANKED_VERSIONS, new YankedVersionsFunction())
         .addSkyFunction(
             SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
             new ModuleExtensionRepoMappingEntriesFunction());

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -87,7 +87,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
-        "//third_party:caffeine",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//third_party:gson",
         "//third_party:guava",
     ],
@@ -139,6 +139,7 @@ java_library(
         "ModuleThreadContext.java",
         "MultipleVersionOverride.java",
         "NonRegistryOverride.java",
+        "RegistryKey.java",
         "RegistryOverride.java",
         "RepoSpecKey.java",
         "SingleExtensionEvalValue.java",
@@ -191,6 +192,7 @@ java_library(
         "ModuleExtensionEvaluationProgress.java",
         "ModuleExtensionRepoMappingEntriesFunction.java",
         "ModuleFileFunction.java",
+        "RegistryFunction.java",
         "RepoSpecFunction.java",
         "Selection.java",
         "SingleExtensionEvalFunction.java",

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java

@@ -55,8 +55,6 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.errorprone.annotations.FormatMethod;
 import java.io.IOException;
-import java.net.URISyntaxException;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -87,7 +85,6 @@ public class ModuleFileFunction implements SkyFunction {
       new Precomputed<>("module_overrides");
 
   private final BazelStarlarkEnvironment starlarkEnv;
-  private final RegistryFactory registryFactory;
   private final Path workspaceRoot;
   private final ImmutableMap<String, NonRegistryOverride> builtinModules;
 
@@ -107,11 +104,9 @@ public class ModuleFileFunction implements SkyFunction {
    */
   public ModuleFileFunction(
       BazelStarlarkEnvironment starlarkEnv,
-      RegistryFactory registryFactory,
       Path workspaceRoot,
       ImmutableMap<String, NonRegistryOverride> builtinModules) {
     this.starlarkEnv = starlarkEnv;
-    this.registryFactory = registryFactory;
     this.workspaceRoot = workspaceRoot;
     this.builtinModules = builtinModules;
   }
@@ -478,14 +473,18 @@ private GetModuleFileResult getModuleFile(
               "unrecognized override type %s for module %s",
               override.getClass().getSimpleName(), key));
     }
-    List<Registry> registryObjects = new ArrayList<>(registries.size());
-    for (String registryUrl : registries) {
-      try {
-        registryObjects.add(registryFactory.getRegistryWithUrl(registryUrl));
-      } catch (URISyntaxException e) {
-        throw errorf(Code.INVALID_REGISTRY_URL, e, "Invalid registry URL");
-      }
+
+    List<RegistryKey> registryKeys =
+        registries.stream().map(RegistryKey::create).collect(ImmutableList.toImmutableList());
+    var registryResult = env.getValuesAndExceptions(registryKeys);
+    if (env.valuesMissing()) {
+      return null;
     }
+    List<Registry> registryObjects =
+        registryKeys.stream()
+            .map(registryResult::get)
+            .map(Registry.class::cast)
+            .collect(ImmutableList.toImmutableList());
 
     // Now go through the list of registries and use the first one that contains the requested
     // module.

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Registry.java

@@ -17,11 +17,12 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.skyframe.SkyValue;
 import java.io.IOException;
 import java.util.Optional;
 
 /** A database where module metadata is stored. */
-public interface Registry {
+public interface Registry extends SkyValue {
 
   /** The URL that uniquely identifies the registry. */
   String getUrl();

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFactory.java

@@ -20,6 +20,10 @@
 /** A factory type for {@link Registry}. */
 public interface RegistryFactory {
 
-  /** Returns the unique registry associated with the given URL. */
-  Registry getRegistryWithUrl(String url) throws URISyntaxException;
+  /**
+   * Creates a registry associated with the given URL.
+   *
+   * <p>Outside of tests, only {@link RegistryFunction} should call this method.
+   */
+  Registry createRegistry(String url) throws URISyntaxException;
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFactoryImpl.java

@@ -15,8 +15,6 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
-import com.github.benmanes.caffeine.cache.Cache;
-import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
 import com.google.devtools.build.lib.vfs.Path;
 import java.net.URI;
@@ -29,7 +27,6 @@ public class RegistryFactoryImpl implements RegistryFactory {
   private final Path workspacePath;
   private final DownloadManager downloadManager;
   private final Supplier<Map<String, String>> clientEnvironmentSupplier;
-  private final Cache<String, Registry> registries = Caffeine.newBuilder().build();
 
   public RegistryFactoryImpl(
       Path workspacePath,
@@ -41,7 +38,7 @@ public RegistryFactoryImpl(
   }
 
   @Override
-  public Registry getRegistryWithUrl(String unresolvedUrl) throws URISyntaxException {
+  public Registry createRegistry(String unresolvedUrl) throws URISyntaxException {
     URI uri = new URI(unresolvedUrl.replace("%workspace%", workspacePath.getPathString()));
     if (uri.getScheme() == null) {
       throw new URISyntaxException(
@@ -55,17 +52,10 @@ public Registry getRegistryWithUrl(String unresolvedUrl) throws URISyntaxExcepti
           "Registry URL path is not valid -- did you mean to use file:///foo/bar "
               + "or file:///c:/foo/bar for Windows?");
     }
-    switch (uri.getScheme()) {
-      case "http":
-      case "https":
-      case "file":
-        return registries.get(
-            unresolvedUrl,
-            unused ->
-                new IndexRegistry(
-                    uri, unresolvedUrl, downloadManager, clientEnvironmentSupplier.get()));
-      default:
-        throw new URISyntaxException(uri.toString(), "Unrecognized registry URL protocol");
-    }
+    return switch (uri.getScheme()) {
+      case "http", "https", "file" ->
+          new IndexRegistry(uri, unresolvedUrl, downloadManager, clientEnvironmentSupplier.get());
+      default -> throw new URISyntaxException(uri.toString(), "Unrecognized registry URL protocol");
+    };
   }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFunction.java

@@ -0,0 +1,57 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package com.google.devtools.build.lib.bazel.bzlmod;
+
+import com.google.devtools.build.lib.server.FailureDetails;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionException;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+import java.net.URISyntaxException;
+import javax.annotation.Nullable;
+
+/** A simple SkyFunction that creates a {@link Registry} with a given URL. */
+public class RegistryFunction implements SkyFunction {
+  private final RegistryFactory registryFactory;
+
+  public RegistryFunction(RegistryFactory registryFactory) {
+    this.registryFactory = registryFactory;
+  }
+
+  @Override
+  @Nullable
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws InterruptedException, RegistryException {
+    RegistryKey key = (RegistryKey) skyKey.argument();
+    try {
+      return registryFactory.createRegistry(key.getUrl());
+    } catch (URISyntaxException e) {
+      throw new RegistryException(
+          ExternalDepsException.withCauseAndMessage(
+              FailureDetails.ExternalDeps.Code.INVALID_REGISTRY_URL,
+              e,
+              "Invalid registry URL: %s",
+              key.getUrl()));
+    }
+  }
+
+  static final class RegistryException extends SkyFunctionException {
+
+    RegistryException(ExternalDepsException cause) {
+      super(cause, Transience.TRANSIENT);
+    }
+  }
+}

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryKey.java

@@ -0,0 +1,46 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package com.google.devtools.build.lib.bazel.bzlmod;
+
+import com.google.auto.value.AutoValue;
+import com.google.devtools.build.lib.skyframe.SkyFunctions;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.devtools.build.skyframe.SkyKey;
+
+/** The key for {@link RegistryFunction}. */
+@AutoCodec
+@AutoValue
+abstract class RegistryKey implements SkyKey {
+  private static final SkyKeyInterner<RegistryKey> interner = SkyKey.newInterner();
+
+  abstract String getUrl();
+
+  @AutoCodec.Instantiator
+  static RegistryKey create(String url) {
+    return interner.intern(new AutoValue_RegistryKey(url));
+  }
+
+  @Override
+  public SkyFunctionName functionName() {
+    return SkyFunctions.REGISTRY;
+  }
+
+  @Override
+  public SkyKeyInterner<RegistryKey> getSkyKeyInterner() {
+    return interner;
+  }
+}

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecFunction.java

@@ -24,42 +24,36 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.IOException;
-import java.net.URISyntaxException;
 import javax.annotation.Nullable;
 
 /**
  * A simple SkyFunction that computes a {@link RepoSpec} for the given {@link InterimModule} by
  * fetching required information from its {@link Registry}.
  */
 public class RepoSpecFunction implements SkyFunction {
-  private final RegistryFactory registryFactory;
-
-  public RepoSpecFunction(RegistryFactory registryFactory) {
-    this.registryFactory = registryFactory;
-  }
 
   @Override
   @Nullable
   public SkyValue compute(SkyKey skyKey, Environment env)
       throws InterruptedException, RepoSpecException {
     RepoSpecKey key = (RepoSpecKey) skyKey.argument();
+
+    Registry registry = (Registry) env.getValue(RegistryKey.create(key.getRegistryUrl()));
+    if (registry == null) {
+      return null;
+    }
+
     try (SilentCloseable c =
         Profiler.instance()
             .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + key.getModuleKey())) {
-      return registryFactory
-          .getRegistryWithUrl(key.getRegistryUrl())
-          .getRepoSpec(key.getModuleKey(), env.getListener());
+      return registry.getRepoSpec(key.getModuleKey(), env.getListener());
     } catch (IOException e) {
       throw new RepoSpecException(
           ExternalDepsException.withCauseAndMessage(
               FailureDetails.ExternalDeps.Code.ERROR_ACCESSING_REGISTRY,
               e,
               "Unable to get module repo spec for %s from registry",
               key.getModuleKey()));
-    } catch (URISyntaxException e) {
-      // This should never happen since we obtain the registry URL from an already constructed
-      // registry.
-      throw new IllegalStateException(e);
     }
   }