Skip to content

Commit

Permalink
Intern module extension repo mapping entries via a SkyFunction
Browse files Browse the repository at this point in the history
Since 74aadb2, the entries of `RepositoryMapping`s are interned to prevent quadratic memory usage for module extensions that create many repos. By using a dedicated `SkyFunction` instead of an interner, the CPU time spent on computing the identical repository mappings for extension repos, which showed up in profiles, is saved in addition to the memory usage.

Closes #20095.

PiperOrigin-RevId: 581177545
Change-Id: Ied5682ba90f5f771fedffbea1491e7b2f6c27145
  • Loading branch information
fmeum authored and copybara-github committed Nov 10, 2023
1 parent e2c0276 commit 07a571f
Show file tree
Hide file tree
Showing 16 changed files with 203 additions and 51 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorValue.AugmentedModule.ResolutionReason;
import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionFunction;
import com.google.devtools.build.lib.bazel.bzlmod.LocalPathOverride;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionRepoMappingEntriesFunction;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileFunction;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleOverride;
import com.google.devtools.build.lib.bazel.bzlmod.NonRegistryOverride;
Expand Down Expand Up @@ -290,7 +291,10 @@ 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.REPO_SPEC, new RepoSpecFunction(registryFactory))
.addSkyFunction(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction());
filesystem = runtime.getFileSystem();

