[improve][meta] Fix invalid use of drain API and race condition in cl…

…osing metadata store (apache#22585) Co-authored-by: Matteo Merli <mmerli@apache.org>
coderzc · Jun 14, 2024 · f7d35e5 · f7d35e5
1 parent c2702e9
commit f7d35e5
Showing 1 changed file with 25 additions and 4 deletions.
diff --git a/.../src/main/java/org/apache/pulsar/metadata/impl/batching/AbstractBatchedMetadataStore.java b/.../src/main/java/org/apache/pulsar/metadata/impl/batching/AbstractBatchedMetadataStore.java
@@ -86,9 +86,13 @@ public void close() throws Exception {
             // Fail all the pending items
             MetadataStoreException ex =
                     new MetadataStoreException.AlreadyClosedException("Metadata store is getting closed");
-            readOps.drain(op -> op.getFuture().completeExceptionally(ex));
-            writeOps.drain(op -> op.getFuture().completeExceptionally(ex));
-
+            MetadataOp op;
+            while ((op = readOps.poll()) != null) {
+                op.getFuture().completeExceptionally(ex);
+            }
+            while ((op = writeOps.poll()) != null) {
+                op.getFuture().completeExceptionally(ex);
+            }
             scheduledTask.cancel(true);
         }
         super.close();
@@ -98,7 +102,13 @@ public void close() throws Exception {
     private void flush() {
         while (!readOps.isEmpty()) {
             List<MetadataOp> ops = new ArrayList<>();
-            readOps.drain(ops::add, maxOperations);
+            for (int i = 0; i < maxOperations; i++) {
+                MetadataOp op = readOps.poll();
+                if (op == null) {
+                    break;
+                }
+                ops.add(op);
+            }
             internalBatchOperation(ops);
         }
 
@@ -167,6 +177,11 @@ public void updateMetadataEventSynchronizer(MetadataEventSynchronizer synchroniz
     }
 
     private void enqueue(MessagePassingQueue<MetadataOp> queue, MetadataOp op) {
+        if (isClosed()) {
+            MetadataStoreException ex = new MetadataStoreException.AlreadyClosedException();
+            op.getFuture().completeExceptionally(ex);
+            return;
+        }
         if (enabled) {
             if (!queue.offer(op)) {
                 // Execute individually if we're failing to enqueue
@@ -182,6 +197,12 @@ private void enqueue(MessagePassingQueue<MetadataOp> queue, MetadataOp op) {
     }
 
     private void internalBatchOperation(List<MetadataOp> ops) {
+        if (isClosed()) {
+            MetadataStoreException ex =
+                    new MetadataStoreException.AlreadyClosedException();
+            ops.forEach(op -> op.getFuture().completeExceptionally(ex));
+            return;
+        }
         long now = System.currentTimeMillis();
         for (MetadataOp op : ops) {
             this.batchMetadataStoreStats.recordOpWaiting(now - op.created());