fix(issue2139): prevent read object info from closed ObjectReader

Signed-off-by: Shichao Nie <[email protected]>
AutoMQ · Nov 8, 2024 · ebd5e51 · ebd5e51
1 parent c098ee3
commit ebd5e51
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java b/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java
@@ -16,15 +16,19 @@
 import com.automq.stream.s3.objects.ObjectAttributes;
 import com.automq.stream.utils.biniarysearch.AbstractOrderedCollection;
 import com.automq.stream.utils.biniarysearch.ComparableItem;
-import io.netty.buffer.ByteBuf;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+
+import io.netty.buffer.ByteBuf;
 
 import static com.automq.stream.s3.ByteBufAlloc.BLOCK_CACHE;
 import static com.automq.stream.s3.CompositeObject.FOOTER_MAGIC;
@@ -41,6 +45,7 @@ public class CompositeObjectReader implements ObjectReader {
     private CompletableFuture<BasicObjectInfo> basicObjectInfoCf;
     private CompletableFuture<Integer> sizeCf;
     private final AtomicInteger refCount = new AtomicInteger(1);
+    private final AtomicBoolean isShutdown = new AtomicBoolean(false);
 
     public CompositeObjectReader(S3ObjectMetadata objectMetadata, RangeReader rangeReader) {
         this.objectMetadata = objectMetadata;
@@ -59,6 +64,9 @@ public String objectKey() {
 
     @Override
     public synchronized CompletableFuture<BasicObjectInfo> basicObjectInfo() {
+        if (isShutdown.get()) {
+            return CompletableFuture.failedFuture(new IllegalStateException("ObjectReader is already shutdown"));
+        }
         if (basicObjectInfoCf == null) {
             this.basicObjectInfoCf = new CompletableFuture<>();
             this.basicObjectInfoCf.exceptionally(ex -> {
@@ -101,6 +109,9 @@ public synchronized CompletableFuture<Integer> size() {
     }
 
     public synchronized void close0() {
+        if (!isShutdown.compareAndSet(false, true)) {
+            return;
+        }
         if (basicObjectInfoCf != null) {
             basicObjectInfoCf.thenAccept(BasicObjectInfo::close);
         }

diff --git a/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java b/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java
@@ -30,6 +30,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -89,6 +90,7 @@ class DefaultObjectReader implements ObjectReader {
         private CompletableFuture<BasicObjectInfo> basicObjectInfoCf;
         private CompletableFuture<Integer> sizeCf;
         private final AtomicInteger refCount = new AtomicInteger(1);
+        private final AtomicBoolean isShutdown = new AtomicBoolean(false);
 
         public DefaultObjectReader(S3ObjectMetadata metadata, ObjectStorage objectStorage) {
             this.metadata = metadata;
@@ -105,6 +107,9 @@ public String objectKey() {
         }
 
         public synchronized CompletableFuture<BasicObjectInfo> basicObjectInfo() {
+            if (isShutdown.get()) {
+                return CompletableFuture.failedFuture(new IllegalStateException("ObjectReader is already shutdown"));
+            }
             if (basicObjectInfoCf == null) {
                 this.basicObjectInfoCf = new CompletableFuture<>();
                 asyncGetBasicObjectInfo();
@@ -186,6 +191,9 @@ public synchronized CompletableFuture<Integer> size() {
         }
 
         public synchronized void close0() {
+            if (!isShutdown.compareAndSet(false, true)) {
+                return;
+            }
             if (basicObjectInfoCf != null) {
                 basicObjectInfoCf.thenAccept(BasicObjectInfo::close);
             }