Skip to content

Commit

Permalink
When we're inlining BZL_LOAD and STARLARK_BUILTINS nodes, cache-a…
Browse files Browse the repository at this point in the history 
…nd-reuse the single `STARLARK_BUILTINS` node value rather than repeatedly computing it inline.

PiperOrigin-RevId: 635950392
Change-Id: If8a132792a6e97c6b3484b295405d59206d1afb1
  • Loading branch information
@haxorz
haxorz authored and Copybara-Service committed May 21, 2024
1 parent 6d840d1 commit 2c0faf9
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 28 deletions.
65 changes: 39 additions & 26 deletions src/main/java/com/google/devtools/build/lib/skyframe/BzlLoadFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import javax.annotation.Nullable;
import net.starlark.java.eval.EvalException;
Expand Down Expand Up @@ -118,20 +119,20 @@ public class BzlLoadFunction implements SkyFunction {
// comment in create() for rationale.
private final ValueGetter getter;

// Handles inlining of BzlLoadFunction calls.
@Nullable private final CachedBzlLoadDataManager cachedBzlLoadDataManager;
// Handles inlining of BzlLoadFunction and StarlarkBuiltinsFunction calls.
@Nullable final InlineCacheManager inlineCacheManager;

private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

private BzlLoadFunction(
RuleClassProvider ruleClassProvider,
BlazeDirectories directories,
ValueGetter getter,
@Nullable CachedBzlLoadDataManager cachedBzlLoadDataManager) {
@Nullable InlineCacheManager inlineCacheManager) {
this.ruleClassProvider = ruleClassProvider;
this.directories = directories;
this.getter = getter;
this.cachedBzlLoadDataManager = cachedBzlLoadDataManager;
this.inlineCacheManager = inlineCacheManager;
}

public static BzlLoadFunction create(
Expand Down Expand Up @@ -167,7 +168,7 @@ public static BzlLoadFunction create(
// (b) The memory overhead of the extra Skyframe node and edge per bzl file is pure
// waste.
new InliningAndCachingGetter(ruleClassProvider, hashFunction, bzlCompileCache),
/* cachedBzlLoadDataManager= */ null);
/* inlineCacheManager= */ null);
}

public static BzlLoadFunction createForInlining(
Expand All @@ -183,7 +184,7 @@ public static BzlLoadFunction createForInlining(
// of a BzlLoadValue inlining cache miss). This is important in the situation where a bzl
// file is loaded by a lot of other bzl files or BUILD files.
RegularSkyframeGetter.INSTANCE,
new CachedBzlLoadDataManager(bzlLoadValueCacheSize));
new InlineCacheManager(bzlLoadValueCacheSize));
}

