From 5f24acc489bfccc21ed11a0a49711e3e5cfacda0 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sat, 11 Jan 2025 16:10:15 +0000
Subject: [PATCH 01/28] Porting changes from STORM-3693

---
 .../src/jvm/org/apache/storm/executor/Executor.java         | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/storm-client/src/jvm/org/apache/storm/executor/Executor.java b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
index b3df9c9e631..2d287c77739 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/Executor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
@@ -285,6 +285,8 @@ public void accept(Object event) {
         int taskId = addressedTuple.getDest();
 
         TupleImpl tuple = (TupleImpl) addressedTuple.getTuple();
+        String streamId = tuple.getSourceStreamId();
+        boolean isSpout = this instanceof SpoutExecutor;
         if (isDebug) {
             LOG.info("Processing received TUPLE: {} for TASK: {} ", tuple, taskId);
         }
@@ -292,6 +294,10 @@ public void accept(Object event) {
         try {
             if (taskId != AddressedTuple.BROADCAST_DEST) {
                 tupleActionFn(taskId, tuple);
+            } else if (isSpout && streamId.equals(Constants.SYSTEM_TICK_STREAM_ID)) {
+                //taskId is irrelevant here. Ensures pending.rotate() is called once per tick.
+                tupleActionFn(taskIds.get(0), tuple);
+
             } else {
                 for (Integer t : taskIds) {
                     tupleActionFn(t, tuple);

From 9227922fa3692802c1d140ab90c200dc537d5301 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:31:10 +0000
Subject: [PATCH 02/28] Refactoring changes proposed on STORM-3693 and adding
 an unit test

---
 .../org/apache/storm/executor/Executor.java   | 10 +--
 .../storm/executor/spout/SpoutExecutor.java   | 38 ++++++++--
 .../storm/executor/SpoutExecutorTest.java     | 75 +++++++++++++++++++
 3 files changed, 110 insertions(+), 13 deletions(-)
 create mode 100644 storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java

diff --git a/storm-client/src/jvm/org/apache/storm/executor/Executor.java b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
index 2d287c77739..562c49102de 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/Executor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
@@ -285,19 +285,17 @@ public void accept(Object event) {
         int taskId = addressedTuple.getDest();
 
         TupleImpl tuple = (TupleImpl) addressedTuple.getTuple();
-        String streamId = tuple.getSourceStreamId();
-        boolean isSpout = this instanceof SpoutExecutor;
         if (isDebug) {
             LOG.info("Processing received TUPLE: {} for TASK: {} ", tuple, taskId);
         }
 
+        acceptTupleAction(taskId, tuple);
+    }
+
+    protected void acceptTupleAction(int taskId, TupleImpl tuple) {
         try {
             if (taskId != AddressedTuple.BROADCAST_DEST) {
                 tupleActionFn(taskId, tuple);
-            } else if (isSpout && streamId.equals(Constants.SYSTEM_TICK_STREAM_ID)) {
-                //taskId is irrelevant here. Ensures pending.rotate() is called once per tick.
-                tupleActionFn(taskIds.get(0), tuple);
-
             } else {
                 for (Integer t : taskIds) {
                     tupleActionFn(t, tuple);
diff --git a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
index 25c84474a49..a4081fe6c42 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
@@ -2,9 +2,9 @@
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
- *
+ * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
- *
+ * <p>
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
@@ -17,6 +17,7 @@
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.apache.storm.Config;
 import org.apache.storm.Constants;
 import org.apache.storm.ICredentialsListener;
@@ -80,7 +81,7 @@ public SpoutExecutor(final WorkerState workerData, final List<Long> executorId,
         this.emittedCount = new MutableLong(0);
         this.emptyEmitStreak = new MutableLong(0);
         this.stats = new SpoutExecutorStats(
-            ConfigUtils.samplingRate(this.getTopoConf()), ObjectReader.getInt(this.getTopoConf().get(Config.NUM_STAT_BUCKETS)));
+                ConfigUtils.samplingRate(this.getTopoConf()), ObjectReader.getInt(this.getTopoConf().get(Config.NUM_STAT_BUCKETS)));
         this.skippedMaxSpoutMs = workerData.getMetricRegistry().rateCounter("__skipped-max-spout-ms", componentId,
                 taskIds.get(0));
         this.skippedInactiveMs = workerData.getMetricRegistry().rateCounter("__skipped-inactive-ms", componentId,
@@ -131,8 +132,8 @@ public void expire(Long key, TupleInfo tupleInfo) {
             }
             ISpout spoutObject = (ISpout) taskData.getTaskObject();
             spoutOutputCollector = new SpoutOutputCollectorImpl(
-                spoutObject, this, taskData, emittedCount,
-                hasAckers, rand, hasEventLoggers, isDebug, pending);
+                    spoutObject, this, taskData, emittedCount,
+                    hasAckers, rand, hasEventLoggers, isDebug, pending);
             SpoutOutputCollector outputCollector = new SpoutOutputCollector(spoutOutputCollector);
             this.outputCollectors.add(outputCollector);
 
@@ -283,6 +284,29 @@ private void inactiveExecute() throws InterruptedException {
         skippedInactiveMs.inc(Time.currentTimeMillis() - start);
     }
 
+    @Override
+    protected void acceptTupleAction(int taskId, TupleImpl tuple) {
+
+        String streamId = tuple.getSourceStreamId();
+
+        try {
+            if (taskId != AddressedTuple.BROADCAST_DEST) {
+                tupleActionFn(taskId, tuple);
+            } else if (streamId.equals(Constants.SYSTEM_TICK_STREAM_ID)) {
+                //taskId is irrelevant here. Ensures pending.rotate() is called once per tick.
+                tupleActionFn(taskIds.get(0), tuple);
+
+            } else {
+                for (Integer t : taskIds) {
+                    tupleActionFn(t, tuple);
+                }
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+
     @Override
     public void tupleActionFn(int taskId, TupleImpl tuple) throws Exception {
         String streamId = tuple.getSourceStreamId();
@@ -370,6 +394,6 @@ public int getSpoutRecvqCheckSkipCount() {
 
     public long getThreadId() {
         return threadId;
-    }   
-    
+    }
+
 }
diff --git a/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
new file mode 100644
index 00000000000..d2871f6db10
--- /dev/null
+++ b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
@@ -0,0 +1,75 @@
+package org.apache.storm.executor;
+
+import org.apache.storm.Constants;
+import org.apache.storm.cluster.IStateStorage;
+import org.apache.storm.daemon.worker.WorkerState;
+import org.apache.storm.executor.spout.SpoutExecutor;
+import org.apache.storm.generated.ComponentCommon;
+import org.apache.storm.metrics2.RateCounter;
+import org.apache.storm.metrics2.StormMetricRegistry;
+import org.apache.storm.task.WorkerTopologyContext;
+import org.apache.storm.tuple.AddressedTuple;
+import org.apache.storm.tuple.TupleImpl;
+import org.apache.storm.utils.RotatingMap;
+import org.apache.storm.utils.Utils;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.commons.util.ReflectionUtils;
+import org.mockito.Mockito;
+
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+
+public class SpoutExecutorTest {
+
+
+    @Test
+    public void testPendingTuplesRotateShouldBeCalledOnlyOnce() throws Exception {
+
+        RateCounter rateCounter = Mockito.mock(RateCounter.class);
+
+        StormMetricRegistry stormMetricRegistry = Mockito.mock(StormMetricRegistry.class);
+        Mockito.when(stormMetricRegistry.rateCounter(anyString(),anyString(),anyInt())).thenReturn(rateCounter);
+
+        Map<String,Object> hashmap = Utils.readDefaultConfig();
+
+        IStateStorage stateStorage = Mockito.mock(IStateStorage.class);
+
+        ComponentCommon componentCommon = Mockito.mock(ComponentCommon.class);
+        Mockito.when(componentCommon.get_json_conf()).thenReturn(null);
+
+        WorkerTopologyContext workerTopologyContext = Mockito.mock(WorkerTopologyContext.class);
+        Mockito.when(workerTopologyContext.getComponentId(anyInt())).thenReturn("1");
+        Mockito.when(workerTopologyContext.getComponentCommon(anyString())).thenReturn(componentCommon);
+
+        WorkerState workerState = Mockito.mock(WorkerState.class);
+        Mockito.when(workerState.getWorkerTopologyContext()).thenReturn(workerTopologyContext);
+        Mockito.when(workerState.getStateStorage()).thenReturn(stateStorage);
+        Mockito.when(workerState.getTopologyConf()).thenReturn(hashmap);
+        Mockito.when(workerState.getMetricRegistry()).thenReturn(stormMetricRegistry);
+
+        SpoutExecutor spoutExecutor = new SpoutExecutor(workerState,List.of(1L,5L),new HashMap<>());
+
+        TupleImpl tuple =   Mockito.mock(TupleImpl.class);
+        Mockito.when(tuple.getSourceStreamId()).thenReturn(Constants.SYSTEM_TICK_STREAM_ID);
+        AddressedTuple addressedTuple =   Mockito.mock(AddressedTuple.class);
+        Mockito.when(addressedTuple.getDest()).thenReturn(AddressedTuple.BROADCAST_DEST);
+        Mockito.when(addressedTuple.getTuple()).thenReturn(tuple);
+
+        RotatingMap rotatingMap = Mockito.mock(RotatingMap.class);
+        Field fieldRotatingMap = ReflectionUtils
+                .findFields(SpoutExecutor.class, f -> f.getName().equals("pending"),
+                        ReflectionUtils.HierarchyTraversalMode.TOP_DOWN)
+                .get(0);
+        fieldRotatingMap.setAccessible(true);
+        fieldRotatingMap.set(spoutExecutor, rotatingMap);
+
+        spoutExecutor.accept(addressedTuple);
+
+        Mockito.verify(rotatingMap,Mockito.times(1)).rotate();
+    }
+}

From f5db47bc18963a78eebf667ac72f21e68c2b37fc Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:44:48 +0000
Subject: [PATCH 03/28] Update Executor.java

---
 storm-client/src/jvm/org/apache/storm/executor/Executor.java | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/storm-client/src/jvm/org/apache/storm/executor/Executor.java b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
index 562c49102de..edd77743fcd 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/Executor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/Executor.java
@@ -2,9 +2,9 @@
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
- * <p>
+ * 
  * http://www.apache.org/licenses/LICENSE-2.0
- * <p>
+ *
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.

From 5ded4cdb67d1463daebe26434571ec09c014ee64 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:45:39 +0000
Subject: [PATCH 04/28] Update SpoutExecutor.java

---
 .../jvm/org/apache/storm/executor/spout/SpoutExecutor.java    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
index a4081fe6c42..191c1cb0a3b 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
@@ -2,9 +2,9 @@
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
- * <p>
+ * 
  * http://www.apache.org/licenses/LICENSE-2.0
- * <p>
+ * 
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.

From e3f76e60323246c2146e5c1c845c733b8434c042 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:47:12 +0000
Subject: [PATCH 05/28] Update SpoutExecutor.java

---
 .../jvm/org/apache/storm/executor/spout/SpoutExecutor.java    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
index 191c1cb0a3b..38adeff1872 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
@@ -2,9 +2,9 @@
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
- * 
+ *
  * http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.

From 677756f19ffdaea11d3cc18c7b618e137d59741c Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:47:50 +0000
Subject: [PATCH 06/28] Update SpoutExecutor.java

---
 .../src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
index 38adeff1872..734fca2a2d6 100644
--- a/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
+++ b/storm-client/src/jvm/org/apache/storm/executor/spout/SpoutExecutor.java
@@ -17,7 +17,6 @@
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.atomic.AtomicBoolean;
-
 import org.apache.storm.Config;
 import org.apache.storm.Constants;
 import org.apache.storm.ICredentialsListener;

From 365958476695eafdd1cf8ab62a01ecf5bf60a441 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 22:49:57 +0000
Subject: [PATCH 07/28] Adding missing license

---
 .../apache/storm/executor/SpoutExecutorTest.java    | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
index d2871f6db10..4bcb2fdd6cc 100644
--- a/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
+++ b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
@@ -1,3 +1,16 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
+ * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+
 package org.apache.storm.executor;
 
 import org.apache.storm.Constants;

From f34416e7f6d99a0fa7224aba0cc37ebbbb6578b1 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sat, 18 Jan 2025 23:26:11 +0000
Subject: [PATCH 08/28] Adding missing license

---
 .../test/jvm/org/apache/storm/executor/SpoutExecutorTest.java    | 1 -
 1 file changed, 1 deletion(-)

diff --git a/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
index 4bcb2fdd6cc..a44597cde9a 100644
--- a/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
+++ b/storm-client/test/jvm/org/apache/storm/executor/SpoutExecutorTest.java
@@ -10,7 +10,6 @@
  * and limitations under the License.
  */
 
-
 package org.apache.storm.executor;
 
 import org.apache.storm.Constants;

From 3455c2a9248043d9a1a979da864c4e335b558e8c Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Wed, 12 Feb 2025 09:38:21 +0000
Subject: [PATCH 09/28] Update IWorkerHook.java

---
 .../jvm/org/apache/storm/hooks/IWorkerHook.java   | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/storm-client/src/jvm/org/apache/storm/hooks/IWorkerHook.java b/storm-client/src/jvm/org/apache/storm/hooks/IWorkerHook.java
index 7533e1ccd6d..fa60975f2c7 100644
--- a/storm-client/src/jvm/org/apache/storm/hooks/IWorkerHook.java
+++ b/storm-client/src/jvm/org/apache/storm/hooks/IWorkerHook.java
@@ -22,18 +22,7 @@
  * when you want to execute operations before topology processing starts, or cleanup operations before your workers shut down.
  */
 public interface IWorkerHook extends Serializable {
-    /**
-     * This method is called when a worker is started.
-     *
-     * @param topoConf The Storm configuration for this worker
-     * @param context  This object can be used to get information about this worker's place within the topology
-     *
-     * @deprecated see {@link IWorkerHook#start(Map, WorkerUserContext)}
-     */
-    default void start(Map<String, Object> topoConf, WorkerTopologyContext context) {
-        // NOOP
-    }
-
+  
     /**
      * This method is called when a worker is started and can be used to do necessary prep-processing and allow initialization of shared
      * application state.
@@ -43,7 +32,7 @@ default void start(Map<String, Object> topoConf, WorkerTopologyContext context)
      * {@link WorkerUserContext#setResource(String, Object)} to set the shared application state.
      */
     default void start(Map<String, Object> topoConf, WorkerUserContext context) {
-        start(topoConf, context);
+         // NOOP
     }
 
     /**

From 357368ce7d4275449321f64fc0817377956b68e4 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sun, 20 Apr 2025 01:04:46 +0100
Subject: [PATCH 10/28] Adding new MetricSet

---
 .../apache/storm/kafka/spout/KafkaSpout.java  |   2 +-
 .../metrics2/KafkaOffsetMetricManager.java    |   7 +
 .../KafkaOffsetPartitionAndTopicMetrics.java  | 211 ++++++++++++++++++
 .../metrics2/KafkaOffsetPartitionMetrics.java |   2 +-
 4 files changed, 220 insertions(+), 2 deletions(-)
 create mode 100644 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
index a328aec5d70..4a7dee0f498 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
@@ -634,7 +634,7 @@ private void refreshAssignment() {
         boolean partitionChanged = topicAssigner.assignPartitions(consumer, assignedPartitions, rebalanceListener);
         if (partitionChanged && canRegisterMetrics()) {
             LOG.info("Partitions assignments has changed, updating metrics.");
-            kafkaOffsetMetricManager.registerMetricsForNewTopicPartitions(assignedPartitions);
+            kafkaOffsetMetricManager.registerPartitionAndTopicLevelMetrics(assignedPartitions);
         }
     }
 
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index 2b150797ffe..3a3a67a3a55 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -75,6 +75,13 @@ public void registerMetricsForNewTopicPartitions(Set<TopicPartition> newAssignme
         }
     }
 
+    public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignment) {
+
+        KafkaOffsetPartitionAndTopicMetrics topicPartitionMetricSet
+                = new KafkaOffsetPartitionAndTopicMetrics<>(offsetManagerSupplier, adminSupplier,newAssignment);
+        topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
+    }
+
     public Map<TopicPartition, KafkaOffsetPartitionMetrics> getTopicPartitionMetricsMap() {
         return topicPartitionMetricsMap;
     }
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
new file mode 100644
index 00000000000..4516ad957ad
--- /dev/null
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -0,0 +1,211 @@
+package org.apache.storm.kafka.spout.metrics2;
+
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Metric;
+import com.codahale.metrics.MetricSet;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ListOffsetsResult;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.RetriableException;
+import org.apache.storm.kafka.spout.internal.OffsetManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Supplier;
+
+/**
+ * This class is used compute the partition and topic level offset metrics.
+ * <p>
+ * Partition level metrics are:
+ * topicName/partition_{number}/earliestTimeOffset //gives beginning offset of the partition
+ * topicName/partition_{number}/latestTimeOffset //gives end offset of the partition
+ * topicName/partition_{number}/latestEmittedOffset //gives latest emitted offset of the partition from the spout
+ * topicName/partition_{number}/latestCompletedOffset //gives latest committed offset of the partition from the spout
+ * topicName/partition_{number}/spoutLag // the delta between the latest Offset and latestCompletedOffset
+ * topicName/partition_{number}/recordsInPartition // total number of records in the partition
+ * </p>
+ * <p>
+ * Topic level metrics are:
+ * topicName/totalEarliestTimeOffset //gives the total beginning offset of all the associated partitions of this spout
+ * topicName/totalLatestTimeOffset //gives the total end offset of all the associated partitions of this spout
+ * topicName/totalLatestEmittedOffset //gives the total latest emitted offset of all the associated partitions of this spout
+ * topicName/totalLatestCompletedOffset //gives the total latest committed offset of all the associated partitions of this spout
+ * topicName/spoutLag // total spout lag of all the associated partitions of this spout
+ * topicName/totalRecordsInPartitions //total number of records in all the associated partitions of this spout
+ * </p>
+ */
+public class KafkaOffsetPartitionAndTopicMetrics <K, V> implements MetricSet {
+
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetPartitionAndTopicMetrics.class);
+    private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
+    private final Supplier<Admin> adminSupplier;
+    private final Set<TopicPartition> assignment;
+    private Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
+
+
+    public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier, Supplier<Admin> adminSupplier, Set<TopicPartition> assignment) {
+        this.offsetManagerSupplier = offsetManagerSupplier;
+        this.adminSupplier = adminSupplier;
+        this.assignment = assignment;
+    }
+
+    @Override
+    public Map<String, Metric> getMetrics() {
+
+        Map<String, Metric> metrics = new HashMap<>();
+
+        for (TopicPartition topicPartition : assignment) {
+
+            String topic=topicPartition.topic();
+            KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
+            if (topicMetrics == null) {
+                topicMetrics = new KafkaOffsetTopicMetrics(topic);
+                topicMetricsMap.put(topic, topicMetrics);
+            }
+
+            String metricPath = topicPartition.topic()  + "/partition_" + topicPartition.partition();
+            KafkaOffsetTopicMetrics finalTopicMetrics = topicMetrics;
+            Gauge<Long> spoutLagGauge = () -> {
+                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
+                if (endOffsets == null || endOffsets.isEmpty()) {
+                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                    return 0L;
+                }
+                // add value to topic level metric
+                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                Long ret = endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
+                finalTopicMetrics.totalSpoutLag += ret;
+                return ret;
+            };
+
+            Gauge<Long> earliestTimeOffsetGauge = () -> {
+                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
+                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                    return 0L;
+                }
+                // add value to topic level metric
+                Long ret = beginningOffsets.get(topicPartition);
+                finalTopicMetrics.totalEarliestTimeOffset += ret;
+                return ret;
+            };
+
+            Gauge<Long> latestTimeOffsetGauge = () -> {
+                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
+                if (endOffsets == null || endOffsets.isEmpty()) {
+                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                    return 0L;
+                }
+                // add value to topic level metric
+                Long ret = endOffsets.get(topicPartition);
+                finalTopicMetrics.totalLatestTimeOffset += ret;
+                return ret;
+            };
+
+            Gauge<Long> latestEmittedOffsetGauge = () -> {
+                // add value to topic level metric
+                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                Long ret = offsetManager.getLatestEmittedOffset();
+                finalTopicMetrics.totalLatestEmittedOffset+=ret;
+                return ret;
+            };
+
+            Gauge<Long> latestCompletedOffsetGauge = () -> {
+                // add value to topic level metric
+                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                Long ret = offsetManager.getCommittedOffset();
+                finalTopicMetrics.totalLatestCompletedOffset+=ret;
+                return ret;
+            };
+
+            Gauge<Long> recordsInPartitionGauge = () -> {
+                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
+                if (endOffsets == null || endOffsets.isEmpty()) {
+                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                    return 0L;
+                }
+                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
+                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                    return 0L;
+                }
+                // add value to topic level metric
+                Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
+                finalTopicMetrics.totalRecordsInPartitions+=ret;
+                return ret;
+            };
+
+            metrics.put(metricPath + "/" + "spoutLag", spoutLagGauge);
+            metrics.put(metricPath + "/" + "earliestTimeOffset", earliestTimeOffsetGauge);
+            metrics.put(metricPath + "/" + "latestTimeOffset", latestTimeOffsetGauge);
+            metrics.put(metricPath + "/" + "latestEmittedOffset", latestEmittedOffsetGauge);
+            metrics.put(metricPath + "/" + "latestCompletedOffset", latestCompletedOffsetGauge);
+            metrics.put(metricPath + "/" + "recordsInPartition", recordsInPartitionGauge);
+
+        }
+
+        metrics.putAll(topicMetricsMap);
+
+        return metrics;
+
+    }
+
+    private Map<TopicPartition, Long> getBeginningOffsets(Set<TopicPartition> topicPartitions) {
+        Admin admin = adminSupplier.get();
+        if (admin == null) {
+            LOG.error("Kafka admin object is null, returning 0.");
+            return Collections.EMPTY_MAP;
+        }
+
+        Map<TopicPartition, Long> beginningOffsets;
+        try {
+            beginningOffsets = getOffsets(admin, topicPartitions, OffsetSpec.earliest());
+        } catch (RetriableException | ExecutionException | InterruptedException e) {
+            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
+            return Collections.EMPTY_MAP;
+        }
+        return beginningOffsets;
+    }
+
+    private Map<TopicPartition, Long> getEndOffsets(Set<TopicPartition> topicPartitions) {
+        Admin admin = adminSupplier.get();
+        if (admin == null) {
+            LOG.error("Kafka admin object is null, returning 0.");
+            return Collections.EMPTY_MAP;
+        }
+
+        Map<TopicPartition, Long> endOffsets;
+        try {
+            endOffsets = getOffsets(admin, topicPartitions, OffsetSpec.latest());
+        } catch (RetriableException | ExecutionException | InterruptedException e) {
+            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
+            return Collections.EMPTY_MAP;
+        }
+        return endOffsets;
+    }
+
+    private static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
+            throws InterruptedException, ExecutionException {
+
+        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
+        for (TopicPartition topicPartition : topicPartitions) {
+            offsetSpecMap.put(topicPartition, offsetSpec);
+        }
+        Map<TopicPartition, Long> ret = new HashMap<>();
+        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
+        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
+        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
+                topicPartitionListOffsetsResultInfoMap.entrySet()) {
+            ret.put(entry.getKey(), entry.getValue().offset());
+        }
+        return ret;
+    }
+}
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
index 6b29eefc919..0bb0d2c2168 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
@@ -101,7 +101,7 @@ public Long getValue() {
                 }
                 // add value to topic level metric
                 Long ret = beginningOffsets.get(topicPartition);
-                topicMetrics.totalEarliestTimeOffset += beginningOffsets.get(topicPartition);
+                topicMetrics.totalEarliestTimeOffset += ret;
                 return ret;
             }
         };

From 7403cebd4cd66dd7d3ade78027734dffff436696 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sun, 20 Apr 2025 19:38:11 +0100
Subject: [PATCH 11/28] Adding unit tests

---
 .../apache/storm/kafka/spout/KafkaSpout.java  |   1 +
 .../metrics2/KafkaOffsetMetricManager.java    |  23 +---
 .../KafkaOffsetPartitionAndTopicMetrics.java  |   9 +-
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 106 ++++++++++++++++++
 4 files changed, 114 insertions(+), 25 deletions(-)
 create mode 100644 external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
index 4a7dee0f498..2b35fdd9272 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
@@ -57,6 +57,7 @@
 import org.apache.storm.kafka.spout.internal.OffsetManager;
 import org.apache.storm.kafka.spout.internal.Timer;
 import org.apache.storm.kafka.spout.metrics2.KafkaOffsetMetricManager;
+import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionAndTopicMetrics;
 import org.apache.storm.kafka.spout.subscription.TopicAssigner;
 import org.apache.storm.spout.SpoutOutputCollector;
 import org.apache.storm.task.TopologyContext;
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index 3a3a67a3a55..c78b5637e89 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -54,31 +54,10 @@ public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> off
         LOG.info("Running KafkaOffsetMetricManager");
     }
 
-    public void registerMetricsForNewTopicPartitions(Set<TopicPartition> newAssignment) {
-        for (TopicPartition topicPartition : newAssignment) {
-            if (!topicPartitionMetricsMap.containsKey(topicPartition)) {
-                LOG.info("Registering metric for topicPartition: {}", topicPartition);
-                // create topic level metrics for given topic if absent
-                String topic = topicPartition.topic();
-                KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
-                if (topicMetrics == null) {
-                    topicMetrics = new KafkaOffsetTopicMetrics(topic);
-                    topicMetricsMap.put(topic, topicMetrics);
-                    topologyContext.registerMetricSet("kafkaOffset", topicMetrics);
-                }
-
-                KafkaOffsetPartitionMetrics topicPartitionMetricSet
-                    = new KafkaOffsetPartitionMetrics<>(offsetManagerSupplier, adminSupplier, topicPartition, topicMetrics);
-                topicPartitionMetricsMap.put(topicPartition, topicPartitionMetricSet);
-                topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
-            }
-        }
-    }
-
     public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignment) {
 
         KafkaOffsetPartitionAndTopicMetrics topicPartitionMetricSet
-                = new KafkaOffsetPartitionAndTopicMetrics<>(offsetManagerSupplier, adminSupplier,newAssignment);
+                = new KafkaOffsetPartitionAndTopicMetrics(offsetManagerSupplier, adminSupplier,newAssignment);
         topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
     }
 
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 4516ad957ad..973ce99f313 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -41,19 +41,20 @@
  * topicName/totalRecordsInPartitions //total number of records in all the associated partitions of this spout
  * </p>
  */
-public class KafkaOffsetPartitionAndTopicMetrics <K, V> implements MetricSet {
+public class KafkaOffsetPartitionAndTopicMetrics implements MetricSet {
 
     private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetPartitionAndTopicMetrics.class);
     private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
     private final Supplier<Admin> adminSupplier;
     private final Set<TopicPartition> assignment;
-    private Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
+    private final Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
 
 
     public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier, Supplier<Admin> adminSupplier, Set<TopicPartition> assignment) {
         this.offsetManagerSupplier = offsetManagerSupplier;
         this.adminSupplier = adminSupplier;
         this.assignment = assignment;
+        this.topicMetricsMap=new HashMap<>();
     }
 
     @Override
@@ -151,7 +152,9 @@ public Map<String, Metric> getMetrics() {
 
         }
 
-        metrics.putAll(topicMetricsMap);
+        for( KafkaOffsetTopicMetrics kafkaOffsetTopicMetrics  : topicMetricsMap.values()){
+            metrics.putAll(kafkaOffsetTopicMetrics.getMetrics());
+        }
 
         return metrics;
 
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
new file mode 100644
index 00000000000..0475155a993
--- /dev/null
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -0,0 +1,106 @@
+package org.apache.storm.kafka.spout.metric2;
+
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Metric;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ListOffsetsResult;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.storm.kafka.spout.internal.OffsetManager;
+import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionAndTopicMetrics;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.util.*;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+public class KafkaOffsetPartitionAndTopicMetricsTest {
+
+    private Set<TopicPartition> assignment;
+    private Admin admin;
+    private HashMap<TopicPartition, OffsetManager> offsetManagers;
+    private ListOffsetsResult listOffsetsResult;
+    private KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> kafkaFuture;
+    private Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap;
+
+    @Test
+    public void registerPartitionAndTopicMetrics() throws ExecutionException, InterruptedException {
+
+        TopicPartition tAp1 = new TopicPartition("topicA",1);
+        TopicPartition tAp2 = new TopicPartition("topicA",2);
+        TopicPartition tBp1 = new TopicPartition("topicB",1);
+        TopicPartition tBp2 = new TopicPartition("topicB",2);
+
+        ListOffsetsResult.ListOffsetsResultInfo tAp1ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400,System.currentTimeMillis(),Optional.empty());
+
+        topicPartitionListOffsetsResultInfoMap = new HashMap<>();
+
+        topicPartitionListOffsetsResultInfoMap.put(tAp1,tAp1ListOffsetsResultInfo);
+        topicPartitionListOffsetsResultInfoMap.put(tAp2,tAp2ListOffsetsResultInfo);
+        topicPartitionListOffsetsResultInfoMap.put(tBp1,tBp1ListOffsetsResultInfo);
+        topicPartitionListOffsetsResultInfoMap.put(tBp2,tBp2ListOffsetsResultInfo);
+
+        kafkaFuture = mock(KafkaFuture.class);
+        when(kafkaFuture.get()).thenReturn(topicPartitionListOffsetsResultInfoMap);
+
+        listOffsetsResult = mock(ListOffsetsResult.class);
+        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
+
+        admin=mock(Admin.class);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
+
+        offsetManagers= new HashMap<>();
+
+        OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
+        when(offsetManagerTaP1.getCommittedOffset()).thenReturn(90L);
+
+        OffsetManager offsetManagerTaP2 = mock(OffsetManager.class);
+        when(offsetManagerTaP2.getCommittedOffset()).thenReturn(170L);
+
+        OffsetManager offsetManagerTbP1 = mock(OffsetManager.class);
+        when(offsetManagerTbP1.getCommittedOffset()).thenReturn(200L);
+
+        OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
+        when(offsetManagerTbP2.getCommittedOffset()).thenReturn(350L);
+
+        offsetManagers.put(tAp1,offsetManagerTaP1);
+        offsetManagers.put(tAp2,offsetManagerTaP2);
+        offsetManagers.put(tBp1,offsetManagerTbP1);
+        offsetManagers.put(tBp2,offsetManagerTbP2);
+
+        assignment=new HashSet<>();
+        assignment.add(tAp1);
+        assignment.add(tAp2);
+        assignment.add(tBp1);
+        assignment.add(tBp2);
+
+
+        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers),() -> admin,assignment);
+        Map<String, Metric>  result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1= (Gauge) result.get("topicA/partition_1/spoutLag");
+        Gauge g2= (Gauge) result.get("topicA/partition_2/spoutLag");
+        Gauge g3= (Gauge) result.get("topicB/partition_1/spoutLag");
+        Gauge g4= (Gauge) result.get("topicB/partition_2/spoutLag");
+        assertEquals(g1.getValue(),10L);
+        assertEquals(g2.getValue(),30L);
+        assertEquals(g3.getValue(),100L);
+        assertEquals(g4.getValue(),50L);
+
+        Gauge gATotal= (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(gATotal.getValue(),40L);
+        Gauge gBTotal= (Gauge) result.get("topicB/totalSpoutLag");
+        assertEquals(gBTotal.getValue(),150L);
+
+
+    }
+}

From 4a5c383a7d59ba7dd4a992e070205ad59c209fc7 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Sun, 20 Apr 2025 19:59:56 +0100
Subject: [PATCH 12/28] Adding unit tests

---
 .../KafkaOffsetPartitionAndTopicMetrics.java    |  6 +++---
 ...KafkaOffsetPartitionAndTopicMetricsTest.java | 17 +++++++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 973ce99f313..a14a44c1d1e 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -47,19 +47,19 @@ public class KafkaOffsetPartitionAndTopicMetrics implements MetricSet {
     private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
     private final Supplier<Admin> adminSupplier;
     private final Set<TopicPartition> assignment;
-    private final Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
-
 
     public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier, Supplier<Admin> adminSupplier, Set<TopicPartition> assignment) {
         this.offsetManagerSupplier = offsetManagerSupplier;
         this.adminSupplier = adminSupplier;
         this.assignment = assignment;
-        this.topicMetricsMap=new HashMap<>();
+
     }
 
     @Override
     public Map<String, Metric> getMetrics() {
 
+        Map<String, KafkaOffsetTopicMetrics> topicMetricsMap=new HashMap<>();
+
         Map<String, Metric> metrics = new HashMap<>();
 
         for (TopicPartition topicPartition : assignment) {
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index 0475155a993..224a2d6cddb 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -101,6 +101,23 @@ public void registerPartitionAndTopicMetrics() throws ExecutionException, Interr
         Gauge gBTotal= (Gauge) result.get("topicB/totalSpoutLag");
         assertEquals(gBTotal.getValue(),150L);
 
+        //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+         g1= (Gauge) result.get("topicA/partition_1/spoutLag");
+         g2= (Gauge) result.get("topicA/partition_2/spoutLag");
+         g3= (Gauge) result.get("topicB/partition_1/spoutLag");
+         g4= (Gauge) result.get("topicB/partition_2/spoutLag");
+        assertEquals(g1.getValue(),10L);
+        assertEquals(g2.getValue(),30L);
+        assertEquals(g3.getValue(),100L);
+        assertEquals(g4.getValue(),50L);
+
+        gATotal= (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(gATotal.getValue(),40L);
+        gBTotal= (Gauge) result.get("topicB/totalSpoutLag");
+        assertEquals(gBTotal.getValue(),150L);
+
 
     }
 }

From 7a4639cd808d07ee4fe0f0ef8b3adb8ed2ea9651 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Tue, 22 Apr 2025 22:20:06 +0100
Subject: [PATCH 13/28] Adding unit tests

---
 .../KafkaOffsetPartitionAndTopicMetrics.java  |  12 +-
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 133 ++++++++++++++----
 2 files changed, 111 insertions(+), 34 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index a14a44c1d1e..a4c3b4f4ddf 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -79,7 +79,7 @@ public Map<String, Metric> getMetrics() {
                     LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
                     return 0L;
                 }
-                // add value to topic level metric
+
                 OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
                 Long ret = endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
                 finalTopicMetrics.totalSpoutLag += ret;
@@ -92,7 +92,7 @@ public Map<String, Metric> getMetrics() {
                     LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
                     return 0L;
                 }
-                // add value to topic level metric
+
                 Long ret = beginningOffsets.get(topicPartition);
                 finalTopicMetrics.totalEarliestTimeOffset += ret;
                 return ret;
@@ -104,14 +104,14 @@ public Map<String, Metric> getMetrics() {
                     LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
                     return 0L;
                 }
-                // add value to topic level metric
+
                 Long ret = endOffsets.get(topicPartition);
                 finalTopicMetrics.totalLatestTimeOffset += ret;
                 return ret;
             };
 
             Gauge<Long> latestEmittedOffsetGauge = () -> {
-                // add value to topic level metric
+
                 OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
                 Long ret = offsetManager.getLatestEmittedOffset();
                 finalTopicMetrics.totalLatestEmittedOffset+=ret;
@@ -119,7 +119,7 @@ public Map<String, Metric> getMetrics() {
             };
 
             Gauge<Long> latestCompletedOffsetGauge = () -> {
-                // add value to topic level metric
+
                 OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
                 Long ret = offsetManager.getCommittedOffset();
                 finalTopicMetrics.totalLatestCompletedOffset+=ret;
@@ -137,7 +137,7 @@ public Map<String, Metric> getMetrics() {
                     LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
                     return 0L;
                 }
-                // add value to topic level metric
+
                 Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
                 finalTopicMetrics.totalRecordsInPartitions+=ret;
                 return ret;
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index 224a2d6cddb..a0424dddda3 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -8,6 +8,7 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.storm.kafka.spout.internal.OffsetManager;
 import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionAndTopicMetrics;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -16,50 +17,51 @@
 import java.util.concurrent.ExecutionException;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.*;
 
 @ExtendWith(MockitoExtension.class)
 public class KafkaOffsetPartitionAndTopicMetricsTest {
 
     private Set<TopicPartition> assignment;
-    private Admin admin;
+    private Admin admin = mock(Admin.class);
     private HashMap<TopicPartition, OffsetManager> offsetManagers;
-    private ListOffsetsResult listOffsetsResult;
-    private KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> kafkaFuture;
-    private Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap;
+    private ListOffsetsResult listOffsetsResultEarliest;
+    private KafkaFuture kafkaFuture = mock(KafkaFuture.class);
+
+    @BeforeEach
+    public void initializeTests() {
+        reset(admin, kafkaFuture);
+
+    }
 
     @Test
-    public void registerPartitionAndTopicMetrics() throws ExecutionException, InterruptedException {
+    public void registerMetricsGetLatest() throws ExecutionException, InterruptedException {
 
         TopicPartition tAp1 = new TopicPartition("topicA",1);
         TopicPartition tAp2 = new TopicPartition("topicA",2);
         TopicPartition tBp1 = new TopicPartition("topicB",1);
         TopicPartition tBp2 = new TopicPartition("topicB",2);
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2ListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400,System.currentTimeMillis(),Optional.empty());
 
-        topicPartitionListOffsetsResultInfoMap = new HashMap<>();
+         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap;
 
-        topicPartitionListOffsetsResultInfoMap.put(tAp1,tAp1ListOffsetsResultInfo);
-        topicPartitionListOffsetsResultInfoMap.put(tAp2,tAp2ListOffsetsResultInfo);
-        topicPartitionListOffsetsResultInfoMap.put(tBp1,tBp1ListOffsetsResultInfo);
-        topicPartitionListOffsetsResultInfoMap.put(tBp2,tBp2ListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1,tAp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2,tAp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1,tBp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2,tBp2LatestListOffsetsResultInfo);
 
-        kafkaFuture = mock(KafkaFuture.class);
-        when(kafkaFuture.get()).thenReturn(topicPartitionListOffsetsResultInfoMap);
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 
-        listOffsetsResult = mock(ListOffsetsResult.class);
-        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
+        listOffsetsResultEarliest = mock(ListOffsetsResult.class);
+        when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
 
         admin=mock(Admin.class);
-        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
-
-        offsetManagers= new HashMap<>();
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
 
         OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
         when(offsetManagerTaP1.getCommittedOffset()).thenReturn(90L);
@@ -73,6 +75,7 @@ public void registerPartitionAndTopicMetrics() throws ExecutionException, Interr
         OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
         when(offsetManagerTbP2.getCommittedOffset()).thenReturn(350L);
 
+        offsetManagers= new HashMap<>();
         offsetManagers.put(tAp1,offsetManagerTaP1);
         offsetManagers.put(tAp2,offsetManagerTaP2);
         offsetManagers.put(tBp1,offsetManagerTbP1);
@@ -104,10 +107,10 @@ public void registerPartitionAndTopicMetrics() throws ExecutionException, Interr
         //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
 
         result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-         g1= (Gauge) result.get("topicA/partition_1/spoutLag");
-         g2= (Gauge) result.get("topicA/partition_2/spoutLag");
-         g3= (Gauge) result.get("topicB/partition_1/spoutLag");
-         g4= (Gauge) result.get("topicB/partition_2/spoutLag");
+        g1= (Gauge) result.get("topicA/partition_1/spoutLag");
+        g2= (Gauge) result.get("topicA/partition_2/spoutLag");
+        g3= (Gauge) result.get("topicB/partition_1/spoutLag");
+        g4= (Gauge) result.get("topicB/partition_2/spoutLag");
         assertEquals(g1.getValue(),10L);
         assertEquals(g2.getValue(),30L);
         assertEquals(g3.getValue(),100L);
@@ -119,5 +122,79 @@ public void registerPartitionAndTopicMetrics() throws ExecutionException, Interr
         assertEquals(gBTotal.getValue(),150L);
 
 
+    }
+
+    @Test
+    public void registerMetricsGetEarliest() throws ExecutionException, InterruptedException {
+
+        TopicPartition tAp1 = new TopicPartition("topicA",1);
+        TopicPartition tAp2 = new TopicPartition("topicA",2);
+        TopicPartition tBp1 = new TopicPartition("topicB",1);
+        TopicPartition tBp2 = new TopicPartition("topicB",2);
+
+        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
+
+         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap;
+
+        topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1,tAp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2,tAp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1,tBp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2,tBp2EarliestListOffsetsResultInfo);
+
+        when(kafkaFuture.get()).thenReturn(topicPartitionEarliestListOffsetsResultInfoMap);
+
+        listOffsetsResultEarliest = mock(ListOffsetsResult.class);
+        when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
+
+        admin=mock(Admin.class);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
+
+        offsetManagers= new HashMap<>();
+
+        assignment=new HashSet<>();
+        assignment.add(tAp1);
+        assignment.add(tAp2);
+        assignment.add(tBp1);
+        assignment.add(tBp2);
+
+        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers),() -> admin,assignment);
+        Map<String, Metric>  result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1= (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
+        Gauge g2= (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
+        Gauge g3= (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
+        Gauge g4= (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
+        assertEquals(g1.getValue(),1L);
+        assertEquals(g2.getValue(),1L);
+        assertEquals(g3.getValue(),1L);
+        assertEquals(g4.getValue(),1L);
+
+        Gauge gATotal= (Gauge) result.get("topicA/totalEarliestTimeOffset");
+        assertEquals(2L,gATotal.getValue());
+        Gauge gBTotal= (Gauge) result.get("topicB/totalEarliestTimeOffset");
+        assertEquals(2L,gBTotal.getValue());
+
+        //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+
+        g1=(Gauge) result.get("topicA/partition_1/earliestTimeOffset");
+        g2=(Gauge) result.get("topicA/partition_2/earliestTimeOffset");
+        g3=(Gauge) result.get("topicB/partition_1/earliestTimeOffset");
+        g4=(Gauge) result.get("topicB/partition_2/earliestTimeOffset");
+        assertEquals(g1.getValue(),1L);
+        assertEquals(g2.getValue(),1L);
+        assertEquals(g3.getValue(),1L);
+        assertEquals(g4.getValue(),1L);
+
+        gATotal= (Gauge) result.get("topicA/totalEarliestTimeOffset");
+        assertEquals(2L,gATotal.getValue());
+        gBTotal= (Gauge) result.get("topicB/totalEarliestTimeOffset");
+        assertEquals(2L,gBTotal.getValue());
+
+
     }
 }

From 1e1edb5b28cc97593197bfb38b5620b7eb3af0b4 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Tue, 22 Apr 2025 22:40:31 +0100
Subject: [PATCH 14/28] Adding unit tests

---
 .../KafkaOffsetPartitionAndTopicMetrics.java  |  14 +-
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 182 +++++++++---------
 2 files changed, 97 insertions(+), 99 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index a4c3b4f4ddf..707796d9d9c 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -58,20 +58,20 @@ public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetMa
     @Override
     public Map<String, Metric> getMetrics() {
 
-        Map<String, KafkaOffsetTopicMetrics> topicMetricsMap=new HashMap<>();
+        Map<String, KafkaOffsetTopicMetrics> topicMetricsMap = new HashMap<>();
 
         Map<String, Metric> metrics = new HashMap<>();
 
         for (TopicPartition topicPartition : assignment) {
 
-            String topic=topicPartition.topic();
+            String topic = topicPartition.topic();
             KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
             if (topicMetrics == null) {
                 topicMetrics = new KafkaOffsetTopicMetrics(topic);
                 topicMetricsMap.put(topic, topicMetrics);
             }
 
-            String metricPath = topicPartition.topic()  + "/partition_" + topicPartition.partition();
+            String metricPath = topicPartition.topic() + "/partition_" + topicPartition.partition();
             KafkaOffsetTopicMetrics finalTopicMetrics = topicMetrics;
             Gauge<Long> spoutLagGauge = () -> {
                 Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
@@ -114,7 +114,7 @@ public Map<String, Metric> getMetrics() {
 
                 OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
                 Long ret = offsetManager.getLatestEmittedOffset();
-                finalTopicMetrics.totalLatestEmittedOffset+=ret;
+                finalTopicMetrics.totalLatestEmittedOffset += ret;
                 return ret;
             };
 
@@ -122,7 +122,7 @@ public Map<String, Metric> getMetrics() {
 
                 OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
                 Long ret = offsetManager.getCommittedOffset();
-                finalTopicMetrics.totalLatestCompletedOffset+=ret;
+                finalTopicMetrics.totalLatestCompletedOffset += ret;
                 return ret;
             };
 
@@ -139,7 +139,7 @@ public Map<String, Metric> getMetrics() {
                 }
 
                 Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
-                finalTopicMetrics.totalRecordsInPartitions+=ret;
+                finalTopicMetrics.totalRecordsInPartitions += ret;
                 return ret;
             };
 
@@ -152,7 +152,7 @@ public Map<String, Metric> getMetrics() {
 
         }
 
-        for( KafkaOffsetTopicMetrics kafkaOffsetTopicMetrics  : topicMetricsMap.values()){
+        for (KafkaOffsetTopicMetrics kafkaOffsetTopicMetrics : topicMetricsMap.values()) {
             metrics.putAll(kafkaOffsetTopicMetrics.getMetrics());
         }
 
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index a0424dddda3..7fcd66681ae 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -37,30 +37,30 @@ public void initializeTests() {
     @Test
     public void registerMetricsGetLatest() throws ExecutionException, InterruptedException {
 
-        TopicPartition tAp1 = new TopicPartition("topicA",1);
-        TopicPartition tAp2 = new TopicPartition("topicA",2);
-        TopicPartition tBp1 = new TopicPartition("topicB",1);
-        TopicPartition tBp2 = new TopicPartition("topicB",2);
+        TopicPartition tAp1 = new TopicPartition("topicA", 1);
+        TopicPartition tAp2 = new TopicPartition("topicA", 2);
+        TopicPartition tBp1 = new TopicPartition("topicB", 1);
+        TopicPartition tBp2 = new TopicPartition("topicB", 2);
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
 
-         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap;
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap;
 
         topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1,tAp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2,tAp2LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1,tBp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2,tBp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1, tAp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2, tAp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1, tBp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2, tBp2LatestListOffsetsResultInfo);
 
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 
         listOffsetsResultEarliest = mock(ListOffsetsResult.class);
         when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
 
-        admin=mock(Admin.class);
+        admin = mock(Admin.class);
         when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
 
         OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
@@ -75,51 +75,51 @@ public void registerMetricsGetLatest() throws ExecutionException, InterruptedExc
         OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
         when(offsetManagerTbP2.getCommittedOffset()).thenReturn(350L);
 
-        offsetManagers= new HashMap<>();
-        offsetManagers.put(tAp1,offsetManagerTaP1);
-        offsetManagers.put(tAp2,offsetManagerTaP2);
-        offsetManagers.put(tBp1,offsetManagerTbP1);
-        offsetManagers.put(tBp2,offsetManagerTbP2);
+        offsetManagers = new HashMap<>();
+        offsetManagers.put(tAp1, offsetManagerTaP1);
+        offsetManagers.put(tAp2, offsetManagerTaP2);
+        offsetManagers.put(tBp1, offsetManagerTbP1);
+        offsetManagers.put(tBp2, offsetManagerTbP2);
 
-        assignment=new HashSet<>();
+        assignment = new HashSet<>();
         assignment.add(tAp1);
         assignment.add(tAp2);
         assignment.add(tBp1);
         assignment.add(tBp2);
 
 
-        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers),() -> admin,assignment);
-        Map<String, Metric>  result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        Gauge g1= (Gauge) result.get("topicA/partition_1/spoutLag");
-        Gauge g2= (Gauge) result.get("topicA/partition_2/spoutLag");
-        Gauge g3= (Gauge) result.get("topicB/partition_1/spoutLag");
-        Gauge g4= (Gauge) result.get("topicB/partition_2/spoutLag");
-        assertEquals(g1.getValue(),10L);
-        assertEquals(g2.getValue(),30L);
-        assertEquals(g3.getValue(),100L);
-        assertEquals(g4.getValue(),50L);
+        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
+        Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1 = (Gauge) result.get("topicA/partition_1/spoutLag");
+        Gauge g2 = (Gauge) result.get("topicA/partition_2/spoutLag");
+        Gauge g3 = (Gauge) result.get("topicB/partition_1/spoutLag");
+        Gauge g4 = (Gauge) result.get("topicB/partition_2/spoutLag");
+        assertEquals(g1.getValue(), 10L);
+        assertEquals(g2.getValue(), 30L);
+        assertEquals(g3.getValue(), 100L);
+        assertEquals(g4.getValue(), 50L);
 
-        Gauge gATotal= (Gauge) result.get("topicA/totalSpoutLag");
-        assertEquals(gATotal.getValue(),40L);
-        Gauge gBTotal= (Gauge) result.get("topicB/totalSpoutLag");
-        assertEquals(gBTotal.getValue(),150L);
+        Gauge gATotal = (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(gATotal.getValue(), 40L);
+        Gauge gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
+        assertEquals(gBTotal.getValue(), 150L);
 
         //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
 
         result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        g1= (Gauge) result.get("topicA/partition_1/spoutLag");
-        g2= (Gauge) result.get("topicA/partition_2/spoutLag");
-        g3= (Gauge) result.get("topicB/partition_1/spoutLag");
-        g4= (Gauge) result.get("topicB/partition_2/spoutLag");
-        assertEquals(g1.getValue(),10L);
-        assertEquals(g2.getValue(),30L);
-        assertEquals(g3.getValue(),100L);
-        assertEquals(g4.getValue(),50L);
+        g1 = (Gauge) result.get("topicA/partition_1/spoutLag");
+        g2 = (Gauge) result.get("topicA/partition_2/spoutLag");
+        g3 = (Gauge) result.get("topicB/partition_1/spoutLag");
+        g4 = (Gauge) result.get("topicB/partition_2/spoutLag");
+        assertEquals(g1.getValue(), 10L);
+        assertEquals(g2.getValue(), 30L);
+        assertEquals(g3.getValue(), 100L);
+        assertEquals(g4.getValue(), 50L);
 
-        gATotal= (Gauge) result.get("topicA/totalSpoutLag");
-        assertEquals(gATotal.getValue(),40L);
-        gBTotal= (Gauge) result.get("topicB/totalSpoutLag");
-        assertEquals(gBTotal.getValue(),150L);
+        gATotal = (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(gATotal.getValue(), 40L);
+        gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
+        assertEquals(gBTotal.getValue(), 150L);
 
 
     }
@@ -127,74 +127,72 @@ public void registerMetricsGetLatest() throws ExecutionException, InterruptedExc
     @Test
     public void registerMetricsGetEarliest() throws ExecutionException, InterruptedException {
 
-        TopicPartition tAp1 = new TopicPartition("topicA",1);
-        TopicPartition tAp2 = new TopicPartition("topicA",2);
-        TopicPartition tBp1 = new TopicPartition("topicB",1);
-        TopicPartition tBp2 = new TopicPartition("topicB",2);
+        TopicPartition tAp1 = new TopicPartition("topicA", 1);
+        TopicPartition tAp2 = new TopicPartition("topicA", 2);
+        TopicPartition tBp1 = new TopicPartition("topicB", 1);
+        TopicPartition tBp2 = new TopicPartition("topicB", 2);
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1,System.currentTimeMillis(),Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
 
-         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap;
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap;
 
         topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1,tAp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2,tAp2EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1,tBp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2,tBp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1, tAp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2, tAp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1, tBp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2, tBp2EarliestListOffsetsResultInfo);
 
         when(kafkaFuture.get()).thenReturn(topicPartitionEarliestListOffsetsResultInfoMap);
 
         listOffsetsResultEarliest = mock(ListOffsetsResult.class);
         when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
 
-        admin=mock(Admin.class);
+        admin = mock(Admin.class);
         when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
 
-        offsetManagers= new HashMap<>();
+        offsetManagers = new HashMap<>();
 
-        assignment=new HashSet<>();
+        assignment = new HashSet<>();
         assignment.add(tAp1);
         assignment.add(tAp2);
         assignment.add(tBp1);
         assignment.add(tBp2);
 
-        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers),() -> admin,assignment);
-        Map<String, Metric>  result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        Gauge g1= (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
-        Gauge g2= (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
-        Gauge g3= (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
-        Gauge g4= (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
-        assertEquals(g1.getValue(),1L);
-        assertEquals(g2.getValue(),1L);
-        assertEquals(g3.getValue(),1L);
-        assertEquals(g4.getValue(),1L);
-
-        Gauge gATotal= (Gauge) result.get("topicA/totalEarliestTimeOffset");
-        assertEquals(2L,gATotal.getValue());
-        Gauge gBTotal= (Gauge) result.get("topicB/totalEarliestTimeOffset");
-        assertEquals(2L,gBTotal.getValue());
+        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
+        Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1 = (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
+        Gauge g2 = (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
+        Gauge g3 = (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
+        Gauge g4 = (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
+        assertEquals(g1.getValue(), 1L);
+        assertEquals(g2.getValue(), 1L);
+        assertEquals(g3.getValue(), 1L);
+        assertEquals(g4.getValue(), 1L);
+
+        Gauge gATotal = (Gauge) result.get("topicA/totalEarliestTimeOffset");
+        assertEquals(2L, gATotal.getValue());
+        Gauge gBTotal = (Gauge) result.get("topicB/totalEarliestTimeOffset");
+        assertEquals(2L, gBTotal.getValue());
 
         //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
 
         result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
 
-        g1=(Gauge) result.get("topicA/partition_1/earliestTimeOffset");
-        g2=(Gauge) result.get("topicA/partition_2/earliestTimeOffset");
-        g3=(Gauge) result.get("topicB/partition_1/earliestTimeOffset");
-        g4=(Gauge) result.get("topicB/partition_2/earliestTimeOffset");
-        assertEquals(g1.getValue(),1L);
-        assertEquals(g2.getValue(),1L);
-        assertEquals(g3.getValue(),1L);
-        assertEquals(g4.getValue(),1L);
-
-        gATotal= (Gauge) result.get("topicA/totalEarliestTimeOffset");
-        assertEquals(2L,gATotal.getValue());
-        gBTotal= (Gauge) result.get("topicB/totalEarliestTimeOffset");
-        assertEquals(2L,gBTotal.getValue());
-
-
+        g1 = (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
+        g2 = (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
+        g3 = (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
+        g4 = (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
+        assertEquals(g1.getValue(), 1L);
+        assertEquals(g2.getValue(), 1L);
+        assertEquals(g3.getValue(), 1L);
+        assertEquals(g4.getValue(), 1L);
+
+        gATotal = (Gauge) result.get("topicA/totalEarliestTimeOffset");
+        assertEquals(2L, gATotal.getValue());
+        gBTotal = (Gauge) result.get("topicB/totalEarliestTimeOffset");
+        assertEquals(2L, gBTotal.getValue());
     }
 }

From 83bbf231cc16f038076568105055bc466edfb64e Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 08:10:59 +0100
Subject: [PATCH 15/28] Adding unit tests

---
 .../apache/storm/kafka/spout/KafkaSpout.java  | 59 ++++++++++---------
 1 file changed, 30 insertions(+), 29 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
index 2b35fdd9272..764ee643b65 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
@@ -39,6 +39,7 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
+
 import org.apache.commons.lang.Validate;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.consumer.Consumer;
@@ -155,7 +156,7 @@ public void open(Map<String, Object> conf, TopologyContext context, SpoutOutputC
 
         tupleListener.open(conf, context);
         this.kafkaOffsetMetricManager
-            = new KafkaOffsetMetricManager<>(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, context);
+                = new KafkaOffsetMetricManager<>(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, context);
 
         LOG.info("Kafka Spout opened with the following configuration: {}", kafkaSpoutConfig);
     }
@@ -184,7 +185,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
             previousAssignment = partitions;
 
             LOG.info("Partitions revoked. [consumer-group={}, consumer={}, topic-partitions={}]",
-                kafkaSpoutConfig.getConsumerGroupId(), consumer, partitions);
+                    kafkaSpoutConfig.getConsumerGroupId(), consumer, partitions);
 
             if (isAtLeastOnceProcessing()) {
                 commitOffsetsForAckedTuples();
@@ -194,7 +195,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
         @Override
         public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
             LOG.info("Partitions reassignment. [task-ID={}, consumer-group={}, consumer={}, topic-partitions={}]",
-                context.getThisTaskId(), kafkaSpoutConfig.getConsumerGroupId(), consumer, partitions);
+                    context.getThisTaskId(), kafkaSpoutConfig.getConsumerGroupId(), consumer, partitions);
 
             initialize(partitions);
             tupleListener.onPartitionsReassigned(partitions);
@@ -222,7 +223,7 @@ private void initialize(Collection<TopicPartition> partitions) {
                 final OffsetAndMetadata committedOffset = consumer.committed(newTp);
                 final long fetchOffset = doSeek(newTp, committedOffset);
                 LOG.debug("Set consumer position to [{}] for topic-partition [{}] with [{}] and committed offset [{}]",
-                    fetchOffset, newTp, firstPollOffsetStrategy, committedOffset);
+                        fetchOffset, newTp, firstPollOffsetStrategy, committedOffset);
                 if (isAtLeastOnceProcessing() && !offsetManagers.containsKey(newTp)) {
                     offsetManagers.put(newTp, new OffsetManager(newTp, fetchOffset));
                 }
@@ -235,13 +236,13 @@ private void initialize(Collection<TopicPartition> partitions) {
          */
         private long doSeek(TopicPartition newTp, OffsetAndMetadata committedOffset) {
             LOG.trace("Seeking offset for topic-partition [{}] with [{}] and committed offset [{}]",
-                newTp, firstPollOffsetStrategy, committedOffset);
+                    newTp, firstPollOffsetStrategy, committedOffset);
 
             if (committedOffset != null) {
                 // offset was previously committed for this consumer group and topic-partition, either by this or another topology.
                 if (commitMetadataManager.isOffsetCommittedByThisTopology(newTp,
-                    committedOffset,
-                    Collections.unmodifiableMap(offsetManagers))) {
+                        committedOffset,
+                        Collections.unmodifiableMap(offsetManagers))) {
                     // Another KafkaSpout instance (of this topology) already committed, therefore FirstPollOffsetStrategy does not apply.
                     consumer.seek(newTp, committedOffset.offset());
                 } else {
@@ -282,7 +283,7 @@ public void nextTuple() {
                     commitOffsetsForAckedTuples();
                 } else if (kafkaSpoutConfig.getProcessingGuarantee() == ProcessingGuarantee.NO_GUARANTEE) {
                     Map<TopicPartition, OffsetAndMetadata> offsetsToCommit =
-                        createFetchedOffsetsMetadata(consumer.assignment());
+                            createFetchedOffsetsMetadata(consumer.assignment());
                     consumer.commitAsync(offsetsToCommit, null);
                     LOG.debug("Committed offsets {} to Kafka", offsetsToCommit);
                 }
@@ -337,7 +338,7 @@ private PollablePartitionsInfo getPollablePartitionsInfo() {
                     pollablePartitions.add(tp);
                 } else {
                     LOG.debug("Not polling on partition [{}]. It has [{}] uncommitted offsets, which exceeds the limit of [{}]. ", tp,
-                        numUncommittedOffsets, maxUncommittedOffsets);
+                            numUncommittedOffsets, maxUncommittedOffsets);
                 }
             }
         }
@@ -346,7 +347,7 @@ private PollablePartitionsInfo getPollablePartitionsInfo() {
 
     private boolean isWaitingToEmit() {
         return waitingToEmit.values().stream()
-            .anyMatch(list -> !list.isEmpty());
+                .anyMatch(list -> !list.isEmpty());
     }
 
     private void setWaitingToEmit(ConsumerRecords<K, V> consumerRecords) {
@@ -366,11 +367,11 @@ private ConsumerRecords<K, V> pollKafkaBroker(PollablePartitionsInfo pollablePar
             ackRetriableOffsetsIfCompactedAway(pollablePartitionsInfo.pollableEarliestRetriableOffsets, consumerRecords);
             final int numPolledRecords = consumerRecords.count();
             LOG.debug("Polled [{}] records from Kafka",
-                numPolledRecords);
+                    numPolledRecords);
             if (kafkaSpoutConfig.getProcessingGuarantee() == KafkaSpoutConfig.ProcessingGuarantee.AT_MOST_ONCE) {
                 //Commit polled records immediately to ensure delivery is at-most-once.
                 Map<TopicPartition, OffsetAndMetadata> offsetsToCommit =
-                    createFetchedOffsetsMetadata(consumer.assignment());
+                        createFetchedOffsetsMetadata(consumer.assignment());
                 consumer.commitSync(offsetsToCommit);
                 LOG.debug("Committed offsets {} to Kafka", offsetsToCommit);
             }
@@ -388,7 +389,7 @@ private void doSeekRetriableTopicPartitions(Map<TopicPartition, Long> pollableEa
     }
 
     private void ackRetriableOffsetsIfCompactedAway(Map<TopicPartition, Long> earliestRetriableOffsets,
-        ConsumerRecords<K, V> consumerRecords) {
+                                                    ConsumerRecords<K, V> consumerRecords) {
         for (Entry<TopicPartition, Long> entry : earliestRetriableOffsets.entrySet()) {
             TopicPartition tp = entry.getKey();
             List<ConsumerRecord<K, V>> records = consumerRecords.records(tp);
@@ -530,7 +531,7 @@ private void commitOffsetsForAckedTuples() {
                      * to the committed offset.
                      */
                     LOG.debug("Consumer fell behind committed offset. Catching up. Position was [{}], skipping to [{}]",
-                        position, committedOffset);
+                            position, committedOffset);
                     consumer.seek(tp, committedOffset);
                 }
                 /**
@@ -540,8 +541,8 @@ private void commitOffsetsForAckedTuples() {
                 if (waitingToEmitForTp != null) {
                     //Discard the pending records that are already committed
                     waitingToEmit.put(tp, waitingToEmitForTp.stream()
-                        .filter(record -> record.offset() >= committedOffset)
-                        .collect(Collectors.toCollection(LinkedList::new)));
+                            .filter(record -> record.offset() >= committedOffset)
+                            .collect(Collectors.toCollection(LinkedList::new)));
                 }
 
                 final OffsetManager offsetManager = offsetManagers.get(tp);
@@ -574,11 +575,11 @@ public void ack(Object messageId) {
 
         if (!emitted.contains(msgId)) {
             LOG.debug("Received ack for message [{}], associated with tuple emitted for a ConsumerRecord that "
-                + "came from a topic-partition that this consumer group instance is no longer tracking "
-                + "due to rebalance/partition reassignment. No action taken.", msgId);
+                    + "came from a topic-partition that this consumer group instance is no longer tracking "
+                    + "due to rebalance/partition reassignment. No action taken.", msgId);
         } else {
             Validate.isTrue(!retryService.isScheduled(msgId), "The message id " + msgId + " is queued for retry while being acked."
-                + " This should never occur barring errors in the RetryService implementation or the spout code.");
+                    + " This should never occur barring errors in the RetryService implementation or the spout code.");
             offsetManagers.get(msgId.getTopicPartition()).addToAckMsgs(msgId);
             emitted.remove(msgId);
         }
@@ -596,11 +597,11 @@ public void fail(Object messageId) {
         final KafkaSpoutMessageId msgId = (KafkaSpoutMessageId) messageId;
         if (!emitted.contains(msgId)) {
             LOG.debug("Received fail for tuple this spout is no longer tracking."
-                + " Partitions may have been reassigned. Ignoring message [{}]", msgId);
+                    + " Partitions may have been reassigned. Ignoring message [{}]", msgId);
             return;
         }
         Validate.isTrue(!retryService.isScheduled(msgId), "The message id " + msgId + " is queued for retry while being failed."
-            + " This should never occur barring errors in the RetryService implementation or the spout code.");
+                + " This should never occur barring errors in the RetryService implementation or the spout code.");
 
         msgId.incrementNumFails();
 
@@ -631,7 +632,7 @@ private void refreshAssignment() {
         List<TopicPartition> allPartitionsSorted = new ArrayList<>(allPartitions);
         Collections.sort(allPartitionsSorted, TopicPartitionComparator.INSTANCE);
         Set<TopicPartition> assignedPartitions = kafkaSpoutConfig.getTopicPartitioner()
-            .getPartitionsForThisTask(allPartitionsSorted, context);
+                .getPartitionsForThisTask(allPartitionsSorted, context);
         boolean partitionChanged = topicAssigner.assignPartitions(consumer, assignedPartitions, rebalanceListener);
         if (partitionChanged && canRegisterMetrics()) {
             LOG.info("Partitions assignments has changed, updating metrics.");
@@ -684,9 +685,9 @@ public void declareOutputFields(OutputFieldsDeclarer declarer) {
     @Override
     public String toString() {
         return "KafkaSpout{"
-            + "offsetManagers =" + offsetManagers
-            + ", emitted=" + emitted
-            + "}";
+                + "offsetManagers =" + offsetManagers
+                + ", emitted=" + emitted
+                + "}";
     }
 
     @Override
@@ -719,8 +720,8 @@ private boolean isPrimitiveOrWrapper(Class<?> type) {
 
     private boolean isWrapper(Class<?> type) {
         return type == Double.class || type == Float.class || type == Long.class
-            || type == Integer.class || type == Short.class || type == Character.class
-            || type == Byte.class || type == Boolean.class || type == String.class;
+                || type == Integer.class || type == Short.class || type == Character.class
+                || type == Byte.class || type == Boolean.class || type == String.class;
     }
 
     private String getTopicsString() {
@@ -736,8 +737,8 @@ private static class PollablePartitionsInfo {
         PollablePartitionsInfo(Set<TopicPartition> pollablePartitions, Map<TopicPartition, Long> earliestRetriableOffsets) {
             this.pollablePartitions = pollablePartitions;
             this.pollableEarliestRetriableOffsets = earliestRetriableOffsets.entrySet().stream()
-                .filter(entry -> pollablePartitions.contains(entry.getKey()))
-                .collect(Collectors.toMap(entry -> entry.getKey(), entry -> entry.getValue()));
+                    .filter(entry -> pollablePartitions.contains(entry.getKey()))
+                    .collect(Collectors.toMap(entry -> entry.getKey(), entry -> entry.getValue()));
         }
 
         public boolean shouldPoll() {

From c07874c4d2b65e6c4b109899c93abf6dc3fa8248 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 08:38:44 +0100
Subject: [PATCH 16/28] Adding unit tests

---
 .../KafkaOffsetPartitionAndTopicMetrics.java  | 36 +++++++++---------
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 38 ++++++++++++++++---
 2 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 707796d9d9c..8cb4cae2d1a 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -55,6 +55,24 @@ public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetMa
 
     }
 
+    private static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
+            throws InterruptedException, ExecutionException {
+
+        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
+        for (TopicPartition topicPartition : topicPartitions) {
+            offsetSpecMap.put(topicPartition, offsetSpec);
+        }
+        Map<TopicPartition, Long> ret = new HashMap<>();
+        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
+        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
+        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
+                topicPartitionListOffsetsResultInfoMap.entrySet()) {
+            ret.put(entry.getKey(), entry.getValue().offset());
+        }
+        return ret;
+    }
+
     @Override
     public Map<String, Metric> getMetrics() {
 
@@ -193,22 +211,4 @@ private Map<TopicPartition, Long> getEndOffsets(Set<TopicPartition> topicPartiti
         }
         return endOffsets;
     }
-
-    private static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
-            throws InterruptedException, ExecutionException {
-
-        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
-        for (TopicPartition topicPartition : topicPartitions) {
-            offsetSpecMap.put(topicPartition, offsetSpec);
-        }
-        Map<TopicPartition, Long> ret = new HashMap<>();
-        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
-        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
-        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
-        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
-                topicPartitionListOffsetsResultInfoMap.entrySet()) {
-            ret.put(entry.getKey(), entry.getValue().offset());
-        }
-        return ret;
-    }
 }
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index 7fcd66681ae..da1493d76f2 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -35,7 +35,7 @@ public void initializeTests() {
     }
 
     @Test
-    public void registerMetricsGetLatest() throws ExecutionException, InterruptedException {
+    public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedException {
 
         TopicPartition tAp1 = new TopicPartition("topicA", 1);
         TopicPartition tAp2 = new TopicPartition("topicA", 2);
@@ -125,7 +125,7 @@ public void registerMetricsGetLatest() throws ExecutionException, InterruptedExc
     }
 
     @Test
-    public void registerMetricsGetEarliest() throws ExecutionException, InterruptedException {
+    public void registerMetricsGetEarliestAndLatest() throws ExecutionException, InterruptedException {
 
         TopicPartition tAp1 = new TopicPartition("topicA", 1);
         TopicPartition tAp2 = new TopicPartition("topicA", 2);
@@ -137,9 +137,7 @@ public void registerMetricsGetEarliest() throws ExecutionException, InterruptedE
         ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
         ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
 
-        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap;
-
-        topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
         topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1, tAp1EarliestListOffsetsResultInfo);
         topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2, tAp2EarliestListOffsetsResultInfo);
         topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1, tBp1EarliestListOffsetsResultInfo);
@@ -194,5 +192,35 @@ public void registerMetricsGetEarliest() throws ExecutionException, InterruptedE
         assertEquals(2L, gATotal.getValue());
         gBTotal = (Gauge) result.get("topicB/totalEarliestTimeOffset");
         assertEquals(2L, gBTotal.getValue());
+
+        //get the latest offsets
+
+        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1, tAp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2, tAp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1, tBp1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2, tBp2LatestListOffsetsResultInfo);
+
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        g1 = (Gauge) result.get("topicA/partition_1/latestTimeOffset");
+        g2 = (Gauge) result.get("topicA/partition_2/latestTimeOffset");
+        g3 = (Gauge) result.get("topicB/partition_1/latestTimeOffset");
+        g4 = (Gauge) result.get("topicB/partition_2/latestTimeOffset");
+        assertEquals(100L, g1.getValue());
+        assertEquals(200L, g2.getValue());
+        assertEquals(300L, g3.getValue());
+        assertEquals(400L, g4.getValue());
+
+        gATotal = (Gauge) result.get("topicA/totalLatestTimeOffset");
+        assertEquals(300L, gATotal.getValue());
+        gBTotal = (Gauge) result.get("topicB/totalLatestTimeOffset");
+        assertEquals(700L, gBTotal.getValue());
     }
 }

From e1dcaa1cd447456f00f69bf90c9da1e2cd317dde Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 09:08:25 +0100
Subject: [PATCH 17/28] Adding unit tests

---
 .../metrics2/KafkaOffsetMetricManager.java    | 16 +++----
 .../KafkaOffsetPartitionAndTopicMetrics.java  |  4 +-
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 48 +++++++++++++++++++
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index c78b5637e89..c0d7631f68c 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- *
+ * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
- *
+ * <p>
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -18,11 +18,6 @@
 
 package org.apache.storm.kafka.spout.metrics2;
 
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-import java.util.function.Supplier;
-
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.storm.kafka.spout.internal.OffsetManager;
@@ -30,6 +25,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+
 /**
  * This class is used to manage both the partition and topic level offset metrics.
  */
@@ -57,7 +57,7 @@ public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> off
     public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignment) {
 
         KafkaOffsetPartitionAndTopicMetrics topicPartitionMetricSet
-                = new KafkaOffsetPartitionAndTopicMetrics(offsetManagerSupplier, adminSupplier,newAssignment);
+                = new KafkaOffsetPartitionAndTopicMetrics(offsetManagerSupplier, adminSupplier, newAssignment);
         topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
     }
 
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 8cb4cae2d1a..fb9d3d31c3e 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -48,7 +48,9 @@ public class KafkaOffsetPartitionAndTopicMetrics implements MetricSet {
     private final Supplier<Admin> adminSupplier;
     private final Set<TopicPartition> assignment;
 
-    public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier, Supplier<Admin> adminSupplier, Set<TopicPartition> assignment) {
+    public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
+                                               Supplier<Admin> adminSupplier,
+                                               Set<TopicPartition> assignment) {
         this.offsetManagerSupplier = offsetManagerSupplier;
         this.adminSupplier = adminSupplier;
         this.assignment = assignment;
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index da1493d76f2..dff2ad1bb94 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -151,7 +151,27 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
         admin = mock(Admin.class);
         when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
 
+        OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
+        when(offsetManagerTaP1.getLatestEmittedOffset()).thenReturn(50L);
+        when(offsetManagerTaP1.getCommittedOffset()).thenReturn(40L);
+
+        OffsetManager offsetManagerTaP2 = mock(OffsetManager.class);
+        when(offsetManagerTaP2.getLatestEmittedOffset()).thenReturn(100L);
+        when(offsetManagerTaP2.getCommittedOffset()).thenReturn(90L);
+
+        OffsetManager offsetManagerTbP1 = mock(OffsetManager.class);
+        when(offsetManagerTbP1.getLatestEmittedOffset()).thenReturn(150L);
+        when(offsetManagerTbP1.getCommittedOffset()).thenReturn(149L);
+
+        OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
+        when(offsetManagerTbP2.getLatestEmittedOffset()).thenReturn(200L);
+        when(offsetManagerTbP2.getCommittedOffset()).thenReturn(200L);
+
         offsetManagers = new HashMap<>();
+        offsetManagers.put(tAp1, offsetManagerTaP1);
+        offsetManagers.put(tAp2, offsetManagerTaP2);
+        offsetManagers.put(tBp1, offsetManagerTbP1);
+        offsetManagers.put(tBp2, offsetManagerTbP2);
 
         assignment = new HashSet<>();
         assignment.add(tAp1);
@@ -222,5 +242,33 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
         assertEquals(300L, gATotal.getValue());
         gBTotal = (Gauge) result.get("topicB/totalLatestTimeOffset");
         assertEquals(700L, gBTotal.getValue());
+
+        g1 = (Gauge) result.get("topicA/partition_1/latestEmittedOffset");
+        g2 = (Gauge) result.get("topicA/partition_2/latestEmittedOffset");
+        g3 = (Gauge) result.get("topicB/partition_1/latestEmittedOffset");
+        g4 = (Gauge) result.get("topicB/partition_2/latestEmittedOffset");
+        assertEquals(50L, g1.getValue());
+        assertEquals(100L, g2.getValue());
+        assertEquals(150L, g3.getValue());
+        assertEquals(200L, g4.getValue());
+
+        gATotal = (Gauge) result.get("topicA/totalLatestEmittedOffset");
+        assertEquals(150L, gATotal.getValue());
+        gBTotal = (Gauge) result.get("topicB/totalLatestEmittedOffset");
+        assertEquals(350L, gBTotal.getValue());
+
+        g1 = (Gauge) result.get("topicA/partition_1/latestCompletedOffset");
+        g2 = (Gauge) result.get("topicA/partition_2/latestCompletedOffset");
+        g3 = (Gauge) result.get("topicB/partition_1/latestCompletedOffset");
+        g4 = (Gauge) result.get("topicB/partition_2/latestCompletedOffset");
+        assertEquals(40L, g1.getValue());
+        assertEquals(90L, g2.getValue());
+        assertEquals(149L, g3.getValue());
+        assertEquals(200L, g4.getValue());
+
+        gATotal = (Gauge) result.get("topicA/totalLatestCompletedOffset");
+        assertEquals(130L, gATotal.getValue());
+        gBTotal = (Gauge) result.get("topicB/totalLatestCompletedOffset");
+        assertEquals(349L, gBTotal.getValue());
     }
 }

From 2a8198ebc0e1673a90d6aec44fff171e3d7b2f83 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 09:15:27 +0100
Subject: [PATCH 18/28] Adding unit tests

---
 .../spout/metrics2/KafkaOffsetMetricManager.java     |  8 ++++----
 .../KafkaOffsetPartitionAndTopicMetrics.java         | 12 ++++++------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index c0d7631f68c..9b6f662f541 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -18,6 +18,10 @@
 
 package org.apache.storm.kafka.spout.metrics2;
 
+import java.util.function.Supplier;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.storm.kafka.spout.internal.OffsetManager;
@@ -25,10 +29,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-import java.util.function.Supplier;
 
 /**
  * This class is used to manage both the partition and topic level offset metrics.
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index fb9d3d31c3e..1b91636d5a2 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -3,6 +3,12 @@
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricSet;
+import java.util.function.Supplier;
+import java.util.Collections;
+import java.util.concurrent.ExecutionException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.ListOffsetsResult;
 import org.apache.kafka.clients.admin.OffsetSpec;
@@ -13,12 +19,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.function.Supplier;
 
 /**
  * This class is used compute the partition and topic level offset metrics.

From 4dcd178de560c594c74a3dd9121c11afede5da44 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 22:52:31 +0100
Subject: [PATCH 19/28] Adding unit tests

---
 .../KafkaOffsetPartitionAndTopicMetrics.java  |  2 -
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 75 ++++++++++++++-----
 2 files changed, 55 insertions(+), 22 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 1b91636d5a2..1d98624098b 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -169,7 +169,6 @@ public Map<String, Metric> getMetrics() {
             metrics.put(metricPath + "/" + "latestEmittedOffset", latestEmittedOffsetGauge);
             metrics.put(metricPath + "/" + "latestCompletedOffset", latestCompletedOffsetGauge);
             metrics.put(metricPath + "/" + "recordsInPartition", recordsInPartitionGauge);
-
         }
 
         for (KafkaOffsetTopicMetrics kafkaOffsetTopicMetrics : topicMetricsMap.values()) {
@@ -177,7 +176,6 @@ public Map<String, Metric> getMetrics() {
         }
 
         return metrics;
-
     }
 
     private Map<TopicPartition, Long> getBeginningOffsets(Set<TopicPartition> topicPartitions) {
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index dff2ad1bb94..02789b90595 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -1,5 +1,13 @@
 package org.apache.storm.kafka.spout.metric2;
 
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Metric;
 import org.apache.kafka.clients.admin.Admin;
@@ -13,9 +21,6 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
 
-import java.util.*;
-import java.util.concurrent.ExecutionException;
-
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.Mockito.*;
 
@@ -25,7 +30,6 @@ public class KafkaOffsetPartitionAndTopicMetricsTest {
     private Set<TopicPartition> assignment;
     private Admin admin = mock(Admin.class);
     private HashMap<TopicPartition, OffsetManager> offsetManagers;
-    private ListOffsetsResult listOffsetsResultEarliest;
     private KafkaFuture kafkaFuture = mock(KafkaFuture.class);
 
     @BeforeEach
@@ -35,7 +39,7 @@ public void initializeTests() {
     }
 
     @Test
-    public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedException {
+    public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionException, InterruptedException {
 
         TopicPartition tAp1 = new TopicPartition("topicA", 1);
         TopicPartition tAp2 = new TopicPartition("topicA", 2);
@@ -47,9 +51,7 @@ public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedE
         ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
         ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
 
-        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap;
-
-        topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
         topicPartitionLatestListOffsetsResultInfoMap.put(tAp1, tAp1LatestListOffsetsResultInfo);
         topicPartitionLatestListOffsetsResultInfoMap.put(tAp2, tAp2LatestListOffsetsResultInfo);
         topicPartitionLatestListOffsetsResultInfoMap.put(tBp1, tBp1LatestListOffsetsResultInfo);
@@ -57,11 +59,11 @@ public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedE
 
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 
-        listOffsetsResultEarliest = mock(ListOffsetsResult.class);
-        when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
+        ListOffsetsResult listOffsetsResult = mock(ListOffsetsResult.class);
+        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
 
         admin = mock(Admin.class);
-        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
 
         OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
         when(offsetManagerTaP1.getCommittedOffset()).thenReturn(90L);
@@ -94,15 +96,15 @@ public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedE
         Gauge g2 = (Gauge) result.get("topicA/partition_2/spoutLag");
         Gauge g3 = (Gauge) result.get("topicB/partition_1/spoutLag");
         Gauge g4 = (Gauge) result.get("topicB/partition_2/spoutLag");
-        assertEquals(g1.getValue(), 10L);
-        assertEquals(g2.getValue(), 30L);
-        assertEquals(g3.getValue(), 100L);
-        assertEquals(g4.getValue(), 50L);
+        assertEquals(10L, g1.getValue());
+        assertEquals(30L, g2.getValue());
+        assertEquals(100L, g3.getValue());
+        assertEquals(50L, g4.getValue());
 
         Gauge gATotal = (Gauge) result.get("topicA/totalSpoutLag");
-        assertEquals(gATotal.getValue(), 40L);
+        assertEquals(40L, gATotal.getValue());
         Gauge gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
-        assertEquals(gBTotal.getValue(), 150L);
+        assertEquals(150L, gBTotal.getValue());
 
         //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
 
@@ -121,6 +123,39 @@ public void registerMetricsGetSpoutLag() throws ExecutionException, InterruptedE
         gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
         assertEquals(gBTotal.getValue(), 150L);
 
+        //get partition records
+
+        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(2, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(3, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(4, System.currentTimeMillis(), Optional.empty());
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1, tAp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2, tAp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1, tBp1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2, tBp2EarliestListOffsetsResultInfo);
+
+        //mock consecutive calls. Each call to the recordsInPartition gauge will call kafkaFuture.get() twice
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
+                topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
+                topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
+                topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap);
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        g1 = (Gauge) result.get("topicA/partition_1/recordsInPartition");
+        g2 = (Gauge) result.get("topicA/partition_2/recordsInPartition");
+        g3 = (Gauge) result.get("topicB/partition_1/recordsInPartition");
+        g4 = (Gauge) result.get("topicB/partition_2/recordsInPartition");
+        assertEquals(99L, g1.getValue());
+        assertEquals(198L, g2.getValue());
+        assertEquals(297L, g3.getValue());
+        assertEquals(396L, g4.getValue());
+
+        gATotal = (Gauge) result.get("topicA/totalRecordsInPartitions");
+        assertEquals(297L, gATotal.getValue());
+        gBTotal = (Gauge) result.get("topicB/totalRecordsInPartitions");
+        assertEquals(693L, gBTotal.getValue());
 
     }
 
@@ -145,11 +180,11 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
 
         when(kafkaFuture.get()).thenReturn(topicPartitionEarliestListOffsetsResultInfoMap);
 
-        listOffsetsResultEarliest = mock(ListOffsetsResult.class);
-        when(listOffsetsResultEarliest.all()).thenReturn(kafkaFuture);
+        ListOffsetsResult listOffsetsResult = mock(ListOffsetsResult.class);
+        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
 
         admin = mock(Admin.class);
-        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResultEarliest);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
 
         OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
         when(offsetManagerTaP1.getLatestEmittedOffset()).thenReturn(50L);

From 08fd1a91c51173673a212ccd2fcab50f378761da Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 23:07:34 +0100
Subject: [PATCH 20/28] Adding unit tests

---
 .../metrics2/KafkaOffsetMetricManager.java    |   2 +-
 .../KafkaOffsetPartitionAndTopicMetrics.java  |   2 +-
 ...fkaOffsetPartitionAndTopicMetricsTest.java | 152 +++++++++---------
 3 files changed, 78 insertions(+), 78 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index 9b6f662f541..ad82ce4ec8b 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -18,10 +18,10 @@
 
 package org.apache.storm.kafka.spout.metrics2;
 
-import java.util.function.Supplier;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Supplier;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.storm.kafka.spout.internal.OffsetManager;
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 1d98624098b..da32549628b 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -3,12 +3,12 @@
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricSet;
-import java.util.function.Supplier;
 import java.util.Collections;
 import java.util.concurrent.ExecutionException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Supplier;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.ListOffsetsResult;
 import org.apache.kafka.clients.admin.OffsetSpec;
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index 02789b90595..79f0c348926 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -41,21 +41,21 @@ public void initializeTests() {
     @Test
     public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionException, InterruptedException {
 
-        TopicPartition tAp1 = new TopicPartition("topicA", 1);
-        TopicPartition tAp2 = new TopicPartition("topicA", 2);
-        TopicPartition tBp1 = new TopicPartition("topicB", 1);
-        TopicPartition tBp2 = new TopicPartition("topicB", 2);
+        TopicPartition topicAPartition1 = new TopicPartition("topicA", 1);
+        TopicPartition topicAPartition2 = new TopicPartition("topicA", 2);
+        TopicPartition topicBPartition1 = new TopicPartition("topicB", 1);
+        TopicPartition topicBPartition2 = new TopicPartition("topicB", 2);
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
 
         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1, tAp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2, tAp2LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1, tBp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2, tBp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition2, topicAPartition2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicBPartition1, topicBPartition1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicBPartition2, topicBPartition2LatestListOffsetsResultInfo);
 
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 
@@ -65,29 +65,29 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
         admin = mock(Admin.class);
         when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
 
-        OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
-        when(offsetManagerTaP1.getCommittedOffset()).thenReturn(90L);
+        OffsetManager offsetManagertopicAPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition1.getCommittedOffset()).thenReturn(90L);
 
-        OffsetManager offsetManagerTaP2 = mock(OffsetManager.class);
-        when(offsetManagerTaP2.getCommittedOffset()).thenReturn(170L);
+        OffsetManager offsetManagertopicAPartition2 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition2.getCommittedOffset()).thenReturn(170L);
 
-        OffsetManager offsetManagerTbP1 = mock(OffsetManager.class);
-        when(offsetManagerTbP1.getCommittedOffset()).thenReturn(200L);
+        OffsetManager offsetManagertopicBPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicBPartition1.getCommittedOffset()).thenReturn(200L);
 
-        OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
-        when(offsetManagerTbP2.getCommittedOffset()).thenReturn(350L);
+        OffsetManager offsetManagertopicBPartition2 = mock(OffsetManager.class);
+        when(offsetManagertopicBPartition2.getCommittedOffset()).thenReturn(350L);
 
         offsetManagers = new HashMap<>();
-        offsetManagers.put(tAp1, offsetManagerTaP1);
-        offsetManagers.put(tAp2, offsetManagerTaP2);
-        offsetManagers.put(tBp1, offsetManagerTbP1);
-        offsetManagers.put(tBp2, offsetManagerTbP2);
+        offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
+        offsetManagers.put(topicAPartition2, offsetManagertopicAPartition2);
+        offsetManagers.put(topicBPartition1, offsetManagertopicBPartition1);
+        offsetManagers.put(topicBPartition2, offsetManagertopicBPartition2);
 
         assignment = new HashSet<>();
-        assignment.add(tAp1);
-        assignment.add(tAp2);
-        assignment.add(tBp1);
-        assignment.add(tBp2);
+        assignment.add(topicAPartition1);
+        assignment.add(topicAPartition2);
+        assignment.add(topicBPartition1);
+        assignment.add(topicBPartition2);
 
 
         KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
@@ -125,16 +125,16 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
 
         //get partition records
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(2, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(3, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(4, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(2, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(3, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(4, System.currentTimeMillis(), Optional.empty());
 
         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1, tAp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2, tAp2EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1, tBp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2, tBp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition2, topicAPartition2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicBPartition1, topicBPartition1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicBPartition2, topicBPartition2EarliestListOffsetsResultInfo);
 
         //mock consecutive calls. Each call to the recordsInPartition gauge will call kafkaFuture.get() twice
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
@@ -162,21 +162,21 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
     @Test
     public void registerMetricsGetEarliestAndLatest() throws ExecutionException, InterruptedException {
 
-        TopicPartition tAp1 = new TopicPartition("topicA", 1);
-        TopicPartition tAp2 = new TopicPartition("topicA", 2);
-        TopicPartition tBp1 = new TopicPartition("topicB", 1);
-        TopicPartition tBp2 = new TopicPartition("topicB", 2);
+        TopicPartition topicAPartition1 = new TopicPartition("topicA", 1);
+        TopicPartition topicAPartition2 = new TopicPartition("topicA", 2);
+        TopicPartition topicBPartition1 = new TopicPartition("topicB", 1);
+        TopicPartition topicBPartition2 = new TopicPartition("topicB", 2);
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
 
         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp1, tAp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tAp2, tAp2EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp1, tBp1EarliestListOffsetsResultInfo);
-        topicPartitionEarliestListOffsetsResultInfoMap.put(tBp2, tBp2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition2, topicAPartition2EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicBPartition1, topicBPartition1EarliestListOffsetsResultInfo);
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicBPartition2, topicBPartition2EarliestListOffsetsResultInfo);
 
         when(kafkaFuture.get()).thenReturn(topicPartitionEarliestListOffsetsResultInfoMap);
 
@@ -186,33 +186,33 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
         admin = mock(Admin.class);
         when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
 
-        OffsetManager offsetManagerTaP1 = mock(OffsetManager.class);
-        when(offsetManagerTaP1.getLatestEmittedOffset()).thenReturn(50L);
-        when(offsetManagerTaP1.getCommittedOffset()).thenReturn(40L);
+        OffsetManager offsetManagertopicAPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition1.getLatestEmittedOffset()).thenReturn(50L);
+        when(offsetManagertopicAPartition1.getCommittedOffset()).thenReturn(40L);
 
-        OffsetManager offsetManagerTaP2 = mock(OffsetManager.class);
-        when(offsetManagerTaP2.getLatestEmittedOffset()).thenReturn(100L);
-        when(offsetManagerTaP2.getCommittedOffset()).thenReturn(90L);
+        OffsetManager offsetManagertopicAPartition2 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition2.getLatestEmittedOffset()).thenReturn(100L);
+        when(offsetManagertopicAPartition2.getCommittedOffset()).thenReturn(90L);
 
-        OffsetManager offsetManagerTbP1 = mock(OffsetManager.class);
-        when(offsetManagerTbP1.getLatestEmittedOffset()).thenReturn(150L);
-        when(offsetManagerTbP1.getCommittedOffset()).thenReturn(149L);
+        OffsetManager offsetManagertopicBPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicBPartition1.getLatestEmittedOffset()).thenReturn(150L);
+        when(offsetManagertopicBPartition1.getCommittedOffset()).thenReturn(149L);
 
-        OffsetManager offsetManagerTbP2 = mock(OffsetManager.class);
-        when(offsetManagerTbP2.getLatestEmittedOffset()).thenReturn(200L);
-        when(offsetManagerTbP2.getCommittedOffset()).thenReturn(200L);
+        OffsetManager offsetManagertopicBPartition2 = mock(OffsetManager.class);
+        when(offsetManagertopicBPartition2.getLatestEmittedOffset()).thenReturn(200L);
+        when(offsetManagertopicBPartition2.getCommittedOffset()).thenReturn(200L);
 
         offsetManagers = new HashMap<>();
-        offsetManagers.put(tAp1, offsetManagerTaP1);
-        offsetManagers.put(tAp2, offsetManagerTaP2);
-        offsetManagers.put(tBp1, offsetManagerTbP1);
-        offsetManagers.put(tBp2, offsetManagerTbP2);
+        offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
+        offsetManagers.put(topicAPartition2, offsetManagertopicAPartition2);
+        offsetManagers.put(topicBPartition1, offsetManagertopicBPartition1);
+        offsetManagers.put(topicBPartition2, offsetManagertopicBPartition2);
 
         assignment = new HashSet<>();
-        assignment.add(tAp1);
-        assignment.add(tAp2);
-        assignment.add(tBp1);
-        assignment.add(tBp2);
+        assignment.add(topicAPartition1);
+        assignment.add(topicAPartition2);
+        assignment.add(topicBPartition1);
+        assignment.add(topicBPartition2);
 
         KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
         Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
@@ -250,16 +250,16 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
 
         //get the latest offsets
 
-        ListOffsetsResult.ListOffsetsResultInfo tAp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tAp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
-        ListOffsetsResult.ListOffsetsResultInfo tBp2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(200, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(300, System.currentTimeMillis(), Optional.empty());
+        ListOffsetsResult.ListOffsetsResultInfo topicBPartition2LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(400, System.currentTimeMillis(), Optional.empty());
 
         Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp1, tAp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tAp2, tAp2LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp1, tBp1LatestListOffsetsResultInfo);
-        topicPartitionLatestListOffsetsResultInfoMap.put(tBp2, tBp2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition2, topicAPartition2LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicBPartition1, topicBPartition1LatestListOffsetsResultInfo);
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicBPartition2, topicBPartition2LatestListOffsetsResultInfo);
 
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 

From 2e2d6fb0d8e5b6e2ad3edfe9311419c11fa90821 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Wed, 23 Apr 2025 23:11:16 +0100
Subject: [PATCH 21/28] Adding unit tests

---
 .../spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index da32549628b..9bf56b44438 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -4,10 +4,10 @@
 import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricSet;
 import java.util.Collections;
-import java.util.concurrent.ExecutionException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.ListOffsetsResult;

From 44f1bdc32c0ac67aa9d8ca5f149a1e6ab789d187 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Thu, 24 Apr 2025 11:46:45 +0100
Subject: [PATCH 22/28] Update KafkaOffsetPartitionAndTopicMetrics.java

---
 .../KafkaOffsetPartitionAndTopicMetrics.java   | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
index 9bf56b44438..6f6d766b491 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
@@ -1,3 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ *   or more contributor license agreements.  See the NOTICE file
+ *   distributed with this work for additional information
+ *   regarding copyright ownership.  The ASF licenses this file
+ *   to you under the Apache License, Version 2.0 (the
+ *   "License"); you may not use this file except in compliance
+ *   with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
 package org.apache.storm.kafka.spout.metrics2;
 
 import com.codahale.metrics.Gauge;

From be536d7fdfa5db6af524db392c4adbb856766fe4 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Thu, 24 Apr 2025 11:47:31 +0100
Subject: [PATCH 23/28] Update KafkaOffsetPartitionAndTopicMetricsTest.java

---
 ...afkaOffsetPartitionAndTopicMetricsTest.java | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
index 79f0c348926..cc774e8b24d 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
@@ -1,3 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ *   or more contributor license agreements.  See the NOTICE file
+ *   distributed with this work for additional information
+ *   regarding copyright ownership.  The ASF licenses this file
+ *   to you under the Apache License, Version 2.0 (the
+ *   "License"); you may not use this file except in compliance
+ *   with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
 package org.apache.storm.kafka.spout.metric2;
 
 

From 08f02575ed7d7498cfbbbd52eccda29f6d0c6d90 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Thu, 24 Apr 2025 12:41:18 +0100
Subject: [PATCH 24/28] Delete
 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java

---
 .../metrics2/KafkaOffsetPartitionMetrics.java | 227 ------------------
 1 file changed, 227 deletions(-)
 delete mode 100644 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
deleted file mode 100644
index 0bb0d2c2168..00000000000
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
+++ /dev/null
@@ -1,227 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.storm.kafka.spout.metrics2;
-
-import com.codahale.metrics.Gauge;
-import com.codahale.metrics.Metric;
-import com.codahale.metrics.MetricSet;
-
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.function.Supplier;
-
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.ListOffsetsResult;
-import org.apache.kafka.clients.admin.OffsetSpec;
-import org.apache.kafka.common.KafkaFuture;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.RetriableException;
-import org.apache.storm.kafka.spout.internal.OffsetManager;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * Partition level metrics.
- * <p>
- * topicName/partition_{number}/earliestTimeOffset //gives beginning offset of the partition
- * topicName/partition_{number}/latestTimeOffset //gives end offset of the partition
- * topicName/partition_{number}/latestEmittedOffset //gives latest emitted offset of the partition from the spout
- * topicName/partition_{number}/latestCompletedOffset //gives latest committed offset of the partition from the spout
- * topicName/partition_{number}/spoutLag // the delta between the latest Offset and latestCompletedOffset
- * topicName/partition_{number}/recordsInPartition // total number of records in the partition
- * </p>
- */
-public class KafkaOffsetPartitionMetrics<K, V> implements MetricSet {
-    private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetPartitionMetrics.class);
-    private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
-    private final Supplier<Admin> adminSupplier;
-
-    private TopicPartition topicPartition;
-    private KafkaOffsetTopicMetrics topicMetrics;
-
-    public KafkaOffsetPartitionMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
-                                       Supplier<Admin> adminSupplier,
-                                       TopicPartition topicPartition,
-                                       KafkaOffsetTopicMetrics topicMetrics) {
-        this.offsetManagerSupplier = offsetManagerSupplier;
-        this.adminSupplier = adminSupplier;
-        this.topicPartition = topicPartition;
-        this.topicMetrics = topicMetrics;
-
-        LOG.info("Running KafkaOffsetMetricSet");
-    }
-
-    @Override
-    public Map<String, Metric> getMetrics() {
-        Map<String, Metric> metrics = new HashMap();
-
-        String metricPath = topicPartition.topic()  + "/partition_" + topicPartition.partition();
-        Gauge<Long> spoutLagGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                // add value to topic level metric
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
-                topicMetrics.totalSpoutLag += ret;
-                return ret;
-            }
-        };
-
-        Gauge<Long> earliestTimeOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
-                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
-                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                // add value to topic level metric
-                Long ret = beginningOffsets.get(topicPartition);
-                topicMetrics.totalEarliestTimeOffset += ret;
-                return ret;
-            }
-        };
-
-        Gauge<Long> latestTimeOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                // add value to topic level metric
-                Long ret = endOffsets.get(topicPartition);
-                topicMetrics.totalLatestTimeOffset += ret;
-                return ret;
-            }
-        };
-
-        Gauge<Long> latestEmittedOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                // add value to topic level metric
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = offsetManager.getLatestEmittedOffset();
-                topicMetrics.totalLatestEmittedOffset += ret;
-                return ret;
-            }
-        };
-
-        Gauge<Long> latestCompletedOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                // add value to topic level metric
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = offsetManager.getCommittedOffset();
-                topicMetrics.totalLatestCompletedOffset += ret;
-                return ret;
-            }
-        };
-
-        Gauge<Long> recordsInPartitionGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
-                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
-                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                // add value to topic level metric
-                Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
-                topicMetrics.totalRecordsInPartitions += ret;
-                return ret;
-            }
-        };
-
-        metrics.put(metricPath + "/" + "spoutLag", spoutLagGauge);
-        metrics.put(metricPath + "/" + "earliestTimeOffset", earliestTimeOffsetGauge);
-        metrics.put(metricPath + "/" + "latestTimeOffset", latestTimeOffsetGauge);
-        metrics.put(metricPath + "/" + "latestEmittedOffset", latestEmittedOffsetGauge);
-        metrics.put(metricPath + "/" + "latestCompletedOffset", latestCompletedOffsetGauge);
-        metrics.put(metricPath + "/" + "recordsInPartition", recordsInPartitionGauge);
-
-        return metrics;
-    }
-
-    private Map<TopicPartition, Long> getBeginningOffsets(Set<TopicPartition> topicPartitions) {
-        Admin admin = adminSupplier.get();
-        if (admin == null) {
-            LOG.error("Kafka admin object is null, returning 0.");
-            return Collections.EMPTY_MAP;
-        }
-
-        Map<TopicPartition, Long> beginningOffsets;
-        try {
-            beginningOffsets = getOffsets(admin, topicPartitions, OffsetSpec.earliest());
-        } catch (RetriableException | ExecutionException | InterruptedException e) {
-            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartition, e);
-            return Collections.EMPTY_MAP;
-        }
-        return beginningOffsets;
-    }
-
-    private Map<TopicPartition, Long> getEndOffsets(Set<TopicPartition> topicPartitions) {
-        Admin admin = adminSupplier.get();
-        if (admin == null) {
-            LOG.error("Kafka admin object is null, returning 0.");
-            return Collections.EMPTY_MAP;
-        }
-
-        Map<TopicPartition, Long> endOffsets;
-        try {
-            endOffsets = getOffsets(admin, topicPartitions, OffsetSpec.latest());
-        } catch (RetriableException | ExecutionException | InterruptedException e) {
-            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartition, e);
-            return Collections.EMPTY_MAP;
-        }
-        return endOffsets;
-    }
-
-    private static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
-        throws InterruptedException, ExecutionException {
-
-        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
-        for (TopicPartition topicPartition : topicPartitions) {
-            offsetSpecMap.put(topicPartition, offsetSpec);
-        }
-        Map<TopicPartition, Long> ret = new HashMap<>();
-        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
-        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
-        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
-        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
-                topicPartitionListOffsetsResultInfoMap.entrySet()) {
-            ret.put(entry.getKey(), entry.getValue().offset());
-        }
-        return ret;
-    }
-}

