Remove and close TCM according to client address

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();

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();
+    }
 }

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaTopicConsumerManagerCache.java

@@ -67,22 +67,36 @@ public void forEach(final Consumer<CompletableFuture<KafkaTopicConsumerManager>>
         });
     }
 
-    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<KafkaTopicConsumerManager> 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) -> {

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();
-    }
 }

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;
             }

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,14 +54,14 @@
 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;
 import org.apache.pulsar.broker.service.persistent.PersistentTopic;
 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<KafkaTopicConsumerManager> 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<KafkaConsumer<String, String>> 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<String, String> producer = new KafkaProducer<>(newKafkaProducerProperties());
+        for (int i = 0; i < numPartitions; i++) {
+            producer.send(new ProducerRecord<>(topic, i, null, "msg-" + i)).get();
+            final ConsumerRecords<String, String> records = consumers.get(i).poll(Duration.ofSeconds(1));
+            assertEquals(records.count(), 1);
+            assertEquals(records.iterator().next().value(), "msg-" + i);
+        }
+
+        final Function<Integer, KafkaTopicConsumerManager> getTcmForPartition = partition -> {
+            final String fullTopicName = new KopTopic(topic).getPartitionName(partition);
+            final List<KafkaTopicConsumerManager> tcmList =
+                    KafkaTopicConsumerManagerCache.getInstance().getTopicConsumerManagers(fullTopicName);
+            return tcmList.isEmpty() ? null : tcmList.get(0);
+        };
+
+        final List<KafkaTopicConsumerManager> 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());
+    }
 }