refactor(s3stream/object-wal): complete appends sequentially (#2550)

refactor(s3stream/object-wal): complete appends sequentially (#2426)

* chore: add todos



* refactor(s3stream/object-wal): sequentially succeed



* refactor(s3stream/object-wal): drop discontinuous objects during recovery



* test: introduce `MockObjectStorage`



* test: test sequentially succeed



* refactor: record endOffset in the object path



* feat: different version of wal object header



* refactor: adapt to the new object header format



* feat: recover from the trim offset



* test: test recover continuous objects from trim offset



* test: test marshal and unmarshal wal object header



* test: fix tests



* test: test recover from discontinuous objects



* test: test recover from v0 and v1 objects



* style: fix lint



---------

Signed-off-by: Ning Yu <[email protected]>

Author: Chillax-0v0

s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java

@@ -36,6 +36,8 @@
 
 import java.io.IOException;
 import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
@@ -67,7 +69,7 @@ public ObjectWALService(Time time, ObjectStorage objectStorage, ObjectWALConfig
     }
 
     // Visible for testing.
-    protected RecordAccumulator accumulator() {
+    RecordAccumulator accumulator() {
         return accumulator;
     }
 
@@ -192,7 +194,8 @@ public static class RecoverIterator implements Iterator<RecoverResult> {
 
         public RecoverIterator(List<RecordAccumulator.WALObject> objectList, ObjectStorage objectStorage,
             int readAheadObjectSize) {
-            this.objectList = objectList;
+            long trimOffset = getTrimOffset(objectList, objectStorage);
+            this.objectList = getContinuousFromTrimOffset(objectList, trimOffset);
             this.objectStorage = objectStorage;
             this.readAheadObjectSize = readAheadObjectSize;
             this.readAheadQueue = new ArrayDeque<>(readAheadObjectSize);
@@ -203,6 +206,63 @@ public RecoverIterator(List<RecordAccumulator.WALObject> objectList, ObjectStora
             }
         }
 
+        /**
+         * Get the latest trim offset from the newest object.
+         */
+        private static long getTrimOffset(List<RecordAccumulator.WALObject> objectList, ObjectStorage objectStorage) {
+            if (objectList.isEmpty()) {
+                return -1;
+            }
+
+            RecordAccumulator.WALObject object = objectList.get(objectList.size() - 1);
+            ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions()
+                .throttleStrategy(ThrottleStrategy.BYPASS)
+                .bucket(object.bucketId());
+            ByteBuf buffer = objectStorage.rangeRead(options, object.path(), 0, object.length()).join();
+            WALObjectHeader header = WALObjectHeader.unmarshal(buffer);
+            buffer.release();
+            return header.trimOffset();
+        }
+
+        // Visible for testing.
+        static List<RecordAccumulator.WALObject> getContinuousFromTrimOffset(List<RecordAccumulator.WALObject> objectList, long trimOffset) {
+            if (objectList.isEmpty()) {
+                return Collections.emptyList();
+            }
+
+            int startIndex = objectList.size();
+            for (int i = 0; i < objectList.size(); i++) {
+                if (objectList.get(i).endOffset() > trimOffset) {
+                    startIndex = i;
+                    break;
+                }
+            }
+            if (startIndex > 0) {
+                for (int i = 0; i < startIndex; i++) {
+                    log.warn("drop trimmed object: {}", objectList.get(i));
+                }
+            }
+            if (startIndex >= objectList.size()) {
+                return Collections.emptyList();
+            }
+
+            int endIndex = startIndex + 1;
+            for (int i = startIndex + 1; i < objectList.size(); i++) {
+                if (objectList.get(i).startOffset() != objectList.get(i - 1).endOffset()) {
+                    break;
+                }
+                endIndex = i + 1;
+            }
+            if (endIndex < objectList.size()) {
+                for (int i = endIndex; i < objectList.size(); i++) {
+                    log.warn("drop discontinuous object: {}", objectList.get(i));
+                }
+            }
+
+            return new ArrayList<>(objectList.subList(startIndex, endIndex));
+        }
+
+
         @Override
         public boolean hasNext() {
             return dataBuffer.isReadable() || !readAheadQueue.isEmpty() || nextIndex < objectList.size();
@@ -215,7 +275,7 @@ private void loadNextBuffer(boolean skipStickyRecord) {
 
             // Check header
             WALObjectHeader header = WALObjectHeader.unmarshal(dataBuffer);
-            dataBuffer.skipBytes(WALObjectHeader.WAL_HEADER_SIZE);
+            dataBuffer.skipBytes(header.size());
 
             if (skipStickyRecord && header.stickyRecordLength() != 0) {
                 dataBuffer.skipBytes((int) header.stickyRecordLength());

s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecordAccumulator.java

@@ -34,16 +34,20 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
 import java.util.PriorityQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentNavigableMap;
 import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Function;
@@ -56,11 +60,13 @@ public class RecordAccumulator implements Closeable {
 
     private static final long DEFAULT_LOCK_WARNING_TIMEOUT = TimeUnit.MILLISECONDS.toNanos(5);
     private static final long DEFAULT_UPLOAD_WARNING_TIMEOUT = TimeUnit.SECONDS.toNanos(5);
+    private static final String OBJECT_PATH_OFFSET_DELIMITER = "-";
+    private static final String OBJECT_PATH_FORMAT = "%s%d" + OBJECT_PATH_OFFSET_DELIMITER + "%d"; // {objectPrefix}/{startOffset}-{endOffset}
     protected final ObjectWALConfig config;
     protected final Time time;
     protected final ObjectStorage objectStorage;
     private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
-    private final ConcurrentNavigableMap<Long /* inclusive first offset */, List<Record>> uploadMap = new ConcurrentSkipListMap<>();
+    private final ConcurrentNavigableMap<Long /* inclusive first offset */, UploadTask> uploadMap = new ConcurrentSkipListMap<>();
     private final ConcurrentNavigableMap<Pair<Long /* epoch */, Long /* exclusive end offset */>, WALObject> previousObjectMap = new ConcurrentSkipListMap<>();
     private final ConcurrentNavigableMap<Long /* exclusive end offset */, WALObject> objectMap = new ConcurrentSkipListMap<>();
     private final String nodePrefix;
@@ -77,6 +83,7 @@ public class RecordAccumulator implements Closeable {
     private volatile long lastUploadTimestamp = System.currentTimeMillis();
     private final AtomicLong nextOffset = new AtomicLong();
     private final AtomicLong flushedOffset = new AtomicLong();
+    private final AtomicLong trimOffset = new AtomicLong(-1);
 
     public RecordAccumulator(Time time, ObjectStorage objectStorage,
         ObjectWALConfig config) {
@@ -100,21 +107,32 @@ public void start() {
                 String path = object.key();
                 String[] parts = path.split("/");
                 try {
-                    long firstOffset = Long.parseLong(parts[parts.length - 1]);
-                    long epoch = Long.parseLong(parts[parts.length - 3]);
+                    WALObject walObject;
 
+                    long epoch = Long.parseLong(parts[parts.length - 3]);
                     // Skip the object if it belongs to a later epoch.
                     if (epoch > config.epoch()) {
                         return;
                     }
 
                     long length = object.size();
-                    long endOffset = firstOffset + length - WALObjectHeader.WAL_HEADER_SIZE;
+
+                    String rawOffset = parts[parts.length - 1];
+                    if (rawOffset.contains(OBJECT_PATH_OFFSET_DELIMITER)) {
+                        // new format: {startOffset}-{endOffset}
+                        long startOffset = Long.parseLong(rawOffset.substring(0, rawOffset.indexOf(OBJECT_PATH_OFFSET_DELIMITER)));
+                        long endOffset = Long.parseLong(rawOffset.substring(rawOffset.indexOf(OBJECT_PATH_OFFSET_DELIMITER) + 1));
+                        walObject = new WALObject(object.bucketId(), path, startOffset, endOffset, length);
+                    } else {
+                        // old format: {startOffset}
+                        long startOffset = Long.parseLong(rawOffset);
+                        walObject = new WALObject(object.bucketId(), path, startOffset, length);
+                    }
 
                     if (epoch != config.epoch()) {
-                        previousObjectMap.put(Pair.of(epoch, endOffset), new WALObject(object.bucketId(), path, firstOffset, length));
+                        previousObjectMap.put(Pair.of(epoch, walObject.endOffset()), walObject);
                     } else {
-                        objectMap.put(endOffset, new WALObject(object.bucketId(), path, firstOffset, length));
+                        objectMap.put(walObject.endOffset(), walObject);
                     }
                     objectDataBytes.addAndGet(length);
                 } catch (NumberFormatException e) {
@@ -196,13 +214,17 @@ public void close() {
         // Wait for all upload tasks to complete.
         if (!pendingFutureMap.isEmpty()) {
             log.info("Wait for {} pending upload tasks to complete.", pendingFutureMap.size());
-            CompletableFuture.allOf(pendingFutureMap.keySet().toArray(new CompletableFuture[0])).join();
+            try {
+                CompletableFuture.allOf(pendingFutureMap.keySet().toArray(new CompletableFuture[0])).get(30, TimeUnit.SECONDS);
+            } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                log.error("Failed to wait for pending upload tasks to complete.", e);
+            }
         }
 
         if (utilityService != null && !utilityService.isShutdown()) {
             utilityService.shutdown();
             try {
-                if (!utilityService.awaitTermination(30, TimeUnit.SECONDS)) {
+                if (!utilityService.awaitTermination(1, TimeUnit.SECONDS)) {
                     log.error("Monitor executor {} did not terminate in time", executorService);
                     utilityService.shutdownNow();
                 }
@@ -246,7 +268,12 @@ public List<WALObject> objectList() throws WALFencedException {
     }
 
     // Visible for testing
-    public ScheduledExecutorService executorService() {
+    String objectPrefix() {
+        return objectPrefix;
+    }
+
+    // Visible for testing
+    ScheduledExecutorService executorService() {
         return executorService;
     }
 
@@ -374,6 +401,7 @@ public CompletableFuture<Void> trim(long offset) throws WALFencedException {
             return CompletableFuture.completedFuture(null);
         }
 
+        trimOffset.set(offset);
         long startTime = time.nanoseconds();
 
         List<ObjectStorage.ObjectPath> deleteObjectList = new ArrayList<>();
@@ -409,7 +437,7 @@ public CompletableFuture<Void> trim(long offset) throws WALFencedException {
 
     // Not thread safe, caller should hold lock.
     // Visible for testing.
-    public void unsafeUpload(boolean force) throws WALFencedException {
+    void unsafeUpload(boolean force) throws WALFencedException {
         if (!force) {
             checkWriteStatus();
         }
@@ -452,7 +480,7 @@ private void unsafeUpload(PriorityQueue<Record> recordQueue) {
 
         while (!recordQueue.isEmpty()) {
             // The retained bytes in the batch must larger than record header size.
-            long retainedBytesInBatch = config.maxBytesInBatch() - dataBuffer.readableBytes() - WALObjectHeader.WAL_HEADER_SIZE;
+            long retainedBytesInBatch = config.maxBytesInBatch() - dataBuffer.readableBytes() - WALObjectHeader.DEFAULT_WAL_HEADER_SIZE;
             if (config.strictBatchLimit() && retainedBytesInBatch <= RecordHeader.RECORD_HEADER_SIZE) {
                 break;
             }
@@ -461,7 +489,7 @@ private void unsafeUpload(PriorityQueue<Record> recordQueue) {
 
             // Records larger than the batch size will be uploaded immediately.
             assert record != null;
-            if (config.strictBatchLimit() && record.record.readableBytes() >= config.maxBytesInBatch() - WALObjectHeader.WAL_HEADER_SIZE) {
+            if (config.strictBatchLimit() && record.record.readableBytes() >= config.maxBytesInBatch() - WALObjectHeader.DEFAULT_WAL_HEADER_SIZE) {
                 dataBuffer.addComponent(true, record.record);
                 recordList.add(record);
                 break;
@@ -490,17 +518,17 @@ assert record != null;
         long endOffset = firstOffset + dataLength;
 
         CompositeByteBuf objectBuffer = ByteBufAlloc.compositeByteBuffer();
-        WALObjectHeader header = new WALObjectHeader(firstOffset, dataLength, stickyRecordLength, config.nodeId(), config.epoch());
+        WALObjectHeader header = new WALObjectHeader(firstOffset, dataLength, stickyRecordLength, config.nodeId(), config.epoch(), trimOffset.get());
         objectBuffer.addComponent(true, header.marshal());
         objectBuffer.addComponent(true, dataBuffer);
 
         // Trigger upload.
         int objectLength = objectBuffer.readableBytes();
-        uploadMap.put(firstOffset, recordList);
+        uploadMap.put(firstOffset, new UploadTask(recordList));
 
         // Enable fast retry.
         ObjectStorage.WriteOptions writeOptions = new ObjectStorage.WriteOptions().enableFastRetry(true);
-        String path = objectPrefix + firstOffset;
+        String path = String.format(OBJECT_PATH_FORMAT, objectPrefix, firstOffset, endOffset);
         CompletableFuture<ObjectStorage.WriteResult> uploadFuture = objectStorage.write(writeOptions, path, objectBuffer);
 
         CompletableFuture<Void> finalFuture = recordUploadMetrics(uploadFuture, startTime, objectLength)
@@ -512,10 +540,16 @@ assert record != null;
                         log.warn("Failed to acquire lock in {}ms, cost: {}ms, operation: upload", TimeUnit.NANOSECONDS.toMillis(DEFAULT_LOCK_WARNING_TIMEOUT), TimeUnit.NANOSECONDS.toMillis(time.nanoseconds() - lockStartTime));
                     }
 
-                    objectMap.put(endOffset, new WALObject(result.bucket(), path, firstOffset, objectLength));
+                    objectMap.put(endOffset, new WALObject(result.bucket(), path, firstOffset, endOffset, objectLength));
                     objectDataBytes.addAndGet(objectLength);
 
-                    List<Record> uploadedRecords = uploadMap.remove(firstOffset);
+                    uploadMap.get(firstOffset).markFinished();
+                    List<UploadTask> finishedTasks = new ArrayList<>();
+                    // Remove consecutive completed tasks from head.
+                    Map.Entry<Long, UploadTask> entry;
+                    while ((entry = uploadMap.firstEntry()) != null && entry.getValue().isFinished()) {
+                        finishedTasks.add(uploadMap.remove(entry.getKey()));
+                    }
 
                     // Update flushed offset
                     if (!uploadMap.isEmpty()) {
@@ -527,7 +561,11 @@ assert record != null;
                     }
 
                     // Release lock and complete future in callback thread.
-                    callbackService.submit(() -> uploadedRecords.forEach(record -> record.future.complete(flushedOffset::get)));
+                    for (UploadTask task : finishedTasks) {
+                        callbackService.submit(() ->
+                            task.records().forEach(record -> record.future.complete(flushedOffset::get))
+                        );
+                    }
                 } finally {
                     lock.writeLock().unlock();
                 }
@@ -536,7 +574,7 @@ assert record != null;
                 bufferedDataBytes.addAndGet(-dataLength);
                 throwable = ExceptionUtils.getRootCause(throwable);
                 if (throwable instanceof WALFencedException) {
-                    List<Record> uploadedRecords = uploadMap.remove(firstOffset);
+                    List<Record> uploadedRecords = uploadMap.remove(firstOffset).records();
                     Throwable finalThrowable = throwable;
                     // Release lock and complete future in callback thread.
                     callbackService.submit(() -> uploadedRecords.forEach(record -> record.future.completeExceptionally(finalThrowable)));
@@ -560,6 +598,27 @@ protected CompletableFuture<ObjectStorage.WriteResult> recordUploadMetrics(
         });
     }
 
+    private static class UploadTask {
+        private final List<Record> records;
+        private boolean finished = false;
+
+        public UploadTask(List<Record> records) {
+            this.records = records;
+        }
+
+        public List<Record> records() {
+            return records;
+        }
+
+        public void markFinished() {
+            this.finished = true;
+        }
+
+        public boolean isFinished() {
+            return finished;
+        }
+    }
+
     protected static class Record {
         public final ByteBuf record;
         public final CompletableFuture<AppendResult.CallbackResult> future;
@@ -577,12 +636,22 @@ public static class WALObject implements Comparable<WALObject> {
         private final short bucketId;
         private final String path;
         private final long startOffset;
+        private final long endOffset;
         private final long length;
 
         public WALObject(short bucketId, String path, long startOffset, long length) {
             this.bucketId = bucketId;
             this.path = path;
             this.startOffset = startOffset;
+            this.endOffset = WALObjectHeader.calculateEndOffsetV0(startOffset, length);
+            this.length = length;
+        }
+
+        public WALObject(short bucketId, String path, long startOffset, long endOffset, long length) {
+            this.bucketId = bucketId;
+            this.path = path;
+            this.startOffset = startOffset;
+            this.endOffset = endOffset;
             this.length = length;
         }
 
@@ -606,5 +675,33 @@ public long startOffset() {
         public long length() {
             return length;
         }
+
+        public long endOffset() {
+            return endOffset;
+        }
+
+        @Override
+        public String toString() {
+            return "WALObject{" +
+                "bucketId=" + bucketId +
+                ", path='" + path + '\'' +
+                ", startOffset=" + startOffset +
+                ", endOffset=" + endOffset +
+                ", length=" + length +
+                '}';
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (!(o instanceof WALObject))
+                return false;
+            WALObject object = (WALObject) o;
+            return bucketId == object.bucketId && startOffset == object.startOffset && endOffset == object.endOffset && length == object.length && Objects.equals(path, object.path);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(bucketId, path, startOffset, endOffset, length);
+        }
     }
 }