From ce3acec2ec307cce137661d38a75f8e7c00da857 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Thu, 24 Apr 2025 13:11:32 +0100
Subject: [PATCH 25/28] Update KafkaOffsetMetricManager.java

---
 .../spout/metrics2/KafkaOffsetMetricManager.java | 16 ++++------------
 1 file changed, 4 insertions(+), 12 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index ad82ce4ec8b..dd96f6d1887 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -38,9 +38,8 @@ public class KafkaOffsetMetricManager<K, V> {
     private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
     private final Supplier<Admin> adminSupplier;
     private TopologyContext topologyContext;
+    private KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics;
 
-    private Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
-    private Map<TopicPartition, KafkaOffsetPartitionMetrics> topicPartitionMetricsMap;
 
     public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
                                     Supplier<Admin> adminSupplier,
@@ -48,9 +47,7 @@ public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> off
         this.offsetManagerSupplier = offsetManagerSupplier;
         this.adminSupplier = adminSupplier;
         this.topologyContext = topologyContext;
-
-        this.topicMetricsMap = new HashMap<>();
-        this.topicPartitionMetricsMap = new HashMap<>();
+        
         LOG.info("Running KafkaOffsetMetricManager");
     }
 
@@ -58,14 +55,9 @@ public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignm
 
         KafkaOffsetPartitionAndTopicMetrics topicPartitionMetricSet
                 = new KafkaOffsetPartitionAndTopicMetrics(offsetManagerSupplier, adminSupplier, newAssignment);
