Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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 <http://www.gnu.org/licenses/>.
*/

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
Comment thread
jeqo marked this conversation as resolved.
}

// 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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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) =>
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -315,15 +410,29 @@ 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,
new FetchRequest.PartitionData(
fetchState.topicId().orElse(Uuid.ZERO_UUID),
fetchState.fetchOffset(),
logStartOffset,
fetchSize,
partitionFetchSize,
Optional.of(fetchState.currentLeaderEpoch()),
lastFetchedEpoch
)
Expand Down
Loading
Loading