From df70c959a659c0e08bf3b46d56fc77244878215a Mon Sep 17 00:00:00 2001 From: Fabian Meumertzheim Date: Wed, 23 Aug 2023 11:52:07 +0200 Subject: [PATCH] Use a SkyFunction --- .../lib/bazel/BazelRepositoryModule.java | 4 +- .../devtools/build/lib/bazel/bzlmod/BUILD | 3 + .../bzlmod/BazelModuleResolutionFunction.java | 160 ++++++++---------- .../build/lib/bazel/bzlmod/IndexRegistry.java | 18 ++ .../build/lib/bazel/bzlmod/RepoSpec.java | 3 +- .../lib/bazel/bzlmod/RepoSpecFunction.java | 61 +++++++ .../build/lib/bazel/bzlmod/RepoSpecKey.java | 41 +++++ .../build/lib/skyframe/SkyFunctions.java | 1 + 8 files changed, 197 insertions(+), 94 deletions(-) create mode 100644 src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecFunction.java create mode 100644 src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecKey.java diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java index 9e952ecd5fdbf3..a730599e7bf2b1 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java +++ b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java @@ -47,6 +47,7 @@ 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.RepoSpec; +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; import com.google.devtools.build.lib.bazel.bzlmod.YankedVersionsUtil; @@ -278,7 +279,8 @@ SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace .addSkyFunction(SkyFunctions.BAZEL_MODULE_INSPECTION, new BazelModuleInspectorFunction()) .addSkyFunction(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction()) .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction) - .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction()); + .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction()) + .addSkyFunction(SkyFunctions.REPO_SPEC, new RepoSpecFunction()); filesystem = runtime.getFileSystem(); credentialModule = Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class)); diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD index 39afe3ea578fca..1969e0dcb6079b 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD +++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD @@ -25,6 +25,7 @@ java_library( ], deps = [ "//src/main/java/com/google/devtools/build/lib/cmdline", + "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects", "//src/main/java/net/starlark/java/eval", "//third_party:auto_value", "//third_party:gson", @@ -131,6 +132,7 @@ java_library( "MultipleVersionOverride.java", "NonRegistryOverride.java", "RegistryOverride.java", + "RepoSpecKey.java", "SingleExtensionEvalValue.java", "SingleExtensionUsagesValue.java", "SingleVersionOverride.java", @@ -173,6 +175,7 @@ java_library( "ModuleExtensionContext.java", "ModuleFileFunction.java", "ModuleFileGlobals.java", + "RepoSpecFunction.java", "Selection.java", "SingleExtensionEvalFunction.java", "SingleExtensionUsagesFunction.java", diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java index 13ab1aa91a04d8..c5067de66d00c7 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java +++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java @@ -15,10 +15,13 @@ package com.google.devtools.build.lib.bazel.bzlmod; +import static com.google.common.collect.ImmutableSet.toImmutableSet; + import com.google.common.base.Strings; import com.google.common.collect.ImmutableCollection; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Maps; import com.google.devtools.build.lib.analysis.BlazeVersionInfo; import com.google.devtools.build.lib.bazel.BazelVersion; @@ -28,7 +31,6 @@ import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.CheckDirectDepsMode; import com.google.devtools.build.lib.events.Event; import com.google.devtools.build.lib.events.EventHandler; -import com.google.devtools.build.lib.events.ExtendedEventHandler; import com.google.devtools.build.lib.profiler.Profiler; import com.google.devtools.build.lib.profiler.ProfilerTask; import com.google.devtools.build.lib.profiler.SilentCloseable; @@ -39,18 +41,11 @@ import com.google.devtools.build.skyframe.SkyFunctionException.Transience; import com.google.devtools.build.skyframe.SkyKey; import com.google.devtools.build.skyframe.SkyValue; -import java.io.IOException; +import com.google.devtools.build.skyframe.SkyframeLookupResult; import java.util.Map; import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import javax.annotation.Nullable; -import static com.google.common.collect.ImmutableMap.toImmutableMap; - /** * Discovers the whole dependency graph and runs selection algorithm on it to produce the pruned * dependency graph and runs checks on it. @@ -71,12 +66,56 @@ public SkyValue compute(SkyKey skyKey, Environment env) if (root == null) { return null; } + + var state = env.getState(ModuleResolutionComputeState::new); + if (state.selectionResult == null) { + state.selectionResult = discoverAndSelect(env, root); + if (state.selectionResult == null) { + return null; + } + } + + ImmutableSet repoSpecKeys = + state.selectionResult.getResolvedDepGraph().values().stream() + // Modules with a null registry have a non-registry override. We don't need to + // fetch or store the repo spec in this case. + .filter(module -> module.getRegistry() != null) + .map(RepoSpecKey::of) + .collect(toImmutableSet()); + SkyframeLookupResult repoSpecResults = env.getValuesAndExceptions(repoSpecKeys); + ImmutableMap.Builder remoteRepoSpecs = ImmutableMap.builder(); + for (RepoSpecKey repoSpecKey : repoSpecKeys) { + RepoSpec repoSpec = (RepoSpec) repoSpecResults.get(repoSpecKey); + if (repoSpec == null) { + return null; + } + remoteRepoSpecs.put(repoSpecKey.getModuleKey(), repoSpec); + } + + ImmutableMap finalDepGraph; + try (SilentCloseable c = + Profiler.instance().profile(ProfilerTask.BZLMOD, "compute final dep graph")) { + finalDepGraph = + computeFinalDepGraph( + state.selectionResult.getResolvedDepGraph(), + root.getOverrides(), + remoteRepoSpecs.buildOrThrow()); + } + + Profiler.instance().profile(ProfilerTask.BZLMOD, "module resolution completed").close(); + + return BazelModuleResolutionValue.create( + finalDepGraph, state.selectionResult.getUnprunedDepGraph()); + } + + private static Selection.Result discoverAndSelect(Environment env, RootModuleFileValue root) + throws BazelModuleResolutionFunctionException, InterruptedException { ImmutableMap initialDepGraph; try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "discovery")) { initialDepGraph = Discovery.run(env, root); - if (initialDepGraph == null) { - return null; - } + } + if (initialDepGraph == null) { + return null; } Selection.Result selectionResult; @@ -109,16 +148,7 @@ public SkyValue compute(SkyKey skyKey, Environment env) checkNoYankedVersions(resolvedDepGraph); } - ImmutableMap finalDepGraph; - try (SilentCloseable c = - Profiler.instance().profile(ProfilerTask.BZLMOD, "compute final dep graph")) { - finalDepGraph = - computeFinalDepGraph(resolvedDepGraph, root.getOverrides(), env.getListener()); - } - - Profiler.instance().profile(ProfilerTask.BZLMOD, "module resolution completed").close(); - - return BazelModuleResolutionValue.create(finalDepGraph, selectionResult.getUnprunedDepGraph()); + return selectionResult; } private static void verifyRootModuleDirectDepsAreAccurate( @@ -237,44 +267,12 @@ private static RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOv .build(); } - @Nullable - private static RepoSpec computeRepoSpec( - InterimModule interimModule, ModuleOverride override, ExtendedEventHandler eventHandler) - throws BazelModuleResolutionFunctionException, InterruptedException { - if (interimModule.getRegistry() == null) { - // This module has a non-registry override. We don't need to store the repo spec in this case. - return null; - } - try { - RepoSpec moduleRepoSpec = - interimModule - .getRegistry() - .getRepoSpec( - interimModule.getKey(), interimModule.getCanonicalRepoName(), eventHandler); - return maybeAppendAdditionalPatches(moduleRepoSpec, override); - } catch (IOException e) { - throw new BazelModuleResolutionFunctionException( - ExternalDepsException.withMessage( - Code.ERROR_ACCESSING_REGISTRY, - "Unable to get module repo spec from registry: %s", - e.getMessage()), - Transience.PERSISTENT); - } - } - /** * Builds a {@link Module} from an {@link InterimModule}, discarding unnecessary fields and adding * extra necessary ones (such as the repo spec). */ static Module moduleFromInterimModule( - InterimModule interim, ModuleOverride override, ExtendedEventHandler eventHandler) - throws BazelModuleResolutionFunctionException, InterruptedException { - RepoSpec repoSpec; - try (SilentCloseable c = - Profiler.instance() - .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + interim.getKey())) { - repoSpec = computeRepoSpec(interim, override, eventHandler); - } + InterimModule interim, ModuleOverride override, RepoSpec remoteRepoSpec) { return Module.builder() .setName(interim.getName()) .setVersion(interim.getVersion()) @@ -283,7 +281,7 @@ static Module moduleFromInterimModule( .setExecutionPlatformsToRegister(interim.getExecutionPlatformsToRegister()) .setToolchainsToRegister(interim.getToolchainsToRegister()) .setDeps(ImmutableMap.copyOf(Maps.transformValues(interim.getDeps(), DepSpec::toModuleKey))) - .setRepoSpec(repoSpec) + .setRepoSpec(maybeAppendAdditionalPatches(remoteRepoSpec, override)) .setExtensionUsages(interim.getExtensionUsages()) .build(); } @@ -291,43 +289,21 @@ static Module moduleFromInterimModule( private static ImmutableMap computeFinalDepGraph( ImmutableMap resolvedDepGraph, ImmutableMap overrides, - ExtendedEventHandler eventHandler) - throws BazelModuleResolutionFunctionException, InterruptedException { - ExecutorService executorService = Executors.newWorkStealingPool(8); - var future = - executorService.submit( - () -> - resolvedDepGraph.entrySet().stream() - .parallel() - .map( - e -> { - try { - return Maps.immutableEntry( - e.getKey(), - moduleFromInterimModule( - e.getValue(), - overrides.get(e.getKey().getName()), - eventHandler)); - } catch (BazelModuleResolutionFunctionException ex) { - throw new CompletionException(ex); - } catch (InterruptedException ex) { - throw new CompletionException(ex); - } - }) - .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue))); - try { - var result = future.get(); - executorService.shutdown(); - return result; - } catch (ExecutionException e) { - if (e.getCause() instanceof BazelModuleResolutionFunctionException) { - throw (BazelModuleResolutionFunctionException) e.getCause(); - } else if (e.getCause() instanceof InterruptedException) { - throw (InterruptedException) e.getCause(); - } else { - throw new IllegalStateException(e); - } + ImmutableMap remoteRepoSpecs) { + ImmutableMap.Builder finalDepGraph = ImmutableMap.builder(); + for (Map.Entry entry : resolvedDepGraph.entrySet()) { + finalDepGraph.put( + entry.getKey(), + moduleFromInterimModule( + entry.getValue(), + overrides.get(entry.getKey().getName()), + remoteRepoSpecs.get(entry.getKey()))); } + return finalDepGraph.buildOrThrow(); + } + + private static class ModuleResolutionComputeState implements Environment.SkyKeyComputeState { + Selection.Result selectionResult; } static class BazelModuleResolutionFunctionException extends SkyFunctionException { diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java index 1716c27d09ea70..900855d9121e33 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java +++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java @@ -17,6 +17,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; +import com.google.common.base.Objects; import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -65,6 +66,23 @@ public IndexRegistry(URI uri, DownloadManager downloadManager, Map "compute repo spec: " + module.getModuleKey())) { + return module + .getRegistry() + .getRepoSpec(module.getModuleKey(), module.getCanonicalRepoName(), 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", + module.getModuleKey())); + } + } + + static final class RepoSpecException extends SkyFunctionException { + + RepoSpecException(ExternalDepsException cause) { + super(cause, Transience.TRANSIENT); + } + } +} diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecKey.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecKey.java new file mode 100644 index 00000000000000..a634a1ff4bc364 --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecKey.java @@ -0,0 +1,41 @@ +package com.google.devtools.build.lib.bazel.bzlmod; + +import com.google.auto.value.AutoValue; +import com.google.devtools.build.lib.cmdline.RepositoryName; +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 RepoSpecFunction}. */ +@AutoCodec +@AutoValue +abstract class RepoSpecKey implements SkyKey { + private static final SkyKeyInterner interner = SkyKey.newInterner(); + + static RepoSpecKey of(InterimModule module) { + return create(module.getKey(), module.getRegistry(), module.getCanonicalRepoName()); + } + + abstract ModuleKey getModuleKey(); + + abstract Registry getRegistry(); + + abstract RepositoryName getCanonicalRepoName(); + + @AutoCodec.Instantiator + static RepoSpecKey create( + ModuleKey moduleKey, Registry registry, RepositoryName canonicalRepoName) { + return interner.intern(new AutoValue_RepoSpecKey(moduleKey, registry, canonicalRepoName)); + } + + @Override + public SkyFunctionName functionName() { + return SkyFunctions.REPO_SPEC; + } + + @Override + public SkyKeyInterner getSkyKeyInterner() { + return interner; + } +} diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java index 1ea238f7007e50..d914bf70940ae3 100644 --- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java +++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java @@ -158,6 +158,7 @@ public final class SkyFunctions { SkyFunctionName.createHermetic("BAZEL_DEP_GRAPH"); public static final SkyFunctionName BAZEL_LOCK_FILE = SkyFunctionName.createHermetic("BAZEL_LOCK_FILE"); + public static final SkyFunctionName REPO_SPEC = SkyFunctionName.createNonHermetic("REPO_SPEC"); public static Predicate isSkyFunction(SkyFunctionName functionName) { return key -> key.functionName().equals(functionName);