Batch Coordinate cache PoC

Author: giuseppelillo

core/src/main/scala/kafka/server/DelayedFetch.scala

@@ -159,10 +159,13 @@ class DelayedFetch(
     }
 
     // Case G
-    if (accumulatedSize >= minBytes.getOrElse(params.minBytes))
+    if (accumulatedSize >= minBytes.getOrElse(params.minBytes)) {
+      DelayedFetchMetrics.enoughBytesMeter.mark()
       forceComplete()
-    else
+    } else {
+      DelayedFetchMetrics.notEnoughBytesMeter.mark()
       false
+    }
   }
 
   /**
@@ -183,10 +186,12 @@ class DelayedFetch(
     if (requests.isEmpty) return Some(0)
 
     val response = try {
-      replicaManager.findDisklessBatches(requests, Int.MaxValue)
+      DelayedFetchMetrics.findBatchesRequestsMeter.mark()
+      replicaManager.findDisklessBatches(requests)
     } catch {
       case e: Throwable =>
         error("Error while trying to find diskless batches on delayed fetch.", e)
+        DelayedFetchMetrics.findBatchesErrorsMeter.mark()
         return None  // Case C
     }
 
@@ -201,6 +206,7 @@ class DelayedFetch(
             val fetchPartitionStatus = fetchPartitionStatusMap.get(topicIdPartition)
             if (fetchPartitionStatus.isEmpty) {
               warn(s"Fetch partition status for $topicIdPartition not found in delayed fetch $this.")
+              DelayedFetchMetrics.findBatchesErrorsMeter.mark()
               return None  // Case C
             }
 
@@ -211,13 +217,20 @@ class DelayedFetch(
             if (fetchOffset.messageOffset > endOffset) {
               // Truncation happened
               debug(s"Satisfying fetch $this since it is fetching later segments of partition $topicIdPartition.")
+              DelayedFetchMetrics.logTruncationsMeter.mark()
               return None  // Case A
             } else if (fetchOffset.messageOffset < endOffset) {
               val bytesAvailable = r.estimatedByteSize(fetchOffset.messageOffset)
               accumulatedSize += bytesAvailable // Case B: accumulate the size of the batches
-            } // Case D: same as fetchOffset == endOffset, no new data available
+            } else {
+              // Case D: same as fetchOffset == endOffset, no new data available
+              DelayedFetchMetrics.noNewDataMeter.mark()
+            }
           }
-        case _ => return None  // Case C
+        case _ => {
+          DelayedFetchMetrics.findBatchesErrorsMeter.mark()
+          return None
+        }  // Case C
       }
     }
 
@@ -297,5 +310,12 @@ object DelayedFetchMetrics {
   private val FetcherTypeKey = "fetcherType"
   val followerExpiredRequestMeter: Meter = metricsGroup.newMeter("ExpiresPerSec", "requests", TimeUnit.SECONDS, Map(FetcherTypeKey -> "follower").asJava)
   val consumerExpiredRequestMeter: Meter = metricsGroup.newMeter("ExpiresPerSec", "requests", TimeUnit.SECONDS, Map(FetcherTypeKey -> "consumer").asJava)
+  // diskless metrics
+  val findBatchesRequestsMeter: Meter = metricsGroup.newMeter("FindBatchesRequestsPerSec", "requests", TimeUnit.SECONDS)
+  val findBatchesErrorsMeter: Meter = metricsGroup.newMeter("FindBatchesErrorsPerSec", "requests", TimeUnit.SECONDS)
+  val enoughBytesMeter: Meter = metricsGroup.newMeter("EnoughBytesPerSec", "requests", TimeUnit.SECONDS)
+  val logTruncationsMeter: Meter = metricsGroup.newMeter("LogTruncationsPerSec", "requests", TimeUnit.SECONDS)
+  val noNewDataMeter: Meter = metricsGroup.newMeter("NoNewDataPerSec", "requests", TimeUnit.SECONDS)
+  val notEnoughBytesMeter: Meter = metricsGroup.newMeter("NotEnoughBytesPerSec", "requests", TimeUnit.SECONDS)
 }
 

core/src/main/scala/kafka/server/ReplicaManager.scala

@@ -19,7 +19,7 @@ package kafka.server
 import com.yammer.metrics.core.Meter
 import io.aiven.inkless.common.SharedState
 import io.aiven.inkless.consume.{FetchHandler, FetchOffsetHandler}
-import io.aiven.inkless.control_plane.{FindBatchRequest, FindBatchResponse, MetadataView}
+import io.aiven.inkless.control_plane.{BatchInfo, FindBatchRequest, FindBatchResponse, MetadataView}
 import io.aiven.inkless.delete.{DeleteRecordsInterceptor, FileCleaner, RetentionEnforcer}
 import io.aiven.inkless.merge.FileMerger
 import io.aiven.inkless.produce.AppendHandler
@@ -1719,12 +1719,33 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
-  def findDisklessBatches(requests: Seq[FindBatchRequest], maxBytes: Int): Option[util.List[FindBatchResponse]] = {
+  private def findDisklessBatchesThroughControlPlane(requests: Seq[FindBatchRequest], maxBytes: Int = Int.MaxValue): Option[util.List[FindBatchResponse]] = {
     inklessSharedState.map { sharedState =>
       sharedState.controlPlane().findBatches(requests.asJava, maxBytes, sharedState.config().maxBatchesPerPartitionToFind())
     }
   }
 
+  def findDisklessBatches(requests: Seq[FindBatchRequest]): Option[util.List[FindBatchResponse]] = {
+    inklessSharedState.flatMap { sharedState =>
+      if (!sharedState.isBatchCoordinatesCacheEnabled) {
+        findDisklessBatchesThroughControlPlane(requests)
+      } else {
+        Some(requests.map { request =>
+          val logFragment = sharedState.batchCoordinateCache().get(request.topicIdPartition(), request.offset())
+          if (logFragment == null) {
+            FindBatchResponse.success(util.List.of(), -1, -1)
+          } else {
+            FindBatchResponse.success(
+              logFragment.batches().stream().map[BatchInfo](batchCoordinate => batchCoordinate.batchInfo()).toList,
+              logFragment.logStartOffset(),
+              logFragment.highWaterMark()
+            )
+          }
+        }.asJava)
+      }
+    }
+  }
+
   def fetchDisklessMessages(params: FetchParams,
                             fetchInfos: Seq[(TopicIdPartition, PartitionData)]): CompletableFuture[Seq[(TopicIdPartition, FetchPartitionData)]] = {
     inklessFetchHandler match {

core/src/test/java/kafka/server/InklessClusterTest.java

@@ -37,6 +37,8 @@
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.test.KafkaClusterTestKit;
 import org.apache.kafka.common.test.TestKitNodes;
 import org.apache.kafka.coordinator.group.GroupCoordinatorConfig;
@@ -55,6 +57,7 @@
 import org.testcontainers.junit.jupiter.Testcontainers;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -189,6 +192,81 @@ public void createDisklessTopic(final boolean idempotenceEnable, final Timestamp
         consumeWithSubscription(timestampType, clientConfigs, topicName, now, numRecords);
     }
 
+    @Test
+    public void test() throws Exception {
+        Map<String, Object> clientConfigs = new HashMap<>();
+        clientConfigs.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
+        clientConfigs.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, String.valueOf(true));
+        clientConfigs.put(ProducerConfig.LINGER_MS_CONFIG, "1000");
+        clientConfigs.put(ProducerConfig.BATCH_SIZE_CONFIG, "100000");
+        clientConfigs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+        clientConfigs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, IntegerSerializer.class.getName());
+        clientConfigs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+        clientConfigs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class.getName());
+        // by default is latest and nothing would get consumed.
+        clientConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, AutoOffsetResetStrategy.EARLIEST.name());
+        String topicName = "inkless-cache-topic";
+        int numRecords = 50;
+
+        try (Admin admin = AdminClient.create(clientConfigs)) {
+            final NewTopic topic = new NewTopic(topicName, 2, (short) 1)
+                .configs(Map.of(
+                    TopicConfig.DISKLESS_ENABLE_CONFIG, "true",
+                    TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, TimestampType.LOG_APPEND_TIME.name
+                ));
+            CreateTopicsResult topics = admin.createTopics(Collections.singletonList(topic));
+            topics.all().get(10, TimeUnit.SECONDS);
+        }
+
+        AtomicInteger recordsProduced = new AtomicInteger();
+        final long now = System.currentTimeMillis();
+        Callback produceCb = (metadata, exception) -> {
+            if (exception != null) {
+                log.error("Failed to send record", exception);
+            } else {
+                log.info("Committed value for topic {} at offset {} at {}", metadata.topic(), metadata.offset(), now);
+                recordsProduced.incrementAndGet();
+            }
+        };
+        var producedRecords1 = new ArrayList<Integer>(numRecords);
+        var producedRecords2 = new ArrayList<Integer>(numRecords);
+        try (Producer<byte[], Integer> producer = new KafkaProducer<>(clientConfigs)) {
+            for (int i = 0; i < numRecords; i++) {
+                producedRecords1.add(i);
+                final ProducerRecord<byte[], Integer> record1 = new ProducerRecord<>(topicName, 0, now, null, i);
+                producedRecords2.add(i + 10000);
+                final ProducerRecord<byte[], Integer> record2 = new ProducerRecord<>(topicName, 1, now, null, i + 10000);
+                producer.send(record1, produceCb);
+                Thread.sleep(100);
+                producer.send(record2, produceCb);
+                Thread.sleep(100);
+            }
+            producer.flush();
+        }
+
+        var consumedRecords1 = new ArrayList<Integer>();
+        var consumedRecords2 = new ArrayList<Integer>();
+        assertEquals(numRecords * 2, recordsProduced.get());
+        try (Consumer<byte[], Integer> consumer = new KafkaConsumer<>(clientConfigs)) {
+            consumer.assign(Collections.singletonList(new TopicPartition(topicName, 0)));
+            ConsumerRecords<byte[], Integer> poll = consumer.poll(Duration.ofSeconds(30));
+            for (ConsumerRecord<byte[], Integer> record : poll) {
+                assertTrue(record.timestamp() > now);
+                consumedRecords1.add(record.value());
+            }
+        }
+        try (Consumer<byte[], Integer> consumer = new KafkaConsumer<>(clientConfigs)) {
+            consumer.assign(Collections.singletonList(new TopicPartition(topicName, 1)));
+            ConsumerRecords<byte[], Integer> poll = consumer.poll(Duration.ofSeconds(30));
+            for (ConsumerRecord<byte[], Integer> record : poll) {
+                assertTrue(record.timestamp() > now);
+                consumedRecords2.add(record.value());
+            }
+        }
+        assertEquals(producedRecords1, consumedRecords1);
+        assertEquals(producedRecords2, consumedRecords2);
+    }
+
     @Test
     public void produceToDisklessAndClassic() throws Exception {
         Map<String, Object> clientConfigs = new HashMap<>();

core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala

@@ -346,7 +346,7 @@ class DelayedFetchTest {
       when(mockResponse.highWatermark()).thenReturn(endOffset) // endOffset < fetchOffset (truncation)
 
       val future = Some(Collections.singletonList(mockResponse))
-      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]], anyInt())).thenReturn(future)
+      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]])).thenReturn(future)
 
       // Mock fetchDisklessMessages for onComplete
       when(replicaManager.fetchParamsWithNewMaxBytes(any[FetchParams], any[Float])).thenAnswer(_.getArgument(0))
@@ -435,7 +435,7 @@ class DelayedFetchTest {
       when(mockResponse.estimatedByteSize(fetchOffset)).thenReturn(estimatedBatchSize)
 
       val future = Some(Collections.singletonList(mockResponse))
-      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]], anyInt())).thenReturn(future)
+      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]])).thenReturn(future)
 
       when(replicaManager.readFromLog(
         fetchParams,
@@ -520,7 +520,7 @@ class DelayedFetchTest {
       when(mockResponse.estimatedByteSize(fetchOffset)).thenReturn(estimatedBatchSize)
 
       val future = Some(Collections.singletonList(mockResponse))
-      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]], anyInt())).thenReturn(future)
+      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]])).thenReturn(future)
 
       when(replicaManager.readFromLog(
         fetchParams,
@@ -602,7 +602,7 @@ class DelayedFetchTest {
       when(mockResponse.estimatedByteSize(fetchOffset)).thenReturn(estimatedBatchSize)
 
       val future = Some(Collections.singletonList(mockResponse))
-      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]], anyInt())).thenReturn(future)
+      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]])).thenReturn(future)
 
       // Mock fetchDisklessMessages for onComplete
       when(replicaManager.fetchParamsWithNewMaxBytes(any[FetchParams], anyFloat())).thenAnswer(_.getArgument(0))
@@ -686,7 +686,7 @@ class DelayedFetchTest {
       when(mockResponse.highWatermark()).thenReturn(600L)
 
       val future = Some(Collections.singletonList(mockResponse))
-      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]], anyInt())).thenReturn(future)
+      when(replicaManager.findDisklessBatches(any[Seq[FindBatchRequest]])).thenReturn(future)
 
       // Mock fetchDisklessMessages for onComplete
       when(replicaManager.fetchParamsWithNewMaxBytes(any[FetchParams], anyFloat())).thenAnswer(_.getArgument(0))

docker/examples/docker-compose-files/inkless/docker-compose.yml

@@ -27,7 +27,10 @@ services:
       # Inkless
       KAFKA_DISKLESS_STORAGE_SYSTEM_ENABLE: "true"
       KAFKA_LOG_DISKLESS_ENABLE: "true"
+      KAFKA_INKLESS_CONSUME_BATCH_COORDINATES_CACHE_ENABLED: "false"
+      KAFKA_INKLESS_CONSUME_BATCH_COORDINATES_CACHE_TTL_MS: "5000"
       KAFKA_INKLESS_CONSUME_CACHE_MAX_COUNT: "100"
+      KAFKA_INKLESS_CONSUME_CACHE_EXPIRATION_MAX_IDLE_SEC: "30"
 
       ## Control Plane
       ### Postgresql
@@ -44,7 +47,7 @@ services:
 
       # JMX
       KAFKA_JMX_PORT: 9999
-      KAFKA_HEAP_OPTS: "-Xmx2G -Xms2G"
+      KAFKA_HEAP_OPTS: "-Xmx8G -Xms8G"
       EXTRA_ARGS: -javaagent:/opt/prometheus-jmx-exporter/jmx_prometheus_javaagent.jar=7070:/opt/prometheus-jmx-exporter/kafka.yml
     volumes:
       - ./../../../extra/prometheus-jmx-exporter:/opt/prometheus-jmx-exporter:Z

docs/inkless/configs.rst

@@ -63,6 +63,21 @@ Under ``inkless.``
   * Valid Values: [0,...]
   * Importance: medium
 
+``consume.batch.coordinates.cache.enabled``
+  If true, batch coordinates cache is enabled.
+
+  * Type: boolean
+  * Default: true
+  * Importance: low
+
+``consume.batch.coordinates.cache.ttl.ms``
+  Time to live in milliseconds for an entry in the Batch Coordinates cache. The time to live must be <= than half of the value of of file.cleaner.interval.ms.
+
+  * Type: int
+  * Default: 5000 (5 seconds)
+  * Valid Values: [1,...]
+  * Importance: low
+
 ``consume.cache.block.bytes``
   The number of bytes to fetch as a single block from object storage when serving fetch requests.
 

storage/inkless/src/main/java/io/aiven/inkless/cache/BatchCoordinateCache.java

@@ -0,0 +1,32 @@
+/*
+ * Inkless
+ * Copyright (C) 2024 - 2025 Aiven OY
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package io.aiven.inkless.cache;
+
+import org.apache.kafka.common.TopicIdPartition;
+
+import java.io.Closeable;
+
+public interface BatchCoordinateCache extends Closeable {
+
+    LogFragment get(TopicIdPartition topicIdPartition, long offset);
+
+    void put(CacheBatchCoordinate cacheBatchCoordinate) throws IllegalStateException;
+
+    int invalidatePartition(TopicIdPartition topicIdPartition);
+
+}

storage/inkless/src/main/java/io/aiven/inkless/cache/BatchCoordinateCacheMetrics.java

@@ -0,0 +1,60 @@
+package io.aiven.inkless.cache;
+
+import org.apache.kafka.server.metrics.KafkaMetricsGroup;
+
+import java.io.Closeable;
+import java.util.concurrent.atomic.LongAdder;
+
+
+public final class BatchCoordinateCacheMetrics implements Closeable {
+    static final String CACHE_HITS = "CacheHits";
+    static final String CACHE_MISSES = "CacheMisses";
+    static final String CACHE_INVALIDATIONS = "CacheInvalidations";
+    static final String CACHE_EVICTIONS = "CacheEvictions";
+    static final String CACHE_SIZE = "CacheSize";
+
+    private final KafkaMetricsGroup metricsGroup = new KafkaMetricsGroup(BatchCoordinateCache.class);
+    private final LongAdder cacheHits = new LongAdder();
+    private final LongAdder cacheMisses = new LongAdder();
+    private final LongAdder cacheInvalidations = new LongAdder();
+    private final LongAdder cacheEvictions = new LongAdder();
+    private final LongAdder cacheSize = new LongAdder();
+
+    public BatchCoordinateCacheMetrics() {
+        metricsGroup.newGauge(CACHE_HITS, cacheHits::intValue);
+        metricsGroup.newGauge(CACHE_MISSES, cacheMisses::intValue);
+        metricsGroup.newGauge(CACHE_INVALIDATIONS, cacheInvalidations::intValue);
+        metricsGroup.newGauge(CACHE_EVICTIONS, cacheEvictions::intValue);
+        metricsGroup.newGauge(CACHE_SIZE, cacheSize::intValue);
+    }
+
+    public void recordCacheHit() {
+        cacheHits.increment();
+    }
+
+    public void recordCacheMiss() {
+        cacheMisses.increment();
+    }
+
+    public void recordCacheInvalidation() {
+        cacheInvalidations.increment();
+    }
+
+    public void recordCacheEviction() {
+        cacheEvictions.increment();
+    }
+
+    public void incrementCacheSize() {
+        cacheSize.increment();
+    }
+
+    public void decreaseCacheSize(int removedEntries) {
+        cacheSize.add(-removedEntries);
+    }
+
+    @Override
+    public void close() {
+        metricsGroup.removeMetric(CACHE_HITS);
+        metricsGroup.removeMetric(CACHE_MISSES);
+    }
+}