fix(issues1784): fix upload burst cause object index out of order (AutoMQ#1785)

Signed-off-by: Robin Han <hanxvdovehx@gmail.com>

Author: superhx

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

@@ -111,10 +111,6 @@ public void burst() {
     }
 
     private CompletableFuture<Void> acquireLimiter(int size) {
-        if (this.burst) {
-            return CompletableFuture.completedFuture(null);
-        }
-
         return limiter.acquire(size);
     }
 

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

@@ -38,7 +38,8 @@ public interface ObjectWriter {
     // TODO: first n bit is the compressed flag
     byte DATA_BLOCK_DEFAULT_FLAG = 0x02;
 
-    static ObjectWriter writer(long objectId, ObjectStorage objectStorage, int blockSizeThreshold, int partSizeThreshold) {
+    static ObjectWriter writer(long objectId, ObjectStorage objectStorage, int blockSizeThreshold,
+        int partSizeThreshold) {
         return new DefaultObjectWriter(objectId, objectStorage, blockSizeThreshold, partSizeThreshold);
     }
 
@@ -79,11 +80,14 @@ class DefaultObjectWriter implements ObjectWriter {
         private IndexBlock indexBlock;
         private long size;
 
+        private long lastStreamId = Constants.NOOP_STREAM_ID;
+        private long lastEndOffset = Constants.NOOP_OFFSET;
+
         /**
          * Create a new object writer.
          *
          * @param objectId           object id
-         * @param objectStorage         S3 operator
+         * @param objectStorage      S3 operator
          * @param blockSizeThreshold the max size of a block
          * @param partSizeThreshold  the max size of a part. If it is smaller than {@link Writer#MIN_PART_SIZE}, it will be set to {@link Writer#MIN_PART_SIZE}.
          */
@@ -98,6 +102,7 @@ public DefaultObjectWriter(long objectId, ObjectStorage objectStorage, int block
         }
 
         public synchronized void write(long streamId, List<StreamRecordBatch> records) {
+            check(streamId, records);
             List<List<StreamRecordBatch>> blocks = groupByBlock(records);
             for (List<StreamRecordBatch> blockRecords : blocks) {
                 DataBlock block = new DataBlock(streamId, blockRecords);
@@ -109,6 +114,32 @@ public synchronized void write(long streamId, List<StreamRecordBatch> records) {
             }
         }
 
+        private void check(long streamId, List<StreamRecordBatch> records) {
+            if (records.isEmpty()) {
+                return;
+            }
+            long recordsEndOffset = records.get(records.size() - 1).getLastOffset();
+            if (lastStreamId == Constants.NOOP_STREAM_ID) {
+                lastStreamId = streamId;
+                lastEndOffset = recordsEndOffset;
+                return;
+            }
+            if (lastStreamId > streamId) {
+                throw new IllegalArgumentException(String.format("The incoming streamId=%s is less than last streamId=%s", streamId, lastStreamId));
+            } else if (lastStreamId == streamId) {
+                long recordsStartOffset = records.get(0).getBaseOffset();
+                if (recordsStartOffset < lastEndOffset) {
+                    throw new IllegalArgumentException(String.format("The incoming streamId=%s startOffset=%s is less than lastEndOffset=%s",
+                        streamId, recordsStartOffset, lastEndOffset));
+                } else {
+                    lastEndOffset = recordsEndOffset;
+                }
+            } else {
+                lastStreamId = streamId;
+                lastEndOffset = recordsEndOffset;
+            }
+        }
+
         private List<List<StreamRecordBatch>> groupByBlock(List<StreamRecordBatch> records) {
             List<List<StreamRecordBatch>> blocks = new LinkedList<>();
             List<StreamRecordBatch> blockRecords = new ArrayList<>(records.size());

s3stream/src/main/java/com/automq/stream/utils/AsyncRateLimiter.java

@@ -36,6 +36,9 @@ public AsyncRateLimiter(double bytesPerSec) {
         tickTask = SCHEDULER.scheduleAtFixedRate(this::tick, 1, 1, TimeUnit.MILLISECONDS);
     }
 
+    /**
+     * Async acquire permit. The returned future will be orderly completed when the permit is acquired.
+     */
     public synchronized CompletableFuture<Void> acquire(int size) {
         if (acquireQueue.isEmpty() && rateLimiter.tryAcquire(size)) {
             return CompletableFuture.completedFuture(null);

s3stream/src/test/java/com/automq/stream/s3/ObjectWriterTest.java

@@ -21,6 +21,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 
@@ -103,6 +104,28 @@ public void testWrite() throws ExecutionException, InterruptedException {
         }
     }
 
+    @Test
+    public void testWrite_check() {
+        S3ObjectMetadata metadata = new S3ObjectMetadata(1, 0, S3ObjectType.STREAM_SET);
+
+        ObjectStorage objectStorage = new MemoryObjectStorage();
+        ObjectWriter objectWriter = ObjectWriter.writer(1, objectStorage, 1024, 1024);
+        objectWriter.write(233, List.of(
+            newRecord(233, 10, 5, 512),
+            newRecord(233, 15, 5, 512)
+        ));
+
+        // write smaller stream
+        Assertions.assertThrowsExactly(IllegalArgumentException.class, () -> objectWriter.write(231, List.of(newRecord(231, 10, 5, 512))));
+
+        objectWriter.write(233, List.of(newRecord(233, 20, 5, 512)));
+
+        // write smaller offset
+        Assertions.assertThrowsExactly(IllegalArgumentException.class, () -> objectWriter.write(233, List.of(newRecord(233, 10, 5, 512))));
+
+        objectWriter.write(234, List.of(newRecord(234, 0, 5, 512)));
+    }
+
     StreamRecordBatch newRecord(long streamId, long offset, int count, int payloadSize) {
         return new StreamRecordBatch(streamId, 0, offset, count, TestUtils.random(payloadSize));
     }