KAFKA-19722: Adding missing metric assigned-partitions for new consumer (#20557)

Adding the missing metric to track the number of partitions assigned.
This metric should be registered whenever the consumer is using a
groupId, and should track the number of partitions from the subscription
state, regardless of the subscription type (manual or automatic).

This PR registers the missing metric as part of the
ConsumerRebalanceMetricsManager setup. This manager is created if there
is a group ID, and reused by the consumer membershipMgr and the
streamsMemberhipMgr, so we ensure that the metric is registered for the
new consumer and streams.

Reviewers: Lucas Brutschy <[email protected]>, TengYao Chi
 <[email protected]>

Author: lianetm

clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManager.java

@@ -163,7 +163,7 @@ public ConsumerMembershipManager(String groupId,
             logContext,
             backgroundEventHandler,
             time,
-            new ConsumerRebalanceMetricsManager(metrics),
+            new ConsumerRebalanceMetricsManager(metrics, subscriptions),
             autoCommitEnabled);
     }
 

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

@@ -296,7 +296,7 @@ public StreamsMembershipManager(final String groupId,
         this.backgroundEventHandler = backgroundEventHandler;
         this.streamsRebalanceData = streamsRebalanceData;
         this.subscriptionState = subscriptionState;
-        metricsManager = new ConsumerRebalanceMetricsManager(metrics);
+        metricsManager = new ConsumerRebalanceMetricsManager(metrics, subscriptionState);
         this.time = time;
     }
 

clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java

