Track getenv calls in module extensions

RELNOTES: Changes to environment variables read via `getenv` now correctly invalidate module extensions.

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

@@ -35,7 +35,7 @@
 @GenerateTypeAdapter
 public abstract class BazelLockFileValue implements SkyValue, Postable {
 
-  public static final int LOCK_FILE_VERSION = 10;
+  public static final int LOCK_FILE_VERSION = 11;
 
   @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 

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

@@ -358,6 +358,22 @@ public RepoRecordedInput.Dirents read(JsonReader jsonReader) throws IOException
             }
           };
 
+  private static final TypeAdapter<RepoRecordedInput.EnvVar>
+      REPO_RECORDED_INPUT_ENV_VAR_TYPE_ADAPTER =
+          new TypeAdapter<>() {
+            @Override
+            public void write(JsonWriter jsonWriter, RepoRecordedInput.EnvVar value)
+                throws IOException {
+              jsonWriter.value(value.toStringInternal());
+            }
+
+            @Override
+            public RepoRecordedInput.EnvVar read(JsonReader jsonReader) throws IOException {
+              return (RepoRecordedInput.EnvVar)
+                  RepoRecordedInput.EnvVar.PARSER.parse(jsonReader.nextString());
+            }
+          };
+
   // This can't reuse the existing type adapter factory for Optional as we need to explicitly
   // serialize null values but don't want to rely on GSON's serializeNulls.
   private static final class OptionalChecksumTypeAdapterFactory implements TypeAdapterFactory {
@@ -441,7 +457,9 @@ private static GsonBuilder newGsonBuilder() {
         .registerTypeAdapter(byte[].class, BYTE_ARRAY_TYPE_ADAPTER)
         .registerTypeAdapter(RepoRecordedInput.File.class, REPO_RECORDED_INPUT_FILE_TYPE_ADAPTER)
         .registerTypeAdapter(
-            RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER);
+            RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER)
+        .registerTypeAdapter(
+            RepoRecordedInput.EnvVar.class, REPO_RECORDED_INPUT_ENV_VAR_TYPE_ADAPTER);
   }
 
   private GsonTypeAdapterUtil() {}

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

