This is an automated email from the ASF dual-hosted git repository.
fresh-borzoni pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git
The following commit(s) were added to refs/heads/main by this push:
new 2f77b22a6 [client] Fix stale metadata on readOnlyGateway by adding
RetryableGatewayClientProxy (#3390)
2f77b22a6 is described below
commit 2f77b22a6d070164355739bb7a738cd51a066a21
Author: Hongshun Wang <[email protected]>
AuthorDate: Wed Jun 3 17:40:55 2026 +0800
[client] Fix stale metadata on readOnlyGateway by adding
RetryableGatewayClientProxy (#3390)
* [client] Fix stale metadata on readOnlyGateway by adding
RetryableGatewayClientProxy
* modified based on advice.
* modified based on advice.
---
.../org/apache/fluss/client/admin/FlussAdmin.java | 30 ++-
.../fluss/client/metadata/MetadataUpdater.java | 60 +++++
.../fluss/client/admin/FlussAdminITCase.java | 27 +++
.../fluss/rpc/RetryableGatewayClientProxy.java | 215 ++++++++++++++++
.../fluss/rpc/RetryableGatewayClientProxyTest.java | 269 +++++++++++++++++++++
5 files changed, 599 insertions(+), 2 deletions(-)
diff --git
a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
index a1d429c99..954247850 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
@@ -48,6 +48,7 @@ import org.apache.fluss.metadata.TableInfo;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.metadata.TableStats;
import org.apache.fluss.rpc.GatewayClientProxy;
+import org.apache.fluss.rpc.RetryableGatewayClientProxy;
import org.apache.fluss.rpc.RpcClient;
import org.apache.fluss.rpc.gateway.AdminGateway;
import org.apache.fluss.rpc.gateway.AdminReadOnlyGateway;
@@ -98,6 +99,7 @@ import org.apache.fluss.rpc.messages.TableExistsResponse;
import org.apache.fluss.rpc.protocol.ApiError;
import org.apache.fluss.security.acl.AclBinding;
import org.apache.fluss.security.acl.AclBindingFilter;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.apache.fluss.utils.concurrent.FutureUtils;
import javax.annotation.Nullable;
@@ -111,6 +113,8 @@ import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
import static
org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAlterDatabaseRequest;
import static
org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAlterTableRequest;
@@ -139,13 +143,33 @@ public class FlussAdmin implements Admin {
private final AdminReadOnlyGateway readOnlyGateway;
private final MetadataUpdater metadataUpdater;
+ /**
+ * Single-thread executor that runs the metadata-refresh callback for
{@link
+ * RetryableGatewayClientProxy}.
+ */
+ private final ExecutorService refreshExecutor =
+ Executors.newFixedThreadPool(
+ 1, new
ExecutorThreadFactory("fluss-admin-metadata-refresh"));
+
public FlussAdmin(RpcClient client, MetadataUpdater metadataUpdater) {
+ // TODO: AdminGateway includes non-idempotent write operations
(createTable, dropTable,
+ // createDatabase, etc.). Wrapping it with
RetryableGatewayClientProxy is unsafe because
+ // a request may succeed on the server while the response is lost
(surfacing as a
+ // RetriableException), causing a duplicate mutation on retry. A
future phase should
+ // introduce idempotent retry semantics (e.g., request-id
deduplication) before enabling
+ // retry on the write gateway.
this.gateway =
GatewayClientProxy.createGatewayProxy(
metadataUpdater::getCoordinatorServer, client,
AdminGateway.class);
- this.readOnlyGateway =
+ AdminGateway rawReadOnlyGateway =
GatewayClientProxy.createGatewayProxy(
metadataUpdater::getRandomTabletServer, client,
AdminGateway.class);
+ this.readOnlyGateway =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ rawReadOnlyGateway,
+ metadataUpdater::refreshClusterUntilAvailable,
+ refreshExecutor,
+ AdminGateway.class);
this.metadataUpdater = metadataUpdater;
}
@@ -739,7 +763,9 @@ public class FlussAdmin implements Admin {
@Override
public void close() {
- // nothing to do yet
+ // Stop the metadata-refresh executor; any in-flight refresh will be
interrupted, which
+ // refreshClusterUntilAvailable handles by surfacing the
InterruptedException upward.
+ refreshExecutor.shutdownNow();
}
private static Map<Integer, GetTableStatsRequest>
prepareTableStatsRequests(
diff --git
a/fluss-client/src/main/java/org/apache/fluss/client/metadata/MetadataUpdater.java
b/fluss-client/src/main/java/org/apache/fluss/client/metadata/MetadataUpdater.java
index 2b7c22a31..eb65f9f91 100644
---
a/fluss-client/src/main/java/org/apache/fluss/client/metadata/MetadataUpdater.java
+++
b/fluss-client/src/main/java/org/apache/fluss/client/metadata/MetadataUpdater.java
@@ -293,6 +293,66 @@ public class MetadataUpdater {
}
}
+ /**
+ * Refreshes the cluster node list until a live server is reached, falling
back to a bootstrap
+ * re-initialization when no tablet server is available.
+ *
+ * <p>This is the recovery entry point for stale server addresses (e.g.,
during rolling upgrades
+ * where every tablet server IP changes). Unlike {@link #updateMetadata},
which only probes a
+ * single server per call, this method loops: each failed server is marked
unavailable so the
+ * next iteration picks a different one, and once all known servers are
exhausted it
+ * re-initializes the cluster from the bootstrap servers. The loop is
therefore bounded by the
+ * number of known servers (at most N attempts + 1 bootstrap) and always
converges.
+ *
+ * <p>No additional backoff is introduced here on purpose: each failed
attempt already pays a
+ * connection timeout, and the terminal bootstrap path ({@link
+ * #tryToInitializeClusterWithRetries}) carries its own exponential
backoff. Keeping this method
+ * a plain "loop until available or bootstrap" avoids over-engineering and
stacking redundant
+ * throttling layers.
+ */
+ public void refreshClusterUntilAvailable() {
+ while (true) {
+ ServerNode serverNode =
+ getOneAvailableTabletServerNode(cluster,
unavailableTabletServerIds);
+ if (serverNode == null) {
+ // All known tablet servers are unavailable, re-initialize the
cluster from the
+ // bootstrap servers as the terminal recovery step and stop.
+ synchronized (this) {
+ LOG.info(
+ "No available tablet server to refresh metadata,
re-initializing cluster using bootstrap server.");
+ cluster = initializeCluster(conf, rpcClient);
+ }
+
unavailableTabletServerIds.removeIf(cluster.getAliveTabletServers()::containsKey);
+ return;
+ }
+
+ try {
+ synchronized (this) {
+ cluster =
+ sendMetadataRequestAndRebuildCluster(
+ cluster, rpcClient, null, null, null,
serverNode);
+ }
+ // Refresh succeeded against a live server, the cluster node
list is now up to date.
+
unavailableTabletServerIds.removeIf(cluster.getAliveTabletServers()::containsKey);
+ return;
+ } catch (Exception e) {
+ Throwable t = stripExecutionException(e);
+ if (t instanceof RetriableException || t instanceof
TimeoutException) {
+ unavailableTabletServerIds.add(serverNode.id());
+ LOG.warn(
+ "tabletServer {} is unavailable for refreshing
metadata for retriable exception, trying the next server. unavailable
tabletServer set {}",
+ serverNode,
+ unavailableTabletServerIds,
+ t);
+ // Continue the loop to try the next available server, or
fall back to
+ // bootstrap when none remain.
+ } else {
+ throw new FlussRuntimeException("Failed to refresh
metadata", t);
+ }
+ }
+ }
+ }
+
/**
* Initialize Cluster. This step just to get the coordinator server
address and alive tablet
* servers according to the config {@link ConfigOptions#BOOTSTRAP_SERVERS}.
diff --git
a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java
b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java
index da737dd07..162f3c158 100644
---
a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java
+++
b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java
@@ -117,6 +117,7 @@ import static
org.apache.fluss.config.ConfigOptions.DATALAKE_FORMAT;
import static org.apache.fluss.config.ConfigOptions.TABLE_DATALAKE_ENABLED;
import static org.apache.fluss.config.ConfigOptions.TABLE_DATALAKE_FORMAT;
import static org.apache.fluss.metadata.DataLakeFormat.PAIMON;
+import static
org.apache.fluss.record.TestData.DATA1_PARTITIONED_TABLE_DESCRIPTOR;
import static org.apache.fluss.record.TestData.DATA1_SCHEMA;
import static org.apache.fluss.testutils.DataTestUtils.row;
import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil;
@@ -1044,6 +1045,32 @@ class FlussAdminITCase extends ClientToServerITCaseBase {
}
}
+ @Test
+ void testListPartitionInfosAfterTabletServerRestart() throws Exception {
+ String dbName = DEFAULT_TABLE_PATH.getDatabaseName();
+ TablePath partitionedTablePath = TablePath.of(dbName,
"test_retry_partitioned_table");
+ admin.createTable(partitionedTablePath,
DATA1_PARTITIONED_TABLE_DESCRIPTOR, true).get();
+
FLUSS_CLUSTER_EXTENSION.waitUntilPartitionAllReady(partitionedTablePath);
+
+ // First query should succeed.
+ List<PartitionInfo> partitionInfosBefore =
+ admin.listPartitionInfos(partitionedTablePath).get();
+ assertThat(partitionInfosBefore).isNotEmpty();
+
+ // Restart all tablet servers (they bind to new ports, making cached
addresses stale).
+ for (int i = 0; i <
FLUSS_CLUSTER_EXTENSION.getTabletServerNodes().size(); i++) {
+ FLUSS_CLUSTER_EXTENSION.stopTabletServer(i);
+ FLUSS_CLUSTER_EXTENSION.startTabletServer(i);
+ }
+ FLUSS_CLUSTER_EXTENSION.waitUntilAllGatewayHasSameMetadata();
+
+ // Second query using the same admin client should succeed after retry
with metadata
+ // refresh (verifies RetryableGatewayClientProxy convergence on stale
addresses).
+ List<PartitionInfo> partitionInfosAfter =
+ admin.listPartitionInfos(partitionedTablePath).get();
+ assertThat(partitionInfosAfter).hasSize(partitionInfosBefore.size());
+ }
+
@Test
void testListPartitionInfosByPartitionSpec() throws Exception {
String dbName = DEFAULT_TABLE_PATH.getDatabaseName();
diff --git
a/fluss-rpc/src/main/java/org/apache/fluss/rpc/RetryableGatewayClientProxy.java
b/fluss-rpc/src/main/java/org/apache/fluss/rpc/RetryableGatewayClientProxy.java
new file mode 100644
index 000000000..cc368e05e
--- /dev/null
+++
b/fluss-rpc/src/main/java/org/apache/fluss/rpc/RetryableGatewayClientProxy.java
@@ -0,0 +1,215 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.fluss.rpc;
+
+import org.apache.fluss.annotation.Internal;
+import org.apache.fluss.exception.RetriableException;
+import org.apache.fluss.utils.ExceptionUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * A proxy that wraps an existing {@link RpcGateway} proxy and adds automatic
retry with metadata
+ * refresh on retriable (network) errors.
+ *
+ * <p>This is designed to solve the stale metadata problem where cached server
addresses become
+ * invalid (e.g., during rolling upgrades in Kubernetes). When an RPC call
fails with a {@link
+ * RetriableException}, this proxy triggers a metadata refresh callback and
retries the request with
+ * potentially updated server addresses.
+ *
+ * <p>The retry flow for a cluster with stale tablet servers:
+ *
+ * <ol>
+ * <li>RPC fails with {@link RetriableException} (e.g., connection refused
to stale IP)
+ * <li>Metadata refresh is triggered, which loops until a live server is
reached or the cluster is
+ * re-initialized from bootstrap servers
+ * <li>The RPC is retried once with the refreshed server addresses
+ * </ol>
+ *
+ * <p>A single retry is sufficient because {@code metadataRefreshAction} is
expected to fully
+ * recover the cluster node list on its own; the retry only validates the
refreshed addresses and
+ * absorbs transient errors right after a server/leader switch.
+ *
+ * <p>Concurrent retriers share a single in-flight refresh: when many RPCs
fail at once (e.g.,
+ * during a rolling upgrade), they all piggyback on the first refresh future
instead of each
+ * scheduling its own redundant refresh. This avoids piling up N refresh tasks
behind the metadata
+ * updater's lock and keeps the data plane's table/partition refreshes from
queueing behind admin
+ * retries.
+ */
+@Internal
+public class RetryableGatewayClientProxy implements InvocationHandler {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(RetryableGatewayClientProxy.class);
+
+ private final Object delegate;
+ private final Runnable metadataRefreshAction;
+ private final Executor refreshExecutor;
+
+ /**
+ * Holds the currently in-flight metadata refresh, if any. Concurrent
retriers piggyback on this
+ * future to coalesce duplicate refreshes; once the future completes the
reference is cleared so
+ * subsequent failures trigger a fresh refresh.
+ */
+ private final AtomicReference<CompletableFuture<Void>> inFlightRefresh =
+ new AtomicReference<>();
+
+ RetryableGatewayClientProxy(
+ Object delegate, Runnable metadataRefreshAction, Executor
refreshExecutor) {
+ this.delegate = delegate;
+ this.metadataRefreshAction = metadataRefreshAction;
+ this.refreshExecutor = refreshExecutor;
+ }
+
+ /**
+ * Creates a retryable proxy wrapping an existing gateway proxy. On {@link
RetriableException},
+ * the proxy will invoke {@code metadataRefreshAction} and retry the
failed RPC call once.
+ *
+ * @param delegate the underlying gateway proxy to wrap
+ * @param metadataRefreshAction callback to refresh metadata (e.g., update
cluster info)
+ * @param refreshExecutor executor on which {@code metadataRefreshAction}
is run; must NOT be a
+ * Netty event loop and ideally should be a dedicated, single-thread
executor (the in-flight
+ * refresh is already coalesced to at most one concurrent task)
+ * @param gatewayClass the gateway interface class
+ * @param <T> the gateway type
+ * @return a retryable gateway proxy
+ */
+ public static <T extends RpcGateway> T createRetryableGatewayProxy(
+ T delegate,
+ Runnable metadataRefreshAction,
+ Executor refreshExecutor,
+ Class<T> gatewayClass) {
+ ClassLoader classLoader = gatewayClass.getClassLoader();
+
+ @SuppressWarnings("unchecked")
+ T proxy =
+ (T)
+ Proxy.newProxyInstance(
+ classLoader,
+ new Class<?>[] {gatewayClass},
+ new RetryableGatewayClientProxy(
+ delegate, metadataRefreshAction,
refreshExecutor));
+ return proxy;
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) throws
Throwable {
+ return invokeWithRetry(method, args, true);
+ }
+
+ @SuppressWarnings("unchecked")
+ private <T> CompletableFuture<T> invokeWithRetry(Method method, Object[]
args, boolean retry) {
+ CompletableFuture<T> future;
+ try {
+ future = (CompletableFuture<T>) method.invoke(delegate, args);
+ } catch (InvocationTargetException e) {
+ CompletableFuture<T> failed = new CompletableFuture<>();
+ failed.completeExceptionally(e.getCause());
+ return failed;
+ } catch (Exception e) {
+ CompletableFuture<T> failed = new CompletableFuture<>();
+ failed.completeExceptionally(e);
+ return failed;
+ }
+
+ CompletableFuture<T> resultFuture = new CompletableFuture<>();
+ future.whenComplete(
+ (result, throwable) -> {
+ if (throwable == null) {
+ resultFuture.complete(result);
+ return;
+ }
+ Throwable cause =
ExceptionUtils.stripCompletionException(throwable);
+ if (!(cause instanceof RetriableException) || !retry) {
+ resultFuture.completeExceptionally(cause);
+ return;
+ }
+ LOG.warn(
+ "RPC call {} failed with retriable error, "
+ + "refreshing metadata and retrying once.",
+ method.getName(),
+ cause);
+ // Coalesce concurrent refreshes so N parallel failing
calls trigger only one
+ // metadata refresh (and one round of MetadataUpdater lock
contention).
+ coalescedRefresh()
+ .thenCompose(
+ ignored ->
+
RetryableGatewayClientProxy.this.<T>invokeWithRetry(
+ method, args, false))
+ .whenComplete(
+ (retryResult, retryError) -> {
+ if (retryError != null) {
+ resultFuture.completeExceptionally(
+
ExceptionUtils.stripCompletionException(
+ retryError));
+ } else {
+ resultFuture.complete(retryResult);
+ }
+ });
+ });
+ return resultFuture;
+ }
+
+ /**
+ * Returns a future that completes when a metadata refresh has finished.
Concurrent callers that
+ * arrive while a refresh is in flight all receive the same future and
therefore wait on a
+ * single shared refresh, instead of each running their own.
+ */
+ private CompletableFuture<Void> coalescedRefresh() {
+ while (true) {
+ CompletableFuture<Void> existing = inFlightRefresh.get();
+ if (existing != null && !existing.isDone()) {
+ return existing;
+ }
+ CompletableFuture<Void> mine = new CompletableFuture<>();
+ if (inFlightRefresh.compareAndSet(existing, mine)) {
+ // Run the metadata refresh on the dedicated executor: the
failed future is
+ // typically completed on a Netty EventLoop (see
ServerConnection#close on a
+ // connection reset), and refreshClusterUntilAvailable can
take the
+ // MetadataUpdater lock, issue further RPCs, and back off with
sleeps in the
+ // bootstrap path -- all of which would freeze every
connection sharing that
+ // EventLoop if run inline. ForkJoinPool.commonPool() is also
unsuitable because
+ // commonPool workers are sized for non-blocking CPU work and
may even fall back
+ // to the caller thread on small containers.
+ CompletableFuture.runAsync(
+ () -> {
+ try {
+ metadataRefreshAction.run();
+ } catch (Exception e) {
+ LOG.warn("Failed to refresh metadata during
retry", e);
+ } finally {
+ // Complete first so piggybackers proceed,
then clear the slot so
+ // future failures start a fresh refresh round.
+ mine.complete(null);
+ inFlightRefresh.compareAndSet(mine, null);
+ }
+ },
+ refreshExecutor);
+ return mine;
+ }
+ }
+ }
+}
diff --git
a/fluss-rpc/src/test/java/org/apache/fluss/rpc/RetryableGatewayClientProxyTest.java
b/fluss-rpc/src/test/java/org/apache/fluss/rpc/RetryableGatewayClientProxyTest.java
new file mode 100644
index 000000000..d4c8f9dc3
--- /dev/null
+++
b/fluss-rpc/src/test/java/org/apache/fluss/rpc/RetryableGatewayClientProxyTest.java
@@ -0,0 +1,269 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.fluss.rpc;
+
+import org.apache.fluss.exception.NetworkException;
+import org.apache.fluss.exception.TableNotExistException;
+import org.apache.fluss.rpc.messages.ApiVersionsRequest;
+import org.apache.fluss.rpc.messages.ApiVersionsResponse;
+import org.apache.fluss.rpc.messages.AuthenticateRequest;
+import org.apache.fluss.rpc.messages.AuthenticateResponse;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Unit tests for {@link RetryableGatewayClientProxy}. */
+class RetryableGatewayClientProxyTest {
+
+ private static final Executor REFRESH_EXECUTOR = ForkJoinPool.commonPool();
+
+ @Test
+ void testSuccessfulCallWithoutRetry() throws Exception {
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+
+ RpcGateway delegate = createGateway(callCount, 0);
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate,
+ refreshCount::incrementAndGet,
+ REFRESH_EXECUTOR,
+ RpcGateway.class);
+
+ CompletableFuture<ApiVersionsResponse> result = proxy.apiVersions(new
ApiVersionsRequest());
+ assertThat(result.get()).isNotNull();
+ assertThat(callCount.get()).isEqualTo(1);
+ assertThat(refreshCount.get()).isEqualTo(0);
+ }
+
+ @Test
+ void testRetryOnNetworkExceptionThenSuccess() throws Exception {
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+
+ // Fail with NetworkException for the first call, then succeed on the
single retry
+ RpcGateway delegate = createGateway(callCount, 1);
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate,
+ refreshCount::incrementAndGet,
+ REFRESH_EXECUTOR,
+ RpcGateway.class);
+
+ CompletableFuture<ApiVersionsResponse> result = proxy.apiVersions(new
ApiVersionsRequest());
+ assertThat(result.get()).isNotNull();
+ // Initial call + 1 retry = 2 total calls
+ assertThat(callCount.get()).isEqualTo(2);
+ // Metadata refresh should be called once before the retry
+ assertThat(refreshCount.get()).isEqualTo(1);
+ }
+
+ @Test
+ void testExhaustsRetriesAndPropagatesError() {
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+
+ // Always fail with NetworkException (more failures than the single
retry can recover)
+ RpcGateway delegate = createGateway(callCount, Integer.MAX_VALUE);
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate,
+ refreshCount::incrementAndGet,
+ REFRESH_EXECUTOR,
+ RpcGateway.class);
+
+ CompletableFuture<ApiVersionsResponse> result = proxy.apiVersions(new
ApiVersionsRequest());
+ assertThatThrownBy(result::get)
+ .isInstanceOf(ExecutionException.class)
+ .rootCause()
+ .isInstanceOf(NetworkException.class)
+ .hasMessageContaining("Simulated network error");
+ // Initial call + 1 retry = 2 total calls
+ assertThat(callCount.get()).isEqualTo(2);
+ // Metadata refresh should be called once before the single retry
+ assertThat(refreshCount.get()).isEqualTo(1);
+ }
+
+ @Test
+ void testNonRetriableExceptionNotRetried() {
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+
+ // Fail with a non-retriable exception
+ RpcGateway delegate =
+ new TestRpcGateway() {
+ @Override
+ public CompletableFuture<ApiVersionsResponse> apiVersions(
+ ApiVersionsRequest request) {
+ callCount.incrementAndGet();
+ CompletableFuture<ApiVersionsResponse> future = new
CompletableFuture<>();
+ future.completeExceptionally(
+ new TableNotExistException("table does not
exist"));
+ return future;
+ }
+ };
+
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate,
+ refreshCount::incrementAndGet,
+ REFRESH_EXECUTOR,
+ RpcGateway.class);
+
+ CompletableFuture<ApiVersionsResponse> result = proxy.apiVersions(new
ApiVersionsRequest());
+ assertThatThrownBy(result::get)
+ .isInstanceOf(ExecutionException.class)
+ .rootCause()
+ .isInstanceOf(TableNotExistException.class)
+ .hasMessageContaining("table does not exist");
+ // Should only be called once - no retries for non-retriable exceptions
+ assertThat(callCount.get()).isEqualTo(1);
+ assertThat(refreshCount.get()).isEqualTo(0);
+ }
+
+ @Test
+ void testMetadataRefreshFailureDoesNotPreventRetry() throws Exception {
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+
+ // Fail first call, succeed second call
+ RpcGateway delegate = createGateway(callCount, 1);
+
+ // Metadata refresh throws exception
+ Runnable failingRefresh =
+ () -> {
+ refreshCount.incrementAndGet();
+ throw new RuntimeException("Simulated refresh failure");
+ };
+
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate, failingRefresh, REFRESH_EXECUTOR,
RpcGateway.class);
+
+ // Should still succeed because the retry goes through even if refresh
fails
+ CompletableFuture<ApiVersionsResponse> result = proxy.apiVersions(new
ApiVersionsRequest());
+ assertThat(result.get()).isNotNull();
+ assertThat(callCount.get()).isEqualTo(2);
+ assertThat(refreshCount.get()).isEqualTo(1);
+ }
+
+ /**
+ * Verifies that when many failing RPCs run concurrently, only a single
metadata refresh is
+ * issued: all retriers piggyback on the same in-flight refresh future.
This prevents the
+ * pile-up of N redundant refresh tasks behind the metadata updater's lock.
+ */
+ @Test
+ void testConcurrentFailingCallsShareSingleRefresh() throws Exception {
+ int n = 8;
+ AtomicInteger callCount = new AtomicInteger(0);
+ AtomicInteger refreshCount = new AtomicInteger(0);
+ CountDownLatch refreshGate = new CountDownLatch(1);
+
+ // First n calls fail, subsequent calls succeed (so each retrier's
retry will succeed).
+ RpcGateway delegate = createGateway(callCount, n);
+
+ // Block the refresh until the test releases it, so all retriers have
a chance to arrive
+ // and either start or piggyback on the in-flight refresh.
+ Runnable blockingRefresh =
+ () -> {
+ refreshCount.incrementAndGet();
+ try {
+ refreshGate.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ };
+
+ RpcGateway proxy =
+ RetryableGatewayClientProxy.createRetryableGatewayProxy(
+ delegate, blockingRefresh, REFRESH_EXECUTOR,
RpcGateway.class);
+
+ // Fire n calls; each fails on its initial invocation and queues for a
refresh.
+ List<CompletableFuture<ApiVersionsResponse>> futures = new
ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ futures.add(proxy.apiVersions(new ApiVersionsRequest()));
+ }
+
+ // Wait until the (single) refresh task has actually started.
+ long deadline = System.currentTimeMillis() + 5000;
+ while (refreshCount.get() == 0 && System.currentTimeMillis() <
deadline) {
+ Thread.sleep(10);
+ }
+ assertThat(refreshCount.get()).as("exactly one refresh should have
started").isEqualTo(1);
+
+ // Release the refresh; all retriers should now proceed with their
single retry.
+ refreshGate.countDown();
+ for (CompletableFuture<ApiVersionsResponse> f : futures) {
+ assertThat(f.get(5, TimeUnit.SECONDS)).isNotNull();
+ }
+
+ // Initial n failing calls + n successful retries = 2n total
invocations.
+ assertThat(callCount.get()).isEqualTo(2 * n);
+ // The critical assertion: n concurrent failures triggered only ONE
refresh.
+ assertThat(refreshCount.get())
+ .as("concurrent retriers must share a single refresh")
+ .isEqualTo(1);
+ }
+
+ /**
+ * Creates a test gateway that fails with {@link NetworkException} for the
first {@code
+ * failCount} invocations, then returns a successful response.
+ */
+ private static RpcGateway createGateway(AtomicInteger callCount, int
failCount) {
+ return new TestRpcGateway() {
+ @Override
+ public CompletableFuture<ApiVersionsResponse>
apiVersions(ApiVersionsRequest request) {
+ int count = callCount.incrementAndGet();
+ CompletableFuture<ApiVersionsResponse> future = new
CompletableFuture<>();
+ if (count <= failCount) {
+ future.completeExceptionally(
+ new NetworkException("Simulated network error on
call " + count));
+ } else {
+ future.complete(new ApiVersionsResponse());
+ }
+ return future;
+ }
+ };
+ }
+
+ /** Base test implementation of {@link RpcGateway} that throws on
unimplemented methods. */
+ private abstract static class TestRpcGateway implements RpcGateway {
+
+ @Override
+ public CompletableFuture<ApiVersionsResponse>
apiVersions(ApiVersionsRequest request) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public CompletableFuture<AuthenticateResponse>
authenticate(AuthenticateRequest request) {
+ throw new UnsupportedOperationException();
+ }
+ }
+}
