Skip to content

refactor(gax): implement clean stateless channel refresh & retry mapping #13342

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
lqiu96 wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

refactor(gax): implement clean stateless channel refresh & retry mapping #13342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
14fda9c
refactor(gax): implement clean stateless channel refresh & retry mapping
lqiu96 Jun 2, 2026
14ea140
refactor(gax): decouple retry package from rpc channel refresh classes
lqiu96 Jun 3, 2026
49b08e8
refactor(gax): replace ExceptionMappingFuture class with inline catch…
lqiu96 Jun 3, 2026
7d02d82
refactor(gax): use ApiFutures.catching in AttemptCallable
lqiu96 Jun 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 106 additions & 1 deletion sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,11 @@
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.Status;
import java.io.FileInputStream;
import java.io.IOException;
import java.security.MessageDigest;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CancellationException;
Expand Down Expand Up @@ -81,7 +85,19 @@ class ChannelPool extends ManagedChannel {
private ScheduledFuture<?> refreshFuture = null;
private ScheduledFuture<?> resizeFuture = null;

private static class DiskCheckResult {
final String fingerprint;
final long timestampNanos;

DiskCheckResult(String fingerprint, long timestampNanos) {
this.fingerprint = fingerprint;
this.timestampNanos = timestampNanos;
}
}

private final AtomicReference<DiskCheckResult> lastDiskCheck = new AtomicReference<>(null);
private final Object entryWriteLock = new Object();
private volatile String activeCertFingerprint = "";
@VisibleForTesting final AtomicReference<ImmutableList<Entry>> entries = new AtomicReference<>();
private final AtomicInteger indexTicker = new AtomicInteger();
private final String authority;
Expand Down Expand Up @@ -134,6 +150,11 @@ static ChannelPool create(
entries.set(initialListBuilder.build());
authority = entries.get().get(0).channel.authority();

String certPath = getWorkloadCertPath();
if (certPath != null) {
this.activeCertFingerprint = getCertificateFingerprint(certPath);
}

if (!settings.isStaticSize()) {
resizeFuture =
backgroundExecutorProvider
Expand Down Expand Up @@ -425,6 +446,74 @@ private void refreshSafely() {
}
}


private static String getWorkloadCertPath() {
String configPath = System.getenv("GOOGLE_API_CERTIFICATE_CONFIG");
if (configPath != null && !configPath.isEmpty()) {
java.io.File configFile = new java.io.File(configPath);
if (configFile.exists() && !configFile.isDirectory()) {
// If explicit config exists, check it
}
}
Comment on lines +451 to +457
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The environment variable GOOGLE_API_CERTIFICATE_CONFIG is checked, but the inner if block is empty, meaning the configuration is completely ignored. If the config path points to a valid certificate file, its path should be returned.

Suggested change
String configPath = System.getenv("GOOGLE_API_CERTIFICATE_CONFIG");
if (configPath != null && !configPath.isEmpty()) {
java.io.File configFile = new java.io.File(configPath);
if (configFile.exists() && !configFile.isDirectory()) {
// If explicit config exists, check it
}
}
String configPath = System.getenv("GOOGLE_API_CERTIFICATE_CONFIG");
if (configPath != null && !configPath.isEmpty()) {
java.io.File configFile = new java.io.File(configPath);
if (configFile.exists() && !configFile.isDirectory()) {
return configFile.getAbsolutePath();
}
}

java.io.File bundleFile = new java.io.File("/var/run/secrets/workload-spiffe-credentials/credentialbundle.pem");
if (bundleFile.exists()) {
return bundleFile.getAbsolutePath();
}
java.io.File certsFile = new java.io.File("/var/run/secrets/workload-spiffe-credentials/certificates.pem");
if (certsFile.exists()) {
return certsFile.getAbsolutePath();
}
return null;
}

private static String getCertificateFingerprint(String certPath) {
try (FileInputStream fis = new FileInputStream(certPath)) {
CertificateFactory cf = CertificateFactory.getInstance("X.509");
X509Certificate cert = (X509Certificate) cf.generateCertificate(fis);
MessageDigest md = MessageDigest.getInstance("SHA-256");
byte[] der = cert.getEncoded();
byte[] digest = md.digest(der);
StringBuilder sb = new StringBuilder();
for (byte b : digest) {
sb.append(String.format("%02x", b));
}
return sb.toString();
} catch (Exception e) {
LOG.log(Level.FINE, "Could not read or parse workload certificate at path " + certPath, e);
return "";
}
}

private String getOrUpdateDiskFingerprint(String certPath) {
long now = System.nanoTime();
DiskCheckResult cached = lastDiskCheck.get();
if (cached != null && (now - cached.timestampNanos < java.util.concurrent.TimeUnit.SECONDS.toNanos(1))) {
return cached.fingerprint;
}

synchronized (lastDiskCheck) {
cached = lastDiskCheck.get();
if (cached != null && (now - cached.timestampNanos < java.util.concurrent.TimeUnit.SECONDS.toNanos(1))) {
return cached.fingerprint;
}
String fingerprint = getCertificateFingerprint(certPath);
lastDiskCheck.set(new DiskCheckResult(fingerprint, System.nanoTime()));
return fingerprint;
}
}

boolean shouldRefresh() {
String certPath = getWorkloadCertPath();
if (certPath == null) {
return false;
}
String currentDiskFingerprint = getOrUpdateDiskFingerprint(certPath);
if (currentDiskFingerprint.isEmpty()) {
return false;
}
return !currentDiskFingerprint.equalsIgnoreCase(activeCertFingerprint);
}

/**
* Replace all of the channels in the channel pool with fresh ones. This is meant to mitigate the
* hourly GFE disconnects by giving clients the ability to prime the channel on reconnect.
Expand All @@ -441,7 +530,23 @@ void refresh() {
// - then thread2 will shut down channel that thread1 will put back into circulation (after it
// replaces the list)
synchronized (entryWriteLock) {
LOG.fine("Refreshing all channels");
String certPath = getWorkloadCertPath();
if (certPath == null) {
return;
}
String currentDiskFingerprint = getOrUpdateDiskFingerprint(certPath);
if (currentDiskFingerprint.isEmpty()) {
return;
}

// Double-check fingerprint inside the lock
if (currentDiskFingerprint.equals(this.activeCertFingerprint)) {
LOG.fine("Channel pool was already refreshed by a concurrent thread, skipping duplicate refresh");
return;
}

this.activeCertFingerprint = currentDiskFingerprint;
LOG.fine("Refreshing all channels with new certificate fingerprint: " + activeCertFingerprint);
ArrayList<Entry> newEntries = new ArrayList<>(entries.get());

for (int i = 0; i < newEntries.size(); i++) {
Expand Down
58 changes: 41 additions & 17 deletions ...latform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ public final class GrpcCallContext implements ApiCallContext {
private final ApiCallContextOptions options;
private final EndpointContext endpointContext;
private final boolean isDirectPath;
@Nullable private final TransportChannel transportChannel;

/** Returns an empty instance with a null channel and default {@link CallOptions}. */
public static GrpcCallContext createDefault() {
Expand All @@ -113,7 +114,8 @@ public static GrpcCallContext createDefault() {
null,
null,
null,
false);
false,
null);
}

/** Returns an instance with the given channel and {@link CallOptions}. */
Expand All @@ -131,7 +133,8 @@ public static GrpcCallContext of(Channel channel, CallOptions callOptions) {
null,
null,
null,
false);
false,
null);
}