credentialModule = Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ java_library(
"ModuleBase.java",
"ModuleExtensionEvalFactors.java",
"ModuleExtensionEvalStarlarkThreadContext.java",
"ModuleExtensionRepoMappingEntriesValue.java",
"ModuleExtensionResolutionEvent.java",
"ModuleFileValue.java",
"ModuleOverride.java",
Expand Down Expand Up @@ -182,6 +183,7 @@ java_library(
"GsonTypeAdapterUtil.java",
"ModuleExtensionContext.java",
"ModuleExtensionEvaluationProgress.java",
"ModuleExtensionRepoMappingEntriesFunction.java",
"ModuleFileFunction.java",
"ModuleFileGlobals.java",
"RepoSpecFunction.java",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
// Copyright 2023 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.common.collect.ImmutableMap;
import com.google.devtools.build.lib.cmdline.RepositoryMapping;
import com.google.devtools.build.lib.cmdline.RepositoryName;
import com.google.devtools.build.skyframe.SkyFunction;
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;
import javax.annotation.Nullable;

/**
* A SkyFunction that computes the {@link RepositoryMapping#entries()} for the repos generated by a
* module extension. While the repos differ in {@link RepositoryMapping#ownerRepo()}, the entries
* are identical and sharing the instance ensures that memory and CPU usage are only linear in the
* number of repos.
*/
public class ModuleExtensionRepoMappingEntriesFunction implements SkyFunction {

@Override
@Nullable
public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedException {
BazelDepGraphValue bazelDepGraphValue =
(BazelDepGraphValue) env.getValue(BazelDepGraphValue.KEY);
if (bazelDepGraphValue == null) {
return null;
}

var moduleExtensionId = ((ModuleExtensionRepoMappingEntriesValue.Key) skyKey).argument();
var extensionEvalValue =
(SingleExtensionEvalValue) env.getValue(SingleExtensionEvalValue.key(moduleExtensionId));
if (extensionEvalValue == null) {
return null;
}

return computeRepoMappingEntries(moduleExtensionId, extensionEvalValue, bazelDepGraphValue);
}

/**
* Calculates repo mappings for a repo generated from a module extension. Such a repo can see all
* repos generated by the same module extension, as well as all repos that the Bazel module
* hosting the extension can see (see above).
*/
private ModuleExtensionRepoMappingEntriesValue computeRepoMappingEntries(
ModuleExtensionId extensionId,
SingleExtensionEvalValue extensionEvalValue,
BazelDepGraphValue bazelDepGraphValue) {
// Find the key of the module containing this extension. This will be used to compute additional
// mappings -- any repo generated by an extension contained in the module "foo" can additionally
// see all repos that "foo" can see.
ModuleKey moduleKey =
bazelDepGraphValue
.getCanonicalRepoNameLookup()
.get(extensionId.getBzlFileLabel().getRepository());
// NOTE(wyv): This means that if "foo" has a bazel_dep with the repo name "bar", and the
// extension generates an internal repo name "bar", then within a repo generated by the
// extension, "bar" will refer to the latter. We should explore a way to differentiate between
// the two to avoid any surprises.
ImmutableMap.Builder<String, RepositoryName> entries = ImmutableMap.builder();
entries.putAll(bazelDepGraphValue.getFullRepoMapping(moduleKey).entries());
entries.putAll(extensionEvalValue.getCanonicalRepoNameToInternalNames().inverse());
return ModuleExtensionRepoMappingEntriesValue.create(entries.buildKeepingLast(), moduleKey);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
// Copyright 2023 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.common.collect.ImmutableMap;
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.AbstractSkyKey;
import com.google.devtools.build.skyframe.SkyFunctionName;
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;

/** The value for {@link ModuleExtensionRepoMappingEntriesFunction}. */
@AutoCodec
@AutoValue
public abstract class ModuleExtensionRepoMappingEntriesValue implements SkyValue {

public abstract ImmutableMap<String, RepositoryName> getEntries();

public abstract ModuleKey getModuleKey();

@AutoCodec.Instantiator
public static ModuleExtensionRepoMappingEntriesValue create(
ImmutableMap<String, RepositoryName> entries, ModuleKey moduleKey) {
return new AutoValue_ModuleExtensionRepoMappingEntriesValue(entries, moduleKey);
}

public static ModuleExtensionRepoMappingEntriesValue.Key key(ModuleExtensionId id) {
return ModuleExtensionRepoMappingEntriesValue.Key.create(id);
}

/**
* The {@link com.google.devtools.build.skyframe.SkyKey} of a {@link
* ModuleExtensionRepoMappingEntriesValue}.
*/
@AutoCodec
public static class Key extends AbstractSkyKey<ModuleExtensionId> {

private static final SkyKeyInterner<ModuleExtensionRepoMappingEntriesValue.Key> interner =
SkyKey.newInterner();

protected Key(ModuleExtensionId arg) {
super(arg);
}

@AutoCodec.Instantiator
static ModuleExtensionRepoMappingEntriesValue.Key create(ModuleExtensionId arg) {
return interner.intern(new ModuleExtensionRepoMappingEntriesValue.Key(arg));
}

@Override
public SkyFunctionName functionName() {
return SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES;
}

@Override
public SkyKeyInterner<ModuleExtensionRepoMappingEntriesValue.Key> getSkyKeyInterner() {
return interner;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,6 @@
import com.google.auto.value.AutoValue;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Interner;
import com.google.devtools.build.lib.concurrent.BlazeInterners;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
Expand All @@ -36,15 +34,6 @@
@AutoValue
public abstract class RepositoryMapping {

// All repositories generated by a module extension have the same repository mapping
// entries. Without interning, if a module extension generates N repositories, each such
// repository would have its own copy with (more than) N entries, resulting in memory usage that
// is quadratic in N.
// Note: We intern the entries, not the RepositoryMapping itself, because the latter also includes
// the owner repo, which differs between extension repos.
private static final Interner<ImmutableMap<String, RepositoryName>> ENTRIES_INTERNER =
BlazeInterners.newWeakInterner();

// Always fallback to the requested name
public static final RepositoryMapping ALWAYS_FALLBACK = createAllowingFallback(ImmutableMap.of());

Expand All @@ -71,8 +60,7 @@ public static RepositoryMapping createAllowingFallback(Map<String, RepositoryNam

private static RepositoryMapping createInternal(
Map<String, RepositoryName> entries, RepositoryName ownerRepo) {
return new AutoValue_RepositoryMapping(
ENTRIES_INTERNER.intern(ImmutableMap.copyOf(entries)), ownerRepo);
return new AutoValue_RepositoryMapping(ImmutableMap.copyOf(entries), ownerRepo);
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@
import com.google.devtools.build.lib.bazel.bzlmod.BazelDepGraphValue;
import com.google.devtools.build.lib.bazel.bzlmod.Module;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionId;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionRepoMappingEntriesValue;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleKey;
import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionEvalValue;
import com.google.devtools.build.lib.bazel.bzlmod.Version;
import com.google.devtools.build.lib.cmdline.LabelConstants;
import com.google.devtools.build.lib.cmdline.RepositoryMapping;
Expand Down Expand Up @@ -141,14 +141,16 @@ public SkyValue compute(SkyKey skyKey, Environment env)
maybeGetModuleExtensionForRepo(repositoryName, bazelDepGraphValue);

if (moduleExtensionId.isPresent()) {
SingleExtensionEvalValue extensionEvalValue =
(SingleExtensionEvalValue)
env.getValue(SingleExtensionEvalValue.key(moduleExtensionId.get()));
if (extensionEvalValue == null) {
var repoMappingEntriesValue =
(ModuleExtensionRepoMappingEntriesValue)
env.getValue(ModuleExtensionRepoMappingEntriesValue.key(moduleExtensionId.get()));
if (repoMappingEntriesValue == null) {
return null;
}
return computeForModuleExtensionRepo(
repositoryName, moduleExtensionId.get(), extensionEvalValue, bazelDepGraphValue);
return RepositoryMappingValue.createForBzlmodRepo(
RepositoryMapping.create(repoMappingEntriesValue.getEntries(), repositoryName),
repoMappingEntriesValue.getModuleKey().getName(),
repoMappingEntriesValue.getModuleKey().getVersion());
}
}

Expand Down Expand Up @@ -192,36 +194,6 @@ private Optional<RepositoryMappingValue> computeForBazelModuleRepo(
module.getVersion()));
}

/**
* Calculates repo mappings for a repo generated from a module extension. Such a repo can see all
* repos generated by the same module extension, as well as all repos that the Bazel module
* hosting the extension can see (see above).
*/
private RepositoryMappingValue computeForModuleExtensionRepo(
RepositoryName repositoryName,
ModuleExtensionId extensionId,
SingleExtensionEvalValue extensionEvalValue,
BazelDepGraphValue bazelDepGraphValue) {
// Find the key of the module containing this extension. This will be used to compute additional
// mappings -- any repo generated by an extension contained in the module "foo" can additionally
// see all repos that "foo" can see.
ModuleKey moduleKey =
bazelDepGraphValue
.getCanonicalRepoNameLookup()
.get(extensionId.getBzlFileLabel().getRepository());
Module module = bazelDepGraphValue.getDepGraph().get(moduleKey);
// NOTE(wyv): This means that if "foo" has a bazel_dep with the repo name "bar", and the
// extension generates an internal repo name "bar", then within a repo generated by the
// extension, "bar" will refer to the latter. We should explore a way to differentiate between
// the two to avoid any surprises.
return RepositoryMappingValue.createForBzlmodRepo(
RepositoryMapping.create(
extensionEvalValue.getCanonicalRepoNameToInternalNames().inverse(), repositoryName)
.withAdditionalMappings(bazelDepGraphValue.getFullRepoMapping(moduleKey)),
module.getName(),
module.getVersion());
}

/**
* Calculate repo mappings for a repo generated from WORKSPACE. Such a repo is not subject to
* strict deps, and can additionally see all repos that the root module can see.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,9 @@ public final class SkyFunctions {
SkyFunctionName.createHermetic("BAZEL_FETCH_ALL");
public static final SkyFunctionName REPO_SPEC = SkyFunctionName.createNonHermetic("REPO_SPEC");

public static final SkyFunctionName MODULE_EXTENSION_REPO_MAPPING_ENTRIES =
SkyFunctionName.createHermetic("MODULE_EXTENSION_REPO_MAPPING_ENTRIES");

public static Predicate<SkyKey> isSkyFunction(SkyFunctionName functionName) {
return key -> key.functionName().equals(functionName);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.google.devtools.build.lib.bazel.bzlmod.BazelLockFileFunction;
import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionFunction;
import com.google.devtools.build.lib.bazel.bzlmod.FakeRegistry;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionRepoMappingEntriesFunction;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileFunction;
import com.google.devtools.build.lib.bazel.bzlmod.NonRegistryOverride;
import com.google.devtools.build.lib.bazel.bzlmod.RepoSpecFunction;
Expand Down Expand Up @@ -176,6 +177,9 @@ public ImmutableMap<SkyFunctionName, SkyFunction> getSkyFunctions(BlazeDirectori
new SingleExtensionEvalFunction(directories, ImmutableMap::of, downloadManager))
.put(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(FakeRegistry.DEFAULT_FACTORY))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,9 @@ public void setup() throws Exception {
.put(SkyFunctions.BAZEL_DEP_GRAPH, new BazelDepGraphFunction())
.put(SkyFunctions.BAZEL_MODULE_RESOLUTION, resolutionFunctionMock)
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,9 @@ public void setup() throws Exception {
new ModuleFileFunction(registryFactory, rootDirectory, ImmutableMap.of()))
.put(SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(rootDirectory))
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,9 @@ public void setup() throws Exception {
.put(SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(rootDirectory))
.put(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction())
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,9 @@ public void setup() throws Exception {
SkyFunctions.MODULE_FILE,
new ModuleFileFunction(registryFactory, workspaceRoot, ImmutableMap.of()))
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,9 @@ private void setUpWithBuiltinModules(ImmutableMap<String, NonRegistryOverride> b
BzlmodRepoRuleValue.BZLMOD_REPO_RULE,
new BzlmodRepoRuleFunction(ruleClassProvider, directories))
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,9 @@ public void setup() throws Exception {
.put(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
.put(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction)
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,9 @@ private void setUpWithBuiltinModules(ImmutableMap<String, NonRegistryOverride> b
BzlmodRepoRuleValue.BZLMOD_REPO_RULE,
new BzlmodRepoRuleFunction(ruleClassProvider, directories))
.put(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory))
.put(
SkyFunctions.MODULE_EXTENSION_REPO_MAPPING_ENTRIES,
new ModuleExtensionRepoMappingEntriesFunction())
.put(
SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE,
new ClientEnvironmentFunction(new AtomicReference<>(ImmutableMap.of())))
Expand Down
Loading

0 comments on commit 07a571f

Please sign in to comment.