Kafka concurrent (#159)

* kafka source init

* kafka consumer wrapper and listener

* unit test and checkstyle

* exception handling

* unit tests for kafka source

* interfaces and wrappers for kafka source

* removed kafka dep

* integration tests for kafka source

* config deserialization

* Fixed deserializer

* Fixed deserializer

* styling

* kafka docker

* pipeline config example

* style

* error checking

* error checking

* integration test kafka source from config

* style

* added parser to kafka source

* example pipeline

* Fail integration test on send fail to kafka server

* requested changes

* requested changes

* configurable backoff time for kafka source

* fixed conf deserializer

* concurrent parsing workers

* multi worker unit test

* queue holds record values instead of records

* style

* mock observer for multithreaded testing

* configurable buffer queue size

* moved fill queue integration test to unit test

* style

* ensure queue fills in queue filled test

* refactor kafka source constructors

Author: joeyjackson

src/main/java/com/arpnetworking/metrics/common/kafka/RunnableConsumerImpl.java

@@ -73,7 +73,6 @@ public void run() {
                 _listener.handle(e);
             }
         }
-        _consumer.close();
     }
 
     @Override

src/main/java/com/arpnetworking/metrics/common/sources/KafkaSource.java

@@ -26,18 +26,21 @@
 import com.arpnetworking.steno.LoggerFactory;
 import net.sf.oval.constraint.CheckWith;
 import net.sf.oval.constraint.CheckWithCheck;