@@ -477,7 +477,7 @@ public synchronized List<TopicPartition> assignedPartitionsList() {
      * Provides the number of assigned partitions in a thread safe manner.
      * @return the number of assigned partitions.
      */
-    synchronized int numAssignedPartitions() {
+    public synchronized int numAssignedPartitions() {
         return this.assignment.size();
     }
 

clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManager.java

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.clients.consumer.internals.metrics;
 
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.Metrics;
@@ -27,7 +29,9 @@
 import org.apache.kafka.common.metrics.stats.Rate;
 import org.apache.kafka.common.metrics.stats.WindowedCount;
 
+import java.util.Collection;
 import java.util.concurrent.TimeUnit;
+import java.util.regex.Pattern;
 
 import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
 import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
@@ -44,11 +48,14 @@ public final class ConsumerRebalanceMetricsManager extends RebalanceMetricsManag
     public final MetricName lastRebalanceSecondsAgo;
     public final MetricName failedRebalanceTotal;
     public final MetricName failedRebalanceRate;
+    public final MetricName assignedPartitionsCount;
     private long lastRebalanceEndMs = -1L;
     private long lastRebalanceStartMs = -1L;
+    private final Metrics metrics;
 
-    public ConsumerRebalanceMetricsManager(Metrics metrics) {
+    public ConsumerRebalanceMetricsManager(Metrics metrics, SubscriptionState subscriptions) {
         super(CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX);
+        this.metrics = metrics;
 
         rebalanceLatencyAvg = createMetric(metrics, "rebalance-latency-avg",
                 "The average time in ms taken for a group to complete a rebalance");
@@ -64,6 +71,9 @@ public ConsumerRebalanceMetricsManager(Metrics metrics) {
                 "The total number of failed rebalance events");
         failedRebalanceRate = createMetric(metrics, "failed-rebalance-rate-per-hour",
                 "The number of failed rebalance events per hour");
+        assignedPartitionsCount = createMetric(metrics, "assigned-partitions",
+                "The number of partitions currently assigned to this consumer");
+        registerAssignedPartitionCount(subscriptions);
 
         successfulRebalanceSensor = metrics.sensor("rebalance-latency");
         successfulRebalanceSensor.add(rebalanceLatencyAvg, new Avg());
@@ -106,4 +116,15 @@ public void maybeRecordRebalanceFailed() {
     public boolean rebalanceStarted() {
         return lastRebalanceStartMs > lastRebalanceEndMs;
     }
+
+    /**
+     * Register metric to track the number of assigned partitions.
+     * It will consider partitions assigned to the consumer
+     * regardless of whether they were assigned via {@link KafkaConsumer#subscribe(Pattern)} or
+     * {@link KafkaConsumer#assign(Collection)}
+     */
+    private void registerAssignedPartitionCount(SubscriptionState subscriptions) {
+        Measurable numParts = (config, now) -> subscriptions.numAssignedPartitions();
+        metrics.addMetric(assignedPartitionsCount, numParts);
+    }
 }

clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java

@@ -155,6 +155,8 @@
 
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.clients.consumer.internals.ClassicKafkaConsumer.DEFAULT_REASON;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
 import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
 import static org.apache.kafka.common.utils.Utils.propsToMap;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
@@ -271,6 +273,35 @@ public void testSubscribingCustomMetricsWithSameNameDoesntAffectConsumerMetrics(
         }
     }
 
+    @ParameterizedTest
+    @EnumSource(GroupProtocol.class)
+    public void testAssignedPartitionsMetrics(GroupProtocol groupProtocol) throws InterruptedException {
+        consumer = newConsumer(groupProtocol, time, mock(KafkaClient.class), subscription,
+                mock(ConsumerMetadata.class), assignor, false, groupInstanceId);
+        Metrics metrics = consumer.metricsRegistry();
+
+        // This metric is added in the background thread for the AsyncConsumer, so waiting on it to avoid flakiness.
+        TestUtils.waitForCondition(() -> getMetric(metrics, "assigned-partitions") != null,
+                "Consumer should register the assigned-partitions metric");
+        assertNotNull(getMetric(metrics, "assigned-partitions"));
+        assertEquals(0.0d, getMetric(metrics, "assigned-partitions").metricValue());
+
+        subscription.assignFromUser(Set.of(tp0));
+        assertEquals(1.0d, getMetric(metrics, "assigned-partitions").metricValue());
+
+        subscription.assignFromUser(Set.of(tp0, tp1));
+        assertEquals(2.0d, getMetric(metrics, "assigned-partitions").metricValue());
+
+        subscription.unsubscribe();
+        subscription.subscribe(Set.of(topic), Optional.empty());
+        subscription.assignFromSubscribed(Set.of(tp0));
+        assertEquals(1.0d, getMetric(metrics, "assigned-partitions").metricValue());
+    }
+
+    private KafkaMetric getMetric(Metrics metrics, String name) {
+        return metrics.metrics().get(metrics.metricName(name, CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX));
+    }
+
     @ParameterizedTest
     @EnumSource(GroupProtocol.class)
     public void testUnsubscribingCustomMetricsWithSameNameDoesntAffectConsumerMetrics(GroupProtocol groupProtocol) {

clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManagerTest.java

@@ -71,6 +71,8 @@
 
 import static org.apache.kafka.clients.consumer.internals.AbstractMembershipManager.TOPIC_PARTITION_COMPARATOR;
 import static org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer.invokeRebalanceCallbacks;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
 import static org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
@@ -125,7 +127,7 @@ public void setup() {
         time = new MockTime(0);
         backgroundEventHandler = new BackgroundEventHandler(backgroundEventQueue, time, mock(AsyncConsumerMetrics.class));
         metrics = new Metrics(time);
-        rebalanceMetricsManager = new ConsumerRebalanceMetricsManager(metrics);
+        rebalanceMetricsManager = new ConsumerRebalanceMetricsManager(metrics, subscriptionState);
 
         when(commitRequestManager.maybeAutoCommitSyncBeforeRebalance(anyLong())).thenReturn(CompletableFuture.completedFuture(null));
     }
@@ -181,6 +183,15 @@ public void testMembershipManagerRackId() {
         assertEquals(Optional.of("rack1"), membershipManager.rackId());
     }
 
+    @Test
+    public void testAssignedPartitionCountMetricRegistered() {
+        MetricName metricName = metrics.metricName(
+                "assigned-partitions",
+                CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX
+        );
+        assertNotNull(metrics.metric(metricName), "Metric assigned-partitions should have been registered");
+    }
+
     @Test
     public void testMembershipManagerInitSupportsEmptyGroupInstanceId() {
         createMembershipManagerJoiningGroup();

clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java

@@ -67,6 +67,7 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -131,6 +132,15 @@ public void setup() {
         verifyInStateUnsubscribed(membershipManager);
     }
 
+    @Test
+    public void testAssignedPartitionCountMetricRegistered() {
+        MetricName metricName = metrics.metricName(
+                "assigned-partitions",
+                CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX
+        );
+        assertNotNull(metrics.metric(metricName), "Metric assigned-partitions should have been registered");
+    }
+
     @Test
     public void testUnexpectedErrorInHeartbeatResponse() {
         final String errorMessage = "Nobody expects the Spanish Inquisition!";

clients/src/test/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManagerTest.java

@@ -0,0 +1,63 @@
+/*
+ * 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.kafka.clients.consumer.internals.metrics;
+
+import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy;
+import org.apache.kafka.clients.consumer.internals.SubscriptionState;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Optional;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.mockito.Mockito.mock;
+
+
+class ConsumerRebalanceMetricsManagerTest {
+
+    private final Time time = new MockTime();
+    private final Metrics metrics = new Metrics(time);
+
+    @Test
+    public void testAssignedPartitionCountMetric() {
+        SubscriptionState subscriptionState = new SubscriptionState(mock(LogContext.class), AutoOffsetResetStrategy.EARLIEST);
+        ConsumerRebalanceMetricsManager consumerRebalanceMetricsManager = new ConsumerRebalanceMetricsManager(metrics, subscriptionState);
+
+        assertNotNull(metrics.metric(consumerRebalanceMetricsManager.assignedPartitionsCount), "Metric assigned-partitions has not been registered as expected");
+
+        // Check for manually assigned partitions
+        subscriptionState.assignFromUser(Set.of(new TopicPartition("topic", 0), new TopicPartition("topic", 1)));
+        assertEquals(2.0d, metrics.metric(consumerRebalanceMetricsManager.assignedPartitionsCount).metricValue());
+        subscriptionState.assignFromUser(Set.of());
+        assertEquals(0.0d, metrics.metric(consumerRebalanceMetricsManager.assignedPartitionsCount).metricValue());
+
+        subscriptionState.unsubscribe();
+        assertEquals(0.0d, metrics.metric(consumerRebalanceMetricsManager.assignedPartitionsCount).metricValue());
+
+        // Check for automatically assigned partitions
+        subscriptionState.subscribe(Set.of("topic"), Optional.empty());
+        subscriptionState.assignFromSubscribed(Set.of(new TopicPartition("topic", 0)));
+        assertEquals(1.0d, metrics.metric(consumerRebalanceMetricsManager.assignedPartitionsCount).metricValue());
+    }
+}