diff --git a/extensions-core/kafka-extraction-namespace/src/main/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactory.java b/extensions-core/kafka-extraction-namespace/src/main/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactory.java index 041a4cecd71c..0c1f2061fa68 100644 --- a/extensions-core/kafka-extraction-namespace/src/main/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactory.java +++ b/extensions-core/kafka-extraction-namespace/src/main/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactory.java @@ -42,6 +42,7 @@ import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.StringDeserializer; import javax.annotation.Nonnull; @@ -52,6 +53,7 @@ import java.util.Collections; import java.util.Map; import java.util.Properties; +import java.util.Set; import java.util.UUID; import java.util.concurrent.CancellationException; import java.util.concurrent.CountDownLatch; @@ -176,6 +178,7 @@ public boolean start() final ListenableFuture future = executorService.submit(() -> { final Consumer consumer = getConsumer(); consumer.subscribe(Collections.singletonList(topic)); + Map startupEndOffsets = null; try { while (!executorService.isShutdown()) { try { @@ -183,7 +186,13 @@ public boolean start() break; } final ConsumerRecords records = consumer.poll(Duration.ofMillis(1000)); - startingReads.countDown(); + if (startingReads.getCount() > 0) { + final Set assignment = consumer.assignment(); + if (!assignment.isEmpty() + && (startupEndOffsets == null || !startupEndOffsets.keySet().equals(assignment))) { + startupEndOffsets = consumer.endOffsets(assignment); + } + } for (final ConsumerRecord record : records) { final String key = record.key(); @@ -204,6 +213,11 @@ public boolean start() doubleEventCount.incrementAndGet(); LOG.trace("Placed key[%s] val[%s]", key, message); } + if (startingReads.getCount() > 0 + && startupEndOffsets != null + && hasReachedEndOffsets(consumer, startupEndOffsets)) { + startingReads.countDown(); + } } catch (Exception e) { LOG.error(e, "Error reading stream for topic [%s]", topic); @@ -251,10 +265,8 @@ public void onFailure(Throwable t) } } catch (InterruptedException | ExecutionException | TimeoutException e) { - executorService.shutdown(); - future.cancel(true); LOG.error(e, "Failed to start kafka extraction factory"); - cacheHandler.close(); + shutdownExecutorAndCloseCache(); return false; } @@ -272,13 +284,7 @@ public boolean close() return !started.get(); } started.set(false); - executorService.shutdown(); - - final ListenableFuture future = this.future; - if (future != null) { - future.cancel(true); - } - cacheHandler.close(); + shutdownExecutorAndCloseCache(); return true; } } @@ -381,6 +387,36 @@ ListenableFuture getFuture() return future; } + boolean awaitExecutorTermination(final long timeout, final TimeUnit unit) throws InterruptedException + { + return executorService.awaitTermination(timeout, unit); + } + + private void shutdownExecutorAndCloseCache() + { + // The executor is single-threaded, so the cache is not closed until the Kafka worker stops using it. + executorService.execute(cacheHandler::close); + executorService.shutdown(); + + final ListenableFuture future = this.future; + if (future != null) { + future.cancel(true); + } + } + + private static boolean hasReachedEndOffsets( + final Consumer consumer, + final Map endOffsets + ) + { + for (final Map.Entry endOffset : endOffsets.entrySet()) { + if (consumer.position(endOffset.getKey()) < endOffset.getValue()) { + return false; + } + } + return true; + } + /** * Check that the user has not set forbidden Kafka consumer props * diff --git a/extensions-core/kafka-extraction-namespace/src/test/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactoryTest.java b/extensions-core/kafka-extraction-namespace/src/test/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactoryTest.java index 24a7b481b229..89d2cdc2ed2d 100644 --- a/extensions-core/kafka-extraction-namespace/src/test/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactoryTest.java +++ b/extensions-core/kafka-extraction-namespace/src/test/java/org/apache/druid/query/lookup/KafkaLookupExtractorFactoryTest.java @@ -26,22 +26,31 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Sets; import com.google.common.primitives.Bytes; +import com.google.common.util.concurrent.Uninterruptibles; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.IAE; +import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.server.lookup.namespace.cache.MockNamespaceExtractionCacheManager; import org.apache.druid.server.lookup.namespace.cache.NamespaceExtractionCacheManager; import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.MockConsumer; -import org.apache.kafka.clients.consumer.OffsetResetStrategy; +import org.apache.kafka.common.TopicPartition; import org.easymock.EasyMock; import org.junit.Assert; import org.junit.Before; import org.junit.Test; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; public class KafkaLookupExtractorFactoryTest @@ -75,6 +84,14 @@ public Object findInjectableValue( }); } + private void verifyCacheManagerAfterExecutorTerminates( + final KafkaLookupExtractorFactory factory + ) throws InterruptedException + { + Assert.assertTrue(factory.awaitExecutorTermination(10, TimeUnit.SECONDS)); + EasyMock.verify(cacheManager); + } + @Test public void testSimpleSerDe() throws Exception { @@ -245,9 +262,13 @@ public void testStopWithoutStart() } @Test - public void testStartStop() + public void testStartStop() throws InterruptedException { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + final MockConsumer kafkaConsumer = new MockConsumer<>("earliest"); + final TopicPartition topicPartition = new TopicPartition(TOPIC, 0); + kafkaConsumer.updateBeginningOffsets(ImmutableMap.of(topicPartition, 0L)); + kafkaConsumer.updateEndOffsets(ImmutableMap.of(topicPartition, 0L)); + kafkaConsumer.schedulePollTask(() -> kafkaConsumer.rebalance(Collections.singletonList(topicPartition))); EasyMock.replay(cacheManager); final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( @@ -268,13 +289,96 @@ Consumer getConsumer() Assert.assertTrue(factory.start()); Assert.assertTrue(factory.close()); Assert.assertTrue(factory.getFuture().isDone()); - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); + } + + @Test + public void testStartWaitsForInitialEndOffsets() throws Exception + { + final MockConsumer kafkaConsumer = new MockConsumer<>("earliest"); + final TopicPartition topicPartition = new TopicPartition(TOPIC, 0); + final CountDownLatch firstPollComplete = new CountDownLatch(1); + final CountDownLatch allowCatchUp = new CountDownLatch(1); + + kafkaConsumer.schedulePollTask(() -> { + kafkaConsumer.updateBeginningOffsets(ImmutableMap.of(topicPartition, 0L)); + kafkaConsumer.updateEndOffsets(ImmutableMap.of(topicPartition, 2L)); + kafkaConsumer.rebalance(Collections.singletonList(topicPartition)); + kafkaConsumer.addRecord(new ConsumerRecord<>(TOPIC, 0, 0L, "key-0", "value-0")); + firstPollComplete.countDown(); + }); + kafkaConsumer.schedulePollTask(() -> { + try { + if (!allowCatchUp.await(10, TimeUnit.SECONDS)) { + throw new RuntimeException("Timed out waiting to finish the startup catch-up"); + } + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + kafkaConsumer.addRecord(new ConsumerRecord<>(TOPIC, 0, 1L, "key-1", "value-1")); + }); + + EasyMock.replay(cacheManager); + final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( + cacheManager, + TOPIC, + ImmutableMap.of("bootstrap.servers", "localhost"), + 10_000L, + false + ) + { + @Override + Consumer getConsumer() + { + return kafkaConsumer; + } + }; + final ExecutorService startExecutor = Execs.singleThreaded("kafka-lookup-start-test"); + final Future startFuture = startExecutor.submit(factory::start); + + try { + Assert.assertTrue(firstPollComplete.await(10, TimeUnit.SECONDS)); + Assert.assertThrows( + "start returned before the consumer reached its initial end offsets", + TimeoutException.class, + () -> startFuture.get(100, TimeUnit.MILLISECONDS) + ); + allowCatchUp.countDown(); + Assert.assertTrue(startFuture.get(10, TimeUnit.SECONDS)); + Assert.assertEquals("value-0", factory.get().apply("key-0")); + Assert.assertEquals("value-1", factory.get().apply("key-1")); + } + finally { + allowCatchUp.countDown(); + factory.close(); + startExecutor.shutdownNow(); + } + verifyCacheManagerAfterExecutorTerminates(factory); } @Test - public void testStartFailsFromTimeout() + public void testStartTimeoutReturnsBeforeConsumerStops() throws Exception { + final CountDownLatch pollStarted = new CountDownLatch(1); + final CountDownLatch allowPollToFinish = new CountDownLatch(1); + final CountDownLatch consumerClosed = new CountDownLatch(1); + final MockConsumer kafkaConsumer = new MockConsumer<>("earliest") + { + @Override + public synchronized void close() + { + super.close(); + consumerClosed.countDown(); + } + }; + kafkaConsumer.schedulePollTask(() -> { + pollStarted.countDown(); + Uninterruptibles.awaitUninterruptibly(allowPollToFinish); + }); + EasyMock.replay(cacheManager); final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( cacheManager, @@ -285,28 +389,37 @@ public void testStartFailsFromTimeout() ) { @Override - Consumer getConsumer() + Consumer getConsumer() { - // Lock up - try { - Thread.currentThread().join(); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - throw new RuntimeException("shouldn't make it here"); + return kafkaConsumer; } }; - Assert.assertFalse(factory.start()); - Assert.assertTrue(factory.getFuture().isDone()); - Assert.assertTrue(factory.getFuture().isCancelled()); + final ExecutorService startExecutor = Execs.singleThreaded("kafka-lookup-timeout-test"); + final Future startFuture = startExecutor.submit(factory::start); + + try { + Assert.assertTrue(pollStarted.await(10, TimeUnit.SECONDS)); + Assert.assertFalse(startFuture.get(500, TimeUnit.MILLISECONDS)); + Assert.assertEquals(1L, consumerClosed.getCount()); + Assert.assertFalse(factory.awaitExecutorTermination(100, TimeUnit.MILLISECONDS)); + + allowPollToFinish.countDown(); + Assert.assertTrue(consumerClosed.await(10, TimeUnit.SECONDS)); + Assert.assertTrue(factory.awaitExecutorTermination(10, TimeUnit.SECONDS)); + Assert.assertTrue(factory.getFuture().isDone()); + Assert.assertTrue(factory.getFuture().isCancelled()); + } + finally { + allowPollToFinish.countDown(); + startExecutor.shutdownNow(); + } EasyMock.verify(cacheManager); } @Test - public void testStartStopStart() + public void testStartStopStart() throws InterruptedException { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + Consumer kafkaConsumer = new MockConsumer<>("earliest"); EasyMock.replay(cacheManager); final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( cacheManager, @@ -323,13 +436,17 @@ Consumer getConsumer() Assert.assertTrue(factory.start()); Assert.assertTrue(factory.close()); Assert.assertFalse(factory.start()); - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); } @Test - public void testStartStartStopStop() + public void testStartStartStopStop() throws InterruptedException { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + final MockConsumer kafkaConsumer = new MockConsumer<>("earliest"); + final TopicPartition topicPartition = new TopicPartition(TOPIC, 0); + kafkaConsumer.updateBeginningOffsets(ImmutableMap.of(topicPartition, 0L)); + kafkaConsumer.updateEndOffsets(ImmutableMap.of(topicPartition, 0L)); + kafkaConsumer.schedulePollTask(() -> kafkaConsumer.rebalance(Collections.singletonList(topicPartition))); EasyMock.replay(cacheManager); final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( cacheManager, @@ -349,7 +466,7 @@ Consumer getConsumer() Assert.assertTrue(factory.start()); Assert.assertTrue(factory.close()); Assert.assertTrue(factory.close()); - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); } @Test