+        
+        this.kafkaOffsetPartitionAndTopicMetrics = topicPartitionMetricSet;
         topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
     }
 
-    public Map<TopicPartition, KafkaOffsetPartitionMetrics> getTopicPartitionMetricsMap() {
-        return topicPartitionMetricsMap;
-    }
-
-    public Map<String, KafkaOffsetTopicMetrics> getTopicMetricsMap() {
-        return topicMetricsMap;
-    }
 }

From 781baa10586c4da09e48e4ecb953892efcb6c3b9 Mon Sep 17 00:00:00 2001
From: Rui Abreu <rui.abreu@gmail.com>
Date: Fri, 25 Apr 2025 08:30:30 +0100
Subject: [PATCH 26/28] Adding unit tests

---
 .../metrics2/KafkaOffsetMetricManager.java    |  5 ++++
 .../metrics2/KafkaOffsetTopicMetrics.java     | 23 +---------------
 .../spout/KafkaSpoutReactivationTest.java     | 26 +++++++++++++++++++
 3 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index dd96f6d1887..e871277084a 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -38,6 +38,7 @@ public class KafkaOffsetMetricManager<K, V> {
     private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
     private final Supplier<Admin> adminSupplier;
     private TopologyContext topologyContext;
+
     private KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics;
 
 
@@ -60,4 +61,8 @@ public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignm
         topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
     }
 