@Override
Expand Down Expand Up @@ -262,7 +263,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
@Nullable
BzlLoadValue computeInline(BzlLoadValue.Key key, InliningState inliningState)
throws BzlLoadFailedException, InterruptedException {
Preconditions.checkNotNull(cachedBzlLoadDataManager);
Preconditions.checkNotNull(inlineCacheManager);
CachedBzlLoadData cachedData = computeInlineCachedData(key, inliningState);
return cachedData != null ? cachedData.getValue() : null;
}
Expand All @@ -283,7 +284,7 @@ private CachedBzlLoadData computeInlineCachedData(
// consistency purposes (see the javadoc of #computeInline).
CachedBzlLoadData cachedData = inliningState.successfulLoads.get(key);
if (cachedData == null) {
cachedData = cachedBzlLoadDataManager.cache.getIfPresent(key);
cachedData = inlineCacheManager.bzlLoadCache.getIfPresent(key);
if (cachedData != null) {
// Found a cache hit from another thread's computation. Register the cache hit's recorded
// deps as if we had requested them directly in the unwrapped environment. We do this for
Expand Down Expand Up @@ -315,7 +316,7 @@ private CachedBzlLoadData computeInlineCachedData(
} finally {
if (cachedData != null) {
inliningState.successfulLoads.put(key, cachedData);
cachedBzlLoadDataManager.cache.put(key, cachedData);
inlineCacheManager.bzlLoadCache.put(key, cachedData);
} else {
inliningState.unsuccessfulLoads.add(key);
// Either propagate an exception or fall through for null return.
Expand Down Expand Up @@ -346,7 +347,7 @@ private CachedBzlLoadData computeInlineForCacheMiss(
// Here we are at the boundary between one CachedBzlLoadData and the next. createChildState()
// unwraps the old recording env and starts a new one for a new node.

InliningState childState = inliningState.createChildState(cachedBzlLoadDataManager);
InliningState childState = inliningState.createChildState(inlineCacheManager);
childState.beginLoad(key); // track for cyclic load() detection
BzlLoadValue value;
try {
Expand All @@ -362,7 +363,7 @@ private CachedBzlLoadData computeInlineForCacheMiss(
}

public void resetInliningCache() {
cachedBzlLoadDataManager.reset();
inlineCacheManager.reset();
}

/**
Expand Down Expand Up @@ -472,8 +473,8 @@ static InliningState create(Environment env) {
* up to log dependency metadata into a CachedBzlLoadData node that is a child of this
* InliningState's node.
*/
private InliningState createChildState(CachedBzlLoadDataManager cachedBzlLoadDataManager) {
CachedBzlLoadData.Builder newBuilder = cachedBzlLoadDataManager.cachedDataBuilder();
private InliningState createChildState(InlineCacheManager inlineCacheManager) {
CachedBzlLoadData.Builder newBuilder = inlineCacheManager.cachedDataBuilder();
RecordingSkyFunctionEnvironment newRecordingEnv =
new RecordingSkyFunctionEnvironment(
recordingEnv.getDelegate(),
Expand Down Expand Up @@ -1513,33 +1514,45 @@ public void doneWithBzlCompileValue(BzlCompileValue.Key key) {
* Per-instance manager for {@link CachedBzlLoadData}, used when {@code BzlLoadFunction} calls are
* inlined.
*/
private static class CachedBzlLoadDataManager {
private final int cacheSize;
private Cache<BzlLoadValue.Key, CachedBzlLoadData> cache;
private CachedBzlLoadDataBuilderFactory cachedDataBuilderFactory =
static class InlineCacheManager {
private final int bzlLoadCacheSize;
private Cache<BzlLoadValue.Key, CachedBzlLoadData> bzlLoadCache;
private CachedBzlLoadDataBuilderFactory cachedBzlLoadDataBuilderFactory =
new CachedBzlLoadDataBuilderFactory();

private CachedBzlLoadDataManager(int cacheSize) {
this.cacheSize = cacheSize;
final AtomicReference<StarlarkBuiltinsValue> builtinsRef = new AtomicReference<>();

private InlineCacheManager(int bzlLoadCacheSize) {
this.bzlLoadCacheSize = bzlLoadCacheSize;
}

private CachedBzlLoadData.Builder cachedDataBuilder() {
return cachedDataBuilderFactory.newCachedBzlLoadDataBuilder();
return cachedBzlLoadDataBuilderFactory.newCachedBzlLoadDataBuilder();
}

private void reset() {
if (cache != null) {
logger.atInfo().log("Starlark inlining cache stats from earlier build: %s", cache.stats());
if (bzlLoadCache != null) {
logger.atInfo().log(
"Starlark inlining cache stats from earlier build: %s", bzlLoadCache.stats());
}
cachedDataBuilderFactory = new CachedBzlLoadDataBuilderFactory();
cachedBzlLoadDataBuilderFactory = new CachedBzlLoadDataBuilderFactory();
Preconditions.checkState(
cacheSize >= 0, "Expected positive Starlark cache size if caching. %s", cacheSize);
cache =
bzlLoadCacheSize >= 0,
"Expected positive Starlark cache size if caching. %s",
bzlLoadCacheSize);
bzlLoadCache =
Caffeine.newBuilder()
.initialCapacity(BlazeInterners.concurrencyLevel())
.maximumSize(cacheSize)
.maximumSize(bzlLoadCacheSize)
.recordStats()
.build();
// Don't reset `builtinsRef`.
//
// All usages of BzlLoadFunction inlining assume builtins can never change (i.e. no usage of
// --experimental_builtins_bzl_path, no inter-invocation flipping of Starlark semantics
// options that'd cause evaluation of builtins to differ). If this assumption ever doesn't
// hold, we'd want to reset `builtinsRef` and we'd also want to inline deps on the logical
// Skyframe subgraph when we get a `builtins` cache hit.
}
}

Expand Down
36 changes: 34 additions & 2 deletions src/main/java/com/google/devtools/build/lib/skyframe/StarlarkBuiltinsFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,14 @@

package com.google.devtools.build.lib.skyframe;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.packages.BazelStarlarkEnvironment;
import com.google.devtools.build.lib.packages.BazelStarlarkEnvironment.InjectionException;
import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
import com.google.devtools.build.lib.skyframe.BzlLoadFunction.InlineCacheManager;
import com.google.devtools.build.skyframe.RecordingSkyFunctionEnvironment;
import com.google.devtools.build.skyframe.SkyFunction;
import com.google.devtools.build.skyframe.SkyFunctionException;
Expand Down Expand Up @@ -60,6 +63,16 @@
*
* <p>This function has a trivial key, so there can only be one value in the build at a time. It has
* a single dependency, on the result of evaluating the exports.bzl file to a {@link BzlLoadValue}.
*
* <p>This function supports a special "inlining" mode, similar to {@link BzlLoadFunction} (see that
* class's javadoc and code comments). Whenever we inline {@link BzlLoadFunction} we also inline
* {@link StarlarkBuiltinsFunction} (and {@link StarlarkBuiltinsFunction}'s calls to {@link
* BzlLoadFunction} are then themselves inlined!). Similar to {@link BzlLoadFunction}'s inlining, we
* cache the result of this computation, and this caching is managed by {@link
* BzlLoadFunction.InlineCacheManager}. But since there's only a single {@link
* StarlarkBuiltinsValue} node and we don't need to worry about that node's value changing at future
* invocations or subsequent versions (see {@link InlineCacheManager#reset} for why), our caching
* strategy is much simpler and we don't need to bother inlining deps of the Skyframe subgraph.
*/
public class StarlarkBuiltinsFunction implements SkyFunction {

Expand Down Expand Up @@ -117,11 +130,30 @@ public static StarlarkBuiltinsValue computeInline(
BazelStarlarkEnvironment bazelStarlarkEnvironment,
BzlLoadFunction bzlLoadFunction)
throws BuiltinsFailedException, InterruptedException {
checkNotNull(bzlLoadFunction.inlineCacheManager);
StarlarkBuiltinsValue cachedBuiltins = bzlLoadFunction.inlineCacheManager.builtinsRef.get();
if (cachedBuiltins != null) {
// See the comment in InlineCacheManager#reset for why it's sound to not inline deps of the
// entire subgraph here.
return cachedBuiltins;
}

// See BzlLoadFunction#computeInline and BzlLoadFunction.InliningState for an explanation of the
// inlining mechanism and its invariants. For our purposes, the Skyframe environment to use
// comes from inliningState.
return computeInternal(
inliningState.getEnvironment(), bazelStarlarkEnvironment, inliningState, bzlLoadFunction);
StarlarkBuiltinsValue computedBuiltins =
computeInternal(
inliningState.getEnvironment(),
bazelStarlarkEnvironment,
inliningState,
bzlLoadFunction);
if (computedBuiltins == null) {
return null;
}
// There's a benign race where multiple threads may try to compute-and-cache the single builtins
// value. Ensure the value computed by winner of that race gets used by everyone.
bzlLoadFunction.inlineCacheManager.builtinsRef.compareAndSet(null, computedBuiltins);
return bzlLoadFunction.inlineCacheManager.builtinsRef.get();
}

// bzlLoadFunction and inliningState are non-null iff using inlining code path.
Expand Down

0 comments on commit 2c0faf9

Please sign in to comment.