Add 'environ' attribute to module extensions and update lockfile

- Add 'environ' to module extensions
- Store the variables and their values in the lockfile
- Compare the variables and values with the lockfile and re-evaluate or error if they differ

PiperOrigin-RevId: 548964193
Change-Id: Ic2d52fe3332e93095c414d8bca4c9b4312bca8c2

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

@@ -114,20 +114,28 @@ public ImmutableList<String> getModuleExtensionDiff(
       ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages,
       ModuleExtensionId extensionId,
       byte[] transitiveDigest,
+      ImmutableMap<String, String> envVariables,
       ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages) {
     ImmutableList.Builder<String> extDiff = new ImmutableList.Builder<>();
     if (lockedExtension == null) {
-      extDiff.add("The module extension '" + extensionId + "' does not exist in the lockfile");
-    } else {
-      if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {
+      return extDiff
+          .add("The module extension '" + extensionId + "' does not exist in the lockfile")
+          .build();
+    }
+    if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {
         extDiff.add(
             "The implementation of the extension '"
                 + extensionId
                 + "' or one of its transitive .bzl files has changed");
-      }
-      if (!extensionUsages.equals(lockedExtensionUsages)) {
-        extDiff.add("The usages of the extension named '" + extensionId + "' has changed");
-      }
+    }
+    if (!extensionUsages.equals(lockedExtensionUsages)) {
+      extDiff.add("The usages of the extension '" + extensionId + "' has changed");
+    }
+    if (!envVariables.equals(lockedExtension.getEnvVariables())) {
+      extDiff.add(
+          "The environment variables the extension '"
+              + extensionId
+              + "' depends on (or their values) have changed");
     }
     return extDiff.build();
   }

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

@@ -31,10 +31,15 @@ public abstract class LockFileModuleExtension implements Postable {
   @SuppressWarnings("mutable")
   public abstract byte[] getBzlTransitiveDigest();
 
+  public abstract ImmutableMap<String, String> getEnvVariables();
+
   public abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
   public static LockFileModuleExtension create(
-      byte[] transitiveDigest, ImmutableMap<String, RepoSpec> generatedRepoSpecs) {
-    return new AutoValue_LockFileModuleExtension(transitiveDigest, generatedRepoSpecs);
+      byte[] transitiveDigest,
+      ImmutableMap<String, String> envVariables,
+      ImmutableMap<String, RepoSpec> generatedRepoSpecs) {
+    return new AutoValue_LockFileModuleExtension(
+        transitiveDigest, envVariables, generatedRepoSpecs);
   }
 }

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

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.Label;
 import java.util.Optional;
@@ -44,6 +45,8 @@ public abstract class ModuleExtension implements StarlarkValue {
 
   public abstract Location getLocation();
 
+  public abstract ImmutableList<String> getEnvVariables();
+
   public static Builder builder() {
     return new AutoValue_ModuleExtension.Builder();
   }
@@ -61,6 +64,8 @@ public abstract static class Builder {
 
     public abstract Builder setTagClasses(ImmutableMap<String, TagClass> value);
 
+    public abstract Builder setEnvVariables(ImmutableList<String> value);
+
     public abstract ModuleExtension build();
   }
 }

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

@@ -32,6 +32,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.rules.repository.NeedsSkyframeRestartException;
+import com.google.devtools.build.lib.rules.repository.RepositoryFunction;
 import com.google.devtools.build.lib.runtime.ProcessWrapper;
 import com.google.devtools.build.lib.runtime.RepositoryRemoteExecutor;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps;
@@ -144,9 +145,16 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           Transience.PERSISTENT);
     }
 
-    // Check the lockfile first for that module extension
+    ModuleExtension extension = (ModuleExtension) exported;
+    ImmutableMap<String, String> extensionEnvVars =
+        RepositoryFunction.getEnvVarValues(env, extension.getEnvVariables());
+    if (extensionEnvVars == null) {
+      return null;
+    }
     byte[] bzlTransitiveDigest =
         BazelModuleContext.of(bzlLoadValue.getModule()).bzlTransitiveDigest();
