From 305cf9c34fc275c13f0bb3f52717b7adb8e5ce2a Mon Sep 17 00:00:00 2001
From: fengyubiao <yubiao.feng@streamnative.io>
Date: Thu, 18 May 2023 11:52:02 +0800
Subject: [PATCH] address comments

---
 .../metadata/impl/AbstractMetadataStore.java  | 15 ++++++-
 .../pulsar/metadata/impl/ZKMetadataStore.java | 45 +++++++++----------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/AbstractMetadataStore.java b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/AbstractMetadataStore.java
index b20b68d8b4da7..4cadf2397a7fa 100644
--- a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/AbstractMetadataStore.java
+++ b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/AbstractMetadataStore.java
@@ -41,6 +41,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
@@ -530,7 +531,19 @@ public void execute(Runnable task, CompletableFuture<?> future) {
         try {
             executor.execute(task);
         } catch (Throwable t) {
-            executor.execute(() -> future.completeExceptionally(t));
+            future.completeExceptionally(t);
+        }
+    }
+
+    /**
+     * Run the task in the executor thread and fail the future if the executor is shutting down.
+     */
+    @VisibleForTesting
+    public void execute(Runnable task, Supplier<List<CompletableFuture<?>>> futures) {
+        try {
+            executor.execute(task);
+        } catch (final Throwable t) {
+            futures.get().forEach(f -> f.completeExceptionally(t));
         }
     }
 
diff --git a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
index 43446ed128ed3..079ae3e2ae5c3 100644
--- a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
+++ b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
@@ -204,30 +204,29 @@ protected void batchOperation(List<MetadataOp> ops) {
                 }
 
                 // Trigger all the futures in the batch
-                for (int i = 0; i < ops.size(); i++) {
-                    OpResult opr = results.get(i);
-                    MetadataOp op = ops.get(i);
-                    execute(() -> {
-                        switch (op.getType()) {
-                            case PUT:
-                                handlePutResult(op.asPut(), opr);
-                                break;
-                            case DELETE:
-                                handleDeleteResult(op.asDelete(), opr);
-                                break;
-                            case GET:
-                                handleGetResult(op.asGet(), opr);
-                                break;
-                            case GET_CHILDREN:
-                                handleGetChildrenResult(op.asGetChildren(), opr);
-                                break;
-
-                            default:
-                                op.getFuture().completeExceptionally(new MetadataStoreException(
-                                        "Operation type not supported in multi: " + op.getType()));
+                execute(() -> {
+                    for (int i = 0; i < ops.size(); i++) {
+                        OpResult opr = results.get(i);
+                        MetadataOp op = ops.get(i);
+                            switch (op.getType()) {
+                                case PUT:
+                                    handlePutResult(op.asPut(), opr);
+                                    break;
+                                case DELETE:
+                                    handleDeleteResult(op.asDelete(), opr);
+                                    break;
+                                case GET:
+                                    handleGetResult(op.asGet(), opr);
+                                    break;
+                                case GET_CHILDREN:
+                                    handleGetChildrenResult(op.asGetChildren(), opr);
+                                    break;
+                                default:
+                                    op.getFuture().completeExceptionally(new MetadataStoreException(
+                                            "Operation type not supported in multi: " + op.getType()));
+                            }
                         }
-                    }, op.getFuture());
-                }
+                }, () -> ops.stream().map(MetadataOp::getFuture).collect(Collectors.toList()));
             }, null);
         } catch (Throwable t) {
             ops.forEach(o -> o.getFuture().completeExceptionally(new MetadataStoreException(t)));