Merge pull request #8 from leancloud/feat/unsubscribe-status

unsubscribed status

Author: ylgrgyq

src/main/java/cn/leancloud/kafka/consumer/AbstractCommitPolicy.java

@@ -75,10 +75,12 @@ void clearCachedCompletedPartitionsRecords(Set<TopicPartition> completedPartitio
         }
     }
 
+    @VisibleForTesting
     Map<TopicPartition, Long> topicOffsetHighWaterMark() {
         return topicOffsetHighWaterMark;
     }
 
+    @VisibleForTesting
     Map<TopicPartition, OffsetAndMetadata> completedTopicOffsets() {
         return completedTopicOffsets;
     }

src/main/java/cn/leancloud/kafka/consumer/AbstractRecommitAwareCommitPolicy.java

@@ -40,6 +40,7 @@ void updateNextRecommitTime() {
         updateNextRecommitTime(System.nanoTime());
     }
 
+    @VisibleForTesting
     long nextRecommitNanos() {
         return nextRecommitNanos;
     }

src/main/java/cn/leancloud/kafka/consumer/AsyncCommitPolicy.java

@@ -46,14 +46,17 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
         return partitions;
     }
 
+    @VisibleForTesting
     int pendingAsyncCommitCount() {
         return pendingAsyncCommitCounter;
     }
 
+    @VisibleForTesting
     boolean forceSync() {
         return forceSync;
     }
 
+    @VisibleForTesting
     void setForceSync(boolean forceSync) {
         this.forceSync = forceSync;
     }

src/main/java/cn/leancloud/kafka/consumer/Fetcher.java

