fix(compaction): prevent double release on compaction shutdown (#2115)

Signed-off-by: Shichao Nie <[email protected]>
AutoMQ · Nov 5, 2024 · 04b53f4 · 04b53f4
1 parent 8338f18
commit 04b53f4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java b/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java
@@ -81,6 +81,13 @@ public CompletableFuture<ByteBuf> getDataCf() {
         return this.dataCf;
     }
 
+    public ByteBuf getAndReleaseData() {
+        if (refCount.getAndDecrement() == 0) {
+            throw new IllegalStateException("Data has already been released");
+        }
+        return this.dataCf.join();
+    }
+
     public void releaseRef() {
         refCount.decrementAndGet();
     }

diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java b/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java
@@ -128,8 +128,7 @@ public CompletableFuture<Void> close() {
     private CompositeByteBuf groupWaitingBlocks() {
         CompositeByteBuf buf = ByteBufAlloc.compositeByteBuffer();
         for (StreamDataBlock block : waitingUploadBlocks) {
-            buf.addComponent(true, block.getDataCf().join());
-            block.releaseRef();
+            buf.addComponent(true, block.getAndReleaseData());
             completedBlocks.add(block);
             nextDataBlockPosition += block.getBlockSize();
         }