From 6b25c887929056bac0e6716d05326ca2b293d71d Mon Sep 17 00:00:00 2001
From: Jason Gustafson <jason@confluent.io>
Date: Thu, 28 May 2015 14:51:08 -0700
Subject: [PATCH] KAFKA-2168; Remove synchronization of KafkaConsumer to enable
 non-blocking close

---
 .../clients/consumer/ClosedConsumerException.java  | 25 +++++++
 .../kafka/clients/consumer/KafkaConsumer.java      | 83 +++++++++++++---------
 .../clients/consumer/internals/Coordinator.java    | 23 +++++-
 .../kafka/clients/consumer/internals/Fetcher.java  | 23 +++++-
 4 files changed, 118 insertions(+), 36 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/ClosedConsumerException.java

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ClosedConsumerException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ClosedConsumerException.java
new file mode 100644
index 0000000..e8440df
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ClosedConsumerException.java
@@ -0,0 +1,25 @@
+/**
+ * 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.kafka.clients.consumer;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * Thrown to indicate that a consumer has been closed and cannot process the given request
+ */
+public class ClosedConsumerException extends KafkaException {
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
index d301be4..b50e9ff 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
@@ -23,6 +23,7 @@ import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -297,18 +298,22 @@ import org.slf4j.LoggerFactory;
  * {@link #seekToBeginning(TopicPartition...)} and {@link #seekToEnd(TopicPartition...)} respectively).
  * 
  * <h3>Multithreaded Processing</h3>
- * 
- * The Kafka consumer is threadsafe but coarsely synchronized. All network I/O happens in the thread of the application
- * making the call. We have intentionally avoided implementing a particular threading model for processing.
+ *
+ * The Kafka consumer is NOT thread-safe, which means that you are responsible for ensuring safe multi-threaded
+ * access. In the absence of external locking, the consumer should be confined to a single thread. It is safe, however
+ * to invoke {@link #close()} from another thread (e.g. in a shutdown hook). When the consumer is closed from another
+ * thread, any active request (along with any future requests) will throw {@link ClosedConsumerException}.
+ *
  * <p>
- * This leaves several options for implementing multi-threaded processing of records.
+ * We have intentionally avoided implementing a particular threading model for processing. This leaves several options
+ * for implementing multi-threaded processing of records.
  * 
  * <h4>1. One Consumer Per Thread</h4>
  * 
- * A simple option is to give each thread it's own consumer instance. Here are the pros and cons of this approach:
+ * A simple option is to give each thread its own consumer instance. Here are the pros and cons of this approach:
  * <ul>
  * <li><b>PRO</b>: It is the easiest to implement
- * <li><b>PRO</b>: It is often the fastest as no inter-thread co-ordination is needed
+ * <li><b>PRO</b>: It is often the fastest as no inter-thread coordination is needed
  * <li><b>PRO</b>: It makes in-order processing on a per-partition basis very easy to implement (each thread just
  * processes messages in the order it receives them).
  * <li><b>CON</b>: More consumers means more TCP connections to the cluster (one per thread). In general Kafka handles
@@ -362,7 +367,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
     private final long autoCommitIntervalMs;
     private final ConsumerRebalanceCallback rebalanceCallback;
     private long lastCommitAttemptMs;
-    private boolean closed = false;
+    private AtomicBoolean closed = new AtomicBoolean(false);
 
     /**
      * A consumer is instantiated by providing a set of key-value pairs as configuration. Valid configuration strings
@@ -541,7 +546,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * then this will give the set of topics currently assigned to the consumer (which may be none if the assignment
      * hasn't happened yet, or the partitions are in the process of getting reassigned).
      */
-    public synchronized Set<TopicPartition> subscriptions() {
+    public Set<TopicPartition> subscriptions() {
         return Collections.unmodifiableSet(this.subscriptions.assignedPartitions());
     }
 
@@ -560,7 +565,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * @param topics A variable list of topics that the consumer wants to subscribe to
      */
     @Override
-    public synchronized void subscribe(String... topics) {
+    public void subscribe(String... topics) {
         ensureNotClosed();
         log.debug("Subscribed to topic(s): {}", Utils.join(topics, ", "));
         for (String topic : topics)
@@ -577,7 +582,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * @param partitions Partitions to incrementally subscribe to
      */
     @Override
-    public synchronized void subscribe(TopicPartition... partitions) {
+    public void subscribe(TopicPartition... partitions) {
         ensureNotClosed();
         log.debug("Subscribed to partitions(s): {}", Utils.join(partitions, ", "));
         for (TopicPartition tp : partitions) {
@@ -592,7 +597,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * 
      * @param topics Topics to unsubscribe from
      */
-    public synchronized void unsubscribe(String... topics) {
+    public void unsubscribe(String... topics) {
         ensureNotClosed();
         log.debug("Unsubscribed from topic(s): {}", Utils.join(topics, ", "));
         // throw an exception if the topic was never subscribed to
@@ -631,7 +636,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      *             offset reset policy has been configured.
      */
     @Override
-    public synchronized ConsumerRecords<K, V> poll(long timeout) {
+    public ConsumerRecords<K, V> poll(long timeout) {
         ensureNotClosed();
         long now = time.milliseconds();
 
@@ -660,6 +665,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
         Cluster cluster = this.metadata.fetch();
         fetcher.initFetches(cluster, now);
         client.poll(timeout, now);
+        ensureNotClosed();
 
         /*
          * initiate a fetch request for any nodes that we just got a response from without blocking
@@ -685,7 +691,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * @param commitType Control whether the commit is blocking
      */
     @Override
-    public synchronized void commit(final Map<TopicPartition, Long> offsets, CommitType commitType) {
+    public void commit(final Map<TopicPartition, Long> offsets, CommitType commitType) {
         ensureNotClosed();
         log.debug("Committing offsets ({}): {} ", commitType.toString().toLowerCase(), offsets);
 
@@ -709,7 +715,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * @param commitType Whether or not the commit should block until it is acknowledged.
      */
     @Override
-    public synchronized void commit(CommitType commitType) {
+    public void commit(CommitType commitType) {
         ensureNotClosed();
         commit(this.subscriptions.allConsumed(), commitType);
     }
@@ -720,7 +726,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * you may lose data if this API is arbitrarily used in the middle of consumption, to reset the fetch offsets
      */
     @Override
-    public synchronized void seek(TopicPartition partition, long offset) {
+    public void seek(TopicPartition partition, long offset) {
         ensureNotClosed();
         log.debug("Seeking to offset {} for partition {}", offset, partition);
         this.subscriptions.seek(partition, offset);
@@ -729,7 +735,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
     /**
      * Seek to the first offset for each of the given partitions
      */
-    public synchronized void seekToBeginning(TopicPartition... partitions) {
+    public void seekToBeginning(TopicPartition... partitions) {
         ensureNotClosed();
         Collection<TopicPartition> parts = partitions.length == 0 ? this.subscriptions.assignedPartitions()
                 : Arrays.asList(partitions);
@@ -742,7 +748,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
     /**
      * Seek to the last offset for each of the given partitions
      */
-    public synchronized void seekToEnd(TopicPartition... partitions) {
+    public void seekToEnd(TopicPartition... partitions) {
         ensureNotClosed();
         Collection<TopicPartition> parts = partitions.length == 0 ? this.subscriptions.assignedPartitions()
                 : Arrays.asList(partitions);
@@ -760,7 +766,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * @throws NoOffsetForPartitionException If a position hasn't been set for a given partition, and no reset policy is
      *             available.
      */
-    public synchronized long position(TopicPartition partition) {
+    public long position(TopicPartition partition) {
         ensureNotClosed();
         if (!this.subscriptions.assignedPartitions().contains(partition))
             throw new IllegalArgumentException("You can only check the position for partitions assigned to this consumer.");
@@ -786,7 +792,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      *             partition.
      */
     @Override
-    public synchronized long committed(TopicPartition partition) {
+    public long committed(TopicPartition partition) {
         ensureNotClosed();
         Set<TopicPartition> partitionsToFetch;
         if (subscriptions.assignedPartitions().contains(partition)) {
@@ -832,25 +838,35 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
     }
 
     @Override
-    public synchronized void close() {
+    public void close() {
         close(false);
     }
 
     private void close(boolean swallowException) {
-        log.trace("Closing the Kafka consumer.");
-        AtomicReference<Throwable> firstException = new AtomicReference<Throwable>();
-        this.closed = true;
-        ClientUtils.closeQuietly(metrics, "consumer metrics", firstException);
-        ClientUtils.closeQuietly(client, "consumer network client", firstException);
-        ClientUtils.closeQuietly(keyDeserializer, "consumer key deserializer", firstException);
-        ClientUtils.closeQuietly(valueDeserializer, "consumer value deserializer", firstException);
-        log.debug("The Kafka consumer has closed.");
-        if (firstException.get() != null && !swallowException) {
-            throw new KafkaException("Failed to close kafka consumer", firstException.get());
+        if (this.closed.compareAndSet(false, true)) {
+            log.trace("Closing the Kafka consumer.");
+            AtomicReference<Throwable> firstException = new AtomicReference<Throwable>();
+
+            // Closing the fetcher and the coordinator will cause any poll loops to
+            // abort with ClosedConsumerException when we wakeup the client.
+            // TODO: This will not be necessary once KAFKA-1894 is resolved
+            ClientUtils.closeQuietly(fetcher, "fetcher", firstException);
+            ClientUtils.closeQuietly(coordinator, "coordinator", firstException);
+
+            // Wakeup the client to abort any pending IO
+            if (client != null) client.wakeup();
+
+            ClientUtils.closeQuietly(metrics, "consumer metrics", firstException);
+            ClientUtils.closeQuietly(client, "consumer network client", firstException);
+            ClientUtils.closeQuietly(keyDeserializer, "consumer key deserializer", firstException);
+            ClientUtils.closeQuietly(valueDeserializer, "consumer value deserializer", firstException);
+            log.debug("The Kafka consumer has closed.");
+            if (firstException.get() != null && !swallowException) {
+                throw new KafkaException("Failed to close kafka consumer", firstException.get());
+            }
         }
     }
 
-
     private boolean shouldAutoCommit(long now) {
         return this.autoCommit && this.lastCommitAttemptMs <= now - this.autoCommitIntervalMs;
     }
@@ -863,6 +879,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
         do {
             long now = time.milliseconds();
             this.client.poll(this.retryBackoffMs, now);
+            ensureNotClosed();
         } while (this.metadata.version() == version);
     }
 
@@ -944,7 +961,7 @@ public class KafkaConsumer<K, V> implements Consumer<K, V> {
      * Check that the consumer hasn't been closed.
      */
     private void ensureNotClosed() {
-        if (this.closed)
-            throw new IllegalStateException("This consumer has already been closed.");
+        if (this.closed.get())
+            throw new ClosedConsumerException();
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Coordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Coordinator.java
index b2764df..effaff6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Coordinator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Coordinator.java
@@ -21,6 +21,7 @@ import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.clients.consumer.ClosedConsumerException;
 import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
@@ -49,17 +50,19 @@ import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * This class manage the coordination process with the consumer coordinator.
  */
-public final class Coordinator {
+public final class Coordinator implements Closeable {
 
     private static final Logger log = LoggerFactory.getLogger(Coordinator.class);
 
@@ -77,6 +80,8 @@ public final class Coordinator {
     private Node consumerCoordinator;
     private String consumerId;
     private int generation;
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
 
     /**
      * Initialize the coordination manager.
@@ -116,6 +121,7 @@ public final class Coordinator {
      * @return The assigned partition info
      */
     public List<TopicPartition> assignPartitions(List<String> subscribedTopics, long now) {
+        ensureNotClosed();
 
         // send a join group request to the coordinator
         log.debug("(Re-)joining group {} with subscribed topics {}", groupId, subscribedTopics);
@@ -158,6 +164,7 @@ public final class Coordinator {
      * @param now The current time
      */
     public void commitOffsets(final Map<TopicPartition, Long> offsets, boolean blocking, long now) {
+        ensureNotClosed();
         if (!offsets.isEmpty()) {
             // create the offset commit request
             Map<TopicPartition, OffsetCommitRequest.PartitionData> offsetData;
@@ -205,6 +212,7 @@ public final class Coordinator {
      * @return The fetched offset values
      */
     public Map<TopicPartition, Long> fetchOffsets(Set<TopicPartition> partitions, long now) {
+        ensureNotClosed();
         log.debug("Fetching committed offsets for partitions: " + Utils.join(partitions, ", "));
 
         while (true) {
@@ -260,6 +268,7 @@ public final class Coordinator {
      * @param now The current time
      */
     public void maybeHeartbeat(long now) {
+        ensureNotClosed();
         if (heartbeat.shouldHeartbeat(now)) {
             HeartbeatRequest req = new HeartbeatRequest(this.groupId, this.generation, this.consumerId);
             this.client.send(initiateCoordinatorRequest(ApiKeys.HEARTBEAT, req.toStruct(), new HeartbeatCompletionHandler(), now));
@@ -308,6 +317,7 @@ public final class Coordinator {
                 } else {
                     log.debug("No connection to coordinator, attempting to connect.");
                     this.client.poll(this.retryBackoffMs, time.milliseconds());
+                    ensureNotClosed();
 
                     // if the coordinator connection has failed, we need to
                     // break the inner loop to re-discover the coordinator
@@ -406,6 +416,8 @@ public final class Coordinator {
         while (node == null || !this.client.ready(node, time.milliseconds())) {
             long now = time.milliseconds();
             this.client.poll(this.retryBackoffMs, now);
+            ensureNotClosed();
+
             node = this.client.leastLoadedNode(now);
 
             // if there is no ready node, backoff before retry
@@ -461,6 +473,15 @@ public final class Coordinator {
         }
     }
 
+    @Override
+    public void close() {
+        this.closed.set(true);
+    }
+
+    private void ensureNotClosed() {
+        if (this.closed.get()) throw new ClosedConsumerException();
+    }
+
     private class HeartbeatCompletionHandler implements RequestCompletionHandler {
         @Override
         public void onComplete(ClientResponse resp) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
index ef9dd52..bb254cb 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
@@ -25,6 +25,7 @@ import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.clients.consumer.ClosedConsumerException;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
@@ -46,6 +47,7 @@ import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.Closeable;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -53,12 +55,13 @@ import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 
 /**
  * This class manage the fetching process with the brokers.
  */
-public class Fetcher<K, V> {
+public class Fetcher<K, V> implements Closeable {
 
     private static final Logger log = LoggerFactory.getLogger(Fetcher.class);
     private static final long EARLIEST_OFFSET_TIMESTAMP = -2L;
@@ -80,7 +83,7 @@ public class Fetcher<K, V> {
     private final AutoOffsetResetStrategy offsetResetStrategy;
     private final Deserializer<K> keyDeserializer;
     private final Deserializer<V> valueDeserializer;
-
+    private final AtomicBoolean closed = new AtomicBoolean(false);
 
     public Fetcher(KafkaClient client,
                    long retryBackoffMs,
@@ -123,6 +126,7 @@ public class Fetcher<K, V> {
      * @param now The current time
      */
     public void initFetches(Cluster cluster, long now) {
+        ensureNotClosed();
         for (ClientRequest request : createFetchRequests(cluster)) {
             Node node = cluster.nodeById(request.request().destination());
             if (client.ready(node, now)) {
@@ -138,6 +142,7 @@ public class Fetcher<K, V> {
      * @return The fetched records per partition
      */
     public Map<TopicPartition, List<ConsumerRecord<K, V>>> fetchedRecords() {
+        ensureNotClosed();
         if (this.subscriptions.partitionAssignmentNeeded()) {
             return Collections.emptyMap();
         } else {
@@ -172,6 +177,7 @@ public class Fetcher<K, V> {
      * @throws org.apache.kafka.clients.consumer.NoOffsetForPartitionException If no offset reset strategy is defined
      */
     public void resetOffset(TopicPartition partition) {
+        ensureNotClosed();
         long timestamp;
         if (this.offsetResetStrategy == AutoOffsetResetStrategy.EARLIEST)
             timestamp = EARLIEST_OFFSET_TIMESTAMP;
@@ -193,6 +199,7 @@ public class Fetcher<K, V> {
      * @return The offset of the message that is published before the given timestamp
      */
     public long offsetBefore(TopicPartition topicPartition, long timestamp) {
+        ensureNotClosed();
         log.debug("Fetching offsets for partition {}.", topicPartition);
         Map<TopicPartition, ListOffsetRequest.PartitionData> partitions = new HashMap<TopicPartition, ListOffsetRequest.PartitionData>(1);
         partitions.put(topicPartition, new ListOffsetRequest.PartitionData(timestamp, 1));
@@ -243,6 +250,7 @@ public class Fetcher<K, V> {
             } else {
                 log.debug("Leader for partition {} is not ready, retry fetching offsets", topicPartition);
                 client.poll(this.retryBackoffMs, now);
+                ensureNotClosed();
             }
         }
     }
@@ -363,9 +371,19 @@ public class Fetcher<K, V> {
         do {
             long now = time.milliseconds();
             this.client.poll(this.retryBackoffMs, now);
+            ensureNotClosed();
         } while (this.metadata.version() == version);
     }
 
+    @Override
+    public void close() {
+        this.closed.set(true);
+    }
+
+    private void ensureNotClosed() {
+        if (this.closed.get()) throw new ClosedConsumerException();
+    }
+
     private static class PartitionRecords<K, V> {
         public long fetchOffset;
         public TopicPartition partition;
@@ -457,4 +475,5 @@ public class Fetcher<K, V> {
             recordsFetched.record(records);
         }
     }
+
 }
-- 
2.3.2 (Apple Git-55)