Keep credentials cached across build commands.

src/main/java/com/google/devtools/build/lib/authandtls/BUILD

@@ -22,6 +22,7 @@ java_library(
         "//src/main/java/com/google/devtools/common/options",
         "//third_party:auth",
         "//third_party:auto_value",
+        "//third_party:caffeine",
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party:netty",

src/main/java/com/google/devtools/build/lib/authandtls/GoogleAuthUtils.java

@@ -14,11 +14,15 @@
 
 package com.google.devtools.build.lib.authandtls;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperCredentials;
 import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperEnvironment;
 import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperProvider;
@@ -48,6 +52,7 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URI;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -248,6 +253,7 @@ public static CallCredentialsProvider newCallCredentialsProvider(@Nullable Crede
    */
   public static Credentials newCredentials(
       CredentialHelperEnvironment credentialHelperEnvironment,
+      Cache<URI, ImmutableMap<String, ImmutableList<String>>> credentialCache,
       CommandLinePathFactory commandLinePathFactory,
       FileSystem fileSystem,
       AuthAndTLSOptions authAndTlsOptions)
@@ -257,12 +263,12 @@ public static Credentials newCredentials(
     Preconditions.checkNotNull(fileSystem);
     Preconditions.checkNotNull(authAndTlsOptions);
 
-    Optional<Credentials> credentials = newGoogleCredentials(authAndTlsOptions);
+    Optional<Credentials> fallbackCredentials = newGoogleCredentials(authAndTlsOptions);
 
-    if (credentials.isEmpty()) {
+    if (fallbackCredentials.isEmpty()) {
       // Fallback to .netrc if it exists.
       try {
-        credentials =
+        fallbackCredentials =
             newCredentialsFromNetrc(credentialHelperEnvironment.getClientEnvironment(), fileSystem);
       } catch (IOException e) {
         // TODO(yannic): Make this fail the build.
@@ -276,8 +282,8 @@ public static Credentials newCredentials(
             commandLinePathFactory,
             authAndTlsOptions.credentialHelpers),
         credentialHelperEnvironment,
-        credentials,
-        authAndTlsOptions.credentialHelperCacheTimeout);
+        credentialCache,
+        fallbackCredentials);
   }
 
   /**

src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper/BUILD

@@ -8,9 +8,23 @@ filegroup(
     visibility = ["//src:__subpackages__"],
 )
 
+java_library(
+    name = "credential_module",
+    srcs = ["CredentialModule.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib:runtime",
+        "//src/main/java/com/google/devtools/build/lib/authandtls",
+        "//third_party:caffeine",
+        "//third_party:guava",
+    ],
+)
+
 java_library(
     name = "credentialhelper",
-    srcs = glob(["*.java"]),
+    srcs = glob(
+        ["*.java"],
+        exclude = ["CredentialModule.java"],
+    ),
     deps = [
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/profiler",

src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialHelperCredentials.java

@@ -14,48 +14,36 @@
 
 package com.google.devtools.build.lib.authandtls.credentialhelper;
 
-import com.github.benmanes.caffeine.cache.CacheLoader;
-import com.github.benmanes.caffeine.cache.Caffeine;
-import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.github.benmanes.caffeine.cache.Cache;
 import com.google.auth.Credentials;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.net.URI;
-import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import javax.annotation.Nullable;
 
 /**
  * Implementation of {@link Credentials} which fetches credentials by invoking a {@code credential
  * helper} as subprocess, falling back to another {@link Credentials} if no suitable helper exists.
  */
 public class CredentialHelperCredentials extends Credentials {
+  private final CredentialHelperProvider credentialHelperProvider;
+  private final CredentialHelperEnvironment credentialHelperEnvironment;
+  private final Cache<URI, ImmutableMap<String, ImmutableList<String>>> credentialCache;
   private final Optional<Credentials> fallbackCredentials;
 
-  private final LoadingCache<URI, GetCredentialsResponse> credentialCache;
-
   public CredentialHelperCredentials(
       CredentialHelperProvider credentialHelperProvider,
       CredentialHelperEnvironment credentialHelperEnvironment,
-      Optional<Credentials> fallbackCredentials,
-      Duration cacheTimeout) {
-    Preconditions.checkNotNull(credentialHelperProvider);
-    Preconditions.checkNotNull(credentialHelperEnvironment);
+      Cache<URI, ImmutableMap<String, ImmutableList<String>>> credentialCache,
+      Optional<Credentials> fallbackCredentials) {
+    this.credentialHelperProvider = Preconditions.checkNotNull(credentialHelperProvider);
+    this.credentialHelperEnvironment = Preconditions.checkNotNull(credentialHelperEnvironment);
+    this.credentialCache = Preconditions.checkNotNull(credentialCache);
     this.fallbackCredentials = Preconditions.checkNotNull(fallbackCredentials);
-    Preconditions.checkNotNull(cacheTimeout);
-    Preconditions.checkArgument(
-        !cacheTimeout.isNegative() && !cacheTimeout.isZero(),
-        "Cache timeout must be greater than 0");
-
-    credentialCache =
-        Caffeine.newBuilder()
-            .expireAfterWrite(cacheTimeout)
-            .build(
-                new CredentialHelperCacheLoader(
-                    credentialHelperProvider, credentialHelperEnvironment));
   }
 
   @Override
@@ -71,9 +59,14 @@ public String getAuthenticationType() {
   public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
     Preconditions.checkNotNull(uri);
 
-    Optional<Map<String, List<String>>> credentials = getRequestMetadataFromCredentialHelper(uri);
-    if (credentials.isPresent()) {
-      return credentials.get();
+    ImmutableMap<String, ImmutableList<String>> credentials;
+    try {
+      credentials = credentialCache.get(uri, this::getCredentialsFromHelper);
+    } catch (Exception e) {
+      throw new IOException(e);
+    }
+    if (credentials != null) {
+      return (Map) credentials;
     }
 
     if (fallbackCredentials.isPresent()) {
@@ -83,13 +76,27 @@ public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException
     return ImmutableMap.of();
   }
 
-  @SuppressWarnings("unchecked") // Map<String, ImmutableList<String>> to Map<String<List<String>>
-  private Optional<Map<String, List<String>>> getRequestMetadataFromCredentialHelper(URI uri) {
+  private ImmutableMap<String, ImmutableList<String>> getCredentialsFromHelper(URI uri) {
     Preconditions.checkNotNull(uri);
 
-    GetCredentialsResponse response = credentialCache.get(uri);
+    Optional<CredentialHelper> maybeCredentialHelper =
+        credentialHelperProvider.findCredentialHelper(uri);
+    if (maybeCredentialHelper.isEmpty()) {
+      return null;
+    }
+    CredentialHelper credentialHelper = maybeCredentialHelper.get();
 
-    return Optional.ofNullable(response).map(value -> (Map) value.getHeaders());
+    GetCredentialsResponse response;
+    try {
+      response = credentialHelper.getCredentials(credentialHelperEnvironment, uri);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+    if (response == null) {
+      return null;
+    }
+
+    return response.getHeaders();
   }
 
   @Override
@@ -110,32 +117,4 @@ public void refresh() throws IOException {
 
     credentialCache.invalidateAll();
   }
-
-  private static final class CredentialHelperCacheLoader
-      implements CacheLoader<URI, GetCredentialsResponse> {
-    private final CredentialHelperProvider credentialHelperProvider;
-    private final CredentialHelperEnvironment credentialHelperEnvironment;
-
-    public CredentialHelperCacheLoader(
-        CredentialHelperProvider credentialHelperProvider,
-        CredentialHelperEnvironment credentialHelperEnvironment) {
-      this.credentialHelperProvider = Preconditions.checkNotNull(credentialHelperProvider);
-      this.credentialHelperEnvironment = Preconditions.checkNotNull(credentialHelperEnvironment);
-    }
-
-    @Nullable
-    @Override
-    public GetCredentialsResponse load(URI uri) throws IOException, InterruptedException {
-      Preconditions.checkNotNull(uri);
-
-      Optional<CredentialHelper> maybeCredentialHelper =
-          credentialHelperProvider.findCredentialHelper(uri);
-      if (maybeCredentialHelper.isEmpty()) {
-        return null;
-      }
-      CredentialHelper credentialHelper = maybeCredentialHelper.get();
-
-      return credentialHelper.getCredentials(credentialHelperEnvironment, uri);
-    }
-  }
 }

src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialModule.java

@@ -0,0 +1,43 @@
+// Copyright 2022 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.authandtls.credentialhelper;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
+import com.google.devtools.build.lib.runtime.BlazeModule;
+import com.google.devtools.build.lib.runtime.CommandEnvironment;
+import java.net.URI;
+import java.time.Duration;
+
+/** A module whose sole purpose is to hold the credential cache which is shared by other modules. */
+public class CredentialModule extends BlazeModule {
+  private final Cache<URI, ImmutableMap<String, ImmutableList<String>>> credentialCache =
+      Caffeine.newBuilder().expireAfterWrite(Duration.ZERO).build();
+
+  /** Returns the credential cache. */
+  public Cache<URI, ImmutableMap<String, ImmutableList<String>>> getCredentialCache() {
+    return credentialCache;
+  }
+
+  @Override
+  public void beforeCommand(CommandEnvironment env) {
+    // Update the cache expiration policy according to the command options.
+    AuthAndTLSOptions authAndTlsOptions = env.getOptions().getOptions(AuthAndTLSOptions.class);
+    credentialCache.policy().expireAfterWrite().get().setExpiresAfter(authAndTlsOptions.credentialHelperCacheTimeout);
+  }
+}

src/main/java/com/google/devtools/build/lib/bazel/BUILD

@@ -26,6 +26,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
+        "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper:credential_module",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:inspection",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:inspection_impl",
@@ -136,6 +137,7 @@ java_library(
         ":spawn_log_module",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_version_info",
+        "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper:credential_module",
         "//src/main/java/com/google/devtools/build/lib/bazel/coverage",
         "//src/main/java/com/google/devtools/build/lib/bazel/debug:workspace-rule-module",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository",

src/main/java/com/google/devtools/build/lib/bazel/Bazel.java

@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.BlazeVersionInfo;
+import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialModule;
 import com.google.devtools.build.lib.runtime.BlazeModule;
 import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import java.io.IOException;
@@ -42,6 +43,8 @@ public final class Bazel {
           // This module needs to be registered before any module providing a SpawnCache
           // implementation.
           com.google.devtools.build.lib.runtime.NoSpawnCacheModule.class,
+          // This module needs to be registered before any module that uses the credential cache.
+          CredentialModule.class,
           com.google.devtools.build.lib.runtime.CommandLogModule.class,
           com.google.devtools.build.lib.runtime.MemoryPressureModule.class,
           com.google.devtools.build.lib.platform.SleepPreventionModule.class,

src/main/java/com/google/devtools/build/lib/buildeventservice/BUILD

@@ -38,9 +38,11 @@ java_library(
         ":buildeventservice-options",
         "//src/main/java/com/google/devtools/build/lib:build-request-options",
         "//src/main/java/com/google/devtools/build/lib:runtime",
+        "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/test_configuration",
         "//src/main/java/com/google/devtools/build/lib/authandtls",
         "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper",
+        "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper:credential_module",
         "//src/main/java/com/google/devtools/build/lib/bugreport",
         "//src/main/java/com/google/devtools/build/lib/buildeventservice/client",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream",

src/main/java/com/google/devtools/build/lib/buildeventservice/BazelBuildEventServiceModule.java

@@ -22,12 +22,16 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
 import com.google.devtools.build.lib.authandtls.GoogleAuthUtils;
 import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperEnvironment;
+import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialModule;
 import com.google.devtools.build.lib.buildeventservice.client.BuildEventServiceClient;
 import com.google.devtools.build.lib.buildeventservice.client.BuildEventServiceGrpcClient;
+import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
+import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
 import io.grpc.ClientInterceptor;
 import io.grpc.ManagedChannel;
 import io.grpc.Metadata;
@@ -68,6 +72,16 @@ static BackendConfig create(
   private BuildEventServiceClient client;
   private BackendConfig config;
 
+  private CredentialModule credentialModule;
+
+  @Override
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
+    Preconditions.checkState(credentialModule == null, "credentialModule must be null");
+    credentialModule =
+        Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class));
+  }
+
   @Override
   protected Class<BuildEventServiceOptions> optionsClass() {
     return BuildEventServiceOptions.class;
@@ -93,10 +107,12 @@ protected BuildEventServiceClient getBesClient(
                   .setClientEnvironment(env.getClientEnv())
                   .setHelperExecutionTimeout(authAndTLSOptions.credentialHelperTimeout)
                   .build(),
+              credentialModule.getCredentialCache(),
               env.getCommandLinePathFactory(),
               env.getRuntime().getFileSystem(),
               newConfig.authAndTLSOptions());
 
+
       config = newConfig;
       client =
           new BuildEventServiceGrpcClient(

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -63,6 +63,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis/platform:platform_utils",
         "//src/main/java/com/google/devtools/build/lib/authandtls",
         "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper",
+        "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper:credential_module",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream",
         "//src/main/java/com/google/devtools/build/lib/clock",