private GrpcCallContext(
Expand All @@ -147,7 +150,8 @@ private GrpcCallContext(
@Nullable RetrySettings retrySettings,
@Nullable Set<StatusCode.Code> retryableCodes,
@Nullable EndpointContext endpointContext,
boolean isDirectPath) {
boolean isDirectPath,
@Nullable TransportChannel transportChannel) {
this.channel = channel;
this.credentials = credentials;
Preconditions.checkNotNull(callOptions);
Expand All @@ -167,6 +171,7 @@ private GrpcCallContext(
this.endpointContext =
endpointContext == null ? EndpointContext.getDefaultInstance() : endpointContext;
this.isDirectPath = isDirectPath;
this.transportChannel = transportChannel;
}

/**
Expand Down Expand Up @@ -208,7 +213,13 @@ public GrpcCallContext withCredentials(Credentials newCredentials) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@Override
public TransportChannel getTransportChannel() {
return transportChannel;
}

@Override
Expand All @@ -232,7 +243,8 @@ public GrpcCallContext withTransportChannel(TransportChannel inputChannel) {
retrySettings,
retryableCodes,
endpointContext,
transportChannel.isDirectPath());
transportChannel.isDirectPath(),
inputChannel);
Comment on lines 243 to +247
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The parameter of withTransportChannel was renamed from transportChannel to inputChannel, but the method body still references transportChannel.isDirectPath(). This resolves to the nullable class field this.transportChannel instead of the new parameter, which will cause a NullPointerException if the field is null, or use the wrong channel's properties if it is not.

Suggested change
retrySettings,
retryableCodes,
endpointContext,
transportChannel.isDirectPath());
transportChannel.isDirectPath(),
inputChannel);
retrySettings,
retryableCodes,
endpointContext,
inputChannel.isDirectPath(),
inputChannel);

}

@Override
Expand All @@ -251,7 +263,8 @@ public GrpcCallContext withEndpointContext(EndpointContext endpointContext) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

/** This method is obsolete. Use {@link #withTimeoutDuration(java.time.Duration)} instead. */
Expand Down Expand Up @@ -286,7 +299,8 @@ public GrpcCallContext withTimeoutDuration(@Nullable java.time.Duration timeout)
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

/** This method is obsolete. Use {@link #getTimeoutDuration()} instead. */
Expand Down Expand Up @@ -335,7 +349,8 @@ public GrpcCallContext withStreamWaitTimeoutDuration(
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

/**
Expand Down Expand Up @@ -370,7 +385,8 @@ public GrpcCallContext withStreamIdleTimeoutDuration(
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@BetaApi("The surface for channel affinity is not stable yet and may change in the future.")
Expand All @@ -388,7 +404,8 @@ public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@BetaApi("The surface for extra headers is not stable yet and may change in the future.")
Expand All @@ -410,7 +427,8 @@ public GrpcCallContext withExtraHeaders(Map<String, List<String>> extraHeaders)
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@Override
Expand All @@ -433,7 +451,8 @@ public GrpcCallContext withRetrySettings(RetrySettings retrySettings) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@Override
Expand All @@ -456,7 +475,8 @@ public GrpcCallContext withRetryableCodes(Set<StatusCode.Code> retryableCodes) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

@Override
Expand Down Expand Up @@ -558,7 +578,8 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
newRetrySettings,
newRetryableCodes,
endpointContext,
newIsDirectPath);
newIsDirectPath,
transportChannel);
}

/** The {@link Channel} set on this context. */
Expand Down Expand Up @@ -641,7 +662,8 @@ public GrpcCallContext withChannel(Channel newChannel) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

/** Returns a new instance with the call options set to the given call options. */
Expand All @@ -659,7 +681,8 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams) {
Expand Down Expand Up @@ -704,7 +727,8 @@ public <T> GrpcCallContext withOption(Key<T> key, T value) {
retrySettings,
retryableCodes,
endpointContext,
isDirectPath);
isDirectPath,
transportChannel);
}

/** {@inheritDoc} */
Expand Down
17 changes: 17 additions & 0 deletions ...rm-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,23 @@ public Channel getChannel() {
return getManagedChannel();
}

@Override
public void refresh() {
Channel channel = getChannel();
if (channel instanceof ChannelPool) {
((ChannelPool) channel).refresh();
}
}

@Override
public boolean shouldRefresh() {
Channel channel = getChannel();
if (channel instanceof ChannelPool) {
return ((ChannelPool) channel).shouldRefresh();
}
return false;
}

@Override
public void shutdown() {
getManagedChannel().shutdown();
Expand Down
Loading
Loading