From de34469a24d7ef7e91a46a1b9d501051c06784c9 Mon Sep 17 00:00:00 2001 From: Googler Date: Tue, 28 Jun 2022 08:08:05 -0700 Subject: [PATCH] Begin to factor BazelStarlarkContext into subclasses BazelStarlarkContext is a hodgepodge of Bazel-specific state needed for or manipulated by various Starlark evaluations. A minor refactoring goal is to break this class up into subclasses, each one representing a different kind of Starlark evaluation that needs different kinds of data. This avoids the proliferation of new @Nullable fields that may or may not be set and used for a given evaluation. This CL begins the process by breaking out the case of initializing a .bzl module (i.e. running its top-level code). The other cases should follow a similar pathway, but I don't intend to pursue that in the immediate future. My motivation for tackling .bzl initialization now is that I want to add a field needed for the .bzl visibility without feeling guilty. I also took this opportunity to clarify that BazelStarlarkContext ought not to be shared or reused between multiple StarlarkThreads, and fixed a location where this was violated by StarlarkCallbackHelper. Work towards #11261. PiperOrigin-RevId: 457732549 Change-Id: I3ee869865c47457275f1b84819f0f154290aa22a --- .../analysis/ConfiguredRuleClassProvider.java | 6 +- .../StarlarkDefinedConfigTransition.java | 2 +- .../starlark/StarlarkRuleClassFunctions.java | 4 - .../lib/packages/BazelStarlarkContext.java | 74 ++++++++++++++----- .../lib/packages/BzlInitThreadContext.java | 72 ++++++++++++++++++ .../lib/packages/StarlarkCallbackHelper.java | 47 ++++++++---- 6 files changed, 163 insertions(+), 42 deletions(-) create mode 100644 src/main/java/com/google/devtools/build/lib/packages/BzlInitThreadContext.java diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java index b3ff10d8883096..4eb4582552be7c 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java @@ -47,7 +47,7 @@ import com.google.devtools.build.lib.cmdline.RepositoryName; import com.google.devtools.build.lib.graph.Digraph; import com.google.devtools.build.lib.graph.Node; -import com.google.devtools.build.lib.packages.BazelStarlarkContext; +import com.google.devtools.build.lib.packages.BzlInitThreadContext; import com.google.devtools.build.lib.packages.NativeAspectClass; import com.google.devtools.build.lib.packages.RuleClass; import com.google.devtools.build.lib.packages.RuleClass.Builder.ThirdPartyLicenseExistencePolicy; @@ -871,12 +871,10 @@ public ImmutableMap getEnvironment() { @Override public void setStarlarkThreadContext(StarlarkThread thread, Label fileLabel) { - new BazelStarlarkContext( - BazelStarlarkContext.Phase.LOADING, + new BzlInitThreadContext( toolsRepository, configurationFragmentMap, new SymbolGenerator<>(fileLabel), - /*analysisRuleLabel=*/ null, networkAllowlistForTests) .storeInThread(thread); } diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java b/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java index 6d564d51f1b12c..8c662bd0dcd489 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java @@ -334,7 +334,7 @@ public ImmutableMap> evaluate( try (Mutability mu = Mutability.create("eval_transition_function")) { StarlarkThread thread = new StarlarkThread(mu, semantics); thread.setPrintHandler(Event.makeDebugPrintHandler(handler)); - // TODO: If the resulting values of Starlark transitions ever evolve to be + // TODO(brandjon): If the resulting values of Starlark transitions ever evolve to be // complex Starlark objects like structs as opposed to the ints, strings, // etc they are today then we need a real symbol generator which is used // to calculate equality between instances of Starlark objects. A candidate diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java index 9b0255bb4ec696..9f174941071bd7 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java @@ -335,10 +335,6 @@ public StarlarkCallable rule( if (implicitOutputs != Starlark.NONE) { if (implicitOutputs instanceof StarlarkFunction) { - // TODO(brandjon): Embedding bazelContext in a callback is not thread safe! Instead - // construct a new BazelStarlarkContext with copies of the relevant fields that are safe to - // share. Maybe create a method on BazelStarlarkContext for safely constructing a child - // context. StarlarkCallbackHelper callback = new StarlarkCallbackHelper( (StarlarkFunction) implicitOutputs, thread.getSemantics(), bazelContext); diff --git a/src/main/java/com/google/devtools/build/lib/packages/BazelStarlarkContext.java b/src/main/java/com/google/devtools/build/lib/packages/BazelStarlarkContext.java index ad5f277a24fd8a..603764c5293b74 100644 --- a/src/main/java/com/google/devtools/build/lib/packages/BazelStarlarkContext.java +++ b/src/main/java/com/google/devtools/build/lib/packages/BazelStarlarkContext.java @@ -27,39 +27,77 @@ import net.starlark.java.eval.Starlark; import net.starlark.java.eval.StarlarkThread; -/** Contextual information associated with each Starlark thread created by Bazel. */ -// TODO(adonovan): rename BazelThreadContext, for symmetry with BazelModuleContext. -// TODO(brandjon): Use composition rather than inheritance for RuleDefinitionEnvironment; clients -// should retrieve the RDE (if it exists) from this class in order to access e.g. the network -// allowlist. The toolsRepository info will be duplicated between this class and RDE but we can -// enforce consistency with a precondition check. -public final class BazelStarlarkContext +/** + * Bazel-specific contextual information associated with a Starlark evaluation thread. + * + *