+    public KafkaOffsetPartitionAndTopicMetrics getKafkaOffsetPartitionAndTopicMetrics() {
+        return kafkaOffsetPartitionAndTopicMetrics;
+    }
+
 }
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
index cd0fbfc6b2c..900bb86ea63 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
@@ -41,7 +41,7 @@ public class KafkaOffsetTopicMetrics implements MetricSet {
 
     private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetTopicMetrics.class);
 
-    private String topic;
+    private final String topic;
     long totalSpoutLag;
     long totalEarliestTimeOffset;
     long totalLatestTimeOffset;
@@ -115,25 +115,4 @@ public Long getValue() {
         metrics.put(topic + "/" + "totalRecordsInPartitions", totalRecordsInPartitionsGauge);
         return metrics;
     }
-
-    private class TopicMetrics {
-        long totalSpoutLag = 0L;
-        long totalEarliestTimeOffset = 0L;
-        long totalLatestTimeOffset = 0L;
-        long totalLatestEmittedOffset = 0L;
-        long totalLatestCompletedOffset = 0L;
-        long totalRecordsInPartitions = 0L;
-
-        public void incrementTotalSpoutLag(long offset) {
-            totalSpoutLag += offset;
-        }
-
-        public void incrementTotalEarliestTimeOffset(long offset) {
-            totalEarliestTimeOffset += offset;
-        }
-
-        public void incrementTotalLatestTimeOffset(long offset) {
-            totalLatestTimeOffset += offset;
-        }
-    }
 }
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
index b4cab2ace70..4e5702a39d9 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
@@ -29,9 +29,12 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Metric;
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -65,6 +68,7 @@ public class KafkaSpoutReactivationTest {
     private final SpoutOutputCollector collector = mock(SpoutOutputCollector.class);
     private final long commitOffsetPeriodMs = 2_000;
     private Consumer<String, String> consumerSpy;
+    private Admin adminSpy;
     private KafkaSpout<String, String> spout;
     private final int maxPollRecords = 10;
 
@@ -78,9 +82,12 @@ public void prepareSpout(int messageCount, FirstPollOffsetStrategy firstPollOffs
                 .build();
         ClientFactory<String, String> clientFactory = new ClientFactoryDefault<>();
         this.consumerSpy = spy(clientFactory.createConsumer(spoutConfig.getKafkaProps()));
+        this.adminSpy = spy(clientFactory.createAdmin(spoutConfig.getKafkaProps()));
         ClientFactory<String, String> clientFactoryMock = mock(ClientFactory.class);
         when(clientFactoryMock.createConsumer(any()))
             .thenReturn(consumerSpy);
+        when(clientFactoryMock.createAdmin(any()))
+                .thenReturn(adminSpy);
         this.spout = new KafkaSpout<>(spoutConfig, clientFactoryMock, new TopicAssigner());
         SingleTopicKafkaUnitSetupHelper.populateTopicData(kafkaUnitExtension.getKafkaUnit(), SingleTopicKafkaSpoutConfiguration.TOPIC, messageCount);
         SingleTopicKafkaUnitSetupHelper.initializeSpout(spout, conf, topologyContext, collector);
@@ -147,4 +154,23 @@ public void testSpoutShouldResumeWhereItLeftOffWithEarliestStrategy() throws Exc
         //With earliest, the spout should also resume where it left off, rather than restart at the earliest offset.
         doReactivationTest(FirstPollOffsetStrategy.EARLIEST);
     }
+
+    @Test
+    public void testSpoutMustHandleGettingMetricsWhileDeactivated() throws Exception {
+        //Storm will try to get metrics from the spout even while deactivated, the spout must be able to handle this
+        prepareSpout(10, FirstPollOffsetStrategy.UNCOMMITTED_EARLIEST);
+
+        for (int i = 0; i < 5; i++) {
+            KafkaSpoutMessageId msgId = emitOne();
+            spout.ack(msgId);
+        }
+
+        spout.deactivate();
+
+        Map<String, Metric> offsetMetric = spout.getKafkaOffsetMetricManager().getKafkaOffsetPartitionAndTopicMetrics().getMetrics();
+        Long partitionLag = (Long) ((Gauge) offsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/partition_0/spoutLag")).getValue();
+        Long spoutLag = (Long) ((Gauge) offsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/totalSpoutLag")).getValue();
+        assertThat(partitionLag, is(5L));
+        assertThat(spoutLag, is(5L));
+    }
 }

From 4d1759ed1935561a1d1180a749b7421bf6c65d42 Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Mon, 9 Jun 2025 21:02:28 +0100
Subject: [PATCH 27/28] Fixing the way Kafka topic level metrics are computed
 (#3)

* Fixing the way Kafka topic level metrics are computed
---
 .../apache/storm/kafka/spout/KafkaSpout.java  |   3 +-
 .../metrics2/KafkaOffsetMetricManager.java    |  41 +++-
 .../KafkaOffsetPartitionAndTopicMetrics.java  | 232 ------------------
 .../metrics2/KafkaOffsetPartitionMetrics.java | 134 ++++++++++
 .../metrics2/KafkaOffsetTopicMetrics.java     | 161 ++++++++----
 .../kafka/spout/metrics2/KafkaOffsetUtil.java |  92 +++++++
 .../spout/KafkaSpoutReactivationTest.java     |  11 +-
 .../KafkaOffsetPartitionMetricsTest.java      | 163 ++++++++++++
 ....java => KafkaOffsetTopicMetricsTest.java} | 212 +++++-----------
 9 files changed, 610 insertions(+), 439 deletions(-)
 delete mode 100644 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
 create mode 100644 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
 create mode 100644 external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetUtil.java
 create mode 100644 external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionMetricsTest.java
 rename external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/{KafkaOffsetPartitionAndTopicMetricsTest.java => KafkaOffsetTopicMetricsTest.java} (57%)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
index 764ee643b65..0fcb22a0446 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java
@@ -58,7 +58,6 @@
 import org.apache.storm.kafka.spout.internal.OffsetManager;
 import org.apache.storm.kafka.spout.internal.Timer;
 import org.apache.storm.kafka.spout.metrics2.KafkaOffsetMetricManager;
-import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionAndTopicMetrics;
 import org.apache.storm.kafka.spout.subscription.TopicAssigner;
 import org.apache.storm.spout.SpoutOutputCollector;
 import org.apache.storm.task.TopologyContext;
@@ -636,7 +635,7 @@ private void refreshAssignment() {
         boolean partitionChanged = topicAssigner.assignPartitions(consumer, assignedPartitions, rebalanceListener);
         if (partitionChanged && canRegisterMetrics()) {
             LOG.info("Partitions assignments has changed, updating metrics.");
-            kafkaOffsetMetricManager.registerPartitionAndTopicLevelMetrics(assignedPartitions);
+            kafkaOffsetMetricManager.registerMetricsForNewTopicPartitions(assignedPartitions);
         }
     }
 
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index e871277084a..57a1d85de17 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -30,6 +30,7 @@
 import org.slf4j.LoggerFactory;
 
 
+
 /**
  * This class is used to manage both the partition and topic level offset metrics.
  */
@@ -39,8 +40,8 @@ public class KafkaOffsetMetricManager<K, V> {
     private final Supplier<Admin> adminSupplier;
     private TopologyContext topologyContext;
 
-    private KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics;
-
+    private Map<String, KafkaOffsetTopicMetrics> topicMetricsMap;
+    private Map<TopicPartition, KafkaOffsetPartitionMetrics> topicPartitionMetricsMap;
 
     public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
                                     Supplier<Admin> adminSupplier,
@@ -48,21 +49,39 @@ public KafkaOffsetMetricManager(Supplier<Map<TopicPartition, OffsetManager>> off
         this.offsetManagerSupplier = offsetManagerSupplier;
         this.adminSupplier = adminSupplier;
         this.topologyContext = topologyContext;
-        
+
+        this.topicMetricsMap = new HashMap<>();
+        this.topicPartitionMetricsMap = new HashMap<>();
         LOG.info("Running KafkaOffsetMetricManager");
     }
 
-    public void registerPartitionAndTopicLevelMetrics(Set<TopicPartition> newAssignment) {
+    public void registerMetricsForNewTopicPartitions(Set<TopicPartition> newAssignment) {
 
-        KafkaOffsetPartitionAndTopicMetrics topicPartitionMetricSet
-                = new KafkaOffsetPartitionAndTopicMetrics(offsetManagerSupplier, adminSupplier, newAssignment);
-        
-        this.kafkaOffsetPartitionAndTopicMetrics = topicPartitionMetricSet;
-        topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
+        for (TopicPartition topicPartition : newAssignment) {
+            if (!topicPartitionMetricsMap.containsKey(topicPartition)) {
+                LOG.info("Registering metric for topicPartition: {}", topicPartition);
+                // create topic level metrics for given topic if absent
+                String topic = topicPartition.topic();
+                KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
+                if (topicMetrics == null) {
+                    topicMetrics = new KafkaOffsetTopicMetrics(topic, offsetManagerSupplier, adminSupplier, newAssignment);
+                    topicMetricsMap.put(topic, topicMetrics);
+                    topologyContext.registerMetricSet("kafkaOffset", topicMetrics);
+                }
+
+                KafkaOffsetPartitionMetrics topicPartitionMetricSet
+                        = new KafkaOffsetPartitionMetrics<>(offsetManagerSupplier, adminSupplier, topicPartition);
+                topicPartitionMetricsMap.put(topicPartition, topicPartitionMetricSet);
+                topologyContext.registerMetricSet("kafkaOffset", topicPartitionMetricSet);
+            }
+        }
     }
 
-    public KafkaOffsetPartitionAndTopicMetrics getKafkaOffsetPartitionAndTopicMetrics() {
-        return kafkaOffsetPartitionAndTopicMetrics;
+    public Map<TopicPartition, KafkaOffsetPartitionMetrics> getTopicPartitionMetricsMap() {
+        return topicPartitionMetricsMap;
     }
 
+    public Map<String, KafkaOffsetTopicMetrics> getTopicMetricsMap() {
+        return topicMetricsMap;
+    }
 }
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
deleted file mode 100644
index 6f6d766b491..00000000000
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionAndTopicMetrics.java
+++ /dev/null
@@ -1,232 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- *   or more contributor license agreements.  See the NOTICE file
- *   distributed with this work for additional information
- *   regarding copyright ownership.  The ASF licenses this file
- *   to you under the Apache License, Version 2.0 (the
- *   "License"); you may not use this file except in compliance
- *   with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- *   Unless required by applicable law or agreed to in writing, software
- *   distributed under the License is distributed on an "AS IS" BASIS,
- *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *   See the License for the specific language governing permissions and
- *   limitations under the License.
- */
-
-package org.apache.storm.kafka.spout.metrics2;
-
-import com.codahale.metrics.Gauge;
-import com.codahale.metrics.Metric;
-import com.codahale.metrics.MetricSet;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.function.Supplier;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.ListOffsetsResult;
-import org.apache.kafka.clients.admin.OffsetSpec;
-import org.apache.kafka.common.KafkaFuture;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.RetriableException;
-import org.apache.storm.kafka.spout.internal.OffsetManager;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-
-/**
- * This class is used compute the partition and topic level offset metrics.
- * <p>
- * Partition level metrics are:
- * topicName/partition_{number}/earliestTimeOffset //gives beginning offset of the partition
- * topicName/partition_{number}/latestTimeOffset //gives end offset of the partition
- * topicName/partition_{number}/latestEmittedOffset //gives latest emitted offset of the partition from the spout
- * topicName/partition_{number}/latestCompletedOffset //gives latest committed offset of the partition from the spout
- * topicName/partition_{number}/spoutLag // the delta between the latest Offset and latestCompletedOffset
- * topicName/partition_{number}/recordsInPartition // total number of records in the partition
- * </p>
- * <p>
- * Topic level metrics are:
- * topicName/totalEarliestTimeOffset //gives the total beginning offset of all the associated partitions of this spout
- * topicName/totalLatestTimeOffset //gives the total end offset of all the associated partitions of this spout
- * topicName/totalLatestEmittedOffset //gives the total latest emitted offset of all the associated partitions of this spout
- * topicName/totalLatestCompletedOffset //gives the total latest committed offset of all the associated partitions of this spout
- * topicName/spoutLag // total spout lag of all the associated partitions of this spout
- * topicName/totalRecordsInPartitions //total number of records in all the associated partitions of this spout
- * </p>
- */
-public class KafkaOffsetPartitionAndTopicMetrics implements MetricSet {
-
-    private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetPartitionAndTopicMetrics.class);
-    private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
-    private final Supplier<Admin> adminSupplier;
-    private final Set<TopicPartition> assignment;
-
-    public KafkaOffsetPartitionAndTopicMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
-                                               Supplier<Admin> adminSupplier,
-                                               Set<TopicPartition> assignment) {
-        this.offsetManagerSupplier = offsetManagerSupplier;
-        this.adminSupplier = adminSupplier;
-        this.assignment = assignment;
-
-    }
-
-    private static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
-            throws InterruptedException, ExecutionException {
-
-        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
-        for (TopicPartition topicPartition : topicPartitions) {
-            offsetSpecMap.put(topicPartition, offsetSpec);
-        }
-        Map<TopicPartition, Long> ret = new HashMap<>();
-        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
-        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
-        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
-        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
-                topicPartitionListOffsetsResultInfoMap.entrySet()) {
-            ret.put(entry.getKey(), entry.getValue().offset());
-        }
-        return ret;
-    }
-
-    @Override
-    public Map<String, Metric> getMetrics() {
-
-        Map<String, KafkaOffsetTopicMetrics> topicMetricsMap = new HashMap<>();
-
-        Map<String, Metric> metrics = new HashMap<>();
-
-        for (TopicPartition topicPartition : assignment) {
-
-            String topic = topicPartition.topic();
-            KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
-            if (topicMetrics == null) {
-                topicMetrics = new KafkaOffsetTopicMetrics(topic);
-                topicMetricsMap.put(topic, topicMetrics);
-            }
-
-            String metricPath = topicPartition.topic() + "/partition_" + topicPartition.partition();
-            KafkaOffsetTopicMetrics finalTopicMetrics = topicMetrics;
-            Gauge<Long> spoutLagGauge = () -> {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
-                finalTopicMetrics.totalSpoutLag += ret;
-                return ret;
-            };
-
-            Gauge<Long> earliestTimeOffsetGauge = () -> {
-                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
-                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
-                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-
-                Long ret = beginningOffsets.get(topicPartition);
-                finalTopicMetrics.totalEarliestTimeOffset += ret;
-                return ret;
-            };
-
-            Gauge<Long> latestTimeOffsetGauge = () -> {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-
-                Long ret = endOffsets.get(topicPartition);
-                finalTopicMetrics.totalLatestTimeOffset += ret;
-                return ret;
-            };
-
-            Gauge<Long> latestEmittedOffsetGauge = () -> {
-
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = offsetManager.getLatestEmittedOffset();
-                finalTopicMetrics.totalLatestEmittedOffset += ret;
-                return ret;
-            };
-
-            Gauge<Long> latestCompletedOffsetGauge = () -> {
-
-                OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
-                Long ret = offsetManager.getCommittedOffset();
-                finalTopicMetrics.totalLatestCompletedOffset += ret;
-                return ret;
-            };
-
-            Gauge<Long> recordsInPartitionGauge = () -> {
-                Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition));
-                if (endOffsets == null || endOffsets.isEmpty()) {
-                    LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-                Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition));
-                if (beginningOffsets == null || beginningOffsets.isEmpty()) {
-                    LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
-                    return 0L;
-                }
-
-                Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
-                finalTopicMetrics.totalRecordsInPartitions += ret;
-                return ret;
-            };
-
-            metrics.put(metricPath + "/" + "spoutLag", spoutLagGauge);
-            metrics.put(metricPath + "/" + "earliestTimeOffset", earliestTimeOffsetGauge);
-            metrics.put(metricPath + "/" + "latestTimeOffset", latestTimeOffsetGauge);
-            metrics.put(metricPath + "/" + "latestEmittedOffset", latestEmittedOffsetGauge);
-            metrics.put(metricPath + "/" + "latestCompletedOffset", latestCompletedOffsetGauge);
-            metrics.put(metricPath + "/" + "recordsInPartition", recordsInPartitionGauge);
-        }
-
-        for (KafkaOffsetTopicMetrics kafkaOffsetTopicMetrics : topicMetricsMap.values()) {
-            metrics.putAll(kafkaOffsetTopicMetrics.getMetrics());
-        }
-
-        return metrics;
-    }
-
-    private Map<TopicPartition, Long> getBeginningOffsets(Set<TopicPartition> topicPartitions) {
-        Admin admin = adminSupplier.get();
-        if (admin == null) {
-            LOG.error("Kafka admin object is null, returning 0.");
-            return Collections.EMPTY_MAP;
-        }
-
-        Map<TopicPartition, Long> beginningOffsets;
-        try {
-            beginningOffsets = getOffsets(admin, topicPartitions, OffsetSpec.earliest());
-        } catch (RetriableException | ExecutionException | InterruptedException e) {
-            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
-            return Collections.EMPTY_MAP;
-        }
-        return beginningOffsets;
-    }
-
-    private Map<TopicPartition, Long> getEndOffsets(Set<TopicPartition> topicPartitions) {
-        Admin admin = adminSupplier.get();
-        if (admin == null) {
-            LOG.error("Kafka admin object is null, returning 0.");
-            return Collections.EMPTY_MAP;
-        }
-
-        Map<TopicPartition, Long> endOffsets;
-        try {
-            endOffsets = getOffsets(admin, topicPartitions, OffsetSpec.latest());
-        } catch (RetriableException | ExecutionException | InterruptedException e) {
-            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
-            return Collections.EMPTY_MAP;
-        }
-        return endOffsets;
-    }
-}
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
new file mode 100644
index 00000000000..df5b0c8694e
--- /dev/null
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetPartitionMetrics.java
@@ -0,0 +1,134 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.storm.kafka.spout.metrics2;
+
+import static org.apache.storm.kafka.spout.metrics2.KafkaOffsetUtil.getBeginningOffsets;
+import static org.apache.storm.kafka.spout.metrics2.KafkaOffsetUtil.getEndOffsets;
+
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Metric;
+import com.codahale.metrics.MetricSet;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Supplier;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.storm.kafka.spout.internal.OffsetManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+
+
+/**
+ * Partition level metrics.
+ * <p>
+ * topicName/partition_{number}/earliestTimeOffset //gives beginning offset of the partition
+ * topicName/partition_{number}/latestTimeOffset //gives end offset of the partition
+ * topicName/partition_{number}/latestEmittedOffset //gives latest emitted offset of the partition from the spout
+ * topicName/partition_{number}/latestCompletedOffset //gives latest committed offset of the partition from the spout
+ * topicName/partition_{number}/spoutLag // the delta between the latest Offset and latestCompletedOffset
+ * topicName/partition_{number}/recordsInPartition // total number of records in the partition
+ * </p>
+ */
+public class KafkaOffsetPartitionMetrics<K, V> implements MetricSet {
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetPartitionMetrics.class);
+    private final Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
+    private final Supplier<Admin> adminSupplier;
+
+    private TopicPartition topicPartition;
+
+    public KafkaOffsetPartitionMetrics(Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier,
+                                       Supplier<Admin> adminSupplier,
+                                       TopicPartition topicPartition) {
+        this.offsetManagerSupplier = offsetManagerSupplier;
+        this.adminSupplier = adminSupplier;
+        this.topicPartition = topicPartition;
+
+        LOG.info("Running KafkaOffsetMetricSet");
+    }
+
+    @Override
+    public Map<String, Metric> getMetrics() {
+        Map<String, Metric> metrics = new HashMap();
+
+        String metricPath = topicPartition.topic() + "/partition_" + topicPartition.partition();
+        Gauge<Long> spoutLagGauge = () -> {
+            Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+            if (endOffsets == null || endOffsets.isEmpty()) {
+                LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                return 0L;
+            }
+            OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+            return endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
+        };
+
+        Gauge<Long> earliestTimeOffsetGauge = () -> {
+            Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition), adminSupplier);
+            if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                return 0L;
+            }
+            return beginningOffsets.get(topicPartition);
+        };
+
+        Gauge<Long> latestTimeOffsetGauge = () -> {
+            Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+            if (endOffsets == null || endOffsets.isEmpty()) {
+                LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                return 0L;
+            }
+            return endOffsets.get(topicPartition);
+        };
+
+        Gauge<Long> latestEmittedOffsetGauge = () -> {
+            OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+            return offsetManager.getLatestEmittedOffset();
+        };
+
+        Gauge<Long> latestCompletedOffsetGauge = () -> {
+            OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+            return offsetManager.getCommittedOffset();
+        };
+
+        Gauge<Long> recordsInPartitionGauge = () -> {
+            Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+            if (endOffsets == null || endOffsets.isEmpty()) {
+                LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                return 0L;
+            }
+            Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition), adminSupplier);
+            if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                return 0L;
+            }
+            return endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
+        };
+
+        metrics.put(metricPath + "/" + "spoutLag", spoutLagGauge);
+        metrics.put(metricPath + "/" + "earliestTimeOffset", earliestTimeOffsetGauge);
+        metrics.put(metricPath + "/" + "latestTimeOffset", latestTimeOffsetGauge);
+        metrics.put(metricPath + "/" + "latestEmittedOffset", latestEmittedOffsetGauge);
+        metrics.put(metricPath + "/" + "latestCompletedOffset", latestCompletedOffsetGauge);
+        metrics.put(metricPath + "/" + "recordsInPartition", recordsInPartitionGauge);
+
+        return metrics;
+    }
+}
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
index 900bb86ea63..8bdcef37485 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetTopicMetrics.java
@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- *
+ * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
- *
+ * <p>
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -18,14 +18,26 @@
 
 package org.apache.storm.kafka.spout.metrics2;
 
