feat(gax): support transparent retries during mTLS certificate rotations

[macastelaz]

Description

This PR introduces robust capability to handle dynamic mTLS certificate rotations transparently, seamlessly retrying in-flight failures and draining transports gracefully without hard-aborting ongoing work.

Changes Made

• Graceful Draining (RefreshingHttpJsonChannel & ChannelPool): Previously, when mTLS certificates rotated on disk, the system would aggressively shut down the obsolete channel. This PR introduces reference-counted outstanding request tracking, allowing all actively executing requests on the retiring channel to complete cleanly before closing the underlying transport.

• Transparent Retries (ApiResultRetryAlgorithm): Handles UNAUTHENTICATED errors cleanly by catching and evaluating them against the current retry strategy, treating certificate-bound 401s as retryable transitions once the channel is refreshed.

• Certificate Source De-duplication (CertificateBasedAccess): Normalizes the static discovery of GOOGLE_API_USE_CLIENT_CERTIFICATE environment overrides and workload certificate checks, addressing concurrent initialization bottlenecks.

Based on PR #13246

[gemini-code-assist]

Code Review

This pull request introduces support for dynamic mTLS certificate rotation across both gRPC and HTTP/JSON transport channels by tracking certificate fingerprints, hot-swapping active channels, and triggering refreshes upon encountering UnauthenticatedException. The code review highlights several critical issues: a bug in ApiResultRetryAlgorithm that bypasses configured retries for non-retryable UnauthenticatedExceptions, a transient IllegalStateException during channel hot-swapping in RefreshingHttpJsonChannel, a backward compatibility break in CertificateBasedAccess when GOOGLE_API_USE_CLIENT_CERTIFICATE is explicitly set to "true", and a memory leak due to unpruned terminated entries in RefreshingHttpJsonChannel. Additionally, the reviewer recommends replacing synchronization on an AtomicReference with a lock-free read and an explicit ReentrantLock to avoid lock contention.

[gemini-code-assist]

The check for UnauthenticatedException is placed before the context.getRetryableCodes() != null check and returns isRetryable() immediately. This completely breaks retries for UnauthenticatedException when UNAUTHENTICATED is in the configured retryable codes but isRetryable() is false (which is the default). We should only return true immediately if isRetryable() is true (indicating a cert rotation retry), and otherwise fall through to the standard context.getRetryableCodes() check.

    if (previousThrowable instanceof UnauthenticatedException
        && ((UnauthenticatedException) previousThrowable).isRetryable()) {
      return true;
    }

[gemini-code-assist]

In getRetainedEntry(), if entry.retain() returns false (because a refresh occurred and oldEntry.shutdownRequested was set to true), the code immediately throws IllegalStateException("Channel has been shut down"). However, the channel pool itself is not shut down; only the old entry is being retired. This will cause transient IllegalStateExceptions for active requests during certificate rotation. We should only throw IllegalStateException if the active entry itself has been shut down (i.e., entry == activeEntry.get()), and otherwise retry the loop with the new active entry.

Suggested change

	private ChannelEntry getRetainedEntry() {
	while (true) {
	ChannelEntry entry = activeEntry.get();
	if (entry.retain()) {
	return entry;
	}
	if (entry.shutdownRequested.get()) {
	throw new IllegalStateException("Channel has been shut down");
	}
	}
	}
	private ChannelEntry getRetainedEntry() {
	while (true) {
	ChannelEntry entry = activeEntry.get();
	if (entry.retain()) {
	return entry;
	}
	if (entry == activeEntry.get()) {
	throw new IllegalStateException("Channel has been shut down");
	}
	}
	}

[gemini-code-assist]

The new implementation of useMtlsClientCertificate() breaks backward compatibility. Previously, if GOOGLE_API_USE_CLIENT_CERTIFICATE was set to "true", it returned true immediately, allowing custom MtlsProviders to be used. Now, it falls through to check for SPIFFE credentials or config files on disk, returning false if they are not found. This will break mTLS for users relying on custom MtlsProviders. We should preserve backward compatibility by returning true immediately if GOOGLE_API_USE_CLIENT_CERTIFICATE is explicitly set to "true".

Suggested change

	// 1. Check the explicit user flag first (Primary override)
	String useClientCertificate = envProvider.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE");
	return "true".equals(useClientCertificate);
	if (useClientCertificate != null && !useClientCertificate.isEmpty()) {
	if ("false".equalsIgnoreCase(useClientCertificate)) {
	return false;
	}
	}
	// 1. Check the explicit user flag first (Primary override)
	String useClientCertificate = envProvider.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE");
	if (useClientCertificate != null && !useClientCertificate.isEmpty()) {
	if ("false".equalsIgnoreCase(useClientCertificate)) {
	return false;
	}
	if ("true".equalsIgnoreCase(useClientCertificate)) {
	return true;
	}
	}

[gemini-code-assist]

allEntries queue accumulates ChannelEntry objects on every refresh(), but they are never removed once they are terminated. This leads to a memory leak over time in long-running processes with frequent certificate rotations. We should prune terminated entries from allEntries during refresh().

  @Override
  public void refresh() {
    synchronized (lock) {
      String certPath = WorkloadCertificateUtils.getWorkloadCertPath();
      if (certPath == null) {
        return;
      }
      String currentDiskFingerprint = getOrUpdateDiskFingerprint(certPath);
      if (currentDiskFingerprint.isEmpty()) {
        return;
      }

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

      this.activeCertFingerprint = currentDiskFingerprint;
      LOG.info("mTLS certificate rotation detected. Triggering HTTP/JSON channel pool refresh.");

      // Prune terminated entries to prevent memory leak
      allEntries.removeIf(entry -> entry.channel.isTerminated());

      ChannelEntry newEntry = new ChannelEntry(channelFactory.get());
      allEntries.add(newEntry);
      ChannelEntry oldEntry = activeEntry.getAndSet(newEntry);

      if (oldEntry != null) {
        oldEntry.requestShutdown();
      }
    }
  }

[gemini-code-assist]

Synchronizing on an AtomicReference (like lastDiskCheck) is a code smell. AtomicReference is designed for lock-free thread-safe operations. In performance-sensitive code, if a lock is needed to prevent concurrent disk reads, we should use an explicit ReentrantLock instead of the synchronized keyword to protect shared state.

Suggested change

	private final AtomicReference<DiskCheckResult> lastDiskCheck = new AtomicReference<>(null);
	private final AtomicReference<DiskCheckResult> lastDiskCheck = new AtomicReference<>(null);
	private final ReentrantLock diskCheckLock = new ReentrantLock();

References

1. In performance-sensitive code, prefer using explicit locks over the 'synchronized' keyword to protect shared state while ensuring thread safety and visibility.

[gemini-code-assist]

To avoid lock contention on the critical path, we should perform a lock-free read on lastDiskCheck first. If the cache is expired, we can then acquire an explicit ReentrantLock to perform the disk read and update the cache, rather than using synchronized.

Suggested change

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

References

1. In performance-sensitive code, prefer using explicit locks over the 'synchronized' keyword to protect shared state while ensuring thread safety and visibility.
2. To avoid lock contention on frequently called monitoring or routing methods (such as getLoad()), maintain state using atomic variables (e.g., AtomicLong, AtomicReference) to allow lock-free reads, rather than acquiring locks on the critical path.