This is stored in the StarlarkThread object as a thread-local. A subclass of this class may be + * used for certain kinds of Starlark evaluations; in that case it is still keyed in the + * thread-locals under {@code BazelStarlarkContext.class}. + * + *

This object is mutable and should not be reused for more than one Starlark thread. + */ +// TODO(b/236456122): rename BazelThreadContext, for symmetry with BazelModuleContext. +// TODO(b/236456122): We should break this class up into subclasses for each kind of evaluation, as +// opposed to storing specialized fields on this class and setting them to null for inapplicable +// environments. Subclasses should define {@code from(StarlarkThread)} and {@code +// fromOrFailFunction(StarlarkThread, String)} methods to be used in place of the {@code +// check*Phase} methods in this class. Kinds of evaluation include: +// - .bzl loading +// - BUILD evaluation (should absorb PackageFactory.PackageContext -- no need to store multiple +// thread-locals in StarlarkThread) +// - WORKSPACE evaluation (shares logic with BUILD) +// - rule and aspect analysis implementation (can/should we store the RuleContext here?) +// - implicit outputs +// - computed defaults +// - transition implementation +// - Args.map_each +// - probably others +// TODO(b/236456122): The inheritance of RuleDefinitionEnvironment should be replaced by +// composition, in an appropriate subclass of this class. Things like the tools repository, network +// allowlist, etc. can be accessed from the RDE. (If any info needs to be duplicated between RDE and +// here, we should assert consistency with a precondition check.) +public class BazelStarlarkContext implements RuleDefinitionEnvironment, StarlarkThread.UncheckedExceptionContext { /** The phase to which this Starlark thread belongs. */ + // TODO(b/236456122): Eliminate. public enum Phase { WORKSPACE, LOADING, ANALYSIS } - /** Return the Bazel information associated with the specified Starlark thread. */ + /** + * Retrieves this context from a Starlark thread, or throws {@link IllegalStateException} if + * unavailable. + */ public static BazelStarlarkContext from(StarlarkThread thread) { - return thread.getThreadLocal(BazelStarlarkContext.class); + BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class); + // ISE rather than NPE for symmetry with subclasses. + Preconditions.checkState( + ctx != null, "Expected BazelStarlarkContext to be available in this Starlark thread"); + return ctx; } - /** Save this BazelStarlarkContext in the specified Starlark thread. */ + /** + * Saves this {@code BazelStarlarkContext} in the specified Starlark thread. Call only once, + * before evaluation begins. + */ public void storeInThread(StarlarkThread thread) { + Preconditions.checkState(thread.getThreadLocal(BazelStarlarkContext.class) == null); thread.setThreadLocal(BazelStarlarkContext.class, this); thread.setUncheckedExceptionContext(this); } + // A generic counter for uniquely identifying symbols created in this Starlark evaluation. + private final SymbolGenerator symbolGenerator; + + // TODO(b/236456122): Migrate the below fields to subclasses. private final Phase phase; // Only necessary for loading phase threads. @Nullable private final RepositoryName toolsRepository; // Only necessary for loading phase threads to construct configuration_field. @Nullable private final ImmutableMap> fragmentNameToClass; - private final SymbolGenerator symbolGenerator; @Nullable private final Label analysisRuleLabel; // TODO(b/192694287): Remove once we migrate all tests from the allowlist @Nullable private final Label networkAllowlistForTests; @@ -75,13 +113,6 @@ public void storeInThread(StarlarkThread thread) { * compared using reference equality. * @param analysisRuleLabel is the label of the rule for an analysis phase (rule or aspect */ - // TODO(adonovan): clearly demarcate which fields are defined in which kinds of threads (loading, - // analysis, workspace, implicit outputs, computed defaults, etc), perhaps by splitting these into - // separate structs, exactly one of which is populated (plus the common fields). And eliminate - // StarlarkUtils.Phase. - // TODO(adonovan): move PackageFactory.PackageContext in here, for loading-phase threads. - // TODO(adonovan): is there any reason not to put the entire RuleContext in this thread, for - // analysis threads? public BazelStarlarkContext( Phase phase, @Nullable RepositoryName toolsRepository, @@ -97,6 +128,11 @@ public BazelStarlarkContext( this.networkAllowlistForTests = networkAllowlistForTests; } + /** Returns the phase associated with this context. */ + public Phase getPhase() { + return phase; + } + /** Returns the name of the tools repository, such as "@bazel_tools". */ @Nullable @Override @@ -138,6 +174,8 @@ public Optional