+import static org.apache.storm.kafka.spout.metrics2.KafkaOffsetUtil.getBeginningOffsets;
+import static org.apache.storm.kafka.spout.metrics2.KafkaOffsetUtil.getEndOffsets;
+
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricSet;
+
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.storm.kafka.spout.internal.OffsetManager;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+
 /**
  * Topic level metrics.
  * <p>
@@ -41,23 +53,19 @@ public class KafkaOffsetTopicMetrics implements MetricSet {
 
     private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetTopicMetrics.class);
 
-    private final String topic;
-    long totalSpoutLag;
-    long totalEarliestTimeOffset;
-    long totalLatestTimeOffset;
-    long totalLatestEmittedOffset;
-    long totalLatestCompletedOffset;
-    long totalRecordsInPartitions;
+    private String topic;
+    Set<TopicPartition> assignment;
+    Supplier<Map<TopicPartition, OffsetManager>> offsetManagerSupplier;
+    Supplier<Admin> adminSupplier;
 
-
-    public KafkaOffsetTopicMetrics(String topic) {
+    public KafkaOffsetTopicMetrics(String topic, Supplier<Map<TopicPartition,
+                                           OffsetManager>> offsetManagerSupplier,
+                                   Supplier<Admin> adminSupplier,
+                                   Set<TopicPartition> newAssignment) {
         this.topic = topic;
-        this.totalSpoutLag = 0L;
-        this.totalEarliestTimeOffset = 0L;
-        this.totalLatestTimeOffset = 0L;
-        this.totalLatestEmittedOffset = 0L;
-        this.totalLatestCompletedOffset = 0L;
-        this.totalRecordsInPartitions = 0L;
+        this.assignment = newAssignment;
+        this.offsetManagerSupplier = offsetManagerSupplier;
+        this.adminSupplier = adminSupplier;
         LOG.info("Create KafkaOffsetTopicMetrics for topic: {}", topic);
     }
 
@@ -65,46 +73,117 @@ public KafkaOffsetTopicMetrics(String topic) {
     public Map<String, Metric> getMetrics() {
         Map<String, Metric> metrics = new HashMap();
 
-        Gauge<Long> totalSpoutLagGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalSpoutLag;
+        Gauge<Long> totalSpoutLagGauge = () -> {
+            Long totalSpoutLag = 0L;
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+                    if (endOffsets == null || endOffsets.isEmpty()) {
+                        LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                        return 0L;
+                    }
+                    // add value to topic level metric
+                    OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                    Long ret = endOffsets.get(topicPartition) - offsetManager.getCommittedOffset();
+                    totalSpoutLag += ret;
+                }
             }
+            return totalSpoutLag;
         };
 
-        Gauge<Long> totalEarliestTimeOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalEarliestTimeOffset;
+        Gauge<Long> totalEarliestTimeOffsetGauge = () -> {
+
+            Long totalEarliestTimeOffset = 0L;
+
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition), adminSupplier);
+                    if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                        LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                        return 0L;
+                    }
+                    // add value to topic level metric
+                    Long ret = beginningOffsets.get(topicPartition);
+                    totalEarliestTimeOffset += ret;
+                }
+
             }
+            return totalEarliestTimeOffset;
         };
 
-        Gauge<Long> totalLatestTimeOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalLatestTimeOffset;
+        Gauge<Long> totalLatestTimeOffsetGauge = () -> {
+
+            Long totalLatestTimeOffset = 0L;
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+                    if (endOffsets == null || endOffsets.isEmpty()) {
+                        LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                        return 0L;
+                    }
+                    // add value to topic level metric
+                    Long ret = endOffsets.get(topicPartition);
+                    totalLatestTimeOffset += ret;
+                }
             }
+            return totalLatestTimeOffset;
         };
 
-        Gauge<Long> totalLatestEmittedOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalLatestEmittedOffset;
+        Gauge<Long> totalLatestEmittedOffsetGauge = () -> {
+
+            Long totalLatestEmittedOffset = 0L;
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                    Long ret = offsetManager.getLatestEmittedOffset();
+                    totalLatestEmittedOffset += ret;
+                }
+
             }
+            return totalLatestEmittedOffset;
         };
 
-        Gauge<Long> totalLatestCompletedOffsetGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalLatestCompletedOffset;
+        Gauge<Long> totalLatestCompletedOffsetGauge = () -> {
+
+            Long totalLatestCompletedOffset = 0L;
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    // add value to topic level metric
+                    OffsetManager offsetManager = offsetManagerSupplier.get().get(topicPartition);
+                    Long ret = offsetManager.getCommittedOffset();
+                    totalLatestCompletedOffset += ret;
+                }
             }
+
+            return totalLatestCompletedOffset;
         };
 
-        Gauge<Long> totalRecordsInPartitionsGauge = new Gauge<Long>() {
-            @Override
-            public Long getValue() {
-                return totalRecordsInPartitions;
+        Gauge<Long> totalRecordsInPartitionsGauge = () -> {
+            Long totalRecordsInPartitions = 0L;
+            for (TopicPartition topicPartition : assignment) {
+                String topicOfPartition = topicPartition.topic();
+                if (topicOfPartition.equals(topic)) {
+                    Map<TopicPartition, Long> endOffsets = getEndOffsets(Collections.singleton(topicPartition), adminSupplier);
+                    if (endOffsets == null || endOffsets.isEmpty()) {
+                        LOG.error("Failed to get endOffsets from Kafka for topic partitions: {}.", topicPartition);
+                        return 0L;
+                    }
+                    Map<TopicPartition, Long> beginningOffsets = getBeginningOffsets(Collections.singleton(topicPartition), adminSupplier);
+                    if (beginningOffsets == null || beginningOffsets.isEmpty()) {
+                        LOG.error("Failed to get beginningOffsets from Kafka for topic partitions: {}.", topicPartition);
+                        return 0L;
+                    }
+                    // add value to topic level metric
+                    Long ret = endOffsets.get(topicPartition) - beginningOffsets.get(topicPartition);
+                    totalRecordsInPartitions += ret;
+                }
             }
+            return totalRecordsInPartitions;
         };
 
         metrics.put(topic + "/" + "totalSpoutLag", totalSpoutLagGauge);
diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetUtil.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetUtil.java
new file mode 100644
index 00000000000..c7f9fabc470
--- /dev/null
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetUtil.java
@@ -0,0 +1,92 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.storm.kafka.spout.metrics2;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Supplier;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ListOffsetsResult;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.RetriableException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class KafkaOffsetUtil {
+
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaOffsetUtil.class);
+
+    public static Map<TopicPartition, Long> getBeginningOffsets(Set<TopicPartition> topicPartitions, Supplier<Admin> adminSupplier) {
+        Admin admin = adminSupplier.get();
+        if (admin == null) {
+            LOG.error("Kafka admin object is null, returning 0.");
+            return Collections.EMPTY_MAP;
+        }
+
+        Map<TopicPartition, Long> beginningOffsets;
+        try {
+            beginningOffsets = getOffsets(admin, topicPartitions, OffsetSpec.earliest());
+        } catch (RetriableException | ExecutionException | InterruptedException e) {
+            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
+            return Collections.EMPTY_MAP;
+        }
+        return beginningOffsets;
+    }
+
+    public static Map<TopicPartition, Long> getEndOffsets(Set<TopicPartition> topicPartitions, Supplier<Admin> adminSupplier) {
+        Admin admin = adminSupplier.get();
+        if (admin == null) {
+            LOG.error("Kafka admin object is null, returning 0.");
+            return Collections.EMPTY_MAP;
+        }
+
+        Map<TopicPartition, Long> endOffsets;
+        try {
+            endOffsets = getOffsets(admin, topicPartitions, OffsetSpec.latest());
+        } catch (RetriableException | ExecutionException | InterruptedException e) {
+            LOG.error("Failed to get offset from Kafka for topic partitions: {}.", topicPartitions, e);
+            return Collections.EMPTY_MAP;
+        }
+        return endOffsets;
+    }
+
+    public static Map<TopicPartition, Long> getOffsets(Admin admin, Set<TopicPartition> topicPartitions, OffsetSpec offsetSpec)
+            throws InterruptedException, ExecutionException {
+
+        Map<TopicPartition, OffsetSpec> offsetSpecMap = new HashMap<>();
+        for (TopicPartition topicPartition : topicPartitions) {
+            offsetSpecMap.put(topicPartition, offsetSpec);
+        }
+        Map<TopicPartition, Long> ret = new HashMap<>();
+        ListOffsetsResult listOffsetsResult = admin.listOffsets(offsetSpecMap);
+        KafkaFuture<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> all = listOffsetsResult.all();
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = all.get();
+        for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry :
+                topicPartitionListOffsetsResultInfoMap.entrySet()) {
+            ret.put(entry.getKey(), entry.getValue().offset());
+        }
+        return ret;
+    }
+}
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
index 4e5702a39d9..b7ae2824b7f 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/KafkaSpoutReactivationTest.java
@@ -164,13 +164,14 @@ public void testSpoutMustHandleGettingMetricsWhileDeactivated() throws Exception
             KafkaSpoutMessageId msgId = emitOne();
             spout.ack(msgId);
         }
-
         spout.deactivate();
 
-        Map<String, Metric> offsetMetric = spout.getKafkaOffsetMetricManager().getKafkaOffsetPartitionAndTopicMetrics().getMetrics();
-        Long partitionLag = (Long) ((Gauge) offsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/partition_0/spoutLag")).getValue();
-        Long spoutLag = (Long) ((Gauge) offsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/totalSpoutLag")).getValue();
+        Map<String, Metric> partitionsOffsetMetric = spout.getKafkaOffsetMetricManager().getTopicPartitionMetricsMap().get(new TopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC ,0)).getMetrics();
+        Long partitionLag = (Long) ((Gauge) partitionsOffsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/partition_0/spoutLag")).getValue();
         assertThat(partitionLag, is(5L));
-        assertThat(spoutLag, is(5L));
+
+        Map<String, Metric> topicOffsetMetric = spout.getKafkaOffsetMetricManager().getTopicMetricsMap().get(SingleTopicKafkaSpoutConfiguration.TOPIC).getMetrics();
+        Long totalSpoutLag = (Long) ((Gauge) topicOffsetMetric.get(SingleTopicKafkaSpoutConfiguration.TOPIC + "/totalSpoutLag")).getValue();
+        assertThat(totalSpoutLag, is(5L));
     }
 }
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionMetricsTest.java
new file mode 100644
index 00000000000..ee085858e86
--- /dev/null
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionMetricsTest.java
@@ -0,0 +1,163 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ *   or more contributor license agreements.  See the NOTICE file
+ *   distributed with this work for additional information
+ *   regarding copyright ownership.  The ASF licenses this file
+ *   to you under the Apache License, Version 2.0 (the
+ *   "License"); you may not use this file except in compliance
+ *   with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
+package org.apache.storm.kafka.spout.metric2;
+
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
+
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Metric;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ListOffsetsResult;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.storm.kafka.spout.internal.OffsetManager;
+import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionMetrics;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@ExtendWith(MockitoExtension.class)
+public class KafkaOffsetPartitionMetricsTest {
+
+    private Set<TopicPartition> assignment;
+    private Admin admin = mock(Admin.class);
+    private HashMap<TopicPartition, OffsetManager> offsetManagers;
+    private KafkaFuture kafkaFuture = mock(KafkaFuture.class);
+
+    @BeforeEach
+    public void initializeTests() {
+        reset(admin, kafkaFuture);
+
+    }
+
+    @Test
+    public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionException, InterruptedException {
+
+        TopicPartition topicAPartition1 = new TopicPartition("topicA", 1);
+
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1LatestListOffsetsResultInfo);
+
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
+
+        ListOffsetsResult listOffsetsResult = mock(ListOffsetsResult.class);
+        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
+
+        admin = mock(Admin.class);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
+
+        OffsetManager offsetManagertopicAPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition1.getCommittedOffset()).thenReturn(90L);
+
+
+        offsetManagers = new HashMap<>();
+        offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
+
+        KafkaOffsetPartitionMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, topicAPartition1);
+        Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1 = (Gauge) result.get("topicA/partition_1/spoutLag");
+        assertEquals(10L, g1.getValue());
+
+
+        //get partition records
+
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1EarliestListOffsetsResultInfo);
+
+        //mock consecutive calls. Each call to the recordsInPartition gauge will call kafkaFuture.get() twice
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap);
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        g1 = (Gauge) result.get("topicA/partition_1/recordsInPartition");
+        assertEquals(99L, g1.getValue());
+    }
+
+    @Test
+    public void registerMetricsGetEarliestAndLatest() throws ExecutionException, InterruptedException {
+
+        TopicPartition topicAPartition1 = new TopicPartition("topicA", 1);
+
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionEarliestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionEarliestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1EarliestListOffsetsResultInfo);
+
+        when(kafkaFuture.get()).thenReturn(topicPartitionEarliestListOffsetsResultInfoMap);
+
+        ListOffsetsResult listOffsetsResult = mock(ListOffsetsResult.class);
+        when(listOffsetsResult.all()).thenReturn(kafkaFuture);
+
+        admin = mock(Admin.class);
+        when(admin.listOffsets(anyMap())).thenReturn(listOffsetsResult);
+
+        OffsetManager offsetManagertopicAPartition1 = mock(OffsetManager.class);
+        when(offsetManagertopicAPartition1.getLatestEmittedOffset()).thenReturn(50L);
+        when(offsetManagertopicAPartition1.getCommittedOffset()).thenReturn(40L);
+
+
+        offsetManagers = new HashMap<>();
+        offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
+
+        assignment = new HashSet<>();
+        assignment.add(topicAPartition1);
+
+        KafkaOffsetPartitionMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, topicAPartition1);
+        Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+        Gauge g1 = (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
+
+        assertEquals(g1.getValue(), 1L);
+
+        //get the latest offsets
+
+        ListOffsetsResult.ListOffsetsResultInfo topicAPartition1LatestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(100, System.currentTimeMillis(), Optional.empty());
+
+
+        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionLatestListOffsetsResultInfoMap = new HashMap<>();
+        topicPartitionLatestListOffsetsResultInfoMap.put(topicAPartition1, topicAPartition1LatestListOffsetsResultInfo);
+
+        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
+
+        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
+
+        g1 = (Gauge) result.get("topicA/partition_1/latestTimeOffset");
+        assertEquals(100L, g1.getValue());
+
+        g1 = (Gauge) result.get("topicA/partition_1/latestEmittedOffset");
+        assertEquals(50L, g1.getValue());
+
+
+        g1 = (Gauge) result.get("topicA/partition_1/latestCompletedOffset");
+        assertEquals(40L, g1.getValue());
+
+    }
+}
diff --git a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetTopicMetricsTest.java
similarity index 57%
rename from external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
rename to external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetTopicMetricsTest.java
index cc774e8b24d..26462503389 100644
--- a/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetPartitionAndTopicMetricsTest.java
+++ b/external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/metric2/KafkaOffsetTopicMetricsTest.java
@@ -1,31 +1,23 @@
-/*
+/**
  * Licensed to the Apache Software Foundation (ASF) under one
- *   or more contributor license agreements.  See the NOTICE file
- *   distributed with this work for additional information
- *   regarding copyright ownership.  The ASF licenses this file
- *   to you under the Apache License, Version 2.0 (the
- *   "License"); you may not use this file except in compliance
- *   with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- *   Unless required by applicable law or agreed to in writing, software
- *   distributed under the License is distributed on an "AS IS" BASIS,
- *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *   See the License for the specific language governing permissions and
- *   limitations under the License.
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package org.apache.storm.kafka.spout.metric2;
 
-
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Metric;
 import org.apache.kafka.clients.admin.Admin;
@@ -33,17 +25,22 @@
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.storm.kafka.spout.internal.OffsetManager;
-import org.apache.storm.kafka.spout.metrics2.KafkaOffsetPartitionAndTopicMetrics;
+import org.apache.storm.kafka.spout.metrics2.KafkaOffsetTopicMetrics;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import java.util.*;
+import java.util.concurrent.ExecutionException;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.Mockito.*;
 
 @ExtendWith(MockitoExtension.class)
-public class KafkaOffsetPartitionAndTopicMetricsTest {
+public class KafkaOffsetTopicMetricsTest {
 
     private Set<TopicPartition> assignment;
     private Admin admin = mock(Admin.class);
@@ -53,7 +50,6 @@ public class KafkaOffsetPartitionAndTopicMetricsTest {
     @BeforeEach
     public void initializeTests() {
         reset(admin, kafkaFuture);
-
     }
 
     @Test
@@ -89,17 +85,9 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
         OffsetManager offsetManagertopicAPartition2 = mock(OffsetManager.class);
         when(offsetManagertopicAPartition2.getCommittedOffset()).thenReturn(170L);
 
-        OffsetManager offsetManagertopicBPartition1 = mock(OffsetManager.class);
-        when(offsetManagertopicBPartition1.getCommittedOffset()).thenReturn(200L);
-
-        OffsetManager offsetManagertopicBPartition2 = mock(OffsetManager.class);
-        when(offsetManagertopicBPartition2.getCommittedOffset()).thenReturn(350L);
-
         offsetManagers = new HashMap<>();
         offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
         offsetManagers.put(topicAPartition2, offsetManagertopicAPartition2);
-        offsetManagers.put(topicBPartition1, offsetManagertopicBPartition1);
-        offsetManagers.put(topicBPartition2, offsetManagertopicBPartition2);
 
         assignment = new HashSet<>();
         assignment.add(topicAPartition1);
@@ -108,40 +96,18 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
         assignment.add(topicBPartition2);
 
 
-        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
-        Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        Gauge g1 = (Gauge) result.get("topicA/partition_1/spoutLag");
-        Gauge g2 = (Gauge) result.get("topicA/partition_2/spoutLag");
-        Gauge g3 = (Gauge) result.get("topicB/partition_1/spoutLag");
-        Gauge g4 = (Gauge) result.get("topicB/partition_2/spoutLag");
-        assertEquals(10L, g1.getValue());
-        assertEquals(30L, g2.getValue());
-        assertEquals(100L, g3.getValue());
-        assertEquals(50L, g4.getValue());
-
-        Gauge gATotal = (Gauge) result.get("topicA/totalSpoutLag");
-        assertEquals(40L, gATotal.getValue());
-        Gauge gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
-        assertEquals(150L, gBTotal.getValue());
-
-        //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
-
-        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        g1 = (Gauge) result.get("topicA/partition_1/spoutLag");
-        g2 = (Gauge) result.get("topicA/partition_2/spoutLag");
-        g3 = (Gauge) result.get("topicB/partition_1/spoutLag");
-        g4 = (Gauge) result.get("topicB/partition_2/spoutLag");
-        assertEquals(g1.getValue(), 10L);
-        assertEquals(g2.getValue(), 30L);
-        assertEquals(g3.getValue(), 100L);
-        assertEquals(g4.getValue(), 50L);
-
-        gATotal = (Gauge) result.get("topicA/totalSpoutLag");
-        assertEquals(gATotal.getValue(), 40L);
-        gBTotal = (Gauge) result.get("topicB/totalSpoutLag");
-        assertEquals(gBTotal.getValue(), 150L);
-
-        //get partition records
+        KafkaOffsetTopicMetrics kafkaOffsetTopicMetricsA = new KafkaOffsetTopicMetrics("topicA", () -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
+        Map<String, Metric> result = kafkaOffsetTopicMetricsA.getMetrics();
+        Gauge g1 = (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(40L, g1.getValue());
+
+        //get again the values from the Gauge. Values cannot change
+        g1 = (Gauge) result.get("topicA/totalSpoutLag");
+        assertEquals(40L, g1.getValue());
+
+        assertNull(result.get("topicB/totalSpoutLag"));
+
+        //get topic records
 
         ListOffsetsResult.ListOffsetsResultInfo topicAPartition1EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(1, System.currentTimeMillis(), Optional.empty());
         ListOffsetsResult.ListOffsetsResultInfo topicAPartition2EarliestListOffsetsResultInfo = new ListOffsetsResult.ListOffsetsResultInfo(2, System.currentTimeMillis(), Optional.empty());
@@ -155,26 +121,20 @@ public void registerMetricsGetSpoutLagAndPartitionRecords() throws ExecutionExce
         topicPartitionEarliestListOffsetsResultInfoMap.put(topicBPartition2, topicBPartition2EarliestListOffsetsResultInfo);
 
         //mock consecutive calls. Each call to the recordsInPartition gauge will call kafkaFuture.get() twice
-        when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
+        when(kafkaFuture.get()).thenReturn(
+                topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
                 topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
                 topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap,
                 topicPartitionLatestListOffsetsResultInfoMap, topicPartitionEarliestListOffsetsResultInfoMap);
 
-        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        g1 = (Gauge) result.get("topicA/partition_1/recordsInPartition");
-        g2 = (Gauge) result.get("topicA/partition_2/recordsInPartition");
-        g3 = (Gauge) result.get("topicB/partition_1/recordsInPartition");
-        g4 = (Gauge) result.get("topicB/partition_2/recordsInPartition");
-        assertEquals(99L, g1.getValue());
-        assertEquals(198L, g2.getValue());
-        assertEquals(297L, g3.getValue());
-        assertEquals(396L, g4.getValue());
-
-        gATotal = (Gauge) result.get("topicA/totalRecordsInPartitions");
+        Gauge gATotal = (Gauge) result.get("topicA/totalRecordsInPartitions");
+        assertEquals(297L, gATotal.getValue());
+
+        //get again the values from the Gauge. Values cannot change
+         gATotal = (Gauge) result.get("topicA/totalRecordsInPartitions");
         assertEquals(297L, gATotal.getValue());
-        gBTotal = (Gauge) result.get("topicB/totalRecordsInPartitions");
-        assertEquals(693L, gBTotal.getValue());
 
+        assertNull(result.get("topicB/totalRecordsInPartitions"));
     }
 
     @Test
@@ -212,19 +172,10 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
         when(offsetManagertopicAPartition2.getLatestEmittedOffset()).thenReturn(100L);
         when(offsetManagertopicAPartition2.getCommittedOffset()).thenReturn(90L);
 
-        OffsetManager offsetManagertopicBPartition1 = mock(OffsetManager.class);
-        when(offsetManagertopicBPartition1.getLatestEmittedOffset()).thenReturn(150L);
-        when(offsetManagertopicBPartition1.getCommittedOffset()).thenReturn(149L);
-
-        OffsetManager offsetManagertopicBPartition2 = mock(OffsetManager.class);
-        when(offsetManagertopicBPartition2.getLatestEmittedOffset()).thenReturn(200L);
-        when(offsetManagertopicBPartition2.getCommittedOffset()).thenReturn(200L);
 
         offsetManagers = new HashMap<>();
         offsetManagers.put(topicAPartition1, offsetManagertopicAPartition1);
         offsetManagers.put(topicAPartition2, offsetManagertopicAPartition2);
-        offsetManagers.put(topicBPartition1, offsetManagertopicBPartition1);
-        offsetManagers.put(topicBPartition2, offsetManagertopicBPartition2);
 
         assignment = new HashSet<>();
         assignment.add(topicAPartition1);
@@ -232,39 +183,18 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
         assignment.add(topicBPartition1);
         assignment.add(topicBPartition2);
 
-        KafkaOffsetPartitionAndTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetPartitionAndTopicMetrics(() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
+        KafkaOffsetTopicMetrics kafkaOffsetPartitionAndTopicMetrics = new KafkaOffsetTopicMetrics("topicA",() -> Collections.unmodifiableMap(offsetManagers), () -> admin, assignment);
         Map<String, Metric> result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        Gauge g1 = (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
-        Gauge g2 = (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
-        Gauge g3 = (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
-        Gauge g4 = (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
-        assertEquals(g1.getValue(), 1L);
-        assertEquals(g2.getValue(), 1L);
-        assertEquals(g3.getValue(), 1L);
-        assertEquals(g4.getValue(), 1L);
 
         Gauge gATotal = (Gauge) result.get("topicA/totalEarliestTimeOffset");
         assertEquals(2L, gATotal.getValue());
-        Gauge gBTotal = (Gauge) result.get("topicB/totalEarliestTimeOffset");
-        assertEquals(2L, gBTotal.getValue());
-
-        //get the metrics a second time. Values should be the same. In particular, the total values for the topic should not accumulate. Each call to getMetrics should reset the total values.
+        assertNull(result.get("topicB/totalEarliestTimeOffset"));
 
-        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-
-        g1 = (Gauge) result.get("topicA/partition_1/earliestTimeOffset");
-        g2 = (Gauge) result.get("topicA/partition_2/earliestTimeOffset");
-        g3 = (Gauge) result.get("topicB/partition_1/earliestTimeOffset");
-        g4 = (Gauge) result.get("topicB/partition_2/earliestTimeOffset");
-        assertEquals(g1.getValue(), 1L);
-        assertEquals(g2.getValue(), 1L);
-        assertEquals(g3.getValue(), 1L);
-        assertEquals(g4.getValue(), 1L);
+        //get the metrics a second time. Values should be the same
 
         gATotal = (Gauge) result.get("topicA/totalEarliestTimeOffset");
         assertEquals(2L, gATotal.getValue());
-        gBTotal = (Gauge) result.get("topicB/totalEarliestTimeOffset");
-        assertEquals(2L, gBTotal.getValue());
+        assertNull(result.get("topicB/totalEarliestTimeOffset"));
 
         //get the latest offsets
 
@@ -281,47 +211,33 @@ public void registerMetricsGetEarliestAndLatest() throws ExecutionException, Int
 
         when(kafkaFuture.get()).thenReturn(topicPartitionLatestListOffsetsResultInfoMap);
 
-        result = kafkaOffsetPartitionAndTopicMetrics.getMetrics();
-        g1 = (Gauge) result.get("topicA/partition_1/latestTimeOffset");
-        g2 = (Gauge) result.get("topicA/partition_2/latestTimeOffset");
-        g3 = (Gauge) result.get("topicB/partition_1/latestTimeOffset");
-        g4 = (Gauge) result.get("topicB/partition_2/latestTimeOffset");
-        assertEquals(100L, g1.getValue());
-        assertEquals(200L, g2.getValue());
-        assertEquals(300L, g3.getValue());
-        assertEquals(400L, g4.getValue());
+        gATotal = (Gauge) result.get("topicA/totalLatestTimeOffset");
+        assertEquals(300L, gATotal.getValue());
+        assertNull((result.get("topicB/totalLatestTimeOffset")));
+
+        gATotal = (Gauge) result.get("topicA/totalLatestEmittedOffset");
+        assertEquals(150L, gATotal.getValue());
+        assertNull( result.get("topicB/totalLatestEmittedOffset"));
+
+
+        gATotal = (Gauge) result.get("topicA/totalLatestCompletedOffset");
+        assertEquals(130L, gATotal.getValue());
+        assertNull(result.get("topiBA/totalLatestCompletedOffset"));
+
+        //get the metrics a second time. Values should be the same
 
         gATotal = (Gauge) result.get("topicA/totalLatestTimeOffset");
         assertEquals(300L, gATotal.getValue());
-        gBTotal = (Gauge) result.get("topicB/totalLatestTimeOffset");
-        assertEquals(700L, gBTotal.getValue());
-
-        g1 = (Gauge) result.get("topicA/partition_1/latestEmittedOffset");
-        g2 = (Gauge) result.get("topicA/partition_2/latestEmittedOffset");
-        g3 = (Gauge) result.get("topicB/partition_1/latestEmittedOffset");
-        g4 = (Gauge) result.get("topicB/partition_2/latestEmittedOffset");
-        assertEquals(50L, g1.getValue());
-        assertEquals(100L, g2.getValue());
-        assertEquals(150L, g3.getValue());
-        assertEquals(200L, g4.getValue());
+        assertNull((result.get("topicB/totalLatestTimeOffset")));
 
         gATotal = (Gauge) result.get("topicA/totalLatestEmittedOffset");
         assertEquals(150L, gATotal.getValue());
-        gBTotal = (Gauge) result.get("topicB/totalLatestEmittedOffset");
-        assertEquals(350L, gBTotal.getValue());
+        assertNull( result.get("topicB/totalLatestEmittedOffset"));
 
-        g1 = (Gauge) result.get("topicA/partition_1/latestCompletedOffset");
-        g2 = (Gauge) result.get("topicA/partition_2/latestCompletedOffset");
-        g3 = (Gauge) result.get("topicB/partition_1/latestCompletedOffset");
-        g4 = (Gauge) result.get("topicB/partition_2/latestCompletedOffset");
-        assertEquals(40L, g1.getValue());
-        assertEquals(90L, g2.getValue());
-        assertEquals(149L, g3.getValue());
-        assertEquals(200L, g4.getValue());
 
         gATotal = (Gauge) result.get("topicA/totalLatestCompletedOffset");
         assertEquals(130L, gATotal.getValue());
-        gBTotal = (Gauge) result.get("topicB/totalLatestCompletedOffset");
-        assertEquals(349L, gBTotal.getValue());
+        assertNull(result.get("topiBA/totalLatestCompletedOffset"));
+
     }
 }

From ed81beb5bedc163a85b6a38596b884934377345c Mon Sep 17 00:00:00 2001
From: reiabreu <rui.abreu@gmail.com>
Date: Wed, 11 Jun 2025 18:29:55 +0100
Subject: [PATCH 28/28] Update KafkaOffsetMetricManager.java

---
 .../storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
index 57a1d85de17..64c42832cf3 100644
--- a/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
+++ b/external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/metrics2/KafkaOffsetMetricManager.java
@@ -64,6 +64,7 @@ public void registerMetricsForNewTopicPartitions(Set<TopicPartition> newAssignme
                 String topic = topicPartition.topic();
                 KafkaOffsetTopicMetrics topicMetrics = topicMetricsMap.get(topic);
                 if (topicMetrics == null) {
+                    LOG.info("Registering metric for topic: {}", topic);
                     topicMetrics = new KafkaOffsetTopicMetrics(topic, offsetManagerSupplier, adminSupplier, newAssignment);
                     topicMetricsMap.put(topic, topicMetrics);
                     topologyContext.registerMetricSet("kafkaOffset", topicMetrics);
