From b3662b7a235fea08c11c28a5381b497961916609 Mon Sep 17 00:00:00 2001
From: houxiaoyu <houxiaoyu@apache.org>
Date: Mon, 4 Dec 2023 10:44:34 +0800
Subject: [PATCH] issue 4109: Fix Flaky-test:
 HandleFailuresTest.testHandleFailureBookieNotInWriteSet (#4110)

Master Issue: https://github.com/apache/bookkeeper/issues/4109
### Motivation

Fix Flaky-test: HandleFailuresTest.testHandleFailureBookieNotInWriteSet

When we call `b1Delay.completeExceptionally(new BKException.BKWriteException())` at line480(code1),  the `preWriteHook` will complete with exception  and then do some actions in the choosen thread(code2), e.g., put the failure bookie to `delayedWriteFailedBookies`(code3).  So the `delayedWriteFailedBookies` update is async.

However, `lh.appendAsync("entry2".getBytes()))`(Line483 in Code4) is invoked in main thread. So when `appendAsync` execute,  the `delayedWriteFailedBookies` could be not updated yet, and `changeInProgress.complete(null)`(code5) will never be invoked.
---
 .../bookkeeper/client/HandleFailuresTest.java      | 14 +++++++++++---
 .../apache/bookkeeper/proto/MockBookieClient.java  |  2 ++
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/HandleFailuresTest.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/HandleFailuresTest.java
index 252fb83330c..8e8f5649dad 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/HandleFailuresTest.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/HandleFailuresTest.java
@@ -31,11 +31,13 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.bookkeeper.client.api.LedgerMetadata;
 import org.apache.bookkeeper.client.api.WriteFlag;
 import org.apache.bookkeeper.common.concurrent.FutureUtils;
 import org.apache.bookkeeper.net.BookieId;
 import org.apache.bookkeeper.net.BookieSocketAddress;
+import org.apache.bookkeeper.proto.MockBookieClient;
 import org.apache.bookkeeper.versioning.Versioned;
 import org.junit.Assert;
 import org.junit.Test;
@@ -480,14 +482,20 @@ public void testHandleFailureBookieNotInWriteSet() throws Exception {
         b1Delay.completeExceptionally(new BKException.BKWriteException());
 
         log.info("write second entry, should have enough bookies, but blocks completion on failure handling");
-        CompletableFuture<?> e2 = lh.appendAsync("entry2".getBytes());
+        AtomicReference<CompletableFuture<?>> e2 = new AtomicReference<>();
+
+        // Execute appendAsync at the same thread of preWriteHook exception thread. So that the
+        // `delayedWriteFailedBookies` could update before appendAsync invoke.
+        ((MockBookieClient) clientCtx.getBookieClient()).getExecutor()
+                .chooseThread(lh.ledgerId)
+                .execute(() -> e2.set(lh.appendAsync("entry2".getBytes())));
         changeInProgress.get();
         assertEventuallyTrue("e2 should eventually complete", () -> lh.pendingAddOps.peek().completed);
-        Assert.assertFalse("e2 shouldn't be completed to client", e2.isDone());
+        Assert.assertFalse("e2 shouldn't be completed to client", e2.get().isDone());
         blockEnsembleChange.complete(null); // allow ensemble change to continue
 
         log.info("e2 should complete");
-        e2.get(10, TimeUnit.SECONDS);
+        e2.get().get(10, TimeUnit.SECONDS);
     }
 
 }
diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/proto/MockBookieClient.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/proto/MockBookieClient.java
index c4344c74d00..2d8315f2f0a 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/proto/MockBookieClient.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/proto/MockBookieClient.java
@@ -32,6 +32,7 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import lombok.Getter;
 import org.apache.bookkeeper.client.BKException;
 import org.apache.bookkeeper.client.api.WriteFlag;
 import org.apache.bookkeeper.common.concurrent.FutureUtils;
@@ -56,6 +57,7 @@
 public class MockBookieClient implements BookieClient {
     static final Logger LOG = LoggerFactory.getLogger(MockBookieClient.class);
 
+    @Getter
     final OrderedExecutor executor;
     final MockBookies mockBookies;
     final Set<BookieId> errorBookies =