diff --git a/core/src/main/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetch.scala b/core/src/main/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetch.scala new file mode 100644 index 00000000000..bc004dfb4e1 --- /dev/null +++ b/core/src/main/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetch.scala @@ -0,0 +1,170 @@ +/* + * Inkless + * Copyright (C) 2024 - 2026 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 . + */ + +package io.aiven.inkless.consolidation + +import io.aiven.inkless.consume.FetchHandler +import io.aiven.inkless.control_plane.FindBatchRequest +import kafka.server.ReplicaManager +import kafka.utils.Logging +import org.apache.kafka.common.TopicIdPartition +import org.apache.kafka.common.protocol.Errors +import org.apache.kafka.common.record.MemoryRecords +import org.apache.kafka.common.requests.FetchRequest +import org.apache.kafka.server.purgatory.DelayedOperation +import org.apache.kafka.server.storage.log.{FetchParams, FetchPartitionData} + +import java.util +import java.util.concurrent.atomic.AtomicBoolean +import scala.jdk.CollectionConverters._ + +/** + * Long-poll wrapper for a consolidation fetch, honoring `minBytes` and `maxWaitMs`. + * + * `tryComplete` runs one cheap metadata probe before parking, completing early only if it already meets `minBytes`. + * It does not re-probe on later wake-ups (none for this delayed operation), so a request below `minBytes` completes + * via the `maxWaitMs` timeout -- trading up to `maxWaitMs` latency for a hard bound on control-plane pressure. + * + * `onComplete` runs the full `FetchHandler.handle` (findBatches + storage fetch) once, then invokes + * the response callback when that future completes. + */ +class DelayedConsolidationFetch( + params: FetchParams, + fetchInfos: util.Map[TopicIdPartition, FetchRequest.PartitionData], + fetchHandler: FetchHandler, + replicaManager: ReplicaManager, + responseCallback: util.Map[TopicIdPartition, FetchPartitionData] => Unit +) extends DelayedOperation(params.maxWaitMs) with Logging { + + private val initialProbeDone = new AtomicBoolean(false) + + override def toString: String = + s"DelayedConsolidationFetch(numPartitions=${fetchInfos.size}, minBytes=${params.minBytes}, maxWaitMs=${params.maxWaitMs})" + + /** + * Cheap probe: sum the bytes [[ReplicaManager.findDisklessBatches]] reports across the requested partitions, + * and force completion if the estimate meets `minBytes`. + * + * With the batch coordinate cache enabled this reads only local, this-broker appends, so it can miss + * data committed on other brokers and under-count. + * This is acceptable here: a miss just keeps the op parked until [[onComplete]] runs the real fetch + * (findBatches through the control plane), which sees the authoritative view. + */ + override def tryComplete(): Boolean = { + if (fetchInfos.isEmpty) return forceComplete() + + // Probe once per op. tryCompleteElseWatch calls tryComplete twice: once before registering the + // watch, then again right after as a race re-check. + // This purgatory has no wake-ups, so that second call is the only repeat, and the CAS collapses it + // to a cheap false, keeping each op at one probe plus the final fetch. + if (!initialProbeDone.compareAndSet(false, true)) return false + + val requests = fetchInfos.entrySet().asScala.toSeq.map { e => + new FindBatchRequest(e.getKey, e.getValue.fetchOffset, e.getValue.maxBytes) + } + + val maybeFindBatchResponses = try { + replicaManager.findDisklessBatches(requests) + } catch { + case e: Throwable => + warn(s"$this partitions=$samplePartitions: error during tryComplete findDisklessBatches; deferring to onComplete", e) + return forceComplete() + } + + val findBatchResponses = maybeFindBatchResponses match { + case Some(r) => r + case None => return forceComplete() // no shared state -- should not happen for consolidation, fail-safe + } + + if (findBatchResponses.size() != fetchInfos.size) { + warn(s"$this partitions=$samplePartitions: findDisklessBatches returned ${findBatchResponses.size()} responses " + + s"for ${fetchInfos.size} requests; deferring to onComplete") + return forceComplete() + } + + var accumulated = 0L + var index = 0 + val it = fetchInfos.entrySet().iterator() + while (it.hasNext && index < findBatchResponses.size()) { + val entry = it.next() + val req = entry.getValue + val resp = findBatchResponses.get(index) + resp.errors() match { + case Errors.NONE => + accumulated += resp.estimatedByteSize(req.fetchOffset) + case _ => + // any error short-circuits the wait -- let onComplete surface it through the real fetch + return forceComplete() + } + index += 1 + } + + // Estimated only: this sums the probe's per-partition byte estimates with no total cap. + // The real fetch in onComplete applies the per-partition and response.max.bytes limits (via findBatches), + // so accumulated here may exceed what the fetch returns -- that is fine for a >= minBytes check. + if (accumulated >= params.minBytes) forceComplete() + else false + } + + override def onExpiration(): Unit = { + debug(s"$this expired at maxWaitMs=${params.maxWaitMs}") + } + + override def onComplete(): Unit = { + try { + fetchHandler.handle(params, fetchInfos).whenComplete { (response, throwable) => + if (throwable == null) { + // handle never completes with a null response (it substitutes an error map on failure). + responseCallback(response) + } else { + warn(s"$this partitions=$samplePartitions: fetchHandler.handle failed in onComplete; returning per-partition errors", throwable) + responseCallback(errorResponse(throwable)) + } + } + } catch { + case e: Throwable => + warn(s"$this partitions=$samplePartitions: fetchHandler.handle failed before completion; returning per-partition errors", e) + responseCallback(errorResponse(e)) + } + } + + private def samplePartitions: String = { + val sample = fetchInfos.keySet().asScala.take(5).map(_.topicPartition).mkString(",") + if (fetchInfos.size > 5) s"[$sample,...]" else s"[$sample]" + } + + private def errorResponse(e: Throwable): util.Map[TopicIdPartition, FetchPartitionData] = { + // Errors.forException unwraps CompletionException/ExecutionException to map the real cause. + val error = Errors.forException(e) + val response = new util.LinkedHashMap[TopicIdPartition, FetchPartitionData](fetchInfos.size) + fetchInfos.keySet().forEach { tp => + response.put(tp, new FetchPartitionData( + error, + -1L, + -1L, + MemoryRecords.EMPTY, + util.Optional.empty(), + util.OptionalLong.empty(), + util.Optional.empty(), + util.OptionalInt.empty(), + false + )) + } + response + } +} diff --git a/core/src/main/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPoint.scala b/core/src/main/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPoint.scala index 8187aed1fde..dc6954032c2 100644 --- a/core/src/main/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPoint.scala +++ b/core/src/main/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPoint.scala @@ -21,27 +21,29 @@ package io.aiven.inkless.consolidation import io.aiven.inkless.consume.{FetchHandler, FetchOffsetHandler} import kafka.server.{KafkaConfig, ReplicaManager, ReplicaQuota} import kafka.utils.Logging -import org.apache.kafka.common.Uuid import org.apache.kafka.common.errors.{KafkaStorageException, UnknownTopicOrPartitionException} -import org.apache.kafka.common.message.{FetchResponseData, OffsetForLeaderEpochRequestData} import org.apache.kafka.common.message.ListOffsetsRequestData.ListOffsetsPartition import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset +import org.apache.kafka.common.message.{FetchResponseData, OffsetForLeaderEpochRequestData} import org.apache.kafka.common.protocol.Errors import org.apache.kafka.common.record.FileRecords.TimestampAndOffset +import org.apache.kafka.common.record.{MemoryRecords, MutableRecordBatch, Records} import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, ListOffsetsRequest, OffsetsForLeaderEpochResponse} -import org.apache.kafka.common.TopicPartition -import org.apache.kafka.common.record.MemoryRecords +import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.metadata.{LeaderAndIsr, PartitionRegistration} import org.apache.kafka.server.common.{MetadataVersion, OffsetAndEpoch} import org.apache.kafka.server.network.BrokerEndPoint -import org.apache.kafka.server.storage.log.{FetchIsolation, FetchParams} +import org.apache.kafka.server.purgatory.TopicPartitionOperationKey +import org.apache.kafka.server.storage.log.{FetchIsolation, FetchParams, FetchPartitionData} import org.apache.kafka.server.{LeaderEndPoint, PartitionFetchState, ReplicaFetch, ResultWithPartitions} import org.apache.kafka.storage.internals.log.OffsetResultHolder.FileRecordsOrError import org.apache.kafka.storage.internals.log.UnifiedLog import java.util +import java.nio.ByteBuffer import java.util.Optional import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentHashMap import scala.collection.mutable import scala.jdk.CollectionConverters._ import scala.util.Try @@ -69,6 +71,7 @@ class DisklessLeaderEndPoint( private val minBytes = brokerConfig.disklessConsolidationFetchMinBytes private val maxBytes = brokerConfig.disklessConsolidationFetchResponseMaxBytes private val fetchSize = brokerConfig.disklessConsolidationFetchMaxBytes + private val unsafeSegmentConfigWarnings = ConcurrentHashMap.newKeySet[TopicPartition]() override def isTruncationOnFetchSupported: Boolean = false @@ -89,24 +92,33 @@ class DisklessLeaderEndPoint( val fetchParams = new FetchParams( FetchRequest.FUTURE_LOCAL_REPLICA_ID, -1, - 0L, - request.minBytes, + // specific consolidation maxWait so the delayed operation parks idle partitions instead of hammering PG. + maxWait.toLong, + // minBytes drives completion only on the initial metadata probe. + minBytes, request.maxBytes, FetchIsolation.LOG_END, Optional.empty() ) - val response = fetchHandler.handle(fetchParams, fetchInfos).get() + val response = awaitDelayedFetch(fetchParams, fetchInfos) response.asScala.map { case (tp, data) => val abortedTransactions = data.abortedTransactions.orElse(null) val lastStableOffset: Long = data.lastStableOffset.orElse(FetchResponse.INVALID_LAST_STABLE_OFFSET) + // Enforce the segment-safe budget reserved in buildFetch, at physical-batch granularity. + val partitionData = fetchInfos.get(tp) + val records = + if (data.error == Errors.NONE && partitionData != null) + clampRecordsToSegment(data.records, partitionData.fetchOffset, partitionData.maxBytes) + else + data.records val fetchResponseData = new FetchResponseData.PartitionData() .setPartitionIndex(tp.topicPartition.partition) .setErrorCode(data.error.code) .setHighWatermark(data.highWatermark) .setLastStableOffset(lastStableOffset) .setAbortedTransactions(abortedTransactions) - .setRecords(data.records) + .setRecords(records) // set local LSO if possible instead of the diskless start offset that is data.logStartOffset replicaManager.getPartitionOrError(tp.topicPartition) match { case Left(error) => @@ -161,6 +173,89 @@ class DisklessLeaderEndPoint( }.toMap.asJava } + /** + * Trim a fetched block so the follower's single append fits `segment.bytes`, restoring the classic + * invariant `returned <= maxBytes + one physical batch (<= max.message.bytes) <= segment.bytes`. + * + * `buildFetch` only reserves the headroom (`maxBytes = segment.bytes - max.message.bytes`); it is not + * enough alone because `find_batches` limits at *coordinate* granularity and always admits the + * coordinate that crosses `maxBytes`, and a `commit_file_v2` coordinate coalesces many physical batches + * into one row whose `byte_size` can far exceed `max.message.bytes` (up to `produce.buffer.max.bytes`). + * That crossing coordinate, or even the first always-admitted one, can then overflow a small + * `segment.bytes` (a supported tuning, e.g. for eager tiering) and fail the append. + * + * Re-applying the limit per physical batch keeps the overshoot to one batch. Do not let the serve path + * emit a larger unit beyond `maxBytes` or this breaks again. Batches with `lastOffset < fetchOffset` are + * dropped: a prior fetch may have truncated a coalesced coordinate mid-run and the follower re-fetches + * into the same coordinate, so its already-appended prefix must not be re-served. At least one batch is + * always emitted for progress (a single batch is bounded by `max.message.bytes`, assumed + * `<= segment.bytes`; the degenerate case is warned about in `buildFetch` and left to fail). + */ + private def clampRecordsToSegment(records: Records, fetchOffset: Long, maxBytes: Int): Records = { + val selected = new util.ArrayList[MutableRecordBatch]() + var accumulated = 0 + var skippedPrefix = false + var stoppedEarly = false + val it = records.batches().iterator() + while (it.hasNext && !stoppedEarly) { + it.next() match { + case batch: MutableRecordBatch => + if (batch.lastOffset() < fetchOffset) { + skippedPrefix = true + } else if (selected.isEmpty || accumulated < maxBytes) { + // First eligible batch always taken; further batches only while bytes *before* them are + // under budget, so overshoot is one batch at most. + selected.add(batch) + accumulated += batch.sizeInBytes() + } else { + stoppedEarly = true + } + case _ => + // Non-MutableRecordBatch is not expected for diskless; serve untrimmed rather than drop data. + return records + } + } + // Nothing trimmed: keep the original to preserve the ConcatenatedRecords zero-copy path. + if (!skippedPrefix && !stoppedEarly) return records + if (selected.isEmpty) return records + val buffer = ByteBuffer.allocate(accumulated) + selected.forEach(batch => batch.writeTo(buffer)) + buffer.flip() + MemoryRecords.readableRecords(buffer) + } + + /** + * Run a [[DelayedConsolidationFetch]] on the broker's `delayedConsolidationFetchPurgatory` + * and block the calling fetcher thread until it completes. + * Park-and-wait when the diskless data is below `minBytes` and expire at `maxWaitMs`. + */ + private def awaitDelayedFetch( + fetchParams: FetchParams, + fetchInfos: util.Map[TopicIdPartition, FetchRequest.PartitionData] + ): util.Map[TopicIdPartition, FetchPartitionData] = { + if (fetchInfos.isEmpty) return util.Map.of() + + val resultFuture = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val delayedFetch = new DelayedConsolidationFetch( + params = fetchParams, + fetchInfos = fetchInfos, + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = response => resultFuture.complete(response) + ) + + val watchKeys = new util.ArrayList[TopicPartitionOperationKey](fetchInfos.size) + fetchInfos.keySet().forEach(tp => watchKeys.add(new TopicPartitionOperationKey(tp.topicPartition))) + + replicaManager.delayedConsolidationFetchPurgatory.tryCompleteElseWatch(delayedFetch, watchKeys) + + // Block until the op completes so only one fetch is in flight per fetcher (backpressure). + // Unbounded on purpose: shutdown must stop the fetchers before the purgatory (see + // ReplicaManager.shutdown), else the reaper never expires a parked op and this + // non-interruptible thread is stranded. + resultFuture.get() + } + override def fetchEarliestOffset(topicPartition: TopicPartition, currentLeaderEpoch: Int): OffsetAndEpoch = listDisklessOffset(topicPartition, currentLeaderEpoch, ListOffsetsRequest.EARLIEST_TIMESTAMP) @@ -315,7 +410,21 @@ class DisklessLeaderEndPoint( partitions.forEach { (topicPartition, fetchState) => if (fetchState.isReadyForFetch && !shouldFollowerThrottle(quota, fetchState, topicPartition)) { try { - val logStartOffset = replicaManager.localLogOrException(topicPartition).logStartOffset + val localLog = replicaManager.localLogOrException(topicPartition) + val logStartOffset = localLog.logStartOffset + val logConfig = localLog.config + // Reserve headroom for the follower append; clampRecordsToSegment enforces it per batch in + // fetch() (find_batches limits at coordinate granularity and a coalesced coordinate can + // overshoot by more than one batch). See clampRecordsToSegment. + val maxOvershootBytes = math.min(logConfig.maxMessageSize, logConfig.segmentSize - 1) + val segmentSafeFetchSize = math.max(1, logConfig.segmentSize - maxOvershootBytes) + val partitionFetchSize = math.min(fetchSize, segmentSafeFetchSize) + if (logConfig.maxMessageSize >= logConfig.segmentSize && unsafeSegmentConfigWarnings.add(topicPartition)) { + logger.warn("Topic-partition {} has max.message.bytes ({}) >= segment.bytes ({}). " + + "Consolidation fetch maxBytes is clamped to {} byte(s) to leave whole-batch overshoot headroom; " + + "increase segment.bytes above max.message.bytes for normal consolidation throughput.", + topicPartition, logConfig.maxMessageSize, logConfig.segmentSize, partitionFetchSize) + } val lastFetchedEpoch = Optional.empty[Integer]() requestMap.put( topicPartition, @@ -323,7 +432,7 @@ class DisklessLeaderEndPoint( fetchState.topicId().orElse(Uuid.ZERO_UUID), fetchState.fetchOffset(), logStartOffset, - fetchSize, + partitionFetchSize, Optional.of(fetchState.currentLeaderEpoch()), lastFetchedEpoch ) diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala index c46b39fd805..43bbaf704db 100644 --- a/core/src/main/scala/kafka/server/ReplicaManager.scala +++ b/core/src/main/scala/kafka/server/ReplicaManager.scala @@ -23,7 +23,7 @@ import io.aiven.inkless.storage_backend.common.ObjectFetcher import io.aiven.inkless.control_plane.{BatchInfo, FindBatchRequest, FindBatchResponse, InitDisklessLogProducerState, RepairDisklessLogRequest} import io.aiven.inkless.delete.{DeleteRecordsInterceptor, FileCleaner, RetentionEnforcer} import io.aiven.inkless.produce.AppendHandler -import io.aiven.inkless.consolidation.{ConsolidatedDisklessLogPruner, ConsolidationFetcherManager, ConsolidationMetrics, ConsolidationReconciler} +import io.aiven.inkless.consolidation.{ConsolidatedDisklessLogPruner, ConsolidationFetcherManager, ConsolidationMetrics, ConsolidationReconciler, DelayedConsolidationFetch} import kafka.cluster.Partition import kafka.log.LogManager import kafka.server.HostedPartition.Online @@ -269,6 +269,21 @@ class ReplicaManager(val config: KafkaConfig, new DelayedOperationPurgatory[DelayedShareFetch]( shareFetchPurgatoryName, delayedShareFetchTimer, config.brokerId, config.shareGroupConfig.shareFetchPurgatoryPurgeIntervalRequests)) + // Hosts long-poll DelayedConsolidationFetch operations created by the consolidation fetcher (DisklessLeaderEndPoint). + // Kept separate from delayedFetchPurgatory because the latter is strictly typed to DelayedFetch. + // Operations perform one initial metadata probe, then complete via maxWaitMs timeout if parked; + // wake-up sites are intentionally not mirrored here to keep control-plane pressure bounded. + // + // purgeInterval = 0 (matches delayedRemoteFetchPurgatory) so completed ops release their captured response + // references immediately for GC. + // Each completed op pins a Map[TopicIdPartition, FetchPartitionData] holding fetched records, + // bounded per partition by diskless.consolidation.fetch.max.bytes and in aggregate by + // diskless.consolidation.fetch.response.max.bytes. + // Without aggressive purging, watch lists accumulate up to ~purgeInterval completed ops worth of records + // and exhaust the heap. + val delayedConsolidationFetchPurgatory = + new DelayedOperationPurgatory[DelayedConsolidationFetch]( + "ConsolidationFetch", config.brokerId, 0) private val _inklessMetadataView: InklessMetadataView = inklessMetadataView.getOrElse(new InklessMetadataView(metadataCache.asInstanceOf[KRaftMetadataCache], () => config.extractLogConfigMap)) private val inklessAppendHandler: Option[AppendHandler] = inklessSharedState.map(new AppendHandler(_)) @@ -2006,6 +2021,16 @@ class ReplicaManager(val config: KafkaConfig, } } + /** + * Best-effort batch-metadata lookup for callers that can tolerate a stale view. + * + * With the batch coordinate cache enabled (default) it consults only this broker's local cache, + * which holds only data this broker produced. + * It cannot see batches committed on other brokers, so a negative answer does not mean "no data + * exists" -- only "none known locally". + * Callers needing an authoritative answer must go to the control plane (see FindBatchesJob). + * Only with the cache disabled does this query the control plane directly. + */ def findDisklessBatches(requests: Seq[FindBatchRequest]): Option[util.List[FindBatchResponse]] = { inklessSharedState.flatMap { sharedState => if (!sharedState.isBatchCoordinateCacheEnabled) { @@ -2014,7 +2039,17 @@ class ReplicaManager(val config: KafkaConfig, Some(requests.map { request => val logFragment = sharedState.batchCoordinateCache().get(request.topicIdPartition(), request.offset()) if (logFragment == null) { - FindBatchResponse.success(util.List.of(), -1, -1) + // Local miss: report empty rather than fall back to PG. + // A cross-broker partition always misses here, so the caller must have its own fallback. + // + // TODO: route by recency once a per-partition HWM reference exists. + // A caught-up fetch expects nothing, so a miss is fine, but a lagging fetch has data to + // catch up on and should go to PG instead of waiting out the timeout. + FindBatchResponse.success( + util.List.of(), + FindBatchResponse.UNKNOWN_OFFSET, + FindBatchResponse.UNKNOWN_OFFSET + ) } else { FindBatchResponse.success( logFragment.batches().stream().map[BatchInfo](batchCoordinate => batchCoordinate.batchInfo(request.topicIdPartition())).toList, @@ -3013,7 +3048,13 @@ class ReplicaManager(val config: KafkaConfig, logDirFailureHandler.shutdown() replicaFetcherManager.shutdown() replicaAlterLogDirsManager.shutdown() + // Stop the consolidation fetchers before their purgatory. A parked fetcher blocks (unbounded) in + // DisklessLeaderEndPoint.awaitDelayedFetch until its op expires, and only the purgatory reaper can + // expire it. Since that thread is non-interruptible, shutting the purgatory down first would strand + // it forever, so the fetcher drain must run while delayedConsolidationFetchPurgatory is still alive. + consolidationFetcherManager.foreach(_.shutdown()) delayedFetchPurgatory.shutdown() + delayedConsolidationFetchPurgatory.shutdown() delayedRemoteFetchPurgatory.shutdown() delayedRemoteListOffsetsPurgatory.shutdown() delayedProducePurgatory.shutdown() @@ -3021,7 +3062,6 @@ class ReplicaManager(val config: KafkaConfig, delayedShareFetchPurgatory.shutdown() if (checkpointHW) checkpointHighWatermarks() - consolidationFetcherManager.foreach(_.shutdown()) consolidationFetchHandler.foreach(_.close()) consolidationMetrics.foreach(_.close()) replicaSelectorPlugin.foreach(_.close) diff --git a/core/src/test/scala/io/aiven/inkless/consolidation/ConsolidationQuotaManagerTest.scala b/core/src/test/scala/io/aiven/inkless/consolidation/ConsolidationQuotaManagerTest.scala index ae0d09c9a21..513d17aede1 100644 --- a/core/src/test/scala/io/aiven/inkless/consolidation/ConsolidationQuotaManagerTest.scala +++ b/core/src/test/scala/io/aiven/inkless/consolidation/ConsolidationQuotaManagerTest.scala @@ -30,7 +30,7 @@ import org.apache.kafka.server.config.ReplicationQuotaManagerConfig import org.apache.kafka.server.network.BrokerEndPoint import org.apache.kafka.server.quota.QuotaType import org.apache.kafka.server.{PartitionFetchState, ReplicaState} -import org.apache.kafka.storage.internals.log.UnifiedLog +import org.apache.kafka.storage.internals.log.{LogConfig, UnifiedLog} import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.{AfterEach, BeforeEach, Test} import org.mockito.Mockito.{mock, when} @@ -70,7 +70,11 @@ class ConsolidationQuotaManagerTest { val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) val replicaManager = mock(classOf[ReplicaManager]) val log = mock(classOf[UnifiedLog]) + val logConfig = mock(classOf[LogConfig]) when(log.logStartOffset).thenReturn(0L) + when(logConfig.segmentSize()).thenReturn(Int.MaxValue) + when(logConfig.maxMessageSize()).thenReturn(1024 * 1024) + when(log.config()).thenReturn(logConfig) when(replicaManager.localLogOrException(tp)).thenReturn(log) val props = TestUtils.createBrokerConfig(brokerEndPoint.id, port = brokerEndPoint.port) diff --git a/core/src/test/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetchTest.scala b/core/src/test/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetchTest.scala new file mode 100644 index 00000000000..ed6fdae075a --- /dev/null +++ b/core/src/test/scala/io/aiven/inkless/consolidation/DelayedConsolidationFetchTest.scala @@ -0,0 +1,339 @@ +/* + * Inkless + * Copyright (C) 2024 - 2026 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 . + */ + +package io.aiven.inkless.consolidation + +import io.aiven.inkless.consume.FetchHandler +import io.aiven.inkless.control_plane.{BatchInfo, BatchMetadata, FindBatchResponse} +import kafka.server.ReplicaManager +import org.apache.kafka.common.{TopicIdPartition, Uuid} +import org.apache.kafka.common.protocol.Errors +import org.apache.kafka.common.record.{MemoryRecords, TimestampType} +import org.apache.kafka.common.requests.FetchRequest +import org.apache.kafka.server.storage.log.{FetchIsolation, FetchParams, FetchPartitionData} +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.function.Executable +import org.mockito.ArgumentMatchers.any +import org.mockito.Mockito.{mock, times, verify, when} + +import java.time.Duration +import java.util +import java.util.concurrent.CompletableFuture +import java.util.concurrent.TimeUnit +import java.util.Optional + +class DelayedConsolidationFetchTest { + private val topicId = Uuid.randomUuid() + private val tip = new TopicIdPartition(topicId, 0, "test-topic") + private val secondTip = new TopicIdPartition(topicId, 1, "test-topic") + + private def newFetchParams(maxWaitMs: Long, minBytes: Int): FetchParams = + new FetchParams( + FetchRequest.FUTURE_LOCAL_REPLICA_ID, + -1, + maxWaitMs, + minBytes, + 1024 * 1024, + FetchIsolation.LOG_END, + Optional.empty() + ) + + private def newFetchInfo(offset: Long, maxBytes: Int): util.Map[TopicIdPartition, FetchRequest.PartitionData] = { + val map = new util.LinkedHashMap[TopicIdPartition, FetchRequest.PartitionData]() + map.put(tip, new FetchRequest.PartitionData(topicId, offset, 0L, maxBytes, Optional.of(0))) + map + } + + private def newFetchInfos(entries: (TopicIdPartition, Long)*): util.Map[TopicIdPartition, FetchRequest.PartitionData] = { + val map = new util.LinkedHashMap[TopicIdPartition, FetchRequest.PartitionData]() + entries.foreach { case (tp, offset) => + map.put(tp, new FetchRequest.PartitionData(tp.topicId, offset, 0L, 1024, Optional.of(0))) + } + map + } + + private def batch(baseOffset: Long, lastOffset: Long, byteSize: Long): BatchInfo = + batchFor(tip, baseOffset, lastOffset, byteSize) + + private def batchFor(topicIdPartition: TopicIdPartition, baseOffset: Long, lastOffset: Long, byteSize: Long): BatchInfo = + new BatchInfo( + baseOffset, + "object-key", + new BatchMetadata( + 2.toByte, topicIdPartition, + baseOffset * 100, byteSize, + baseOffset, lastOffset, + 0L, 0L, + TimestampType.CREATE_TIME + ) + ) + + private def emptyFetchPartitionData: FetchPartitionData = + new FetchPartitionData( + Errors.NONE, + 0L, 0L, + MemoryRecords.EMPTY, + Optional.empty(), + java.util.OptionalLong.empty(), + Optional.empty(), + java.util.OptionalInt.empty(), + false + ) + + @Test + def tryCompleteCompletesWhenAccumulatedAtLeastMinBytes(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + // control plane returns 2 batches of 500 bytes each, totaling 1000 bytes -- above minBytes + val batches = util.List.of(batch(0, 4, 500), batch(5, 9, 500)) + val response = util.List.of(FindBatchResponse.success(batches, 0L, 100L)) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(response)) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData))) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 100), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + assertTrue(op.tryComplete(), "tryComplete should succeed when bytes >= minBytes") + assertTrue(op.isCompleted, "operation must be completed after tryComplete returns true") + assertNotNull(captured.get(1, TimeUnit.SECONDS)) + verify(fetchHandler, times(1)).handle(any(), any()) + } + + @Test + def tryCompleteCompletesWhenAccumulatedAcrossPartitionsReachesMinBytes(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + val firstResponse = FindBatchResponse.success(util.List.of(batchFor(tip, 0, 4, 100)), 0L, 100L) + val secondResponse = FindBatchResponse.success(util.List.of(batchFor(secondTip, 0, 4, 125)), 0L, 100L) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(util.List.of(firstResponse, secondResponse))) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData, secondTip, emptyFetchPartitionData))) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 200), + fetchInfos = newFetchInfos(tip -> 0L, secondTip -> 0L), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + assertTrue(op.tryComplete(), "tryComplete should use aggregate bytes across partitions") + assertEquals(2, captured.get(1, TimeUnit.SECONDS).size) + verify(replicaManager, times(1)).findDisklessBatches(any()) + } + + @Test + def tryCompleteForcesCompletionWhenFindBatchResponseCountDiffers(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + val firstResponse = FindBatchResponse.success(util.List.of(batchFor(tip, 0, 4, 100)), 0L, 100L) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(util.List.of(firstResponse))) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData, secondTip, emptyFetchPartitionData))) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 200), + fetchInfos = newFetchInfos(tip -> 0L, secondTip -> 0L), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + assertTrue(op.tryComplete(), "response-count mismatch should defer to the authoritative fetch") + assertEquals(2, captured.get(1, TimeUnit.SECONDS).size) + verify(fetchHandler, times(1)).handle(any(), any()) + } + + @Test + def tryCompleteParksWhenBelowMinBytes(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + // Empty control-plane response -- accumulated = 0, below minBytes + val response = util.List.of(FindBatchResponse.success(util.List.of(), 0L, 100L)) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(response)) + + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = _ => () + ) + + assertFalse(op.tryComplete(), "tryComplete should park when bytes < minBytes") + assertFalse(op.tryComplete(), "duplicate purgatory registration check should not re-probe the control plane") + assertFalse(op.isCompleted, "operation must remain in purgatory") + verify(replicaManager, times(1)).findDisklessBatches(any()) + verify(fetchHandler, times(0)).handle(any(), any()) + } + + @Test + def expirationCompletesParkedFetchViaFinalFetch(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + val response = util.List.of(FindBatchResponse.success(util.List.of(), 0L, 100L)) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(response)) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData))) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + assertFalse(op.tryComplete(), "initial probe should park below minBytes") + op.run() + + assertEquals(1, captured.get(1, TimeUnit.SECONDS).size) + verify(replicaManager, times(1)).findDisklessBatches(any()) + verify(fetchHandler, times(1)).handle(any(), any()) + } + + @Test + def tryCompleteForcesCompletionOnError(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + // control plane returns OFFSET_OUT_OF_RANGE -- should short-circuit the wait + val response = util.List.of(FindBatchResponse.offsetOutOfRange(0L, 100L)) + when(replicaManager.findDisklessBatches(any())).thenReturn(Some(response)) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData))) + + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = _ => () + ) + + assertTrue(op.tryComplete(), "tryComplete should force completion on control-plane error") + assertTrue(op.isCompleted) + } + + @Test + def tryCompleteForcesCompletionWhenFindDisklessBatchesThrows(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + when(replicaManager.findDisklessBatches(any())).thenThrow(new RuntimeException("control plane down")) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of(tip, emptyFetchPartitionData))) + + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = _ => () + ) + + // Should not propagate the exception -- defer to onComplete + assertTrue(op.tryComplete()) + assertTrue(op.isCompleted) + } + + @Test + def emptyFetchInfosCompletesImmediately(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + when(fetchHandler.handle(any(), any())) + .thenReturn(CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, FetchPartitionData]())) + + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = new util.LinkedHashMap[TopicIdPartition, FetchRequest.PartitionData](), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = _ => () + ) + + assertTrue(op.tryComplete()) + } + + @Test + def onCompleteSurfacesPerPartitionErrorsOnFailure(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + val failed = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + failed.completeExceptionally(new RuntimeException("storage down")) + when(fetchHandler.handle(any(), any())).thenReturn(failed) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + op.forceComplete() + val response = captured.get(1, TimeUnit.SECONDS) + assertNotNull(response) + assertFalse(response.isEmpty, "callback should fire with per-partition errors when handle() fails") + assertEquals(Errors.UNKNOWN_SERVER_ERROR, response.get(tip).error) + } + + @Test + def onCompleteDoesNotBlockOnIncompleteFetchFuture(): Unit = { + val replicaManager = mock(classOf[ReplicaManager]) + val fetchHandler = mock(classOf[FetchHandler]) + + val pending = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + when(fetchHandler.handle(any(), any())).thenReturn(pending) + + val captured = new CompletableFuture[util.Map[TopicIdPartition, FetchPartitionData]]() + val op = new DelayedConsolidationFetch( + params = newFetchParams(maxWaitMs = 1000, minBytes = 1), + fetchInfos = newFetchInfo(0L, 1024), + fetchHandler = fetchHandler, + replicaManager = replicaManager, + responseCallback = r => captured.complete(r) + ) + + assertTimeoutPreemptively(Duration.ofSeconds(1), new Executable { + override def execute(): Unit = assertTrue(op.forceComplete()) + }) + assertFalse(captured.isDone, "callback must wait for the async fetch result") + + pending.complete(util.Map.of(tip, emptyFetchPartitionData)) + assertNotNull(captured.get(1, TimeUnit.SECONDS)) + } +} diff --git a/core/src/test/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPointTest.scala b/core/src/test/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPointTest.scala index ef767ca67eb..efdba146d44 100644 --- a/core/src/test/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPointTest.scala +++ b/core/src/test/scala/io/aiven/inkless/consolidation/DisklessLeaderEndPointTest.scala @@ -18,7 +18,7 @@ package io.aiven.inkless.consolidation -import io.aiven.inkless.consume.{FetchHandler, FetchOffsetHandler} +import io.aiven.inkless.consume.{ConcatenatedRecords, FetchHandler, FetchOffsetHandler} import kafka.cluster.Partition import kafka.server.{KafkaConfig, QuotaFactory, ReplicaManager, ReplicaQuota} import kafka.utils.TestUtils @@ -27,18 +27,20 @@ import org.apache.kafka.server.config.ServerConfigs import org.apache.kafka.common.message.FetchResponseData import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset +import org.apache.kafka.common.compress.Compression import org.apache.kafka.common.protocol.Errors import org.apache.kafka.common.record.FileRecords.TimestampAndOffset -import org.apache.kafka.common.record.MemoryRecords +import org.apache.kafka.common.record.{MemoryRecords, Records, SimpleRecord} import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, ListOffsetsRequest, OffsetsForLeaderEpochResponse} import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.metadata.LeaderAndIsr import org.apache.kafka.server.common.{MetadataVersion, OffsetAndEpoch} import org.apache.kafka.server.network.BrokerEndPoint +import org.apache.kafka.server.purgatory.DelayedOperationPurgatory import org.apache.kafka.server.storage.log.FetchPartitionData import org.apache.kafka.server.{PartitionFetchState, ReplicaState} import org.apache.kafka.storage.internals.log.OffsetResultHolder.FileRecordsOrError -import org.apache.kafka.storage.internals.log.UnifiedLog +import org.apache.kafka.storage.internals.log.{LogConfig, UnifiedLog} import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.Test import org.mockito.ArgumentCaptor @@ -46,6 +48,7 @@ import org.mockito.ArgumentMatchers.{any, eq => eqTo} import org.mockito.Mockito.{doNothing, mock, verify, when} import java.util +import java.nio.ByteBuffer import java.util.Optional import java.util.OptionalInt import java.util.OptionalLong @@ -65,6 +68,33 @@ class DisklessLeaderEndPointTest { KafkaConfig.fromProps(props) } + /** + * A [[ReplicaManager]] mock whose consolidation purgatory completes the delayed fetch inline, so + * [[DisklessLeaderEndPoint.fetch]] runs the mocked [[FetchHandler]] synchronously. + * forceComplete() runs onComplete() directly (bypassing tryComplete/findDisklessBatches), which + * keeps these tests focused on fetch response mapping. The delayed-op mechanics are covered by + * DelayedConsolidationFetchTest. + */ + private def replicaManagerMock(): ReplicaManager = { + val replicaManager = mock(classOf[ReplicaManager]) + val purgatory = mock(classOf[DelayedOperationPurgatory[DelayedConsolidationFetch]]) + when(purgatory.tryCompleteElseWatch(any(), any())).thenAnswer { invocation => + invocation.getArgument(0, classOf[DelayedConsolidationFetch]).forceComplete() + } + when(replicaManager.delayedConsolidationFetchPurgatory).thenReturn(purgatory) + replicaManager + } + + private def unifiedLogMock(logStartOffset: Long, segmentSize: Int, maxMessageSize: Int): UnifiedLog = { + val log = mock(classOf[UnifiedLog]) + val config = mock(classOf[LogConfig]) + when(log.logStartOffset).thenReturn(logStartOffset) + when(config.segmentSize()).thenReturn(segmentSize) + when(config.maxMessageSize()).thenReturn(maxMessageSize) + when(log.config()).thenReturn(config) + log + } + private def newEndPoint( fetchHandler: FetchHandler, fetchOffsetHandler: FetchOffsetHandler, @@ -113,7 +143,7 @@ class DisklessLeaderEndPointTest { ): DisklessLeaderEndPoint = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -140,9 +170,8 @@ class DisklessLeaderEndPointTest { def testBuildFetchProducesReplicaFetch(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) - val log = mock(classOf[UnifiedLog]) - when(log.logStartOffset).thenReturn(11L) + val replicaManager = replicaManagerMock() + val log = unifiedLogMock(logStartOffset = 11L, segmentSize = Int.MaxValue, maxMessageSize = 1024 * 1024) when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) @@ -168,7 +197,7 @@ class DisklessLeaderEndPointTest { def testFetchMapsFetchHandlerResponseToPartitionData(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -192,6 +221,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val result = endPoint.fetch(fetchBuilder).asScala @@ -210,7 +240,7 @@ class DisklessLeaderEndPointTest { def testFetchUsesInvalidLastStableOffsetWhenOptionalEmpty(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -232,6 +262,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -256,7 +287,7 @@ class DisklessLeaderEndPointTest { private def verifyListOffsetTimestamp(expectedTimestamp: Long, invoke: DisklessLeaderEndPoint => OffsetAndEpoch): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -294,7 +325,7 @@ class DisklessLeaderEndPointTest { private def listOffsetEndPointWithPlaceholderEpoch(offset: Long, seal: Long, disklessLeaderEpoch: Int): DisklessLeaderEndPoint = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -346,7 +377,7 @@ class DisklessLeaderEndPointTest { def testListDisklessOffsetThrowsWhenTopicNotDiskless(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) when(fetchOffsetHandler.createJob()).thenReturn(job) @@ -363,7 +394,7 @@ class DisklessLeaderEndPointTest { def testListDisklessOffsetPropagatesHolderException(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -385,7 +416,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsReturnsEmptyForEmptyInput(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) assertTrue(endPoint.fetchEpochEndOffsets(util.Map.of()).isEmpty) @@ -395,7 +426,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsUndefinedEpoch(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) when(fetchOffsetHandler.createJob()).thenReturn(job) doNothing().when(job).start() @@ -421,7 +452,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsUnknownTopicPartition(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) when(fetchOffsetHandler.createJob()).thenReturn(job) @@ -447,7 +478,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsSuccess(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -481,7 +512,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsBelowDisklessEpochReturnsSeal(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) when(fetchOffsetHandler.createJob()).thenReturn(job) @@ -516,7 +547,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsAtOrAboveDisklessEpochReturnsDisklessLeo(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -553,7 +584,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsBornDisklessReturnsDisklessLeo(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -587,7 +618,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsHolderException(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val holder = new FileRecordsOrError( @@ -615,7 +646,7 @@ class DisklessLeaderEndPointTest { def testFetchEpochEndOffsetsFutureGetThrows(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val job = mock(classOf[FetchOffsetHandler.Job]) val failed = new CompletableFuture[FileRecordsOrError]() @@ -641,7 +672,7 @@ class DisklessLeaderEndPointTest { def testBuildFetchReturnsEmptyWhenQuotaExceeded(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val quota = mock(classOf[ReplicaQuota]) when(quota.isQuotaExceeded).thenReturn(true) @@ -665,7 +696,7 @@ class DisklessLeaderEndPointTest { def testBuildFetchMarksPartitionWithKafkaStorageException(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() when(replicaManager.localLogOrException(topicPartition)).thenThrow(new KafkaStorageException("bad log")) val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) @@ -688,7 +719,7 @@ class DisklessLeaderEndPointTest { def testBuildFetchMarksPartitionWithUnknownTopicOrPartitionException(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() when(replicaManager.localLogOrException(topicPartition)).thenThrow(new UnknownTopicOrPartitionException("deleted")) val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) @@ -711,9 +742,8 @@ class DisklessLeaderEndPointTest { def testBuildFetchSkipsPartitionWhenFollowerShouldThrottle(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) - val log = mock(classOf[UnifiedLog]) - when(log.logStartOffset).thenReturn(0L) + val replicaManager = replicaManagerMock() + val log = unifiedLogMock(logStartOffset = 0L, segmentSize = Int.MaxValue, maxMessageSize = 1024 * 1024) when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) val quota = mock(classOf[ReplicaQuota]) @@ -744,7 +774,7 @@ class DisklessLeaderEndPointTest { // the local log (0L), not UNKNOWN_OFFSET. val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -766,6 +796,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -781,7 +812,7 @@ class DisklessLeaderEndPointTest { // the endpoint must redirect just as it does for the NONE+empty-batch case. val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -816,7 +847,7 @@ class DisklessLeaderEndPointTest { def testFetchOverlaysPartitionErrorAndUnknownLogStartWhenLookupFailsAndDisklessWasOk(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() when(replicaManager.getPartitionOrError(topicPartition)).thenReturn(Left(Errors.NOT_LEADER_OR_FOLLOWER)) val fetchData = new FetchPartitionData( @@ -834,6 +865,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -845,7 +877,7 @@ class DisklessLeaderEndPointTest { def testFetchKeepsDisklessErrorWhenLookupFailsButDisklessAlreadyFailed(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() when(replicaManager.getPartitionOrError(topicPartition)).thenReturn(Left(Errors.KAFKA_STORAGE_ERROR)) val fetchData = new FetchPartitionData( @@ -863,6 +895,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -874,7 +907,7 @@ class DisklessLeaderEndPointTest { def testFetchSetsUnknownServerErrorAndUnknownLogStartWhenLocalLogStartAheadOfHighWatermark(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) val localLog = mock(classOf[UnifiedLog]) when(partition.localLogOrException).thenReturn(localLog) @@ -896,6 +929,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -907,7 +941,7 @@ class DisklessLeaderEndPointTest { def testFetchSetsUnknownLogStartWhenLocalLogUnavailable(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) when(partition.localLogOrException).thenThrow(new NotLeaderOrFollowerException("no local log")) when(replicaManager.getPartitionOrError(topicPartition)).thenReturn(Right(partition)) @@ -927,6 +961,7 @@ class DisklessLeaderEndPointTest { val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) @@ -1044,10 +1079,14 @@ class DisklessLeaderEndPointTest { @Test def testFetchDoesNotSignalWhenRequestOmitsOffset(): Unit = { // A request that does not include this partition leaves requestedOffset unresolved (-1): no signal. + // The request carries a different partition (so the fetch is not empty), while the mocked handler + // still returns topicPartition, exercising the getOrElse(-1) fallback in fetch(). val endPoint = consolidatedPrefixEndPoint(localLogStartOffset = 0L, disklessStart = 100L, highWatermark = 200L, remoteLogEnabled = true) - val emptyRequest = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() - val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, emptyRequest).setMaxBytes(100_000) + val otherPartition = new TopicPartition("other-topic", 0) + val request = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + request.put(otherPartition, new FetchRequest.PartitionData(Uuid.randomUuid(), 0L, 0L, 1024 * 1024, Optional.of(Int.box(0)))) + val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, request).setMaxBytes(100_000) val pd = endPoint.fetch(fetchBuilder).get(topicPartition) assertEquals(Errors.NONE.code, pd.errorCode) @@ -1063,10 +1102,9 @@ class DisklessLeaderEndPointTest { // OFFSET_MOVED_TO_TIERED_STORAGE for offset 0. val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) + val replicaManager = replicaManagerMock() val partition = mock(classOf[Partition]) - val localLog = mock(classOf[UnifiedLog]) - when(localLog.logStartOffset).thenReturn(0L) + val localLog = unifiedLogMock(logStartOffset = 0L, segmentSize = Int.MaxValue, maxMessageSize = 1024 * 1024) when(localLog.remoteLogEnabled()).thenReturn(true) when(partition.localLogOrException).thenReturn(localLog) // buildFetch resolves the request log start offset via localLogOrException(tp)... @@ -1109,9 +1147,8 @@ class DisklessLeaderEndPointTest { def testConsolidationFetchConfigsAreUsed(): Unit = { val fetchHandler = mock(classOf[FetchHandler]) val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) - val replicaManager = mock(classOf[ReplicaManager]) - val log = mock(classOf[UnifiedLog]) - when(log.logStartOffset).thenReturn(0L) + val replicaManager = replicaManagerMock() + val log = unifiedLogMock(logStartOffset = 0L, segmentSize = 2 * 1024 * 1024, maxMessageSize = 1024 * 1024) when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) val props = TestUtils.createBrokerConfig(brokerEndPoint.id, port = brokerEndPoint.port) @@ -1150,7 +1187,83 @@ class DisklessLeaderEndPointTest { assertEquals(4096, fetchRequest.minBytes) val partitionData = fetchRequest.fetchData(util.Map.of(topicId, topicPartition.topic)) assertFalse(partitionData.isEmpty, "fetchData should contain the partition") - assertEquals(20971520, partitionData.values().iterator().next().maxBytes) + assertEquals(1024 * 1024, partitionData.values().iterator().next().maxBytes) + } + + @Test + def testConsolidationFetchLeavesHeadroomForWholeBatchOvershoot(): Unit = { + val segmentBytes = 4096 + val maxMessageBytes = 512 + val fetchHandler = mock(classOf[FetchHandler]) + val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) + val replicaManager = replicaManagerMock() + // findBatches returns whole batches, so it may exceed the request maxBytes by one batch. + // Follower append still requires the full returned records block to fit in one segment. + val log = unifiedLogMock(logStartOffset = 0L, segmentSize = segmentBytes, maxMessageSize = maxMessageBytes) + when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) + + val props = TestUtils.createBrokerConfig(brokerEndPoint.id, port = brokerEndPoint.port) + props.setProperty(ServerConfigs.DISKLESS_CONSOLIDATION_FETCH_MAX_BYTES_CONFIG, segmentBytes.toString) + val config = KafkaConfig.fromProps(props) + val endPoint = new DisklessLeaderEndPoint( + brokerEndPoint, + fetchHandler, + fetchOffsetHandler, + replicaManager, + config, + QuotaFactory.UNBOUNDED_QUOTA, + () => MetadataVersion.LATEST_PRODUCTION, + () => 7L + ) + + val fetchState = new PartitionFetchState( + Optional.of(topicId), + 0L, + Optional.empty(), + 1, + Optional.empty(), + ReplicaState.FETCHING, + Optional.empty() + ) + + val result = endPoint.buildFetch(util.Map.of(topicPartition, fetchState)) + assertTrue(result.result.isPresent) + + val fetchRequest = result.result.get.fetchRequest.build() + val partitionData = fetchRequest.fetchData(util.Map.of(topicId, topicPartition.topic)) + val maxBytes = partitionData.values().iterator().next().maxBytes + // Leave one max-message-sized batch of headroom below segment.bytes. + assertEquals(segmentBytes - maxMessageBytes, maxBytes) + assertTrue(maxBytes + maxMessageBytes <= segmentBytes) + } + + @Test + def testConsolidationFetchClampsUnsafeSegmentConfigToOneByte(): Unit = { + val segmentBytes = 1024 * 1024 + val maxMessageBytes = segmentBytes + 12 + val fetchHandler = mock(classOf[FetchHandler]) + val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) + val replicaManager = replicaManagerMock() + val log = unifiedLogMock(logStartOffset = 0L, segmentSize = segmentBytes, maxMessageSize = maxMessageBytes) + when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) + + val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) + val fetchState = new PartitionFetchState( + Optional.of(topicId), + 0L, + Optional.empty(), + 1, + Optional.empty(), + ReplicaState.FETCHING, + Optional.empty() + ) + + val result = endPoint.buildFetch(util.Map.of(topicPartition, fetchState)) + assertTrue(result.result.isPresent) + + val fetchRequest = result.result.get.fetchRequest.build() + val partitionData = fetchRequest.fetchData(util.Map.of(topicId, topicPartition.topic)) + assertEquals(1, partitionData.values().iterator().next().maxBytes) } @Test @@ -1158,13 +1271,192 @@ class DisklessLeaderEndPointTest { val props = TestUtils.createBrokerConfig(brokerEndPoint.id, port = brokerEndPoint.port) val config = KafkaConfig.fromProps(props) - assertEquals(1 * 1024 * 1024, config.disklessConsolidationFetchMaxBytes) - assertEquals(10 * 1024 * 1024, config.disklessConsolidationFetchResponseMaxBytes) - assertEquals(1, config.disklessConsolidationFetchMinBytes) - assertEquals(500, config.disklessConsolidationFetchMaxWaitMs) + assertEquals(10 * 1024 * 1024, config.disklessConsolidationFetchMaxBytes) + assertEquals(64 * 1024 * 1024, config.disklessConsolidationFetchResponseMaxBytes) + assertEquals(8 * 1024 * 1024, config.disklessConsolidationFetchMinBytes) + assertEquals(1000, config.disklessConsolidationFetchMaxWaitMs) assertEquals(1, config.disklessConsolidationNumFetchers) assertEquals(0, config.disklessConsolidationFindBatchesMaxPerPartition) assertEquals(4, config.disklessConsolidationFetchMetadataThreadPoolSize) assertEquals(8, config.disklessConsolidationFetchDataThreadPoolSize) } + + /** Single-batch MemoryRecords starting at `baseOffset` with `numRecords` records of `valueSize` bytes each. */ + private def singleBatch(baseOffset: Long, numRecords: Int, valueSize: Int): MemoryRecords = { + val value = new Array[Byte](valueSize) + val records = (0 until numRecords).map(_ => new SimpleRecord(value)).toArray + MemoryRecords.withRecords(baseOffset, Compression.NONE, records: _*) + } + + /** + * One coordinate holding `batchCount` byte-contiguous single-record physical batches (offsets + * `baseOffset .. baseOffset + batchCount - 1`), modelling a `commit_file_v2` coalesced run: multiple + * physical batches materialized into a single MemoryRecords. + */ + private def coalescedCoordinate(baseOffset: Long, batchCount: Int, valueSize: Int): MemoryRecords = { + val batches = (0 until batchCount).map(i => singleBatch(baseOffset + i, 1, valueSize)) + val buffer = ByteBuffer.allocate(batches.map(_.sizeInBytes()).sum) + batches.foreach(b => buffer.put(b.buffer().duplicate())) + buffer.flip() + MemoryRecords.readableRecords(buffer) + } + + private def baseOffsets(records: Records): Seq[Long] = + records.batches().iterator().asScala.map(_.baseOffset()).toSeq + + /** + * Runs [[DisklessLeaderEndPoint.fetch]] with the fetch handler returning `records`, a per-partition + * request carrying `fetchOffset`/`maxBytes`, and a healthy local log (so no tiered redirect fires and + * the records survive to the response). + */ + private def fetchClampedRecords(records: Records, fetchOffset: Long, maxBytes: Int): Records = { + val fetchHandler = mock(classOf[FetchHandler]) + val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) + val replicaManager = replicaManagerMock() + val partition = mock(classOf[Partition]) + val localLog = mock(classOf[UnifiedLog]) + when(partition.localLogOrException).thenReturn(localLog) + when(localLog.logStartOffset).thenReturn(0L) + when(replicaManager.getPartitionOrError(topicPartition)).thenReturn(Right(partition)) + + val fetchData = new FetchPartitionData( + Errors.NONE, + 1000L, + 0L, + records, + Optional.empty(), + OptionalLong.empty(), + Optional.empty(), + OptionalInt.empty(), + false + ) + when(fetchHandler.handle(any(), any())).thenReturn(CompletableFuture.completedFuture(Map(topicIdPartition -> fetchData).asJava)) + + val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) + val partitionData = new util.HashMap[TopicPartition, FetchRequest.PartitionData]() + partitionData.put(topicPartition, new FetchRequest.PartitionData(topicId, fetchOffset, 0L, maxBytes, Optional.of(Int.box(0)))) + val fetchBuilder = FetchRequest.Builder.forReplica(12, 1, 1L, 100, 1, partitionData).setMaxBytes(100_000_000) + endPoint.fetch(fetchBuilder).get(topicPartition).records.asInstanceOf[Records] + } + + @Test + def testFetchClampStopsAtBatchBoundaryOnceOverBudget(): Unit = { + // Three equal-sized physical batches at offsets 0,1,2. With maxBytes just under the size of the + // first two batches, the first is taken (progress), the second is the one that crosses the budget + // and is still included (bounded overshoot), and the third is dropped. + val batch0 = singleBatch(baseOffset = 0L, numRecords = 1, valueSize = 256) + val size = batch0.sizeInBytes() + val records = new ConcatenatedRecords(util.List.of( + batch0, + singleBatch(baseOffset = 1L, numRecords = 1, valueSize = 256), + singleBatch(baseOffset = 2L, numRecords = 1, valueSize = 256) + )) + + val clamped = fetchClampedRecords(records, fetchOffset = 0L, maxBytes = 2 * size - 1) + + assertEquals(Seq(0L, 1L), baseOffsets(clamped)) + assertEquals(2 * size, clamped.sizeInBytes()) + } + + @Test + def testFetchClampSkipsBatchesBelowFetchOffset(): Unit = { + // A prior fetch truncated a coalesced coordinate, so this re-fetch starts mid-run at offset 1. + // Batches ending before the fetch offset must be dropped so the follower does not re-append them. + val records = new ConcatenatedRecords(util.List.of( + singleBatch(baseOffset = 0L, numRecords = 1, valueSize = 64), + singleBatch(baseOffset = 1L, numRecords = 1, valueSize = 64), + singleBatch(baseOffset = 2L, numRecords = 1, valueSize = 64) + )) + + val clamped = fetchClampedRecords(records, fetchOffset = 1L, maxBytes = 10 * 1024 * 1024) + + assertEquals(Seq(1L, 2L), baseOffsets(clamped)) + } + + @Test + def testFetchClampAlwaysEmitsOneBatchForProgress(): Unit = { + // maxBytes smaller than a single batch: the first batch is still emitted (progress) and the next + // batch is dropped, so a small budget cannot stall the fetcher. + val batch0 = singleBatch(baseOffset = 0L, numRecords = 1, valueSize = 512) + val records = new ConcatenatedRecords(util.List.of( + batch0, + singleBatch(baseOffset = 1L, numRecords = 1, valueSize = 512) + )) + + val clamped = fetchClampedRecords(records, fetchOffset = 0L, maxBytes = 1) + + assertEquals(Seq(0L), baseOffsets(clamped)) + assertEquals(batch0.sizeInBytes(), clamped.sizeInBytes()) + } + + @Test + def testFetchClampIsNoOpWhenBlockFitsBudget(): Unit = { + // Everything fits under maxBytes and starts at the fetch offset: the original records instance is + // returned unchanged to preserve the ConcatenatedRecords zero-copy send path. + val records = new ConcatenatedRecords(util.List.of( + singleBatch(baseOffset = 0L, numRecords = 1, valueSize = 64), + singleBatch(baseOffset = 1L, numRecords = 1, valueSize = 64) + )) + + val clamped = fetchClampedRecords(records, fetchOffset = 0L, maxBytes = 10 * 1024 * 1024) + + assertSame(records, clamped) + } + + @Test + def testFetchClampKeepsCoalescedBlockWithinSegmentEndToEnd(): Unit = { + // With segment.bytes tuned small (supported, e.g. for eager tiering), the control plane can admit two + // coalesced coordinates (each a multi-batch commit_file_v2 run) because its per-partition byte limit + // is applied at coordinate granularity and always includes the coordinate that crosses the limit. + // Their combined size then exceeds segment.bytes, which would make the follower's single append fail + // with RecordBatchTooLargeException. buildFetch reserves maxBytes = segment.bytes - max.message.bytes; + // the per-batch clamp in fetch() must trim the block back within the segment, splitting a coalesced + // coordinate mid-run if needed. + val segmentBytes = 4096 + val maxMessageBytes = 1024 + + val fetchHandler = mock(classOf[FetchHandler]) + val fetchOffsetHandler = mock(classOf[FetchOffsetHandler]) + val replicaManager = replicaManagerMock() + val log = unifiedLogMock(logStartOffset = 0L, segmentSize = segmentBytes, maxMessageSize = maxMessageBytes) + when(replicaManager.localLogOrException(topicPartition)).thenReturn(log) + val partition = mock(classOf[Partition]) + when(partition.localLogOrException).thenReturn(log) + when(replicaManager.getPartitionOrError(topicPartition)).thenReturn(Right(partition)) + + // Coordinate A (offsets 0..4) and B (offsets 5..8), each physical batch < max.message.bytes. + val coordA = coalescedCoordinate(baseOffset = 0L, batchCount = 5, valueSize = 500) + val coordB = coalescedCoordinate(baseOffset = 5L, batchCount = 4, valueSize = 500) + val combined = new ConcatenatedRecords(util.List.of(coordA, coordB)) + assertTrue(combined.sizeInBytes() > segmentBytes, + s"precondition: the unclamped block (${combined.sizeInBytes()}) must overflow the segment ($segmentBytes)") + + val fetchData = new FetchPartitionData( + Errors.NONE, 1000L, 0L, combined, + Optional.empty(), OptionalLong.empty(), Optional.empty(), OptionalInt.empty(), false + ) + when(fetchHandler.handle(any(), any())).thenReturn(CompletableFuture.completedFuture(Map(topicIdPartition -> fetchData).asJava)) + + val endPoint = newEndPoint(fetchHandler, fetchOffsetHandler, replicaManager) + val fetchState = new PartitionFetchState( + Optional.of(topicId), 0L, Optional.empty(), 4, ReplicaState.FETCHING, Optional.empty() + ) + + // buildFetch reserves the segment-safe budget, then fetch() enforces it per batch on the same request. + val replicaFetch = endPoint.buildFetch(util.Map.of(topicPartition, fetchState)) + assertTrue(replicaFetch.result.isPresent) + val requestedMaxBytes = replicaFetch.result.get.partitionData.get(topicPartition).maxBytes + assertEquals(segmentBytes - maxMessageBytes, requestedMaxBytes) + + val clamped = endPoint.fetch(replicaFetch.result.get.fetchRequest).get(topicPartition).records.asInstanceOf[Records] + + // Solved: the block the follower appends now fits the segment (no RecordBatchTooLargeException)... + assertTrue(clamped.sizeInBytes() <= segmentBytes, + s"clamped block (${clamped.sizeInBytes()}) must fit the segment ($segmentBytes)") + // ...within the classic one-batch overshoot bound... + assertTrue(clamped.sizeInBytes() <= requestedMaxBytes + maxMessageBytes) + // ...and it was actually trimmed (a coalesced coordinate split mid-run), starting at the fetch offset. + assertTrue(clamped.sizeInBytes() < combined.sizeInBytes()) + assertEquals(0L, baseOffsets(clamped).head) + } } diff --git a/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java b/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java index 705054dd619..e22c73d1a08 100644 --- a/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java +++ b/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java @@ -162,12 +162,14 @@ public class ServerConfigs { "Kafka tiered log storage on the configured topic."; public static final String DISKLESS_CONSOLIDATION_FETCH_MAX_BYTES_CONFIG = "diskless.consolidation.fetch.max.bytes"; - public static final int DISKLESS_CONSOLIDATION_FETCH_MAX_BYTES_DEFAULT = 1024 * 1024; // 1MB + public static final int DISKLESS_CONSOLIDATION_FETCH_MAX_BYTES_DEFAULT = 10 * 1024 * 1024; // 10MB public static final String DISKLESS_CONSOLIDATION_FETCH_MAX_BYTES_DOC = "The maximum number of bytes per partition the consolidation " + - "fetcher will request per iteration. Larger values fetch more batches per iteration, reducing control-plane query frequency."; + "fetcher will request per iteration. Larger values fetch more batches per iteration, reducing control-plane query frequency. " + + "For consolidated topics, keep topic-level segment.bytes greater than max.message.bytes; otherwise the fetcher must clamp " + + "the per-partition request size to leave headroom for a whole-batch control-plane overshoot."; public static final String DISKLESS_CONSOLIDATION_FETCH_RESPONSE_MAX_BYTES_CONFIG = "diskless.consolidation.fetch.response.max.bytes"; - public static final int DISKLESS_CONSOLIDATION_FETCH_RESPONSE_MAX_BYTES_DEFAULT = 10 * 1024 * 1024; // 10MB + public static final int DISKLESS_CONSOLIDATION_FETCH_RESPONSE_MAX_BYTES_DEFAULT = 64 * 1024 * 1024; // 64MB public static final String DISKLESS_CONSOLIDATION_FETCH_RESPONSE_MAX_BYTES_DOC = "The maximum total bytes the consolidation fetcher " + "will accept across all partitions in a single fetch response. Caps aggregate network I/O per request when multiple " + "partitions are fetched from the same leader."; @@ -195,12 +197,12 @@ public class ServerConfigs { "to concurrently fetch data files from remote storage. Lower than the consumer fetch pool since consolidation is background work."; public static final String DISKLESS_CONSOLIDATION_FETCH_MIN_BYTES_CONFIG = "diskless.consolidation.fetch.min.bytes"; - public static final int DISKLESS_CONSOLIDATION_FETCH_MIN_BYTES_DEFAULT = 1; + public static final int DISKLESS_CONSOLIDATION_FETCH_MIN_BYTES_DEFAULT = 8 * 1024 * 1024; // 8MB public static final String DISKLESS_CONSOLIDATION_FETCH_MIN_BYTES_DOC = "The minimum number of bytes the consolidation fetcher will wait for " + "before returning. Combined with diskless.consolidation.fetch.max.wait.ms, controls how often the fetcher wakes up when there is little new data."; public static final String DISKLESS_CONSOLIDATION_FETCH_MAX_WAIT_MS_CONFIG = "diskless.consolidation.fetch.max.wait.ms"; - public static final int DISKLESS_CONSOLIDATION_FETCH_MAX_WAIT_MS_DEFAULT = 500; + public static final int DISKLESS_CONSOLIDATION_FETCH_MAX_WAIT_MS_DEFAULT = 1000; public static final String DISKLESS_CONSOLIDATION_FETCH_MAX_WAIT_MS_DOC = "The maximum time the consolidation fetcher will wait " + "for minBytes of data before returning. Higher values reduce iteration frequency when there is little new data."; public static final String DISKLESS_CONSOLIDATION_FETCH_RATE_LIMIT_BYTES_PER_SECOND_CONFIG = "diskless.consolidation.fetch.rate.limit.bytes.per.second";