From b7a96bd2f61f6984708bf1f3f1f1692c9d3e767f Mon Sep 17 00:00:00 2001 From: sakshichitnis27 <156598682+sakshichitnis27@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:15:26 +0000 Subject: [PATCH 1/3] fix: Wait for Kafka lookup startup catch-up --- .../lookup/KafkaLookupExtractorFactory.java | 29 ++++++- .../KafkaLookupExtractorFactoryTest.java | 87 ++++++++++++++++++- 2 files changed, 113 insertions(+), 3 deletions(-) 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..b89f565f74c2 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); @@ -381,6 +395,19 @@ ListenableFuture getFuture() return future; } + 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..53892f0aefb1 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 @@ -28,20 +28,29 @@ import com.google.common.primitives.Bytes; 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 @@ -247,7 +256,11 @@ public void testStopWithoutStart() @Test public void testStartStop() { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + final MockConsumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.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( @@ -271,6 +284,72 @@ Consumer getConsumer() EasyMock.verify(cacheManager); } + @Test + public void testStartWaitsForInitialEndOffsets() throws Exception + { + final MockConsumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.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(); + } + EasyMock.verify(cacheManager); + } + @Test public void testStartFailsFromTimeout() @@ -329,7 +408,11 @@ Consumer getConsumer() @Test public void testStartStartStopStop() { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + final MockConsumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.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, From 945415261c328604140edd9988e2c538004fa9fc Mon Sep 17 00:00:00 2001 From: sakshichitnis27 Date: Mon, 3 Aug 2026 15:36:22 +0000 Subject: [PATCH 2/3] Wait for Kafka lookup worker before closing cache --- .../lookup/KafkaLookupExtractorFactory.java | 38 ++++++++--- .../KafkaLookupExtractorFactoryTest.java | 65 ++++++++++++++----- 2 files changed, 75 insertions(+), 28 deletions(-) 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 b89f565f74c2..d8fd8cb12986 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 @@ -265,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; } @@ -286,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; } } @@ -395,6 +387,32 @@ ListenableFuture getFuture() return future; } + private void shutdownExecutorAndCloseCache() + { + executorService.shutdown(); + + final ListenableFuture future = this.future; + if (future != null) { + future.cancel(true); + } + + boolean interrupted = false; + while (!executorService.isTerminated()) { + try { + executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS); + } + catch (InterruptedException e) { + interrupted = true; + } + } + + cacheHandler.close(); + + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + private static boolean hasReachedEndOffsets( final Consumer consumer, final Map endOffsets 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 53892f0aefb1..7c082d70a8aa 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,6 +26,7 @@ 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; @@ -34,7 +35,6 @@ 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; @@ -256,7 +256,7 @@ public void testStopWithoutStart() @Test public void testStartStop() { - final MockConsumer 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)); @@ -287,7 +287,7 @@ Consumer getConsumer() @Test public void testStartWaitsForInitialEndOffsets() throws Exception { - final MockConsumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + 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); @@ -352,8 +352,25 @@ Consumer getConsumer() @Test - public void testStartFailsFromTimeout() + public void testStartTimeoutWaitsForConsumerToStop() 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, @@ -364,28 +381,40 @@ 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.assertThrows( + "start returned before the Kafka consumer stopped", + TimeoutException.class, + () -> startFuture.get(500, TimeUnit.MILLISECONDS) + ); + Assert.assertEquals(1L, consumerClosed.getCount()); + + allowPollToFinish.countDown(); + Assert.assertFalse(startFuture.get(10, TimeUnit.SECONDS)); + Assert.assertTrue(consumerClosed.await(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() { - Consumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + Consumer kafkaConsumer = new MockConsumer<>("earliest"); EasyMock.replay(cacheManager); final KafkaLookupExtractorFactory factory = new KafkaLookupExtractorFactory( cacheManager, @@ -408,7 +437,7 @@ Consumer getConsumer() @Test public void testStartStartStopStop() { - final MockConsumer 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)); From d7d6f3b25a1c49106ca092227beffa2dcb61b9b1 Mon Sep 17 00:00:00 2001 From: sakshichitnis27 Date: Wed, 5 Aug 2026 12:56:53 +0000 Subject: [PATCH 3/3] fix: Avoid blocking Kafka lookup startup timeout --- .../lookup/KafkaLookupExtractorFactory.java | 23 ++++--------- .../KafkaLookupExtractorFactoryTest.java | 33 +++++++++++-------- 2 files changed, 26 insertions(+), 30 deletions(-) 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 d8fd8cb12986..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 @@ -387,30 +387,21 @@ 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); } - - boolean interrupted = false; - while (!executorService.isTerminated()) { - try { - executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS); - } - catch (InterruptedException e) { - interrupted = true; - } - } - - cacheHandler.close(); - - if (interrupted) { - Thread.currentThread().interrupt(); - } } private static boolean hasReachedEndOffsets( 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 7c082d70a8aa..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 @@ -84,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 { @@ -254,7 +262,7 @@ public void testStopWithoutStart() } @Test - public void testStartStop() + public void testStartStop() throws InterruptedException { final MockConsumer kafkaConsumer = new MockConsumer<>("earliest"); final TopicPartition topicPartition = new TopicPartition(TOPIC, 0); @@ -281,7 +289,7 @@ Consumer getConsumer() Assert.assertTrue(factory.start()); Assert.assertTrue(factory.close()); Assert.assertTrue(factory.getFuture().isDone()); - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); } @Test @@ -347,12 +355,12 @@ Consumer getConsumer() factory.close(); startExecutor.shutdownNow(); } - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); } @Test - public void testStartTimeoutWaitsForConsumerToStop() throws Exception + public void testStartTimeoutReturnsBeforeConsumerStops() throws Exception { final CountDownLatch pollStarted = new CountDownLatch(1); final CountDownLatch allowPollToFinish = new CountDownLatch(1); @@ -391,16 +399,13 @@ Consumer getConsumer() try { Assert.assertTrue(pollStarted.await(10, TimeUnit.SECONDS)); - Assert.assertThrows( - "start returned before the Kafka consumer stopped", - TimeoutException.class, - () -> startFuture.get(500, TimeUnit.MILLISECONDS) - ); + Assert.assertFalse(startFuture.get(500, TimeUnit.MILLISECONDS)); Assert.assertEquals(1L, consumerClosed.getCount()); + Assert.assertFalse(factory.awaitExecutorTermination(100, TimeUnit.MILLISECONDS)); allowPollToFinish.countDown(); - Assert.assertFalse(startFuture.get(10, TimeUnit.SECONDS)); Assert.assertTrue(consumerClosed.await(10, TimeUnit.SECONDS)); + Assert.assertTrue(factory.awaitExecutorTermination(10, TimeUnit.SECONDS)); Assert.assertTrue(factory.getFuture().isDone()); Assert.assertTrue(factory.getFuture().isCancelled()); } @@ -412,7 +417,7 @@ Consumer getConsumer() } @Test - public void testStartStopStart() + public void testStartStopStart() throws InterruptedException { Consumer kafkaConsumer = new MockConsumer<>("earliest"); EasyMock.replay(cacheManager); @@ -431,11 +436,11 @@ 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 { final MockConsumer kafkaConsumer = new MockConsumer<>("earliest"); final TopicPartition topicPartition = new TopicPartition(TOPIC, 0); @@ -461,7 +466,7 @@ Consumer getConsumer() Assert.assertTrue(factory.start()); Assert.assertTrue(factory.close()); Assert.assertTrue(factory.close()); - EasyMock.verify(cacheManager); + verifyCacheManagerAfterExecutorTerminates(factory); } @Test