Add the track of the completion of the topology refresh

redis · Dec 31, 2024 · ed56fa4 · ed56fa4
1 parent 5952b0b
commit ed56fa4
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/src/main/java/io/lettuce/core/cluster/ClusterTopologyRefreshScheduler.java b/src/main/java/io/lettuce/core/cluster/ClusterTopologyRefreshScheduler.java
@@ -22,6 +22,7 @@
 import static io.lettuce.core.event.cluster.AdaptiveRefreshTriggeredEvent.*;
 
 import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -112,6 +113,10 @@ public boolean isTopologyRefreshInProgress() {
         return clusterTopologyRefreshTask.get();
     }
 
+    public CompletableFuture<Void> getTopologyRefreshCompletionFuture() {
+        return clusterTopologyRefreshTask.getCompletionFuture();
+    }
+
     @Override
     public void run() {
 
@@ -316,11 +321,16 @@ private static class ClusterTopologyRefreshTask extends AtomicBoolean implements
         private static final long serialVersionUID = -1337731371220365694L;
 
         private final Supplier<CompletionStage<?>> reloadTopologyAsync;
+        private final CompletableFuture<Void> completionFuture = new CompletableFuture<>();
 
         ClusterTopologyRefreshTask(Supplier<CompletionStage<?>> reloadTopologyAsync) {
             this.reloadTopologyAsync = reloadTopologyAsync;
         }
 
+        public CompletableFuture<Void> getCompletionFuture() {
+            return completionFuture;
+        }
+
         public void run() {
 
             if (compareAndSet(false, true)) {
@@ -343,12 +353,16 @@ void doRun() {
 
                     if (throwable != null) {
                         logger.warn("Cannot refresh Redis Cluster topology", throwable);
+                        completionFuture.completeExceptionally(throwable);
+                    } else {
+                        completionFuture.complete(null);
                     }
 
                     set(false);
                 });
             } catch (Exception e) {
                 logger.warn("Cannot refresh Redis Cluster topology", e);
+                completionFuture.completeExceptionally(e);
             }
         }
 

diff --git a/src/main/java/io/lettuce/core/cluster/RedisClusterClient.java b/src/main/java/io/lettuce/core/cluster/RedisClusterClient.java
@@ -1153,7 +1153,13 @@ public CompletableFuture<Void> shutdownAsync(long quietPeriod, long timeout, Tim
 
         suspendTopologyRefresh();
 
-        return super.shutdownAsync(quietPeriod, timeout, timeUnit);
+        CompletableFuture<Void> topologyRefreshFuture = topologyRefreshScheduler.getTopologyRefreshCompletionFuture();
+
+        return topologyRefreshFuture.thenCompose(ignore -> super.shutdownAsync(quietPeriod, timeout, timeUnit))
+                .exceptionally(ex -> {
+                    System.err.println("Error during topology refresh or shutdown: " + ex.getMessage());
+                    return null;
+        });
     }
 
     // -------------------------------------------------------------------------