+import net.sf.oval.constraint.Min;
 import net.sf.oval.constraint.NotNull;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.KafkaException;
 
 import java.time.Duration;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
 /**
- * Produce instances of {@code Record} from the values of entries
+ * Produce instances of {@link com.arpnetworking.metrics.mad.model.Record} from the values of entries
  * from a Kafka topic. The key from the entries gets discarded
  *
  * @param <T> the type of data created by the source
@@ -52,28 +55,47 @@ public final class KafkaSource<T, V> extends BaseSource {
     private final Consumer<?, V> _consumer;
     private final RunnableConsumer _runnableConsumer;
     private final ExecutorService _consumerExecutor;
+    private final ExecutorService _parserExecutor;
     private final Parser<T, V> _parser;
     private final Logger _logger;
     private final Duration _shutdownAwaitTime;
     private final Duration _backoffTime;
+    private final Integer _numWorkerThreads;
+    private final BlockingQueue<V> _buffer;
+    private final ParsingWorker _parsingWorker = new ParsingWorker();
 
     @Override
     public void start() {
         _consumerExecutor.execute(_runnableConsumer);
+        for (int i = 0; i < _numWorkerThreads; i++) {
+            _parserExecutor.execute(_parsingWorker);
+        }
     }
 
     @Override
     public void stop() {
         _runnableConsumer.stop();
-
         _consumerExecutor.shutdown();
         try {
             _consumerExecutor.awaitTermination(_shutdownAwaitTime.toMillis(), TimeUnit.MILLISECONDS);
         } catch (final InterruptedException e) {
-            LOGGER.warn()
+            _logger.warn()
                     .setMessage("Unable to shutdown kafka consumer executor")
                     .setThrowable(e)
                     .log();
+        } finally {
+            _consumer.close();
+        }
+
+        _parsingWorker.stop();
+        _parserExecutor.shutdown();
+        try {
+            _parserExecutor.awaitTermination(_shutdownAwaitTime.toMillis(), TimeUnit.MILLISECONDS);
+        } catch (final InterruptedException e) {
+            _logger.warn()
+                    .setMessage("Unable to shutdown parsing worker executor")
+                    .setThrowable(e)
+                    .log();
         }
     }
 
@@ -97,39 +119,88 @@ public String toString() {
 
     @SuppressWarnings("unused")
     private KafkaSource(final Builder<T, V> builder) {
-        this(builder, LOGGER);
+        this(builder, LOGGER, new ArrayBlockingQueue<>(builder._bufferSize));
     }
 
+    // NOTE: Package private for testing
     /* package private */ KafkaSource(final Builder<T, V> builder, final Logger logger) {
+        this(builder, logger, new ArrayBlockingQueue<>(builder._bufferSize));
+    }
+
+    // NOTE: Package private for testing
+    /* package private */ KafkaSource(final Builder<T, V> builder, final BlockingQueue<V> buffer) {
+        this(builder, LOGGER, buffer);
+    }
+
+    private KafkaSource(final Builder<T, V> builder, final Logger logger, final BlockingQueue<V> buffer) {
         super(builder);
-        _logger = logger;
         _consumer = builder._consumer;
         _parser = builder._parser;
         _runnableConsumer = new RunnableConsumerImpl.Builder<V>()
                 .setConsumer(builder._consumer)
                 .setListener(new LogConsumerListener())
                 .setPollTime(builder._pollTime)
                 .build();
+        _numWorkerThreads = builder._numWorkerThreads;
         _consumerExecutor = Executors.newSingleThreadExecutor(runnable -> new Thread(runnable, "KafkaConsumer"));
+        _parserExecutor = Executors.newFixedThreadPool(_numWorkerThreads);
         _shutdownAwaitTime = builder._shutdownAwaitTime;
         _backoffTime = builder._backoffTime;
+        _logger = logger;
+        _buffer = buffer;
+    }
+
+    private class ParsingWorker implements Runnable {
+        private volatile boolean _isRunning = true;
+
+        @Override
+        public void run() {
+            while (_isRunning || !_buffer.isEmpty()) { // Empty the queue before stopping the workers
+                final V value = _buffer.poll();
+                if (value != null) {
+                    final T record;
+                    try {
+                        record = _parser.parse(value);
+                    } catch (final ParsingException e) {
+                        _logger.error()
+                                .setMessage("Failed to parse data")
+                                .setThrowable(e)
+                                .log();
+                        return;
+                    }
+                    KafkaSource.this.notify(record);
+                } else {
+                    // Queue is empty
+                    try {
+                        Thread.sleep(_backoffTime.toMillis());
+                    } catch (final InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        stop();
+                    }
+                }
+            }
+        }
+
+        public void stop() {
+            _isRunning = false;
+        }
     }
 
     private class LogConsumerListener implements ConsumerListener<V> {
 
         @Override
         public void handle(final ConsumerRecord<?, V> consumerRecord) {
-            final T record;
             try {
-                record = _parser.parse(consumerRecord.value());
-            } catch (final ParsingException e) {
-                _logger.error()
-                        .setMessage("Failed to parse data")
+                _buffer.put(consumerRecord.value());
+            } catch (final InterruptedException e) {
+                _logger.info()
+                        .setMessage("Consumer thread interrupted")
+                        .addData("source", KafkaSource.this)
+                        .addData("action", "stopping")
                         .setThrowable(e)
                         .log();
-                return;
+                _runnableConsumer.stop();
             }
-            KafkaSource.this.notify(record);
         }
 
         @Override
@@ -181,7 +252,7 @@ private void backoff(final Throwable throwable) {
     }
 
     /**
-     * Builder pattern class for {@code KafkaSource}.
+     * Builder pattern class for {@link KafkaSource}.
      *
      * @param <T> the type of data created by the source
      * @param <V> the type of data of value in kafka {@code ConsumerRecords}
@@ -231,11 +302,11 @@ public Builder<T, V> setPollTime(final Duration pollTime) {
         }
 
         /**
-         * Sets the amount of time the {@link KafkaSource} will wait to shutdown the {@code RunnableConsumer} thread.
+         * Sets the amount of time the {@link KafkaSource} will wait to shutdown the {@link RunnableConsumer} thread.
          * Default is 10 seconds. Cannot be null or negative.
          *
          * @param shutdownAwaitTime The {@code Duration} the {@link KafkaSource} will wait to shutdown
-         *                          the {@code RunnableConsumer} thread.
+         *                          the {@link RunnableConsumer} thread.
          * @return This instance of {@link KafkaSource.Builder}.
          */
         public Builder<T, V> setShutdownAwaitTime(final Duration shutdownAwaitTime) {
@@ -256,6 +327,30 @@ public Builder<T, V> setBackoffTime(final Duration backoffTime) {
             return this;
         }
 
+        /**
+         * Sets the number of threads that will parse {@code ConsumerRecord}s from the {@code Consumer}.
+         * Default is 1. Must be greater than or equal to 1.
+         *
+         * @param numWorkerThreads The number of parsing worker threads.
+         * @return This instance of {@link KafkaSource.Builder}.
+         */
+        public Builder<T, V> setNumWorkerThreads(final Integer numWorkerThreads) {
+            _numWorkerThreads = numWorkerThreads;
+            return this;
+        }
+
+        /**
+         * Sets the size of the buffer to hold {@code ConsumerRecord}s from the {@code Consumer} before they are parsed.
+         * Default is 1000. Must be greater than or equal to 1.
+         *
+         * @param bufferSize The size of the buffer.
+         * @return This instance of {@link KafkaSource.Builder}.
+         */
+        public Builder<T, V> setBufferSize(final Integer bufferSize) {
+            _bufferSize = bufferSize;
+            return this;
+        }
+
         @Override
         protected Builder<T, V> self() {
             return this;
@@ -274,6 +369,12 @@ protected Builder<T, V> self() {
         @NotNull
         @CheckWith(value = PositiveDuration.class, message = "Backoff time must be positive.")
         private Duration _backoffTime = Duration.ofSeconds(1);
+        @NotNull
+        @Min(1)
+        private Integer _numWorkerThreads = 1;
+        @NotNull
+        @Min(1)
+        private Integer _bufferSize = 1000;
 
         private static class PositiveDuration implements CheckWithCheck.SimpleCheck {
             @Override

src/test/java/com/arpnetworking/metrics/common/integration/KafkaIT.java

@@ -43,6 +43,7 @@
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
 import java.io.IOException;
@@ -56,6 +57,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
 
 /**
  * Tests to check integration with a kafka topic.
@@ -69,7 +71,8 @@ public class KafkaIT {
     private static final String KEY_DESERIALIZER = "org.apache.kafka.common.serialization.IntegerDeserializer";
     private static final String VALUE_DESERIALIZER = "org.apache.kafka.common.serialization.StringDeserializer";
     private static final Duration POLL_DURATION = Duration.ofSeconds(1);
-    private static final int TIMEOUT = 5000;
+    private static final int TIMEOUT = 10000;
+    private static final int NUM_RECORDS = 500;
 
     private Map<String, Object> _consumerProps;
     private KafkaConsumer<Integer, String> _consumer;
@@ -82,19 +85,7 @@ public void setUp() throws TimeoutException {
         // Create kafka topic
         _topicName = createTopicName();
         createTopic(_topicName);
-
-        // Create and send producer records
-        _producerRecords = createProducerRecords(_topicName, 3);
-        sendRecords(_producerRecords);
-
-        // Create consumer props
-        _consumerProps = Maps.newHashMap();
-        _consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
-        _consumerProps.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer0");
-        _consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, KEY_DESERIALIZER);
-        _consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, VALUE_DESERIALIZER);
-        _consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
-        _consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group" + _topicName);
+        setupKafka();
     }
 
     @After
@@ -189,6 +180,32 @@ public void testKafkaSourceFromConfig() throws IOException {
         }
     }
 
+    @Test
+    public void testKafkaSourceMultipleWorkerThreads() {
+        // Create Kafka Source
+        _consumer = new KafkaConsumer<>(_consumerProps);
+        _consumer.subscribe(Collections.singletonList(_topicName));
+        _source = new KafkaSource.Builder<String, String>()
+                .setName("KafkaSource")
+                .setParser(new StringParser())
+                .setConsumer(_consumer)
+                .setPollTime(POLL_DURATION)
+                .setNumWorkerThreads(4)
+                .build();
+
+        // Observe records
+        final Observer observer = Mockito.mock(Observer.class);
+        _source.attach(observer);
+        _source.start();
+
+        final ArgumentCaptor<Object> captor = ArgumentCaptor.forClass(Object.class);
+        Mockito.verify(observer, Mockito.timeout(TIMEOUT).times(NUM_RECORDS)).notify(Mockito.any(), captor.capture());
+        Assert.assertEquals(
+                _producerRecords.stream().map(ProducerRecord::value).sorted().collect(Collectors.toList()),
+                captor.getAllValues().stream().sorted().collect(Collectors.toList())
+        );
+    }
+
     private String createTopicName() {
         return "topic_" + UUID.randomUUID().toString();
     }
@@ -252,6 +269,21 @@ private void sendRecords(final List<ProducerRecord<Integer, String>> records) {
         producer.close();
     }
 
+    private void setupKafka() {
+        // Create and send producer records
+        _producerRecords = createProducerRecords(_topicName, NUM_RECORDS);
+        sendRecords(_producerRecords);
+
+        // Create consumer props
+        _consumerProps = Maps.newHashMap();
+        _consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+        _consumerProps.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer0");
+        _consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, KEY_DESERIALIZER);
+        _consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, VALUE_DESERIALIZER);
+        _consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        _consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "group" + _topicName);
+    }
+
     /**
      * Class needed for generic object mapping with {@code ObjectMapper}.
      *