Track getenv calls in module extensions

Fixes #22404

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

Closes #22498.

PiperOrigin-RevId: 637058337
Change-Id: Id4aaa4155a728452472eedae4a59c8d29456e512

Author: fmeum

MODULE.bazel.lock

@@ -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/BazelLockFileValue.java

@@ -232,22 +232,34 @@ public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> typeToken) {
           if (elementTypeAdapter == null) {
             return null;
           }
-          return (TypeAdapter<T>) new OptionalTypeAdapter<>(elementTypeAdapter);
+          // Explicit nulls for Optional.empty are required for env variable tracking, but are too
+          // noisy and unnecessary for other types.
+          return (TypeAdapter<T>)
+              new OptionalTypeAdapter<>(
+                  elementTypeAdapter, /* serializeNulls= */ elementType.equals(String.class));
         }
       };
 
   private static final class OptionalTypeAdapter<T> extends TypeAdapter<Optional<T>> {
     private final TypeAdapter<T> elementTypeAdapter;
+    private final boolean serializeNulls;
 
-    public OptionalTypeAdapter(TypeAdapter<T> elementTypeAdapter) {
+    public OptionalTypeAdapter(TypeAdapter<T> elementTypeAdapter, boolean serializeNulls) {
       this.elementTypeAdapter = elementTypeAdapter;
+      this.serializeNulls = serializeNulls;
     }
 
     @Override
     public void write(JsonWriter jsonWriter, Optional<T> t) throws IOException {
       Preconditions.checkNotNull(t);
       if (t.isEmpty()) {
-        jsonWriter.nullValue();
+        boolean oldSerializeNulls = jsonWriter.getSerializeNulls();
+        jsonWriter.setSerializeNulls(serializeNulls);
+        try {
+          jsonWriter.nullValue();
+        } finally {
+          jsonWriter.setSerializeNulls(oldSerializeNulls);
+        }
       } else {
         elementTypeAdapter.write(jsonWriter, t.get());
       }
@@ -358,6 +370,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 +469,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/GsonTypeAdapterUtil.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, Optional<String>> getEnvVariables();
 
   public abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
@@ -78,7 +78,8 @@ 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, Optional<String>> value);
 
     public abstract Builder setGeneratedRepoSpecs(ImmutableMap<String, RepoSpec> value);
 

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

@@ -25,6 +25,7 @@
 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.ImmutableTable;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
@@ -230,6 +231,15 @@ 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) {
+      var envVariables =
+          ImmutableMap.<RepoRecordedInput.EnvVar, Optional<String>>builder()
+              // The environment variable dependencies statically declared via the 'environ'
+              // attribute.
+              .putAll(RepoRecordedInput.EnvVar.wrap(extension.getStaticEnvVars()))
+              // The environment variable dependencies dynamically declared via the 'getenv' method.
+              .putAll(moduleExtensionResult.getRecordedEnvVarInputs())
+              .buildKeepingLast();
+
       lockFileInfo =
           Optional.of(
               new LockFileModuleExtension.WithFactors(
@@ -241,7 +251,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(ImmutableSortedMap.copyOf(envVariables))
                       .setGeneratedRepoSpecs(generatedRepoSpecs)
                       .setModuleExtensionMetadata(moduleExtensionMetadata)
                       .setRecordedRepoMappingEntries(
@@ -284,7 +294,11 @@ private SingleExtensionValue tryGettingValueFromLockFile(
                 + extensionId
                 + "' or one of its transitive .bzl files has changed");
       }
-      if (!extension.getEnvVars().equals(lockedExtension.getEnvVariables())) {
+      if (didRecordedInputsChange(
+          env,
+          directories,
+          // didRecordedInputsChange expects possibly null String values.
+          Maps.transformValues(lockedExtension.getEnvVariables(), v -> v.orElse(null)))) {
         diffRecorder.record(
             "The environment variables the extension '"
                 + extensionId
@@ -415,7 +429,7 @@ private static boolean didRepoMappingsChange(
   private static boolean didRecordedInputsChange(
       Environment env,
       BlazeDirectories directories,
-      ImmutableMap<? extends RepoRecordedInput, String> recordedInputs)
+      Map<? extends RepoRecordedInput, String> recordedInputs)
       throws InterruptedException, NeedsSkyframeRestartException {
     boolean upToDate = RepoRecordedInput.areAllValuesUpToDate(env, directories, recordedInputs);
     if (env.valuesMissing()) {
@@ -523,15 +537,15 @@ private BzlLoadValue loadBzlFile(
    * <p>The general idiom is to "load" such a {@link RunnableExtension} object by getting as much
    * information about it as needed to determine whether it can be reused from the lockfile (hence
    * methods such as {@link #getEvalFactors()}, {@link #getBzlTransitiveDigest()}, {@link
-   * #getEnvVars()}). Then the {@link #run} method can be called if it's determined that we can't
-   * reuse the cached results in the lockfile and have to re-run this extension.
+   * #getStaticEnvVars()}). Then the {@link #run} method can be called if it's determined that we
+   * can't reuse the cached results in the lockfile and have to re-run this extension.
    */
   private interface RunnableExtension {
     ModuleExtensionEvalFactors getEvalFactors();
 
     byte[] getBzlTransitiveDigest();
 
-    ImmutableMap<String, String> getEnvVars();
+    ImmutableMap<String, Optional<String>> getStaticEnvVars();
 
     @Nullable
     RunModuleExtensionResult run(
@@ -682,7 +696,7 @@ public byte[] getBzlTransitiveDigest() {
     }
 
     @Override
-    public ImmutableMap<String, String> getEnvVars() {
+    public ImmutableMap<String, Optional<String>> getStaticEnvVars() {
       return ImmutableMap.of();
     }
 
@@ -774,6 +788,7 @@ public RunModuleExtensionResult run(
         generatedRepoSpecs.put(name, repoSpec);
       }
       return RunModuleExtensionResult.create(
+          ImmutableMap.of(),
           ImmutableMap.of(),
           ImmutableMap.of(),
           generatedRepoSpecs.buildOrThrow(),
@@ -820,7 +835,7 @@ private RegularRunnableExtension loadRegularRunnableExtension(
           Transience.PERSISTENT);
     }
 
-    ImmutableMap<String, String> envVars =
+    ImmutableMap<String, Optional<String>> envVars =
         RepositoryFunction.getEnvVarValues(env, ImmutableSet.copyOf(extension.getEnvVariables()));
     if (envVars == null) {
       return null;
@@ -831,15 +846,15 @@ private RegularRunnableExtension loadRegularRunnableExtension(
   private final class RegularRunnableExtension implements RunnableExtension {
     private final BzlLoadValue bzlLoadValue;
     private final ModuleExtension extension;
-    private final ImmutableMap<String, String> envVars;
+    private final ImmutableMap<String, Optional<String>> staticEnvVars;
 
     RegularRunnableExtension(
         BzlLoadValue bzlLoadValue,
         ModuleExtension extension,
-        ImmutableMap<String, String> envVars) {
+        ImmutableMap<String, Optional<String>> staticEnvVars) {
       this.bzlLoadValue = bzlLoadValue;
       this.extension = extension;
-      this.envVars = envVars;
+      this.staticEnvVars = staticEnvVars;
     }
 
     @Override
@@ -850,8 +865,8 @@ public ModuleExtensionEvalFactors getEvalFactors() {
     }
 
     @Override
-    public ImmutableMap<String, String> getEnvVars() {
-      return envVars;
+    public ImmutableMap<String, Optional<String>> getStaticEnvVars() {
+      return staticEnvVars;
     }
 
     @Override
@@ -952,6 +967,7 @@ public RunModuleExtensionResult run(
       return RunModuleExtensionResult.create(
           moduleContext.getRecordedFileInputs(),
           moduleContext.getRecordedDirentsInputs(),
+          moduleContext.getRecordedEnvVarInputs(),
           threadContext.getGeneratedRepoSpecs(),
           moduleExtensionMetadata,
           repoMappingRecorder.recordedEntries());
@@ -1016,6 +1032,8 @@ abstract static class RunModuleExtensionResult {
 
     abstract ImmutableMap<RepoRecordedInput.Dirents, String> getRecordedDirentsInputs();
 
+    abstract ImmutableMap<RepoRecordedInput.EnvVar, Optional<String>> getRecordedEnvVarInputs();
+
     abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
     abstract Optional<ModuleExtensionMetadata> getModuleExtensionMetadata();
@@ -1025,12 +1043,14 @@ abstract static class RunModuleExtensionResult {
     static RunModuleExtensionResult create(
         ImmutableMap<RepoRecordedInput.File, String> recordedFileInputs,
         ImmutableMap<RepoRecordedInput.Dirents, String> recordedDirentsInputs,
+        ImmutableMap<RepoRecordedInput.EnvVar, Optional<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/bzlmod/SingleExtensionEvalFunction.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, Optional<String>> getRecordedEnvVarInputs()
+      throws InterruptedException {
+    // getEnvVarValues doesn't return null since the Skyframe dependencies have already been
+    // established by getenv calls.
+    return RepoRecordedInput.EnvVar.wrap(
+        ImmutableSortedMap.copyOf(RepositoryFunction.getEnvVarValues(env, accumulatedEnvKeys)));
   }
 
   protected void checkInOutputDirectory(String operation, StarlarkPath path)

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Table;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
@@ -361,16 +362,14 @@ 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(
+          Maps.transformValues(
+              starlarkRepositoryContext.getRecordedEnvVarInputs(), v -> v.orElse(null)));
 
       for (Table.Cell<RepositoryName, String, RepositoryName> repoMappings :
           repoMappingRecorder.recordedEntries().cellSet()) {

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.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,13 @@ public RepoRecordedInput parse(String s) {
 
     final String name;
 
-    public EnvVar(String name) {
+    public static ImmutableMap<EnvVar, Optional<String>> wrap(
+        Map<String, Optional<String>> envVars) {
+      return envVars.entrySet().stream()
+          .collect(toImmutableMap(e -> new EnvVar(e.getKey()), Map.Entry::getValue));
+    }
+
+    private EnvVar(String name) {
       this.name = name;
     }