TODO(adonovan): eliminate the need for this class by making the Starlark calls in the same - * Starlark thread that instantiated the rule. */ +// TODO(brandjon): Consider eliminating this class by executing the callback in the same thread as +// the caller, i.e. in the thread evaluating a BUILD file. This might not be possible for implicit +// outputs without some refactoring to cache the result of the computation (currently RuleContext +// seems to reinvoke the callback). public final class StarlarkCallbackHelper { private final StarlarkFunction callback; @@ -48,15 +49,18 @@ public final class StarlarkCallbackHelper { // Alternatively (or additionally), we could put PackageContext // into BazelStarlarkContext so there's only a single blob of state. private final StarlarkSemantics starlarkSemantics; - private final BazelStarlarkContext context; + // The Bazel context that was in effect at the place where the StarlarkCallbackHelper was defined. + // This will be used as a template for creating a new context each time the callback is invoked + // (since BazelStarlarkContexts cannot be safely shared or reused). + private final BazelStarlarkContext contextTemplate; public StarlarkCallbackHelper( StarlarkFunction callback, StarlarkSemantics starlarkSemantics, - BazelStarlarkContext context) { + BazelStarlarkContext contextTemplate) { this.callback = callback; this.starlarkSemantics = starlarkSemantics; - this.context = context; + this.contextTemplate = contextTemplate; } public ImmutableList getParameterNames() { @@ -65,17 +69,30 @@ public ImmutableList getParameterNames() { // TODO(adonovan): opt: all current callers are forced to construct a temporary Structure. // Instead, make them supply a map. - public Object call(EventHandler eventHandler, Structure ctx, Object... arguments) + public Object call(EventHandler eventHandler, Structure struct, Object... arguments) throws EvalException, InterruptedException { try (Mutability mu = Mutability.create("callback", callback)) { StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics); thread.setPrintHandler(Event.makeDebugPrintHandler(eventHandler)); + // TODO(b/236456122): If we don't eliminate StarlarkCallbackHelper entirely, we can give it + // its own BazelStarlarkContext subclass. + BazelStarlarkContext context = + new BazelStarlarkContext( + contextTemplate.getPhase(), + contextTemplate.getToolsRepository(), + contextTemplate.getFragmentNameToClass(), + // TODO(brandjon): In principle, if we're creating a new symbol generator here, we + // should have a unique owner object to associate it with for distinguishing + // reference-equality objects. But I don't think implicit outputs or computed defaults + // care about identity. + new SymbolGenerator<>(new Object()), + // Analysis rule label should be null for computed defaults and implicit outputs, but + // may as well copy the field for symmetry. + contextTemplate.getAnalysisRuleLabel(), + contextTemplate.getNetworkAllowlistForTests().orElse(null)); context.storeInThread(thread); return Starlark.call( - thread, - callback, - buildArgumentList(ctx, arguments), - /*kwargs=*/ ImmutableMap.of()); + thread, callback, buildArgumentList(struct, arguments), /*kwargs=*/ ImmutableMap.of()); } catch (ClassCastException | IllegalArgumentException e) { // TODO(adonovan): investigate throw new EvalException(e); } @@ -83,18 +100,18 @@ public Object call(EventHandler eventHandler, Structure ctx, Object... arguments /** * Creates a list of actual arguments that contains the given arguments and all attribute values - * required from the specified context. + * required from the specified structure. */ - private ImmutableList buildArgumentList(Structure ctx, Object... arguments) + private ImmutableList buildArgumentList(Structure struct, Object... arguments) throws EvalException { ImmutableList.Builder builder = ImmutableList.builder(); ImmutableList names = getParameterNames(); int requiredParameters = names.size() - arguments.length; for (int pos = 0; pos < requiredParameters; ++pos) { String name = names.get(pos); - Object value = ctx.getValue(name); + Object value = struct.getValue(name); if (value == null) { - throw new IllegalArgumentException(ctx.getErrorMessageForUnknownField(name)); + throw new IllegalArgumentException(struct.getErrorMessageForUnknownField(name)); } builder.add(value); }