+
+    // Check the lockfile first for that module extension
     LockfileMode lockfileMode = BazelLockFileFunction.LOCKFILE_MODE.get(env);
     if (!lockfileMode.equals(LockfileMode.OFF)) {
       BazelLockFileValue lockfile = (BazelLockFileValue) env.getValue(BazelLockFileValue.KEY);
@@ -155,14 +163,18 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
       SingleExtensionEvalValue singleExtensionEvalValue =
           tryGettingValueFromLockFile(
-              extensionId, usagesValue, bzlTransitiveDigest, lockfileMode, lockfile);
+              extensionId,
+              extensionEnvVars,
+              usagesValue,
+              bzlTransitiveDigest,
+              lockfileMode,
+              lockfile);
       if (singleExtensionEvalValue != null) {
         return singleExtensionEvalValue;
       }
     }
 
     // Run that extension!
-    ModuleExtension extension = (ModuleExtension) exported;
     ImmutableMap<String, RepoSpec> generatedRepoSpecs =
         runModuleExtension(
             extensionId, extension, usagesValue, bzlLoadValue.getModule(), starlarkSemantics, env);
@@ -177,14 +189,16 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           .post(
               ModuleExtensionResolutionEvent.create(
                   extensionId,
-                  LockFileModuleExtension.create(bzlTransitiveDigest, generatedRepoSpecs)));
+                  LockFileModuleExtension.create(
+                      bzlTransitiveDigest, extensionEnvVars, generatedRepoSpecs)));
     }
     return createSingleExtentionValue(generatedRepoSpecs, usagesValue);
   }
 
   @Nullable
   private SingleExtensionEvalValue tryGettingValueFromLockFile(
       ModuleExtensionId extensionId,
+      ImmutableMap<String, String> envVariables,
       SingleExtensionUsagesValue usagesValue,
       byte[] bzlTransitiveDigest,
       LockfileMode lockfileMode,
@@ -205,7 +219,8 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
     // If we have the extension, check if the implementation and usage haven't changed
     if (lockedExtension != null
         && Arrays.equals(bzlTransitiveDigest, lockedExtension.getBzlTransitiveDigest())
-        && usagesValue.getExtensionUsages().equals(lockedExtensionUsages)) {
+        && usagesValue.getExtensionUsages().equals(lockedExtensionUsages)
+        && envVariables.equals(lockedExtension.getEnvVariables())) {
       return createSingleExtentionValue(lockedExtension.getGeneratedRepoSpecs(), usagesValue);
     } else if (lockfileMode.equals(LockfileMode.ERROR)) {
       ImmutableList<String> extDiff =
@@ -214,6 +229,7 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
               lockedExtensionUsages,
               extensionId,
               bzlTransitiveDigest,
+              envVariables,
               usagesValue.getExtensionUsages());
       throw new SingleExtensionEvalFunctionException(
           ExternalDepsException.withMessage(

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java

@@ -291,6 +291,7 @@ public Object moduleExtension(
       StarlarkCallable implementation,
       Dict<?, ?> tagClasses, // Dict<String, TagClass>
       Object doc, // <String> or Starlark.NONE
+      Sequence<?> environ, // <String>
       StarlarkThread thread)
       throws EvalException {
     return ModuleExtension.builder()
@@ -300,6 +301,7 @@ public Object moduleExtension(
         .setDoc(Starlark.toJavaOptional(doc, String.class))
         .setDefiningBzlFileLabel(
             BzlInitThreadContext.fromOrFailFunction(thread, "module_extension").getBzlFile())
+        .setEnvVariables(ImmutableList.copyOf(Sequence.cast(environ, String.class, "environ")))
         .setLocation(thread.getCallerLocation())
         .build();
   }

src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java

@@ -20,6 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.BaseEncoding;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
@@ -51,6 +52,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -188,9 +190,9 @@ public abstract RepositoryDirectoryValue.Builder fetch(
       throws InterruptedException, RepositoryFunctionException;
 
   @SuppressWarnings("unchecked")
-  private static Iterable<String> getEnviron(Rule rule) {
+  private static Collection<String> getEnviron(Rule rule) {
     if (rule.isAttrDefined("$environ", Type.STRING_LIST)) {
-      return (Iterable<String>) rule.getAttr("$environ");
+      return (Collection<String>) rule.getAttr("$environ");
     }
     return ImmutableList.of();
   }
@@ -201,7 +203,7 @@ private static Iterable<String> getEnviron(Rule rule) {
    * is needed.
    */
   public boolean verifyMarkerData(Rule rule, Map<String, String> markerData, Environment env)
-      throws InterruptedException, RepositoryFunctionException {
+      throws InterruptedException {
     return verifyEnvironMarkerData(markerData, env, getEnviron(rule))
         && verifyMarkerDataForFiles(rule, markerData, env)
         && verifySemanticsMarkerData(markerData, env);
@@ -306,33 +308,37 @@ public static RootedPath getRootedPathFromLabel(Label label, Environment env)
   protected Map<String, String> declareEnvironmentDependencies(
       Map<String, String> markerData, Environment env, Iterable<String> keys)
       throws InterruptedException {
-    Map<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
+    ImmutableMap<String, String> envDep = getEnvVarValues(env, keys);
+    if (envDep == null) {
+      return null;
+    }
+    // Add the dependencies to the marker file
+    keys.forEach(key -> markerData.put("ENV:" + key, envDep.get(key)));
+    return envDep;
+  }
 
-    // Returns true if there is a null value and we need to wait for some dependencies.
+  @Nullable
+  public static ImmutableMap<String, String> getEnvVarValues(Environment env, Iterable<String> keys)
+      throws InterruptedException {
+    ImmutableMap<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
     if (environ == null) {
       return null;
     }
-
     Map<String, String> repoEnvOverride = PrecomputedValue.REPO_ENV.get(env);
     if (repoEnvOverride == null) {
       return null;
     }
 
     // Only depend on --repo_env values that are specified in the "environ" attribute.
-    Map<String, String> repoEnv = new LinkedHashMap<String, String>(environ);
+    ImmutableMap.Builder<String, String> repoEnv = ImmutableMap.builder();
+    repoEnv.putAll(environ);
     for (String key : keys) {
       String value = repoEnvOverride.get(key);
       if (value != null) {
         repoEnv.put(key, value);
       }
     }
-
-    // Add the dependencies to the marker file
-    for (Map.Entry<String, String> value : repoEnv.entrySet()) {
-      markerData.put("ENV:" + value.getKey(), value.getValue());
-    }
-
-    return repoEnv;
+    return repoEnv.buildKeepingLast();
   }
 
   /**
@@ -341,9 +347,9 @@ protected Map<String, String> declareEnvironmentDependencies(
    * Environment)} function to verify the values for environment variables.
    */
   protected boolean verifyEnvironMarkerData(
-      Map<String, String> markerData, Environment env, Iterable<String> keys)
+      Map<String, String> markerData, Environment env, Collection<String> keys)
       throws InterruptedException {
-    Map<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
+    ImmutableMap<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
     if (env.valuesMissing()) {
       return false; // Returns false so caller knows to return immediately
     }
@@ -364,17 +370,18 @@ protected boolean verifyEnvironMarkerData(
 
     // Verify that all environment variable in the marker file are also in keys
     for (String key : markerData.keySet()) {
-      if (key.startsWith("ENV:") && !repoEnv.containsKey(key.substring(4))) {
+      if (key.startsWith("ENV:") && !keys.contains(key.substring(4))) {
         return false;
       }
     }
+
     // Now verify the values of the marker data
-    for (Map.Entry<String, String> value : repoEnv.entrySet()) {
-      if (!markerData.containsKey("ENV:" + value.getKey())) {
+    for (String key : keys) {
+      if (!markerData.containsKey("ENV:" + key)) {
         return false;
       }
-      String markerValue = markerData.get("ENV:" + value.getKey());
-      if (!Objects.equals(markerValue, value.getValue())) {
+      String markerValue = markerData.get("ENV:" + key);
+      if (!Objects.equals(markerValue, repoEnv.get(key))) {
         return false;
       }
     }
@@ -457,8 +464,7 @@ protected static String getPathAttr(Rule rule) throws RepositoryFunctionExceptio
   }
 
   @VisibleForTesting
-  protected static PathFragment getTargetPath(String userDefinedPath, Path workspace)
-      throws RepositoryFunctionException {
+  protected static PathFragment getTargetPath(String userDefinedPath, Path workspace) {
     PathFragment pathFragment = PathFragment.create(userDefinedPath);
     return workspace.getRelative(pathFragment).asFragment();
   }

src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.bugreport.BugReport;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.skyframe.AbstractSkyKey;
@@ -23,8 +24,6 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.SkyframeLookupResult;
-import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -45,7 +44,7 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
     return env.getValue(ClientEnvironmentFunction.key(key));
   }
 
-  /** @return the SkyKey to invoke this function for the environment variable {@code variable}. */
+  /** Returns the SkyKey to invoke this function for the environment variable {@code variable}. */
   public static Key key(String variable) {
     return Key.create(variable);
   }
@@ -81,8 +80,8 @@ public SkyKeyInterner<Key> getSkyKeyInterner() {
    * if and only if some dependencies from Skyframe still need to be resolved.
    */
   @Nullable
-  public static Map<String, String> getEnvironmentView(Environment env, Iterable<String> keys)
-      throws InterruptedException {
+  public static ImmutableMap<String, String> getEnvironmentView(
+      Environment env, Iterable<String> keys) throws InterruptedException {
     ImmutableList.Builder<SkyKey> skyframeKeysBuilder = ImmutableList.builder();
     for (String key : keys) {
       skyframeKeysBuilder.add(key(key));
@@ -92,8 +91,8 @@ public static Map<String, String> getEnvironmentView(Environment env, Iterable<S
     if (env.valuesMissing()) {
       return null;
     }
-    // To return the initial order and support null values, we use a LinkedHashMap.
-    LinkedHashMap<String, String> result = new LinkedHashMap<>();
+
+    ImmutableMap.Builder<String, String> result = ImmutableMap.builder();
     for (SkyKey key : skyframeKeys) {
       ClientEnvironmentValue value = (ClientEnvironmentValue) values.get(key);
       if (value == null) {
@@ -102,8 +101,10 @@ public static Map<String, String> getEnvironmentView(Environment env, Iterable<S
                 "ClientEnvironmentValue " + key + " was missing, this should never happen"));
         return null;
       }
-      result.put(key.argument().toString(), value.getValue());
+      if (value.getValue() != null) {
+        result.put(key.argument().toString(), value.getValue());
+      }
     }
-    return Collections.unmodifiableMap(result);
+    return result.buildOrThrow();
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java

@@ -30,9 +30,10 @@ public static Key key(String keyString) {
     return ClientEnvironmentFunction.Key.create(keyString);
   }
 
+  /** The Skyframe key for the client environment function. */
   @AutoCodec.VisibleForSerialization
   @AutoCodec
-  static class Key extends AbstractSkyKey<String> {
+  public static class Key extends AbstractSkyKey<String> {
     private static final SkyKeyInterner<Key> interner = SkyKey.newInterner();
 
     private Key(String arg) {
@@ -64,7 +65,7 @@ public ClientEnvironmentFunction(AtomicReference<Map<String, String>> clientEnv)
 
   @Nullable
   @Override
-  public SkyValue compute(SkyKey key, Environment env) throws InterruptedException {
+  public SkyValue compute(SkyKey key, Environment env) {
     return new ClientEnvironmentValue(clientEnv.get().get((String) key.argument()));
   }
 }