diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java index abb1f9281a..6739bb19e2 100644 --- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java +++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java @@ -472,7 +472,8 @@ public void close() { kopEventManager.close(); KafkaTopicManager.LOOKUP_CACHE.clear(); KopBrokerLookupManager.clear(); - KafkaTopicManager.closeKafkaTopicConsumerManagers(); + KafkaTopicManager.cancelCursorExpireTask(); + KafkaTopicConsumerManagerCache.getInstance().close(); KafkaTopicManager.getReferences().clear(); KafkaTopicManager.getTopics().clear(); statsProvider.stop(); diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManager.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManager.java index cc1cec1eb0..e5884ca6e3 100644 --- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManager.java +++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManager.java @@ -265,4 +265,9 @@ public ManagedLedger getManagedLedger() { public int getNumCreatedCursors() { return numCreatedCursors; } + + @VisibleForTesting + public boolean isClosed() { + return closed.get(); + } } diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerCache.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerCache.java index 3e76f62530..478e76db9d 100644 --- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerCache.java +++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerCache.java @@ -67,22 +67,36 @@ public void forEach(final Consumer> }); } - public void removeAndClose(final String fullTopicName) { - // The TCM future could be completed with null, so we should process this case + private static void closeTcmFuture(final CompletableFuture tcmFuture) { + // Use thenAccept to avoid blocking + tcmFuture.thenAccept(tcm -> { + if (tcm != null) { + tcm.close(); + } + }); + } + + public void removeAndCloseByTopic(final String fullTopicName) { Optional.ofNullable(cache.remove(fullTopicName)).ifPresent(map -> map.forEach((remoteAddress, future) -> { if (log.isDebugEnabled()) { log.debug("[{}][{}] Remove and close TCM", fullTopicName, remoteAddress); } - // Use thenAccept to avoid blocking - future.thenAccept(tcm -> { - if (tcm != null) { - tcm.close(); - } - }); + closeTcmFuture(future); })); } + public void removeAndCloseByAddress(final SocketAddress remoteAddress) { + cache.forEach((fullTopicName, internalMap) -> { + Optional.ofNullable(internalMap.remove(remoteAddress)).ifPresent(future -> { + if (log.isDebugEnabled()) { + log.debug("[{}][{}] Remove and close TCM", fullTopicName, remoteAddress); + } + closeTcmFuture(future); + }); + }); + } + public void close() { cache.forEach((fullTopicName, internalMap) -> { internalMap.forEach((remoteAddress, future) -> { diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicManager.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicManager.java index df8c4db836..74d90613d7 100644 --- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicManager.java +++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicManager.java @@ -101,6 +101,15 @@ private static void initializeCursorExpireTask(final ScheduledExecutorService ex } } + public static void cancelCursorExpireTask() { + synchronized (KafkaTopicManager.class) { + if (cursorExpireTask != null) { + cursorExpireTask.cancel(true); + cursorExpireTask = null; + } + } + } + // update Ctx information, since at internalServerCnx create time there is no ctx passed into kafkaRequestHandler. public void setRemoteAddress(SocketAddress remoteAddress) { internalServerCnx.updateCtx(remoteAddress); @@ -326,7 +335,7 @@ public void close() { } try { - closeKafkaTopicConsumerManagers(); + TCM_CACHE.removeAndCloseByAddress(remoteAddress); topics.keySet().forEach(topicName -> { if (log.isDebugEnabled()) { @@ -377,20 +386,10 @@ public static void deReference(String topicName) { try { removeTopicManagerCache(topicName); - TCM_CACHE.removeAndClose(topicName); + TCM_CACHE.removeAndCloseByTopic(topicName); removePersistentTopicAndReferenceProducer(topicName); } catch (Exception e) { log.error("Failed to close reference for individual topic {}. exception:", topicName, e); } } - - public static void closeKafkaTopicConsumerManagers() { - synchronized (KafkaTopicManager.class) { - if (cursorExpireTask != null) { - cursorExpireTask.cancel(true); - cursorExpireTask = null; - } - } - TCM_CACHE.close(); - } } diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/MessageFetchContext.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/MessageFetchContext.java index f7a67e2d60..788b409f64 100644 --- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/MessageFetchContext.java +++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/MessageFetchContext.java @@ -301,7 +301,7 @@ private void handlePartitionData(final TopicPartition topicPartition, statsLogger.getPrepareMetadataStats().registerFailedEvent( MathUtils.elapsedNanos(startPrepareMetadataNanos), TimeUnit.NANOSECONDS); // remove null future cache - KafkaTopicConsumerManagerCache.getInstance().removeAndClose(fullTopicName); + KafkaTopicConsumerManagerCache.getInstance().removeAndCloseByTopic(fullTopicName); addErrorPartitionResponse(topicPartition, Errors.NOT_LEADER_FOR_PARTITION); return; } @@ -335,7 +335,7 @@ private void handlePartitionData(final TopicPartition topicPartition, // tcm is closed, just return a NONE error because the channel may be still active log.warn("[{}] KafkaTopicConsumerManager is closed, remove TCM of {}", requestHandler.ctx, fullTopicName); - KafkaTopicConsumerManagerCache.getInstance().removeAndClose(fullTopicName); + KafkaTopicConsumerManagerCache.getInstance().removeAndCloseByTopic(fullTopicName); addErrorPartitionResponse(topicPartition, Errors.NONE); return; } diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerTest.java index d811c3fc62..0084c20985 100644 --- a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerTest.java +++ b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerTest.java @@ -16,7 +16,11 @@ import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; import io.netty.channel.Channel; @@ -26,6 +30,7 @@ import io.streamnative.pulsar.handlers.kop.stats.NullStatsLogger; import io.streamnative.pulsar.handlers.kop.utils.KopTopic; import java.time.Duration; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Properties; @@ -36,6 +41,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; import lombok.Cleanup; @@ -48,6 +54,7 @@ import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.IntegerSerializer; import org.apache.kafka.common.serialization.StringSerializer; import org.apache.pulsar.broker.protocol.ProtocolHandler; @@ -55,7 +62,6 @@ import org.apache.pulsar.client.admin.PulsarAdminException; import org.apache.pulsar.common.policies.data.TopicStats; import org.apache.pulsar.policies.data.loadbalancer.LocalBrokerData; -import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -369,7 +375,7 @@ public void testOnlyOneCursorCreated() throws Exception { final List tcmList = KafkaTopicConsumerManagerCache.getInstance().getTopicConsumerManagers(partitionName); - Assert.assertFalse(tcmList.isEmpty()); + assertFalse(tcmList.isEmpty()); // Only 1 cursor should be created for a consumer even if there were a lot of FETCH requests // This check is to ensure that KafkaTopicConsumerManager#add is called in FETCH request handler assertEquals(tcmList.get(0).getCreatedCursors().size(), 1); @@ -434,4 +440,60 @@ public void testCursorCountForMultiGroups() throws Exception { assertEquals(tcmList.get(i).getNumCreatedCursors(), 0); } } + + // KafkaTopicManager#close should only remove TCM cache for the specific address + @Test(timeOut = 20000) + public void testTopicManagerClose() throws Exception { + final String topic = "test-topic-manager-close"; + final int numPartitions = 2; + admin.topics().createPartitionedTopic(topic, numPartitions); + + final List> consumers = new ArrayList<>(); + for (int i = 0; i < numPartitions; i++) { + consumers.add(new KafkaConsumer<>(newKafkaConsumerProperties())); + consumers.get(i).assign(Collections.singleton(new TopicPartition(topic, i))); + } + + final KafkaProducer producer = new KafkaProducer<>(newKafkaProducerProperties()); + for (int i = 0; i < numPartitions; i++) { + producer.send(new ProducerRecord<>(topic, i, null, "msg-" + i)).get(); + final ConsumerRecords records = consumers.get(i).poll(Duration.ofSeconds(1)); + assertEquals(records.count(), 1); + assertEquals(records.iterator().next().value(), "msg-" + i); + } + + final Function getTcmForPartition = partition -> { + final String fullTopicName = new KopTopic(topic).getPartitionName(partition); + final List tcmList = + KafkaTopicConsumerManagerCache.getInstance().getTopicConsumerManagers(fullTopicName); + return tcmList.isEmpty() ? null : tcmList.get(0); + }; + + final List originalTcmList = new ArrayList<>(); + for (int i = 0; i < numPartitions; i++) { + final KafkaTopicConsumerManager tcm = getTcmForPartition.apply(i); + assertNotNull(tcm); + assertFalse(tcm.isClosed()); + originalTcmList.add(tcm); + } + + producer.close(); // trigger KafkaTopicManager#close but the TCM cache was not affected + assertSame(getTcmForPartition.apply(0), originalTcmList.get(0)); + assertFalse(originalTcmList.get(0).isClosed()); + assertSame(getTcmForPartition.apply(1), originalTcmList.get(1)); + assertFalse(originalTcmList.get(1).isClosed()); + + consumers.get(1).close(); // trigger KafkaTopicManager#close, only the partition 1 related cache was removed + assertSame(getTcmForPartition.apply(0), originalTcmList.get(0)); + assertFalse(originalTcmList.get(0).isClosed()); + // The tcm of partition 1 was closed and it was removed from cache + assertNull(getTcmForPartition.apply(1)); + assertTrue(originalTcmList.get(1).isClosed()); + + consumers.get(0).close(); // Now all TCM cache was cleared + assertNull(getTcmForPartition.apply(0)); + assertNull(getTcmForPartition.apply(1)); + assertTrue(originalTcmList.get(0).isClosed()); + assertTrue(originalTcmList.get(1).isClosed()); + } }