This is an automated email from the ASF dual-hosted git repository.
FrankChen021 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 017c18ab6db fix: Wait for Kafka lookup startup catch-up (#19858)
017c18ab6db is described below
commit 017c18ab6dbf334f1ec58318faef016817e302e0
Author: sakshichitnis27 <[email protected]>
AuthorDate: Sat Sep 26 05:00:53 2026 +0100
fix: Wait for Kafka lookup startup catch-up (#19858)
* fix: Wait for Kafka lookup startup catch-up
* Wait for Kafka lookup worker before closing cache
* fix: Avoid blocking Kafka lookup startup timeout
* Fix deprecated MockConsumer usage in Kafka lookup tests
* Replace remaining deprecated MockConsumer usages
---
.../query/lookup/KafkaLookupExtractorFactory.java | 58 ++++++--
.../lookup/KafkaLookupExtractorFactoryTest.java | 163 ++++++++++++++++++---
2 files changed, 187 insertions(+), 34 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 81c80de1125..2c19e8ece15 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
@@ -43,6 +43,7 @@ import org.apache.kafka.clients.consumer.ConsumerConfig;
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.time.Duration;
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 class KafkaLookupExtractorFactory implements
LookupExtractorFactory
final ListenableFuture<?> future = executorService.submit(() -> {
final Consumer<String, String> consumer = getConsumer();
consumer.subscribe(Collections.singletonList(topic));
+ Map<TopicPartition, Long> startupEndOffsets = null;
try {
while (!executorService.isShutdown()) {
try {
@@ -183,7 +186,13 @@ public class KafkaLookupExtractorFactory implements
LookupExtractorFactory
break;
}
final ConsumerRecords<String, String> records =
consumer.poll(Duration.ofMillis(1000));
- startingReads.countDown();
+ if (startingReads.getCount() > 0) {
+ final Set<TopicPartition> assignment = consumer.assignment();
+ if (!assignment.isEmpty()
+ && (startupEndOffsets == null ||
!startupEndOffsets.keySet().equals(assignment))) {
+ startupEndOffsets = consumer.endOffsets(assignment);
+ }
+ }
for (final ConsumerRecord<String, String> record : records) {
final String key = record.key();
@@ -204,6 +213,11 @@ public class KafkaLookupExtractorFactory implements
LookupExtractorFactory
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 class KafkaLookupExtractorFactory implements
LookupExtractorFactory
}
}
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 class KafkaLookupExtractorFactory implements
LookupExtractorFactory
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 @@ public class KafkaLookupExtractorFactory implements
LookupExtractorFactory
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<String, String> consumer,
+ final Map<TopicPartition, Long> endOffsets
+ )
+ {
+ for (final Map.Entry<TopicPartition, Long> 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 c9466d28c7f..d81ad1fbf89 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.fasterxml.jackson.databind.ObjectMapper;
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.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.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 class KafkaLookupExtractorFactoryTest
});
}
+ private void verifyCacheManagerAfterExecutorTerminates(
+ final KafkaLookupExtractorFactory factory
+ ) throws InterruptedException
+ {
+ Assertions.assertTrue(factory.awaitExecutorTermination(10,
TimeUnit.SECONDS));
+ EasyMock.verify(cacheManager);
+ }
+
@Test
public void testSimpleSerDe() throws Exception
{
@@ -245,9 +262,13 @@ public class KafkaLookupExtractorFactoryTest
}
@Test
- public void testStartStop()
+ public void testStartStop() throws InterruptedException
{
- Consumer<String, String> kafkaConsumer = new
MockConsumer<>(OffsetResetStrategy.EARLIEST);
+ final MockConsumer<String, String> 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 @@ public class KafkaLookupExtractorFactoryTest
Assertions.assertTrue(factory.start());
Assertions.assertTrue(factory.close());
Assertions.assertTrue(factory.getFuture().isDone());
- EasyMock.verify(cacheManager);
+ verifyCacheManagerAfterExecutorTerminates(factory);
+ }
+
+ @Test
+ public void testStartWaitsForInitialEndOffsets() throws Exception
+ {
+ final MockConsumer<String, String> 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<String, String> getConsumer()
+ {
+ return kafkaConsumer;
+ }
+ };
+ final ExecutorService startExecutor =
Execs.singleThreaded("kafka-lookup-start-test");
+ final Future<Boolean> startFuture = startExecutor.submit(factory::start);
+
+ try {
+ Assertions.assertTrue(firstPollComplete.await(10, TimeUnit.SECONDS));
+ Assertions.assertThrows(
+ TimeoutException.class,
+ () -> startFuture.get(100, TimeUnit.MILLISECONDS),
+ "start returned before the consumer reached its initial end offsets"
+ );
+ allowCatchUp.countDown();
+ Assertions.assertTrue(startFuture.get(10, TimeUnit.SECONDS));
+ Assertions.assertEquals("value-0", factory.get().apply("key-0"));
+ Assertions.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<String, String> 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 class KafkaLookupExtractorFactoryTest
)
{
@Override
- Consumer getConsumer()
+ Consumer<String, String> getConsumer()
{
- // Lock up
- try {
- Thread.currentThread().join();
- }
- catch (InterruptedException e) {
- throw new RuntimeException(e);
- }
- throw new RuntimeException("shouldn't make it here");
+ return kafkaConsumer;
}
};
- Assertions.assertFalse(factory.start());
- Assertions.assertTrue(factory.getFuture().isDone());
- Assertions.assertTrue(factory.getFuture().isCancelled());
+ final ExecutorService startExecutor =
Execs.singleThreaded("kafka-lookup-timeout-test");
+ final Future<Boolean> startFuture = startExecutor.submit(factory::start);
+
+ try {
+ Assertions.assertTrue(pollStarted.await(10, TimeUnit.SECONDS));
+ Assertions.assertFalse(startFuture.get(500, TimeUnit.MILLISECONDS));
+ Assertions.assertEquals(1L, consumerClosed.getCount());
+ Assertions.assertFalse(factory.awaitExecutorTermination(100,
TimeUnit.MILLISECONDS));
+
+ allowPollToFinish.countDown();
+ Assertions.assertTrue(consumerClosed.await(10, TimeUnit.SECONDS));
+ Assertions.assertTrue(factory.awaitExecutorTermination(10,
TimeUnit.SECONDS));
+ Assertions.assertTrue(factory.getFuture().isDone());
+ Assertions.assertTrue(factory.getFuture().isCancelled());
+ }
+ finally {
+ allowPollToFinish.countDown();
+ startExecutor.shutdownNow();
+ }
EasyMock.verify(cacheManager);
}
@Test
- public void testStartStopStart()
+ public void testStartStopStart() throws InterruptedException
{
- Consumer<String, String> kafkaConsumer = new
MockConsumer<>(OffsetResetStrategy.EARLIEST);
+ Consumer<String, String> kafkaConsumer = new MockConsumer<>("earliest");
EasyMock.replay(cacheManager);
final KafkaLookupExtractorFactory factory = new
KafkaLookupExtractorFactory(
cacheManager,
@@ -323,13 +436,17 @@ public class KafkaLookupExtractorFactoryTest
Assertions.assertTrue(factory.start());
Assertions.assertTrue(factory.close());
Assertions.assertFalse(factory.start());
- EasyMock.verify(cacheManager);
+ verifyCacheManagerAfterExecutorTerminates(factory);
}
@Test
- public void testStartStartStopStop()
+ public void testStartStartStopStop() throws InterruptedException
{
- Consumer<String, String> kafkaConsumer = new
MockConsumer<>(OffsetResetStrategy.EARLIEST);
+ final MockConsumer<String, String> 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 @@ public class KafkaLookupExtractorFactoryTest
Assertions.assertTrue(factory.start());
Assertions.assertTrue(factory.close());
Assertions.assertTrue(factory.close());
- EasyMock.verify(cacheManager);
+ verifyCacheManagerAfterExecutorTerminates(factory);
}
@Test
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
