Fix partitionsAutoUpdateFuture never complete (apache#14625)

yangl · Mar 17, 2022 · b06dac6 · b06dac6
1 parent 50a7e50
commit b06dac6
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 2 deletions.
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/SimpleProducerConsumerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/SimpleProducerConsumerTest.java
@@ -39,6 +39,7 @@
 import com.google.common.collect.Sets;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import io.netty.util.Timeout;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -4351,4 +4352,52 @@ public void testShareConsumerWithMessageListener() throws Exception {
             assertEquals(resultSet.size(), total);
         });
     }
+
+    @Test
+    public void testPartitionsAutoUpdate() throws Exception {
+        log.info("-- Starting {} test --", methodName);
+
+        int numPartitions = 3;
+        TopicName topicName = TopicName.get("persistent://my-property/my-ns/partitionsAutoUpdate-1");
+        admin.topics().createPartitionedTopic(topicName.toString(), numPartitions);
+
+        int operationTimeout = 2000; // MILLISECONDS
+        @Cleanup final PulsarClient client = PulsarClient.builder()
+                .serviceUrl(lookupUrl.toString())
+                .operationTimeout(operationTimeout, TimeUnit.MILLISECONDS)
+                .build();
+
+        ProducerBuilder<byte[]> producerBuilder = client.newProducer()
+                .topic(topicName.toString()).sendTimeout(1, TimeUnit.SECONDS);
+
+        @Cleanup
+        PartitionedProducerImpl<byte[]> partitionedProducer =
+                (PartitionedProducerImpl<byte[]>) producerBuilder.autoUpdatePartitions(true).create();
+
+        // Trigger the Connection refused exception
+        stopBroker();
+
+        log.info("trigger partitionsAutoUpdateTimerTask run failed for producer");
+        Timeout timeout = partitionedProducer.getPartitionsAutoUpdateTimeout();
+        timeout.task().run(timeout);
+        Awaitility.await().untilAsserted(() -> {
+            assertNotNull(partitionedProducer.getPartitionsAutoUpdateFuture());
+            assertTrue(partitionedProducer.getPartitionsAutoUpdateFuture().isCompletedExceptionally());
+            assertTrue(FutureUtil.getException(partitionedProducer.getPartitionsAutoUpdateFuture()).get().getMessage()
+                    .contains("Connection refused:"));
+        });
+
+        startBroker();
+
+        log.info("trigger partitionsAutoUpdateTimerTask run successful for producer");
+        timeout = partitionedProducer.getPartitionsAutoUpdateTimeout();
+        timeout.task().run(timeout);
+        Awaitility.await().untilAsserted(() -> {
+            assertNotNull(partitionedProducer.getPartitionsAutoUpdateFuture());
+            assertTrue(partitionedProducer.getPartitionsAutoUpdateFuture().isDone());
+            assertFalse(partitionedProducer.getPartitionsAutoUpdateFuture().isCompletedExceptionally());
+        });
+
+        log.info("-- Exiting {} test --", methodName);
+    }
 }
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/PartitionedProducerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/PartitionedProducerImpl.java
@@ -385,7 +385,7 @@ public CompletableFuture<Void> onTopicsExtended(Collection<String> topicsExtende
 
                 if (log.isDebugEnabled()) {
                     log.debug("[{}] partitions number. old: {}, new: {}",
-                        topic, oldPartitionNumber, currentPartitionNumber);
+                            topic, oldPartitionNumber, currentPartitionNumber);
                 }
 
                 if (oldPartitionNumber == currentPartitionNumber) {
@@ -437,10 +437,14 @@ public CompletableFuture<Void> onTopicsExtended(Collection<String> topicsExtende
                     }
                 } else {
                     log.error("[{}] not support shrink topic partitions. old: {}, new: {}",
-                        topic, oldPartitionNumber, currentPartitionNumber);
+                            topic, oldPartitionNumber, currentPartitionNumber);
                     future.completeExceptionally(new NotSupportedException("not support shrink topic partitions"));
                 }
                 return future;
+            }).exceptionally(throwable -> {
+                log.error("[{}] Auto getting partitions failed", topic, throwable);
+                future.completeExceptionally(throwable);
+                return null;
             });
 
             return future;
@@ -476,6 +480,11 @@ public void run(Timeout timeout) throws Exception {
         }
     };
 
+    @VisibleForTesting
+    public CompletableFuture<Void> getPartitionsAutoUpdateFuture() {
+        return partitionsAutoUpdateFuture;
+    }
+
     @VisibleForTesting
     public Timeout getPartitionsAutoUpdateTimeout() {
         return partitionsAutoUpdateTimeout;