Skip to content

Commit

Permalink
Keep credentials cached across build commands.
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
tjgq committed Nov 22, 2022
1 parent 581796a commit 4e469f7
Show file tree
Hide file tree
Showing 19 changed files with 264 additions and 132 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand All @@ -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.
Expand All @@ -276,8 +282,8 @@ public static Credentials newCredentials(
commandLinePathFactory,
authAndTlsOptions.credentialHelpers),
credentialHelperEnvironment,
credentials,
authAndTlsOptions.credentialHelperCacheTimeout);
credentialCache,
fallbackCredentials);
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()) {
Expand All @@ -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
Expand All @@ -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);
}
}
}
Original file line number Diff line number Diff line change
@@ -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);
}
}
2 changes: 2 additions & 0 deletions src/main/java/com/google/devtools/build/lib/bazel/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
3 changes: 3 additions & 0 deletions src/main/java/com/google/devtools/build/lib/bazel/Bazel.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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(
Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/google/devtools/build/lib/remote/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading

0 comments on commit 4e469f7

Please sign in to comment.