@@ -24,6 +24,7 @@ final class Fetcher<K, V> implements Runnable, Closeable {
     private final Map<ConsumerRecord<K, V>, Future<ConsumerRecord<K, V>>> pendingFutures;
     private final CommitPolicy<K, V> policy;
     private final long gracefulShutdownMillis;
+    private final CompletableFuture<UnsubscribedStatus> unsubscribeStatusFuture;
     private volatile boolean closed;
 
     Fetcher(LcKafkaConsumerBuilder<K, V> consumerBuilder) {
@@ -44,13 +45,15 @@ final class Fetcher<K, V> implements Runnable, Closeable {
         this.service = new ExecutorCompletionService<>(workerPool);
         this.policy = policy;
         this.gracefulShutdownMillis = gracefulShutdownMillis;
+        this.unsubscribeStatusFuture = new CompletableFuture<>();
     }
 
     @Override
     public void run() {
         logger.debug("Fetcher thread started.");
         final long pollTimeout = this.pollTimeout;
         final Consumer<K, V> consumer = this.consumer;
+        UnsubscribedStatus unsubscribedStatus = UnsubscribedStatus.CLOSED;
         while (true) {
             try {
                 final ConsumerRecords<K, V> records = consumer.poll(pollTimeout);
@@ -72,13 +75,14 @@ public void run() {
                     break;
                 }
             } catch (Exception ex) {
+                unsubscribedStatus = UnsubscribedStatus.ERROR;
                 close();
                 logger.error("Fetcher quit with unexpected exception. Will rebalance after poll timeout.", ex);
                 break;
             }
         }
 
-        gracefulShutdown();
+        gracefulShutdown(unsubscribedStatus);
     }
 
     @Override
@@ -87,6 +91,11 @@ public void close() {
         consumer.wakeup();
     }
 
+    CompletableFuture<UnsubscribedStatus> unsubscribeStatusFuture() {
+        return unsubscribeStatusFuture;
+    }
+
+    @VisibleForTesting
     Map<ConsumerRecord<K, V>, Future<ConsumerRecord<K, V>>> pendingFutures() {
         return pendingFutures;
     }
@@ -134,7 +143,7 @@ private void tryCommitRecordOffsets() {
         }
     }
 
-    private void gracefulShutdown() {
+    private void gracefulShutdown(UnsubscribedStatus unsubscribedStatus) {
         final long start = System.currentTimeMillis();
         long remain = gracefulShutdownMillis;
         try {
@@ -163,6 +172,8 @@ private void gracefulShutdown() {
 
         pendingFutures.clear();
 
+        unsubscribeStatusFuture.complete(unsubscribedStatus);
+
         logger.debug("Fetcher thread exit.");
     }
 }

src/main/java/cn/leancloud/kafka/consumer/LcKafkaConsumer.java

@@ -7,7 +7,7 @@
 import java.io.Closeable;
 import java.util.Collection;
 import java.util.Map;
-import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Pattern;
@@ -67,11 +67,15 @@ int code() {
      * Subscribe some Kafka topics to consume records from them.
      *
      * @param topics the topics to consume.
+     * @return a {@link CompletableFuture} which will be complete when the internal
+     * {@link org.apache.kafka.clients.consumer.KafkaConsumer} unsubscribed all the topics that have subscribed to
+     * due to an error occurred or {@link LcKafkaConsumer#close()} was called. If it is due to an error occurred, this
+     * {@code LcKafkaConsumer} will be closed before the returned {@code CompletableFuture} complete.
      * @throws IllegalStateException    if this {@code LcKafkaConsumer} has closed or subscribed to some topics
      * @throws NullPointerException     if the input {@code topics} is null
      * @throws IllegalArgumentException if the input {@code topics} is empty or contains null or empty topic
      */
-    public synchronized void subscribe(Collection<String> topics) {
+    public synchronized CompletableFuture<UnsubscribedStatus> subscribe(Collection<String> topics) {
         requireNonNull(topics, "topics");
 
         if (topics.isEmpty()) {
@@ -91,6 +95,8 @@ public synchronized void subscribe(Collection<String> topics) {
         fetcherThread.start();
 
         state = State.SUBSCRIBED;
+
+        return setupUnsubscribedFuture();
     }
 
     /**
@@ -100,10 +106,14 @@ public synchronized void subscribe(Collection<String> topics) {
      * the max metadata age, the consumer will refresh metadata more often and check for matching topics.
      *
      * @param pattern {@link Pattern} to subscribe to.
-     * @throws IllegalStateException    if this {@code LcKafkaConsumer} has closed or subscribed to some topics
-     * @throws NullPointerException     if the input {@code pattern} is null
+     * @return a {@link CompletableFuture} which will be complete when the internal
+     * {@link org.apache.kafka.clients.consumer.KafkaConsumer} unsubscribed all the topics that have subscribed to
+     * due to an error occurred or {@link LcKafkaConsumer#close()} was called. If it is due to an error occurred, this
+     * {@code LcKafkaConsumer} will be closed before the returned {@code CompletableFuture} complete.
+     * @throws IllegalStateException if this {@code LcKafkaConsumer} has closed or subscribed to some topics
+     * @throws NullPointerException  if the input {@code pattern} is null
      */
-    public synchronized void subscribe(Pattern pattern) {
+    public synchronized CompletableFuture<UnsubscribedStatus> subscribe(Pattern pattern) {
         requireNonNull(pattern, "pattern");
 
         ensureInInit();
@@ -114,6 +124,8 @@ public synchronized void subscribe(Pattern pattern) {
         fetcherThread.start();
 
         state = State.SUBSCRIBED;
+
+        return setupUnsubscribedFuture();
     }
 
     /**
@@ -139,9 +151,11 @@ public void close() {
             state = State.CLOSED;
         }
 
-        fetcher.close();
         try {
-            fetcherThread.join();
+            if (Thread.currentThread() != fetcherThread) {
+                fetcher.close();
+                fetcherThread.join();
+            }
             consumer.close();
             if (shutdownWorkerPoolOnStop) {
                 workerPool.shutdown();
@@ -152,21 +166,37 @@ public void close() {
         }
     }
 
+    @VisibleForTesting
     boolean subscribed() {
         return state.code() > State.INIT.code();
     }
 
+    @VisibleForTesting
     boolean closed() {
         return state == State.CLOSED;
     }
 
+    @VisibleForTesting
     CommitPolicy<K, V> policy() {
         return policy;
     }
 
+    private CompletableFuture<UnsubscribedStatus> setupUnsubscribedFuture() {
+        assert !fetcher.unsubscribeStatusFuture().isDone();
+
+        // use a new CompletableFuture to ensure that close() should be called first, then do the pipeline bind to the
+        // returned CompletableFuture
+        final CompletableFuture<UnsubscribedStatus> ret = new CompletableFuture<>();
+        fetcher.unsubscribeStatusFuture().thenAccept(status -> {
+            close();
+            ret.complete(status);
+        });
+        return ret;
+    }
+
     private void ensureInInit() {
         if (subscribed() || closed()) {
-            throw new IllegalStateException("client is in " + state + " state. expect: " + State.INIT);
+            throw new IllegalStateException("consumer is closed or have subscribed to some topics or pattern");
         }
     }
 

src/main/java/cn/leancloud/kafka/consumer/LcKafkaConsumerBuilder.java

@@ -484,7 +484,7 @@ private Consumer<K, V> buildConsumer(boolean autoCommit) {
         return new KafkaConsumer<>(configs, keyDeserializer, valueDeserializer);
     }
 
-    Duration getRecommitInterval() {
+    private Duration getRecommitInterval() {
         if (recommitInterval == null) {
             logger.warn("Recommit interval is not set for a non-automatic commit consumer, the default " +
                     "interval of 1 hour will be used.");

src/main/java/cn/leancloud/kafka/consumer/PartialAsyncCommitPolicy.java

@@ -47,10 +47,12 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
         }
     }
 
+    @VisibleForTesting
     int pendingAsyncCommitCount() {
         return pendingAsyncCommitCounter;
     }
 
+    @VisibleForTesting
     boolean forceSync() {
         return forceSync;
     }

src/main/java/cn/leancloud/kafka/consumer/UnsubscribedStatus.java

@@ -0,0 +1,22 @@
+package cn.leancloud.kafka.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+/**
+ * A status of {@link LcKafkaConsumer} when it unsubscribed to all the topics that have subscribed to.
+ */
+public enum UnsubscribedStatus {
+    /**
+     * Indicate that the unsubscription is due to {@link LcKafkaConsumer#close()} was called.
+     * Usually, this is an expected status.
+     */
+    CLOSED,
+
+    /**
+     * Indicate that the unsubscription is due to an error occurred, maybe an exception was thrown from
+     * {@link ConsumerRecordHandler#handleRecord(ConsumerRecord)}, or some unrecoverable error happened
+     * within {@link org.apache.kafka.clients.consumer.KafkaConsumer}. You can check the specific exception
+     * in log.
+     */
+    ERROR
+}

src/main/java/cn/leancloud/kafka/consumer/VisibleForTesting.java

@@ -0,0 +1,9 @@
+package cn.leancloud.kafka.consumer;
+
+/**
+ * Indicates that the visibility of a type or member has been relaxed to make the code testable. Idea borrowed from
+ * Guava
+ *
+ */
+public @interface VisibleForTesting {}
+

src/test/java/cn/leancloud/kafka/consumer/FetcherTest.java

@@ -12,6 +12,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -62,6 +63,7 @@ public void tearDown() throws Exception {
     public void testGracefulShutdown() throws Exception {
         executorService = mock(ExecutorService.class);
         fetcher = new Fetcher<>(consumer, pollTimeout, consumerRecordHandler, executorService, policy, 0);
+        final CompletableFuture<UnsubscribedStatus> unsubscribedStatusFuture = fetcher.unsubscribeStatusFuture();
         fetcherThread = new Thread(fetcher);
 
         doNothing().when(executorService).execute(any(Runnable.class));
@@ -76,11 +78,13 @@ public void testGracefulShutdown() throws Exception {
         assertThat(fetcher.pendingFutures()).isEmpty();
         verify(policy, times(1)).partialCommit();
         assertThat(consumer.subscription()).isEmpty();
+        assertThat(unsubscribedStatusFuture).isCompletedWithValue(UnsubscribedStatus.CLOSED);
     }
 
     @Test
     public void testHandleMsgFailed() throws Exception {
         fetcher = new Fetcher<>(consumer, pollTimeout, consumerRecordHandler, executorService, policy, 10_000);
+        final CompletableFuture<UnsubscribedStatus> unsubscribedStatusFuture = fetcher.unsubscribeStatusFuture();
         fetcherThread = new Thread(fetcher);
 
         assignPartitions(consumer, toPartitions(singletonList(0)), 0L);
@@ -96,6 +100,7 @@ public void testHandleMsgFailed() throws Exception {
         verify(policy, never()).completeRecord(any());
         verify(policy, never()).tryCommit(anyBoolean());
         verify(consumerRecordHandler, times(1)).handleRecord(defaultTestingRecord);
+        assertThat(unsubscribedStatusFuture).isCompletedWithValue(UnsubscribedStatus.ERROR);
     }
 
     @Test