KAFKA-18878: Added share session cache and delayed share fetch metrics (KIP-1103) (apache#19059)

The PR implements the ShareSessionCache and DelayedShareFetchMetrics as
defined in KIP-1103.

Reviewers: Andrew Schofield <[email protected]>

Author: apoorvmittal10

core/src/main/java/kafka/server/share/DelayedShareFetch.java

@@ -25,6 +25,7 @@
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.FetchRequest;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.metrics.KafkaMetricsGroup;
 import org.apache.kafka.server.purgatory.DelayedOperation;
 import org.apache.kafka.server.share.SharePartitionKey;
 import org.apache.kafka.server.share.fetch.DelayedShareFetchGroupKey;
@@ -36,6 +37,8 @@
 import org.apache.kafka.storage.internals.log.LogOffsetMetadata;
 import org.apache.kafka.storage.internals.log.LogOffsetSnapshot;
 
+import com.yammer.metrics.core.Meter;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -45,6 +48,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
@@ -61,6 +65,8 @@ public class DelayedShareFetch extends DelayedOperation {
 
     private static final Logger log = LoggerFactory.getLogger(DelayedShareFetch.class);
 
+    private static final String EXPIRES_PER_SEC = "ExpiresPerSec";
+
     private final ShareFetch shareFetch;
     private final ReplicaManager replicaManager;
     private final BiConsumer<SharePartitionKey, Throwable> exceptionHandler;
@@ -70,6 +76,10 @@ public class DelayedShareFetch extends DelayedOperation {
     // The topic partitions that need to be completed for the share fetch request are given by sharePartitions.
     // sharePartitions is a subset of shareFetchData. The order of insertion/deletion of entries in sharePartitions is important.
     private final LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions;
+    /**
+     * Metric for the rate of expired delayed fetch requests.
+     */
+    private final Meter expiredRequestMeter;
     // Tracks the start time to acquire any share partition for a fetch request.
     private long acquireStartTimeMs;
     private LinkedHashMap<TopicIdPartition, Long> partitionsAcquired;
@@ -124,10 +134,14 @@ public DelayedShareFetch(
         this.shareGroupMetrics = shareGroupMetrics;
         this.time = time;
         this.acquireStartTimeMs = time.hiResClockMs();
+        // Register metrics for DelayedShareFetch.
+        KafkaMetricsGroup metricsGroup = new KafkaMetricsGroup("kafka.server", "DelayedShareFetchMetrics");
+        this.expiredRequestMeter = metricsGroup.newMeter(EXPIRES_PER_SEC, "requests", TimeUnit.SECONDS);
     }
 
     @Override
     public void onExpiration() {
+        expiredRequestMeter.mark();
     }
 
     /**
@@ -514,4 +528,9 @@ void releasePartitionLocks(Set<TopicIdPartition> topicIdPartitions) {
     Lock lock() {
         return lock;
     }
+
+    // Visible for testing.
+    Meter expiredRequestMeter() {
+        return expiredRequestMeter;
+    }
 }

core/src/test/java/kafka/server/share/DelayedShareFetchTest.java

@@ -154,6 +154,7 @@ public void testDelayedShareFetchTryCompleteReturnsFalseDueToNonAcquirablePartit
         // Metrics shall not be recorded as no partition is acquired.
         assertNull(shareGroupMetrics.topicPartitionsAcquireTimeMs(groupId));
         assertNull(shareGroupMetrics.topicPartitionsFetchRatio(groupId));
+        assertEquals(0, delayedShareFetch.expiredRequestMeter().count());
 
         delayedShareFetch.lock().unlock();
     }
@@ -1118,6 +1119,23 @@ public void testPartitionMaxBytesFromUniformStrategyInCombineLogReadResponse() {
             true);
     }
 
+    @Test
+    public void testOnCompleteExecutionOnTimeout() {
+        ShareFetch shareFetch = new ShareFetch(
+            FETCH_PARAMS, "grp", Uuid.randomUuid().toString(),
+            new CompletableFuture<>(), new LinkedHashMap<>(), BATCH_SIZE, MAX_FETCH_RECORDS,
+            BROKER_TOPIC_STATS);
+        DelayedShareFetch delayedShareFetch = DelayedShareFetchBuilder.builder()
+            .withShareFetchData(shareFetch)
+            .build();
+        assertFalse(delayedShareFetch.isCompleted());
+        assertFalse(shareFetch.isCompleted());
+        // Call run to execute onComplete and onExpiration.
+        delayedShareFetch.run();
+        assertTrue(shareFetch.isCompleted());
+        assertEquals(1, delayedShareFetch.expiredRequestMeter().count());
+    }
+
     static void mockTopicIdPartitionToReturnDataEqualToMinBytes(ReplicaManager replicaManager, TopicIdPartition topicIdPartition, int minBytes) {
         LogOffsetMetadata hwmOffsetMetadata = new LogOffsetMetadata(1, 1, minBytes);
         LogOffsetSnapshot endOffsetSnapshot = new LogOffsetSnapshot(1, mock(LogOffsetMetadata.class),

core/src/test/java/kafka/server/share/SharePartitionTest.java

@@ -48,7 +48,6 @@
 import org.apache.kafka.coordinator.group.GroupConfig;
 import org.apache.kafka.coordinator.group.GroupConfigManager;
 import org.apache.kafka.coordinator.group.ShareGroupAutoOffsetResetStrategy;
-import org.apache.kafka.server.metrics.KafkaYammerMetrics;
 import org.apache.kafka.server.share.acknowledge.ShareAcknowledgementBatch;
 import org.apache.kafka.server.share.fetch.ShareAcquiredRecords;
 import org.apache.kafka.server.share.metrics.SharePartitionMetrics;
@@ -67,8 +66,6 @@
 import org.apache.kafka.storage.internals.log.OffsetResultHolder;
 import org.apache.kafka.test.TestUtils;
 
-import com.yammer.metrics.core.Gauge;
-
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -91,6 +88,7 @@
 
 import static kafka.server.share.SharePartition.EMPTY_MEMBER_ID;
 import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.memoryRecordsBuilder;
+import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.yammerMetricValue;
 import static org.apache.kafka.test.TestUtils.assertFutureThrows;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -6736,19 +6734,6 @@ public void mockPersisterReadStateMethod(Persister persister) {
         Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult));
     }
 
-    private Number yammerMetricValue(String name) {
-        try {
-            Gauge gauge = (Gauge) KafkaYammerMetrics.defaultRegistry().allMetrics().entrySet().stream()
-                .filter(e -> e.getKey().getMBeanName().contains(name))
-                .findFirst()
-                .orElseThrow()
-                .getValue();
-            return (Number) gauge.value();
-        } catch (Exception e) {
-            return 0;
-        }
-    }
-
     private static class SharePartitionBuilder {
 
         private int defaultAcquisitionLockTimeoutMs = 30000;

server/src/main/java/org/apache/kafka/server/share/session/ShareSessionCache.java

@@ -20,11 +20,15 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.requests.ShareRequestMetadata;
 import org.apache.kafka.common.utils.ImplicitLinkedHashCollection;
+import org.apache.kafka.server.metrics.KafkaMetricsGroup;
 import org.apache.kafka.server.share.CachedSharePartition;
 
+import com.yammer.metrics.core.Meter;
+
 import java.util.HashMap;
 import java.util.Map;
 import java.util.TreeMap;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Caches share sessions.
@@ -37,6 +41,17 @@
  * must never be acquired while an individual ShareSession lock is already held.
  */
 public class ShareSessionCache {
+    // Visible for testing.
+    static final String SHARE_SESSIONS_COUNT = "ShareSessionsCount";
+    // Visible for testing.
+    static final String SHARE_PARTITIONS_COUNT = "SharePartitionsCount";
+    private static final String SHARE_SESSION_EVICTIONS_PER_SEC = "ShareSessionEvictionsPerSec";
+
+    /**
+     * Metric for the rate of eviction of share sessions.
+     */
+    private final Meter evictionsMeter;
+
     private final int maxEntries;
     private final long evictionMs;
     private long numPartitions = 0;
@@ -47,9 +62,15 @@ public class ShareSessionCache {
     // Maps last used times to sessions.
     private final TreeMap<LastUsedKey, ShareSession> lastUsed = new TreeMap<>();
 
+    @SuppressWarnings("this-escape")
     public ShareSessionCache(int maxEntries, long evictionMs) {
         this.maxEntries = maxEntries;
         this.evictionMs = evictionMs;
+        // Register metrics for ShareSessionCache.
+        KafkaMetricsGroup metricsGroup = new KafkaMetricsGroup("kafka.server", "ShareSessionCache");
+        metricsGroup.newGauge(SHARE_SESSIONS_COUNT, this::size);
+        metricsGroup.newGauge(SHARE_PARTITIONS_COUNT, this::totalPartitions);
+        this.evictionsMeter = metricsGroup.newMeter(SHARE_SESSION_EVICTIONS_PER_SEC, "evictions", TimeUnit.SECONDS);
     }
 
     /**
@@ -136,6 +157,7 @@ public synchronized boolean tryEvict(long now) {
         } else if (now - lastUsedEntry.getKey().lastUsedMs() > evictionMs) {
             ShareSession session = lastUsedEntry.getValue();
             remove(session);
+            evictionsMeter.mark();
             return true;
         }
         return false;
@@ -159,4 +181,9 @@ public synchronized ShareSessionKey maybeCreateSession(String groupId, Uuid memb
         }
         return null;
     }
+
+    // Visible for testing.
+    Meter evictionsMeter() {
+        return evictionsMeter;
+    }
 }

server/src/test/java/org/apache/kafka/server/share/fetch/ShareFetchTestUtils.java

@@ -23,8 +23,11 @@
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MemoryRecordsBuilder;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.server.metrics.KafkaYammerMetrics;
 import org.apache.kafka.test.TestUtils;
 
+import com.yammer.metrics.core.Gauge;
+
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.LinkedHashMap;
@@ -145,4 +148,32 @@ public static ShareAcquiredRecords createShareAcquiredRecords(AcquiredRecords ac
             List.of(acquiredRecords), (int) (acquiredRecords.lastOffset() - acquiredRecords.firstOffset() + 1)
         );
     }
+
+    /**
+     * Fetch the gauge value from the yammer metrics.
+     *
+     * @param name The name of the metric.
+     * @return The gauge value as a number.
+     */
+    public static Number yammerMetricValue(String name) {
+        try {
+            Gauge gauge = (Gauge) KafkaYammerMetrics.defaultRegistry().allMetrics().entrySet().stream()
+                .filter(e -> e.getKey().getMBeanName().contains(name))
+                .findFirst()
+                .orElseThrow()
+                .getValue();
+            return (Number) gauge.value();
+        } catch (Exception e) {
+            return 0;
+        }
+    }
+
+    /**
+     * Clear all the yammer metrics.
+     */
+    public static void clearYammerMetrics() {
+        KafkaYammerMetrics.defaultRegistry().allMetrics().keySet().forEach(
+            metricName -> KafkaYammerMetrics.defaultRegistry().removeMetric(metricName)
+        );
+    }
 }