[BUGFIX] Fix loadTxnLogTopics might crash in K8S environment

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

@@ -16,7 +16,6 @@
 import static com.google.common.base.Preconditions.checkState;
 import static io.streamnative.pulsar.handlers.kop.KopServerStats.SERVER_SCOPE;
 import static io.streamnative.pulsar.handlers.kop.utils.TopicNameUtils.getKafkaTopicNameFromPulsarTopicName;
-import static org.apache.pulsar.common.naming.TopicName.PARTITIONED_TOPIC_SUFFIX;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
@@ -34,15 +33,10 @@
 import io.streamnative.pulsar.handlers.kop.utils.MetadataUtils;
 import io.streamnative.pulsar.handlers.kop.utils.timer.SystemTimer;
 import java.net.InetSocketAddress;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.NonNull;
@@ -59,14 +53,12 @@
 import org.apache.pulsar.broker.namespace.NamespaceBundleOwnershipListener;
 import org.apache.pulsar.broker.protocol.ProtocolHandler;
 import org.apache.pulsar.broker.service.BrokerService;
-import org.apache.pulsar.client.admin.Lookup;
 import org.apache.pulsar.client.admin.PulsarAdmin;
 import org.apache.pulsar.client.admin.PulsarAdminException;
 import org.apache.pulsar.common.naming.NamespaceBundle;
 import org.apache.pulsar.common.naming.NamespaceName;
 import org.apache.pulsar.common.naming.TopicName;
 import org.apache.pulsar.common.policies.data.ClusterData;
-import org.apache.pulsar.common.util.FutureUtil;
 
 /**
  * Kafka Protocol Handler load and run by Pulsar Service.
@@ -474,13 +466,7 @@ public void start(BrokerService service) {
         kopEventManager.start();
 
         if (kafkaConfig.isEnableTransactionCoordinator()) {
-            TransactionCoordinator transactionCoordinator =
-                    getTransactionCoordinator(kafkaConfig.getKafkaMetadataTenant());
-            try {
-                loadTxnLogTopics(kafkaConfig.getKafkaMetadataTenant(), transactionCoordinator);
-            } catch (Exception e) {
-                log.error("Failed to load transaction log", e);
-            }
+            getTransactionCoordinator(kafkaConfig.getKafkaMetadataTenant());
         }
 
         Configuration conf = new PropertiesConfiguration();
@@ -680,39 +666,6 @@ public TransactionCoordinator initTransactionCoordinator(String tenant, PulsarAd
         return transactionCoordinator;
     }
 
-    /**
-     * This method discovers ownership of offset topic partitions and attempts to load transaction topics
-     * assigned to this broker.
-     */
-    private void loadTxnLogTopics(String tenant, TransactionCoordinator txnCoordinator) throws Exception {
-        Lookup lookupService = brokerService.pulsar().getAdminClient().lookups();
-        String currentBroker = brokerService.pulsar().getBrokerServiceUrl();
-        String topicBase = MetadataUtils.constructTxnLogTopicBaseName(tenant, kafkaConfig);
-        int numPartitions = kafkaConfig.getTxnLogTopicNumPartitions();
-
-        Map<String, List<Integer>> mapBrokerToPartition = new HashMap<>();
-
-        for (int i = 0; i < numPartitions; i++) {
-            String broker = lookupService.lookupTopic(topicBase + PARTITIONED_TOPIC_SUFFIX + i);
-            mapBrokerToPartition.putIfAbsent(broker, new ArrayList<>());
-            mapBrokerToPartition.get(broker).add(i);
-        }
-
-        mapBrokerToPartition.forEach(
-                (key, value) -> log.info("Discovered broker: {} owns txn log topic partitions: {} ", key, value));
-
-        List<Integer> partitionsOwnedByCurrentBroker = mapBrokerToPartition.get(currentBroker);
-
-        if (null != partitionsOwnedByCurrentBroker && !partitionsOwnedByCurrentBroker.isEmpty()) {
-            List<CompletableFuture<Void>> lists = partitionsOwnedByCurrentBroker.stream().map(
-                    txnCoordinator::handleTxnImmigration).collect(Collectors.toList());
-
-            FutureUtil.waitForAll(lists).get();
-        } else {
-            log.info("Current broker: {} does not own any of the txn log topic partitions", currentBroker);
-        }
-    }
-
     public static @NonNull LookupClient getLookupClient(final PulsarService pulsarService) {
         return LOOKUP_CLIENT_MAP.computeIfAbsent(pulsarService, ignored -> new LookupClient(pulsarService));
     }

tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionStateManagerTest.java

@@ -18,6 +18,7 @@
 import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 
+import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import io.streamnative.pulsar.handlers.kop.KafkaProtocolHandler;
 import io.streamnative.pulsar.handlers.kop.KopProtocolHandlerTestBase;
@@ -26,14 +27,17 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.pulsar.client.admin.PulsarAdminException;
 import org.apache.pulsar.common.policies.data.BundlesData;
 import org.apache.pulsar.common.policies.data.RetentionPolicies;
 import org.apache.pulsar.common.policies.data.TenantInfoImpl;
+import org.apache.pulsar.common.util.FutureUtil;
 import org.awaitility.Awaitility;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
@@ -50,11 +54,17 @@ public class TransactionStateManagerTest extends KopProtocolHandlerTestBase {
     private static final Short producerEpoch = 0;
     private static final Integer transactionTimeoutMs = 1000;
 
+    private int txnLogTopicNumPartitions;
+    private TransactionCoordinator transactionCoordinator;
+
     @BeforeClass
     @Override
     protected void setup() throws Exception {
         this.conf.setEnableTransactionCoordinator(true);
+        this.txnLogTopicNumPartitions = this.conf.getTxnLogTopicNumPartitions();
+
         internalSetup();
+        transactionCoordinator = getTxnCoordinator();
 
         TenantInfoImpl tenantInfo =
                 new TenantInfoImpl(Sets.newHashSet("appid1", "appid2"), Sets.newHashSet("test"));
@@ -103,6 +113,8 @@ public void txnLogStoreAndTCImmigrationTest() throws Exception {
         transactionStates.put(4L, TransactionState.PREPARE_ABORT);
         transactionStates.put(5L, TransactionState.COMPLETE_ABORT);
 
+        // Make sure transaction state log already loaded.
+        this.loadTxnImmigration();
         TransactionStateManager transactionStateManager = getTxnManager();
 
         CountDownLatch countDownLatch = new CountDownLatch(pidMappings.size());
@@ -223,9 +235,24 @@ private void waitTCImmigrationComplete() throws PulsarAdminException {
                 .untilAsserted(() -> assertFalse(txnStateManager.isLoading()));
     }
 
-    private TransactionStateManager getTxnManager() {
+    private TransactionCoordinator getTxnCoordinator() {
         return ((KafkaProtocolHandler) this.pulsar.getProtocolHandlers().protocol("kafka"))
-                .getTransactionCoordinator("public").getTxnManager();
+                .getTransactionCoordinator("public");
     }
 
+    private TransactionStateManager getTxnManager() {
+        return getTxnCoordinator().getTxnManager();
+    }
+
+    private void loadTxnImmigration() {
+        List<CompletableFuture<Void>> allFuture = Lists.newArrayList();
+        for (int i = 0; i < txnLogTopicNumPartitions; i++) {
+            allFuture.add(transactionCoordinator.handleTxnImmigration(i));
+        }
+        try {
+            FutureUtil.waitForAll(allFuture).get();
+        } catch (InterruptedException | ExecutionException e) {
+            log.error("Load txn immigration failed.", e);
+        }
+    }
 }