Skip to content

Commit 89428d6

Browse files
committed
feat(gax): support transparent retries during mTLS certificate rotations
- Add CertificateBasedAccess and WorkloadCertificateUtils for SPIFFE and custom certificate loading - Implement RefreshingHttpJsonChannel and ChannelPool mTLS certificate fingerprint tracking and rotation - Enable transparent retries for retryable UnauthenticatedExceptions in ApiResultRetryAlgorithm and AttemptCallable - Add override delegation for getEndpoint, getHttpTransport, and getExecutor to preserve SLF4J MDC logging in Showcase tests
1 parent a86cc11 commit 89428d6

33 files changed

Lines changed: 1988 additions & 230 deletions

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java

Lines changed: 120 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
import com.google.api.core.InternalApi;
3333
import com.google.api.gax.core.FixedExecutorProvider;
34+
import com.google.api.gax.rpc.mtls.WorkloadCertificateUtils;
3435
import com.google.common.annotations.VisibleForTesting;
3536
import com.google.common.base.Preconditions;
3637
import com.google.common.collect.ImmutableList;
@@ -72,18 +73,36 @@
7273
@NullMarked
7374
class ChannelPool extends ManagedChannel {
7475
static final String CHANNEL_POOL_CONSECUTIVE_RESIZING_WARNING =
75-
"The gRPC ChannelPool used in the client has been flagged to be repeatedly resizing (5+ times). See https://github.com/googleapis/google-cloud-java/blob/main/docs/grpc_channel_pool_guide.md for more information about this behavior.";
76+
"The gRPC ChannelPool used in the client has been flagged to be repeatedly resizing (5+"
77+
+ " times). See"
78+
+ " https://github.com/googleapis/google-cloud-java/blob/main/docs/grpc_channel_pool_guide.md"
79+
+ " for more information about this behavior.";
7680
@VisibleForTesting static final Logger LOG = Logger.getLogger(ChannelPool.class.getName());
7781
private static final java.time.Duration REFRESH_PERIOD = java.time.Duration.ofMinutes(50);
7882

7983
private final ChannelPoolSettings settings;
8084
private final ChannelFactory channelFactory;
8185
private final FixedExecutorProvider backgroundExecutorProvider;
86+
private final String workloadCertPath;
8287

8388
private @Nullable ScheduledFuture<?> refreshFuture = null;
8489
private @Nullable ScheduledFuture<?> resizeFuture = null;
8590

91+
private static class DiskCheckResult {
92+
final String fingerprint;
93+
final long timestampNanos;
94+
95+
DiskCheckResult(String fingerprint, long timestampNanos) {
96+
this.fingerprint = fingerprint;
97+
this.timestampNanos = timestampNanos;
98+
}
99+
}
100+
101+
private volatile DiskCheckResult lastDiskCheck = null;
102+
private final java.util.concurrent.locks.ReentrantLock diskCheckLock =
103+
new java.util.concurrent.locks.ReentrantLock();
86104
private final Object entryWriteLock = new Object();
105+
private volatile String activeCertFingerprint = "";
87106
@VisibleForTesting final AtomicReference<ImmutableList<Entry>> entries = new AtomicReference<>();
88107
private final AtomicInteger indexTicker = new AtomicInteger();
89108
private final String authority;
@@ -100,14 +119,15 @@ class ChannelPool extends ManagedChannel {
100119
static ChannelPool create(
101120
ChannelPoolSettings settings,
102121
ChannelFactory channelFactory,
103-
@Nullable ScheduledExecutorService backgroundExecutor)
122+
@Nullable ScheduledExecutorService backgroundExecutor,
123+
@Nullable String workloadCertPath)
104124
throws IOException {
105125

106126
FixedExecutorProvider executorProvider =
107127
backgroundExecutor == null
108128
? FixedExecutorProvider.create(Executors.newSingleThreadScheduledExecutor(), true)
109129
: FixedExecutorProvider.create(backgroundExecutor, false);
110-
return new ChannelPool(settings, channelFactory, executorProvider);
130+
return new ChannelPool(settings, channelFactory, executorProvider, workloadCertPath);
111131
}
112132

113133
/**
@@ -121,11 +141,13 @@ static ChannelPool create(
121141
ChannelPool(
122142
ChannelPoolSettings settings,
123143
ChannelFactory channelFactory,
124-
FixedExecutorProvider executorProvider)
144+
FixedExecutorProvider executorProvider,
145+
@Nullable String workloadCertPath)
125146
throws IOException {
126147
this.settings = settings;
127148
this.channelFactory = channelFactory;
128149
this.backgroundExecutorProvider = executorProvider;
150+
this.workloadCertPath = workloadCertPath;
129151

130152
ImmutableList.Builder<Entry> initialListBuilder = ImmutableList.builder();
131153

@@ -136,6 +158,11 @@ static ChannelPool create(
136158
entries.set(initialListBuilder.build());
137159
authority = entries.get().get(0).channel.authority();
138160

161+
if (workloadCertPath != null) {
162+
this.activeCertFingerprint =
163+
WorkloadCertificateUtils.getCertificateFingerprint(workloadCertPath);
164+
}
165+
139166
if (!settings.isStaticSize()) {
140167
resizeFuture =
141168
backgroundExecutorProvider
@@ -421,12 +448,54 @@ private void expand(int desiredSize) {
421448

422449
private void refreshSafely() {
423450
try {
424-
refresh();
451+
synchronized (entryWriteLock) {
452+
if (workloadCertPath != null) {
453+
String currentDiskFingerprint = getOrUpdateDiskFingerprint(workloadCertPath);
454+
if (!currentDiskFingerprint.isEmpty()) {
455+
this.activeCertFingerprint = currentDiskFingerprint;
456+
}
457+
}
458+
refreshAll();
459+
}
425460
} catch (Exception e) {
426-
LOG.log(Level.WARNING, "Failed to pre-emptively refresh channnels", e);
461+
LOG.log(Level.WARNING, "Failed to pre-emptively refresh channels", e);
427462
}
428463
}
429464

465+
private String getOrUpdateDiskFingerprint(String certPath) {
466+
long now = System.nanoTime();
467+
DiskCheckResult cached = lastDiskCheck;
468+
if (cached != null
469+
&& (now - cached.timestampNanos < java.util.concurrent.TimeUnit.SECONDS.toNanos(1))) {
470+
return cached.fingerprint;
471+
}
472+
473+
diskCheckLock.lock();
474+
try {
475+
cached = lastDiskCheck;
476+
if (cached != null
477+
&& (now - cached.timestampNanos < java.util.concurrent.TimeUnit.SECONDS.toNanos(1))) {
478+
return cached.fingerprint;
479+
}
480+
String fingerprint = WorkloadCertificateUtils.getCertificateFingerprint(certPath);
481+
lastDiskCheck = new DiskCheckResult(fingerprint, System.nanoTime());
482+
return fingerprint;
483+
} finally {
484+
diskCheckLock.unlock();
485+
}
486+
}
487+
488+
boolean shouldRefresh() {
489+
if (workloadCertPath == null) {
490+
return false;
491+
}
492+
String currentDiskFingerprint = getOrUpdateDiskFingerprint(workloadCertPath);
493+
if (currentDiskFingerprint.isEmpty()) {
494+
return false;
495+
}
496+
return !currentDiskFingerprint.equalsIgnoreCase(activeCertFingerprint);
497+
}
498+
430499
/**
431500
* Replace all of the channels in the channel pool with fresh ones. This is meant to mitigate the
432501
* hourly GFE disconnects by giving clients the ability to prime the channel on reconnect.
@@ -443,7 +512,35 @@ void refresh() {
443512
// - then thread2 will shut down channel that thread1 will put back into circulation (after it
444513
// replaces the list)
445514
synchronized (entryWriteLock) {
446-
LOG.fine("Refreshing all channels");
515+
if (workloadCertPath == null) {
516+
return;
517+
}
518+
String currentDiskFingerprint = getOrUpdateDiskFingerprint(workloadCertPath);
519+
if (currentDiskFingerprint.isEmpty()) {
520+
return;
521+
}
522+
523+
// Double-check fingerprint inside the lock
524+
if (currentDiskFingerprint.equalsIgnoreCase(this.activeCertFingerprint)) {
525+
LOG.fine(
526+
"Channel pool was already refreshed by a concurrent thread, skipping duplicate"
527+
+ " refresh");
528+
return;
529+
}
530+
531+
this.activeCertFingerprint = currentDiskFingerprint;
532+
refreshAll();
533+
}
534+
}
535+
536+
@InternalApi("Visible for testing")
537+
void refreshAll() {
538+
synchronized (entryWriteLock) {
539+
LOG.fine(
540+
"Refreshing all channels"
541+
+ (activeCertFingerprint == null
542+
? ""
543+
: " with certificate fingerprint: " + activeCertFingerprint));
447544
ArrayList<Entry> newEntries = new ArrayList<>(entries.get());
448545

449546
for (int i = 0; i < newEntries.size(); i++) {
@@ -621,7 +718,13 @@ public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
621718
}
622719
}
623720

624-
/** ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion. */
721+
/**
722+
* ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion.
723+
*
724+
* <p>Contract: Exactly one call to {@link #start(Listener, Metadata)} or explicit release via
725+
* {@link #cancel(String, Throwable)} is required to balance reference counts. Early cancellation
726+
* before {@code start()} safely decrements the reference count via atomic compare-and-set.
727+
*/
625728
static class ReleasingClientCall<ReqT, RespT> extends SimpleForwardingClientCall<ReqT, RespT> {
626729
private @Nullable CancellationException cancellationException;
627730
final Entry entry;
@@ -636,6 +739,9 @@ public ReleasingClientCall(ClientCall<ReqT, RespT> delegate, Entry entry) {
636739
@Override
637740
public void start(Listener<RespT> responseListener, Metadata headers) {
638741
if (cancellationException != null) {
742+
if (wasReleased.compareAndSet(false, true)) {
743+
entry.release();
744+
}
639745
throw new IllegalStateException("Call is already cancelled", cancellationException);
640746
}
641747
try {
@@ -646,7 +752,8 @@ public void onClose(Status status, Metadata trailers) {
646752
if (!wasClosed.compareAndSet(false, true)) {
647753
LOG.log(
648754
Level.WARNING,
649-
"Call is being closed more than once. Please make sure that onClose() is not being manually called.");
755+
"Call is being closed more than once. Please make sure that onClose() is not"
756+
+ " being manually called.");
650757
return;
651758
}
652759
try {
@@ -657,7 +764,8 @@ public void onClose(Status status, Metadata trailers) {
657764
} else {
658765
LOG.log(
659766
Level.WARNING,
660-
"Entry was released before the call is closed. This may be due to an exception on start of the call.");
767+
"Entry was released before the call is closed. This may be due to an"
768+
+ " exception on start of the call.");
661769
}
662770
}
663771
}
@@ -670,7 +778,8 @@ public void onClose(Status status, Metadata trailers) {
670778
} else {
671779
LOG.log(
672780
Level.WARNING,
673-
"The entry is already released. This indicates that onClose() has already been called previously");
781+
"The entry is already released. This indicates that onClose() has already been called"
782+
+ " previously");
674783
}
675784
throw e;
676785
}

0 commit comments

Comments
 (0)