feat(s3stream): support call #close(force=true) after #close(force=false) (#1854)

Signed-off-by: Robin Han <[email protected]>

Author: superhx

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

@@ -81,6 +81,7 @@ public class S3Stream implements Stream {
     private final AsyncNetworkBandwidthLimiter networkOutboundLimiter;
     private long startOffset;
     private CompletableFuture<Void> lastPendingTrim = CompletableFuture.completedFuture(null);
+    private CompletableFuture<Void> closeCf;
 
     public S3Stream(long streamId, long epoch, long startOffset, long nextOffset, Storage storage,
         StreamManager streamManager) {
@@ -348,6 +349,15 @@ public CompletableFuture<Void> close(boolean force) {
             if (force) {
                 pendingRequests.forEach(cf -> cf.completeExceptionally(new StreamClientException(ErrorCode.UNEXPECTED, "FORCE_CLOSE")));
             }
+            // Supported repeated call #close.
+            // And it supports the following scenario:
+            // 1. #close(false): try gracefully close the stream. (await the pending request complete, upload data to object storage and close stream)
+            // 2. #close(true): try force close the stream. (complete the pending request with exception, upload data to object storage and close stream)
+            // Even #close(false) is called, we still could use #close(true) to force close the stream.
+            if (closeCf != null) {
+                return closeCf;
+            }
+
             CompletableFuture<Void> awaitPendingRequestsCf = CompletableFuture.allOf(pendingRequests.toArray(new CompletableFuture[0]));
             CompletableFuture<Void> closeCf = new CompletableFuture<>();
 
@@ -366,6 +376,7 @@ public CompletableFuture<Void> close(boolean force) {
                 S3StreamMetricsManager.removePendingStreamFetchLatencySupplier(streamId);
             });
 
+            this.closeCf = closeCf;
             return closeCf;
         } finally {
             writeLock.unlock();

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

@@ -83,11 +83,12 @@ public class S3StreamClient implements StreamClient {
     private final ReentrantLock lock = new ReentrantLock();
 
     final Map<Long, CompletableFuture<Stream>> openingStreams = new ConcurrentHashMap<>();
-    final Map<Long, CompletableFuture<Stream>> closingStreams = new ConcurrentHashMap<>();
+    final Map<Long, StreamWrapper> closingStreams = new ConcurrentHashMap<>();
 
     private final List<StreamLifeCycleListener> streamLifeCycleListeners = new CopyOnWriteArrayList<>();
 
     private boolean closed;
+    private boolean forceCloseMark;
 
     @SuppressWarnings("unused")
     public S3StreamClient(StreamManager streamManager, Storage storage, ObjectManager objectManager,
@@ -201,17 +202,23 @@ public void shutdown() {
         }
 
         TimerUtil timerUtil = new TimerUtil();
-        closeStreams(false);
+        closeStreams();
         LOGGER.info("S3StreamClient shutdown, cost {}ms", timerUtil.elapsedAs(TimeUnit.MILLISECONDS));
     }
 
-    private void closeStreams(boolean force) {
+    private void closeStreams() {
+        if (forceCloseMark) {
+            runInLock(() -> closingStreams.forEach((streamId, stream) -> {
+                LOGGER.info("force close closing stream, streamId={}", streamId);
+                stream.close(true);
+            }));
+        }
         for (; ; ) {
             lock.lock();
             try {
                 openedStreams.forEach((streamId, stream) -> {
                     LOGGER.info("trigger stream close, streamId={}", streamId);
-                    stream.close(force);
+                    stream.close(forceCloseMark);
                 });
                 if (openedStreams.isEmpty() && openingStreams.isEmpty() && closingStreams.isEmpty()) {
                     LOGGER.info("all streams are closed");
@@ -227,11 +234,10 @@ private void closeStreams(boolean force) {
 
     public void forceClose() {
         markClosed();
-        closeStreams(true);
+        runInLock(() -> forceCloseMark = true);
+        closeStreams();
     }
 
-
-
     private void checkState() {
         if (closed) {
             throw new IllegalStateException("S3StreamClient is already closed");
@@ -321,15 +327,19 @@ public CompletableFuture<Void> close() {
         public CompletableFuture<Void> close(boolean force) {
             return runInLock(() -> {
                 CompletableFuture<Stream> cf = new CompletableFuture<>();
-                openedStreams.remove(streamId(), this);
-                closingStreams.put(streamId(), cf);
-                return stream.close(force).whenComplete((v, e) -> runInLock(() -> {
-                    cf.complete(StreamWrapper.this);
-                    closingStreams.remove(streamId(), cf);
-                    for (StreamLifeCycleListener listener : streamLifeCycleListeners) {
-                        listener.onStreamClose(streamId());
-                    }
-                }));
+                long streamId = streamId();
+                if (openedStreams.remove(streamId, this)) {
+                    closingStreams.put(streamId, this);
+                    return stream.close(force).whenComplete((v, e) -> runInLock(() -> {
+                        cf.complete(StreamWrapper.this);
+                        closingStreams.remove(streamId, this);
+                        for (StreamLifeCycleListener listener : streamLifeCycleListeners) {
+                            listener.onStreamClose(streamId);
+                        }
+                    }));
+                } else {
+                    return stream.close(force);
+                }
             });
         }
 
@@ -338,10 +348,10 @@ public CompletableFuture<Void> destroy() {
             return runInLock(() -> {
                 CompletableFuture<Stream> cf = new CompletableFuture<>();
                 openedStreams.remove(streamId(), this);
-                closingStreams.put(streamId(), cf);
+                closingStreams.put(streamId(), this);
                 return stream.destroy().whenComplete((v, e) -> runInLock(() -> {
                     cf.complete(StreamWrapper.this);
-                    closingStreams.remove(streamId(), cf);
+                    closingStreams.remove(streamId(), this);
                 }));
             });
         }