@@ -48,7 +48,7 @@ public static Builder builder() {
 
   public abstract ImmutableMap<RepoRecordedInput.Dirents, String> getRecordedDirentsInputs();
 
-  public abstract ImmutableMap<String, String> getEnvVariables();
+  public abstract ImmutableMap<RepoRecordedInput.EnvVar, String> getEnvVariables();
 
   public abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
@@ -78,7 +78,7 @@ public abstract Builder setRecordedFileInputs(
     public abstract Builder setRecordedDirentsInputs(
         ImmutableMap<RepoRecordedInput.Dirents, String> value);
 
-    public abstract Builder setEnvVariables(ImmutableMap<String, String> value);
+    public abstract Builder setEnvVariables(ImmutableMap<RepoRecordedInput.EnvVar, String> value);
 
     public abstract Builder setGeneratedRepoSpecs(ImmutableMap<String, RepoSpec> value);
 

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

@@ -230,6 +230,13 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     // result is taken from the lockfile, we can already populate the lockfile info. This is
     // necessary to prevent the extension from rerunning when only the imports change.
     if (lockfileMode == LockfileMode.UPDATE || lockfileMode == LockfileMode.REFRESH) {
+      ImmutableMap.Builder<RepoRecordedInput.EnvVar, String> combinedEnvVars =
+          ImmutableMap.builder();
+      // The environment variable dependencies statically declared via the 'environ' attribute.
+      combinedEnvVars.putAll(RepoRecordedInput.EnvVar.wrap(extension.getEnvVars()));
+      // The environment variable dependencies dynamically declared via the 'getenv' method.
+      combinedEnvVars.putAll(moduleExtensionResult.getRecordedEnvVarInputs());
+
       lockFileInfo =
           Optional.of(
               new LockFileModuleExtension.WithFactors(
@@ -241,7 +248,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                               GsonTypeAdapterUtil.SINGLE_EXTENSION_USAGES_VALUE_GSON, usagesValue))
                       .setRecordedFileInputs(moduleExtensionResult.getRecordedFileInputs())
                       .setRecordedDirentsInputs(moduleExtensionResult.getRecordedDirentsInputs())
-                      .setEnvVariables(extension.getEnvVars())
+                      .setEnvVariables(combinedEnvVars.buildKeepingLast())
                       .setGeneratedRepoSpecs(generatedRepoSpecs)
                       .setModuleExtensionMetadata(moduleExtensionMetadata)
                       .setRecordedRepoMappingEntries(
@@ -284,7 +291,7 @@ private SingleExtensionValue tryGettingValueFromLockFile(
                 + extensionId
                 + "' or one of its transitive .bzl files has changed");
       }
-      if (!extension.getEnvVars().equals(lockedExtension.getEnvVariables())) {
+      if (didRecordedInputsChange(env, directories, lockedExtension.getEnvVariables())) {
         diffRecorder.record(
             "The environment variables the extension '"
                 + extensionId
@@ -769,6 +776,7 @@ public RunModuleExtensionResult run(
         generatedRepoSpecs.put(name, repoSpec);
       }
       return RunModuleExtensionResult.create(
+          ImmutableMap.of(),
           ImmutableMap.of(),
           ImmutableMap.of(),
           generatedRepoSpecs.buildOrThrow(),
@@ -947,6 +955,7 @@ public RunModuleExtensionResult run(
       return RunModuleExtensionResult.create(
           moduleContext.getRecordedFileInputs(),
           moduleContext.getRecordedDirentsInputs(),
+          moduleContext.getRecordedEnvVarInputs(),
           threadContext.getGeneratedRepoSpecs(),
           moduleExtensionMetadata,
           repoMappingRecorder.recordedEntries());
@@ -1011,6 +1020,8 @@ abstract static class RunModuleExtensionResult {
 
     abstract ImmutableMap<RepoRecordedInput.Dirents, String> getRecordedDirentsInputs();
 
+    abstract ImmutableMap<RepoRecordedInput.EnvVar, String> getRecordedEnvVarInputs();
+
     abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
     abstract Optional<ModuleExtensionMetadata> getModuleExtensionMetadata();
@@ -1020,12 +1031,14 @@ abstract static class RunModuleExtensionResult {
     static RunModuleExtensionResult create(
         ImmutableMap<RepoRecordedInput.File, String> recordedFileInputs,
         ImmutableMap<RepoRecordedInput.Dirents, String> recordedDirentsInputs,
+        ImmutableMap<RepoRecordedInput.EnvVar, String> recordedEnvVarInputs,
         ImmutableMap<String, RepoSpec> generatedRepoSpecs,
         Optional<ModuleExtensionMetadata> moduleExtensionMetadata,
         ImmutableTable<RepositoryName, String, RepositoryName> recordedRepoMappingEntries) {
       return new AutoValue_SingleExtensionEvalFunction_RunModuleExtensionResult(
           recordedFileInputs,
           recordedDirentsInputs,
+          recordedEnvVarInputs,
           generatedRepoSpecs,
           moduleExtensionMetadata,
           recordedRepoMappingEntries);

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

@@ -22,7 +22,6 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Maps;
 import com.google.common.util.concurrent.Futures;
@@ -214,9 +213,12 @@ public ImmutableMap<Dirents, String> getRecordedDirentsInputs() {
     return ImmutableMap.copyOf(recordedDirentsInputs);
   }
 
-  /** Returns set of environment variable keys encountered so far. */
-  public ImmutableSet<String> getAccumulatedEnvKeys() {
-    return ImmutableSet.copyOf(accumulatedEnvKeys);
+  public ImmutableMap<RepoRecordedInput.EnvVar, String> getRecordedEnvVarInputs()
+      throws InterruptedException {
+    // getEnvVarValues doesn't return null since the Skyframe dependencies have already been
+    // established by getenv calls.
+    return RepoRecordedInput.EnvVar.wrap(
+        RepositoryFunction.getEnvVarValues(env, accumulatedEnvKeys));
   }
 
   protected void checkInOutputDirectory(String operation, StarlarkPath path)

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

@@ -361,16 +361,12 @@ private RepositoryDirectoryValue.Builder fetchInternal(
         env.getListener().handle(Event.debug(defInfo));
       }
 
-      // Modify marker data to include the files/dirents used by the rule's implementation function.
+      // Modify marker data to include the files/dirents/env vars used by the rule's implementation
+      // function.
       recordedInputValues.putAll(starlarkRepositoryContext.getRecordedFileInputs());
       recordedInputValues.putAll(starlarkRepositoryContext.getRecordedDirentsInputs());
       recordedInputValues.putAll(starlarkRepositoryContext.getRecordedDirTreeInputs());
-
-      // Ditto for environment variables accessed via `getenv`.
-      for (String envKey : starlarkRepositoryContext.getAccumulatedEnvKeys()) {
-        recordedInputValues.put(
-            new RepoRecordedInput.EnvVar(envKey), clientEnvironment.get(envKey));
-      }
+      recordedInputValues.putAll(starlarkRepositoryContext.getRecordedEnvVarInputs());
 
       for (Table.Cell<RepositoryName, String, RepositoryName> repoMappings :
           repoMappingRecorder.recordedEntries().cellSet()) {

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

@@ -15,11 +15,13 @@
 package com.google.devtools.build.lib.rules.repository;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Splitter;
+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;
@@ -491,7 +493,7 @@ public boolean isUpToDate(
 
   /** Represents an environment variable accessed during the repo fetch. */
   public static final class EnvVar extends RepoRecordedInput {
-    static final Parser PARSER =
+    public static final Parser PARSER =
         new Parser() {
           @Override
           public String getPrefix() {
@@ -506,7 +508,12 @@ public RepoRecordedInput parse(String s) {
 
     final String name;
 
-    public EnvVar(String name) {
+    public static ImmutableMap<EnvVar, String> wrap(Map<String, String> envVars) {
+      return envVars.entrySet().stream()
+          .collect(toImmutableMap(e -> new EnvVar(e.getKey()), Map.Entry::getValue));
+    }
+
+    private EnvVar(String name) {
       this.name = name;
     }
 

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.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -262,9 +263,7 @@ protected Map<String, String> declareEnvironmentDependencies(
       return null;
     }
     // Add the dependencies to the marker file
-    for (String key : keys) {
-      recordedInputValues.put(new RepoRecordedInput.EnvVar(key), envDep.get(key));
-    }
+    recordedInputValues.putAll(RepoRecordedInput.EnvVar.wrap(envDep));
     return envDep;
   }
 
@@ -289,7 +288,7 @@ public static ImmutableMap<String, String> getEnvVarValues(Environment env, Set<
         repoEnv.put(key, value);
       }
     }
-    return repoEnv.buildKeepingLast();
+    return ImmutableSortedMap.copyOf(repoEnv.buildKeepingLast());
   }
 
   /**