From 33260e906bd63f0f14e172b654e562718457cf74 Mon Sep 17 00:00:00 2001
From: Yunze Xu <xyzinfernity@163.com>
Date: Wed, 7 Jul 2021 21:59:27 +0800
Subject: [PATCH] Fix group id z-node not updated for the same client id

---
 .../handlers/kop/utils/ZooKeeperUtils.java    | 16 +++++++-
 .../handlers/kop/MetricsProviderTest.java     | 39 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/utils/ZooKeeperUtils.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/utils/ZooKeeperUtils.java
index dcf13f2778..d62fde7255 100644
--- a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/utils/ZooKeeperUtils.java
+++ b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/utils/ZooKeeperUtils.java
@@ -35,7 +35,21 @@ public static void tryCreatePath(ZooKeeper zooKeeper, String path, byte[] data)
             if (log.isDebugEnabled()) {
                 log.debug("Created ZK path: {} data: {}", path, new String(data, StandardCharsets.UTF_8));
             }
-        } catch (KeeperException | InterruptedException e) {
+        } catch (KeeperException e) {
+            if (!e.code().equals(KeeperException.Code.NODEEXISTS)) {
+                log.error("Failed to create ZooKeeper node {}: {}", path, e.getMessage());
+            } else {
+                // update the group id
+                try {
+                    zooKeeper.setData(path, data, -1);
+                    if (log.isDebugEnabled()) {
+                        log.debug("Updated ZK path: {} data: {}", path, new String(data, StandardCharsets.UTF_8));
+                    }
+                } catch (KeeperException | InterruptedException setDataException) {
+                    log.error("Failed to set path '{}''s data to {}", path, data, setDataException);
+                }
+            }
+        } catch (InterruptedException e) {
             log.error("Failed to create ZooKeeper node {}: {}", path, e.getMessage());
         }
     }
diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/MetricsProviderTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/MetricsProviderTest.java
index b76457ca04..fe787aa2ca 100644
--- a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/MetricsProviderTest.java
+++ b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/MetricsProviderTest.java
@@ -16,8 +16,11 @@
 import java.io.BufferedReader;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.Collections;
 import java.util.List;
+import java.util.Properties;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -28,8 +31,10 @@
 import org.apache.http.client.HttpClient;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.impl.client.HttpClientBuilder;
+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.clients.producer.ProducerRecord;
 import org.apache.kafka.common.TopicPartition;
 import org.testng.Assert;
@@ -192,4 +197,38 @@ public void testMetricsProvider() throws Exception {
         Assert.assertTrue(sb.toString().contains("kop_server_BYTES_IN{partition=\"0\","
                 + "topic=\"kopKafkaProducePulsarMetrics1\"} 1170"));
     }
+
+    @Test(timeOut = 20000)
+    public void testUpdateGroupId() throws Exception {
+        final String topic = "testUpdateGroupId";
+        final String clientId = "my-client";
+        final String group1 = "my-group-1";
+        final String group2 = "my-group-2";
+
+        tryConsume(topic, clientId, group1);
+        List<String> children = mockZooKeeper.getChildren(conf.getGroupIdZooKeeperPath(), false);
+        Assert.assertEquals(children.size(), 1);
+        Assert.assertEquals(children.get(0), "127.0.0.1-" + clientId);
+        byte[] data = mockZooKeeper.getData(conf.getGroupIdZooKeeperPath() + "/" + children.get(0), false, null);
+        Assert.assertEquals(new String(data, StandardCharsets.UTF_8), group1);
+
+        // Create a consumer with the same hostname and client id, the existed z-node will be updated
+        tryConsume(topic, clientId, group2);
+        children = mockZooKeeper.getChildren(conf.getGroupIdZooKeeperPath(), false);
+        Assert.assertEquals(children.size(), 1);
+        Assert.assertEquals(children.get(0), "127.0.0.1-" + clientId);
+        data = mockZooKeeper.getData(conf.getGroupIdZooKeeperPath() + "/" + children.get(0), false, null);
+        Assert.assertEquals(new String(data, StandardCharsets.UTF_8), group2);
+    }
+
+    private void tryConsume(final String topic, final String clientId, final String groupId) {
+        final Properties props = newKafkaConsumerProperties();
+        props.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
+        props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+
+        final KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
+        consumer.subscribe(Collections.singleton(topic));
+        consumer.poll(Duration.ofSeconds(3));
+        consumer.close();
+    }
 }