fix(log): Prevent potential offset overflow in ElasticLogSegment (#2720)

* fix(log): Prevent potential offset overflow in ElasticLogSegment

This commit addresses an issue where a log segment could accommodate more than Integer.MAX_VALUE records, leading to a potential integer overflow when calculating relative offsets.

The root cause was that the check `offset - baseOffset <= Integer.MAX_VALUE` allowed a relative offset to be exactly `Integer.MAX_VALUE`. Since offsets are 0-based, this allows for `Integer.MAX_VALUE + 1` records, which cannot be represented by a standard Integer.

This fix implements the following changes:
1.  In `ElasticLogSegment`, the offset validation is changed from `<=` to `< Integer.MAX_VALUE` to ensure the relative offset strictly fits within an Integer's bounds.
2.  In `LogCleaner`, a new segment grouping method `groupSegmentsBySizeV2` is introduced for `ElasticUnifiedLog`. This method uses the same stricter offset check to prevent incorrectly grouping segments that would exceed the offset limit.
3.  The corresponding unit tests in `LogCleanerTest` have been updated to reflect these new boundaries and validate the fix.

Fixes: #2718

* fix(logCleaner): unify segment grouping logic

* fix(logCleaner): extract offset range check for segment grouping to prevent overflow in ElasticLogSegment

* style(logCleaner): fix indentation in segment grouping condition for readability

* style(logCleaner): fix line break in offset range check for readability

* chore: add AutoMQ inject

* style(logCleaner): remove unnecessary blank line after segment grouping

* fix(stream): validate record batch count to prevent negative values in append

Author: Gezi-lzq

core/src/main/scala/kafka/log/LogCleaner.scala

@@ -17,31 +17,31 @@
 
 package kafka.log
 
-import java.io.{File, IOException}
-import java.nio._
-import java.util.Date
-import java.util.concurrent.TimeUnit
 import kafka.common._
-import kafka.log.LogCleaner.{CleanerRecopyPercentMetricName, DeadThreadCountMetricName, MaxBufferUtilizationPercentMetricName, MaxCleanTimeMetricName, MaxCompactionDelayMetricsName}
+import kafka.log.LogCleaner._
 import kafka.log.streamaspect.ElasticLogSegment
 import kafka.server.{BrokerReconfigurable, KafkaConfig}
 import kafka.utils.{Logging, Pool}
-import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.config.ConfigException
 import org.apache.kafka.common.errors.{CorruptRecordException, KafkaStorageException}
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.{BufferSupplier, Time}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.server.config.ServerConfigs
 import org.apache.kafka.server.metrics.KafkaMetricsGroup
 import org.apache.kafka.server.util.ShutdownableThread
-import org.apache.kafka.storage.internals.log.{AbortedTxn, CleanerConfig, LastRecord, LogDirFailureChannel, LogSegment, LogSegmentOffsetOverflowException, OffsetMap, SkimpyOffsetMap, TransactionIndex}
+import org.apache.kafka.storage.internals.log._
 import org.apache.kafka.storage.internals.utils.Throttler
 
-import scala.jdk.CollectionConverters._
+import java.io.{File, IOException}
+import java.nio._
+import java.util.Date
+import java.util.concurrent.TimeUnit
 import scala.collection.mutable.ListBuffer
 import scala.collection.{Iterable, Seq, Set, mutable}
+import scala.jdk.CollectionConverters._
 import scala.util.control.ControlThrowable
 
 /**
@@ -977,6 +977,18 @@ private[log] class Cleaner(val id: Int,
   private[log] def groupSegmentsBySize(segments: Iterable[LogSegment], maxSize: Int, maxIndexSize: Int, firstUncleanableOffset: Long): List[Seq[LogSegment]] = {
     var grouped = List[List[LogSegment]]()
     var segs = segments.toList
+
+    // AutoMQ inject start
+    def isOffsetRangeValid(group: List[LogSegment]) = {
+      val offsetRange = lastOffsetForFirstSegment(segs, firstUncleanableOffset) - group.last.baseOffset
+      // For ElasticLogSegment, use a stricter offset range check (`< Int.MaxValue`) to prevent a potential overflow
+      // issue as described in https://github.com/AutoMQ/automq/issues/2717.
+      // For other segment types, the original less-strict check (`<= Int.MaxValue`) is retained.
+      if (group.last.isInstanceOf[ElasticLogSegment]) offsetRange < Int.MaxValue
+      else offsetRange <= Int.MaxValue
+    }
+    // AutoMQ inject end
+
     while (segs.nonEmpty) {
       var group = List(segs.head)
       var logSize = segs.head.size.toLong
@@ -990,7 +1002,9 @@ private[log] class Cleaner(val id: Int,
             //if first segment size is 0, we don't need to do the index offset range check.
             //this will avoid empty log left every 2^31 message.
             (segs.head.size == 0 ||
-              lastOffsetForFirstSegment(segs, firstUncleanableOffset) - group.last.baseOffset <= Int.MaxValue)) {
+              // AutoMQ inject start
+              isOffsetRangeValid(group))) {
+              // AutoMQ inject end
         group = segs.head :: group
         logSize += segs.head.size
         indexSize += offsetIndexSize(segs.head)

core/src/main/scala/kafka/log/streamaspect/ElasticLogSegment.java

@@ -39,6 +39,7 @@
 import org.apache.kafka.storage.internals.log.LogFileUtils;
 import org.apache.kafka.storage.internals.log.LogOffsetMetadata;
 import org.apache.kafka.storage.internals.log.LogSegment;
+import org.apache.kafka.storage.internals.log.LogSegmentOffsetOverflowException;
 import org.apache.kafka.storage.internals.log.OffsetIndex;
 import org.apache.kafka.storage.internals.log.OffsetPosition;
 import org.apache.kafka.storage.internals.log.ProducerAppendInfo;
@@ -194,9 +195,26 @@ public int size() {
         return log.sizeInBytes();
     }
 
+    /**
+     * Checks that the argument offset can be represented as an integer offset relative to the baseOffset.
+     * This method is similar in purpose to {@see org.apache.kafka.storage.internals.log.LogSegment#canConvertToRelativeOffset}.
+     * <p>
+     * The implementation is inspired by {@see org.apache.kafka.storage.internals.log.AbstractIndex#canAppendOffset},
+     * but uses {@code < Integer.MAX_VALUE} instead of {@code <= Integer.MAX_VALUE} to address an offset overflow issue.
+     *
+     * @param offset The offset to check.
+     * @return true if the offset can be converted, false otherwise.
+     * @see <a href="https://github.com/AutoMQ/automq/issues/2718">Issue #2718</a>
+     */
     private boolean canConvertToRelativeOffset(long offset) {
         long relativeOffset = offset - baseOffset;
-        return relativeOffset >= 0 && relativeOffset <= Integer.MAX_VALUE;
+        // Note: The check is `relativeOffset < Integer.MAX_VALUE` instead of `<=` to avoid overflow.
+        // See https://github.com/AutoMQ/automq/issues/2718 for details.
+        return relativeOffset >= 0 && relativeOffset < Integer.MAX_VALUE;
+    }
+    private void ensureOffsetInRange(long offset) throws IOException {
+        if (!canConvertToRelativeOffset(offset))
+            throw new LogSegmentOffsetOverflowException(this, offset);
     }
 
     @Override
@@ -216,6 +234,8 @@ public void append(
                 meta.firstBatchTimestamp(largestTimestampMs);
             }
 
+            ensureOffsetInRange(largestOffset);
+
             // append the messages
             long appendedBytes = log.append(records, largestOffset + 1);
             if (LOGGER.isTraceEnabled()) {

core/src/test/scala/unit/kafka/log/LogCleanerTest.scala

@@ -30,7 +30,7 @@ import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.coordinator.transaction.TransactionLogConfigs
 import org.apache.kafka.server.metrics.{KafkaMetricsGroup, KafkaYammerMetrics}
 import org.apache.kafka.server.util.MockTime
-import org.apache.kafka.storage.internals.log.{AbortedTxn, AppendOrigin, CleanerConfig, LogAppendInfo, LogConfig, LogDirFailureChannel, LogFileUtils, LogSegment, LogSegments, LogStartOffsetIncrementReason, OffsetMap, ProducerStateManager, ProducerStateManagerConfig}
+import org.apache.kafka.storage.internals.log._
 import org.apache.kafka.storage.internals.utils.Throttler
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, Test}
@@ -1437,8 +1437,20 @@ class LogCleanerTest extends Logging {
     //create 3 segments
     for (i <- 0 until 3) {
       log.appendAsLeader(TestUtils.singletonRecords(value = v, key = k), leaderEpoch = 0)
-      //0 to Int.MaxValue is Int.MaxValue+1 message, -1 will be the last message of i-th segment
-      val records = messageWithOffset(k, v, (i + 1L) * (Int.MaxValue + 1L) -1 )
+
+      // AutoMQ inject start
+      val records = if (log.isInstanceOf[ElasticUnifiedLog]) {
+        // Create a sparse segment by appending a record with a large offset.
+        // A segment can contain up to Int.MaxValue messages (see https://github.com/AutoMQ/automq/issues/2717).
+        // The offset `(i + 1L) * Int.MaxValue - 1` ensures that this is the last message of the i-th segment,
+        // creating a large gap to the next segment's base offset. This helps test segment grouping with sparse offsets.
+        messageWithOffset(k, v, (i + 1L) * Int.MaxValue - 1)
+      } else {
+        //0 to Int.MaxValue is Int.MaxValue+1 message, -1 will be the last message of i-th segment
+        messageWithOffset(k, v, (i + 1L) * (Int.MaxValue + 1L) -1 )
+      }
+
+      // AutoMQ inject end
       log.appendAsFollower(records)
       assertEquals(i + 1, log.numberOfSegments)
     }
@@ -1469,6 +1481,13 @@ class LogCleanerTest extends Logging {
     //trigger a clean and 2 empty segments should cleaned to 1
     cleaner.clean(LogToClean(log.topicPartition, log, 0, firstUncleanableOffset))
     assertEquals(totalSegments - 1, log.numberOfSegments)
+
+    // AutoMQ inject start
+
+    // after clean, the 2nd and 3rd segment should be none empty
+    assertEquals(2, log.logSegments.asScala.takeRight(2).count(_.size > 0))
+
+    // AutoMQ inject end
   }
 
   /**
@@ -1492,12 +1511,21 @@ class LogCleanerTest extends Logging {
       log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0)
 
     // forward offset and append message to next segment at offset Int.MaxValue
-    val records = messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 1)
+    // AutoMQ inject start
+    val records = if (log.isInstanceOf[ElasticUnifiedLog]) {
+      messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 2)
+    } else {
+      messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 1)
+    }
+    // AutoMQ inject end
+
     log.appendAsFollower(records)
     log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0)
 
     // AutoMQ inject start
-    if (!log.isInstanceOf[ElasticUnifiedLog]) {
+    if (log.isInstanceOf[ElasticUnifiedLog]) {
+      assertEquals(Int.MaxValue - 1, log.activeSegment.readNextOffset() - 1)
+    } else {
       assertEquals(Int.MaxValue, log.activeSegment.offsetIndex.lastOffset)
     }
     // AutoMQ inject end

s3stream/src/main/java/com/automq/stream/s3/S3Stream.java

@@ -183,6 +183,9 @@ public CompletableFuture<AppendResult> append(AppendContext context, RecordBatch
         if (snapshotRead()) {
             return FutureUtil.failedFuture(new IllegalStateException("Append operation is not support for readonly stream"));
         }
+        if (recordBatch.count() < 0) {
+            return FutureUtil.failedFuture(new IllegalArgumentException("record batch count is negative"));
+        }
         long startTimeNanos = System.nanoTime();
         readLock.lock();
         try {