Keep credentials cached across build commands.

When using a credential helper, the lifetime of the credential cache is currently tied to an individual command, which causes the helper to be called for every command resulting in poor incremental build latency for builds using a non-trivial helper. Since the cache must be shared by RemoteModule and BazelBuildServiceModule, I've introduced a new CredentialModule whose sole purpose is to provide access to it. Closes bazelbuild#16822. PiperOrigin-RevId: 491598103 Change-Id: Ib668954b635a0e9498f0a7418707d6a2dfae0265
sluongng · Nov 29, 2022 · 2baa5a4 · 2baa5a4
1 parent 31e4bf4
commit 2baa5a4
Show file tree

Hide file tree

Showing 22 changed files with 328 additions and 146 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/authandtls/AuthAndTLSOptions.java b/src/main/java/com/google/devtools/build/lib/authandtls/AuthAndTLSOptions.java
@@ -175,7 +175,11 @@ public class AuthAndTLSOptions extends OptionsBase {
       converter = DurationConverter.class,
       documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
       effectTags = {OptionEffectTag.UNKNOWN},
-      help = "Configures the duration for which credentials from Credential Helpers are cached.")
+      help =
+          "Configures the duration for which credentials from Credential Helpers are cached.\n\n"
+              + "Invoking with a different value will adjust the lifetime of preexisting entries;"
+              + " pass zero to clear the cache. A clean command always clears the cache, regardless"
+              + " of this flag.")
   public Duration credentialHelperCacheTimeout;
 
   /** One of the values of the `--credential_helper` flag. */

diff --git a/src/main/java/com/google/devtools/build/lib/authandtls/BUILD b/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",

diff --git a/src/main/java/com/google/devtools/build/lib/authandtls/GoogleAuthUtils.java b/src/main/java/com/google/devtools/build/lib/authandtls/GoogleAuthUtils.java
@@ -14,11 +14,14 @@
 
 package com.google.devtools.build.lib.authandtls;
 
+import com.github.benmanes.caffeine.cache.Cache;
 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 +51,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 +252,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 +262,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 +281,8 @@ public static Credentials newCredentials(
             commandLinePathFactory,
             authAndTlsOptions.credentialHelpers),
         credentialHelperEnvironment,
-        credentials,
-        authAndTlsOptions.credentialHelperCacheTimeout);
+        credentialCache,
+        fallbackCredentials);
   }
 
   /**

diff --git a/src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper/BUILD b/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",

diff --git a/...main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialHelper.java b/...main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialHelper.java
@@ -67,7 +67,7 @@ public Path getPath() {
    * @return The response from the subprocess.
    */
   public GetCredentialsResponse getCredentials(CredentialHelperEnvironment environment, URI uri)
-      throws InterruptedException, IOException {
+      throws IOException {
     Preconditions.checkNotNull(environment);
     Preconditions.checkNotNull(uri);
 
@@ -81,7 +81,16 @@ public GetCredentialsResponse getCredentials(CredentialHelperEnvironment environ
           GSON.toJson(GetCredentialsRequest.newBuilder().setUri(uri).build(), stdin);
         }
 
-        process.waitFor();
+        try {
+          process.waitFor();
+        } catch (InterruptedException e) {
+          throw new CredentialHelperException(
+              String.format(
+                  Locale.US,
+                  "Failed to get credentials for '%s' from helper '%s': process was interrupted",
+                  uri,
+                  path));
+        }
 
         if (process.timedout()) {
           throw new CredentialHelperException(

diff --git a/...om/google/devtools/build/lib/authandtls/credentialhelper/CredentialHelperCredentials.java b/...om/google/devtools/build/lib/authandtls/credentialhelper/CredentialHelperCredentials.java
@@ -14,15 +14,13 @@
 
 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;
@@ -33,29 +31,34 @@
  * 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;
+  /** Wraps around an {@link IOException} so we can smuggle it through {@link Cache#get}. */
+  public static final class WrappedIOException extends RuntimeException {
+    private final IOException wrapped;
+
+    WrappedIOException(IOException e) {
+      super(e);
+      this.wrapped = e;
+    }
+
+    IOException getWrapped() {
+      return wrapped;
+    }
+  }
 
   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
@@ -68,12 +71,18 @@ public String getAuthenticationType() {
   }
 
   @Override
+  @SuppressWarnings("unchecked") // Map<String, ImmutableList<String>> to Map<String<List<String>>
   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 (WrappedIOException e) {
+      throw e.getWrapped();
+    }
+    if (credentials != null) {
+      return (Map) credentials;
     }
 
     if (fallbackCredentials.isPresent()) {
@@ -83,13 +92,28 @@ 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) {
+  @Nullable
+  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();
+
+    GetCredentialsResponse response;
+    try {
+      response = credentialHelper.getCredentials(credentialHelperEnvironment, uri);
+    } catch (IOException e) {
+      throw new WrappedIOException(e);
+    }
+    if (response == null) {
+      return null;
+    }
 
-    return Optional.ofNullable(response).map(value -> (Map) value.getHeaders());
+    return response.getHeaders();
   }
 
   @Override
@@ -110,32 +134,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);
-    }
-  }
 }
diff --git a/...main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialModule.java b/...main/java/com/google/devtools/build/lib/authandtls/credentialhelper/CredentialModule.java
@@ -0,0 +1,52 @@
+// 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);
+
+    // Clear the cache on clean.
+    if (env.getCommand().name().equals("clean")) {
+      credentialCache.invalidateAll();
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/BUILD
@@ -136,6 +136,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",

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/Bazel.java b/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,

diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BUILD b/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",