From 5363380d576a4962a8e83b0dabc1d4bf117efc4d Mon Sep 17 00:00:00 2001
From: Caideyipi <87789683+Caideyipi@users.noreply.github.com>
Date: Wed, 3 Jun 2026 15:23:43 +0800
Subject: [PATCH] Fix flaky stamped lock test

---
 .../lock/StampedWriterPreferredLockTest.java  | 42 ++++++++++++++-----
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/metadata/mtree/lock/StampedWriterPreferredLockTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/metadata/mtree/lock/StampedWriterPreferredLockTest.java
index 2c9fdbdfd3d54..b559aaf81529a 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/metadata/mtree/lock/StampedWriterPreferredLockTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/metadata/mtree/lock/StampedWriterPreferredLockTest.java
@@ -25,6 +25,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -143,22 +144,41 @@ public void testThreadReadLock() {
   }
 
   @Test
-  public void testAcquireReadLockWhileWriting() {
+  public void testAcquireReadLockWhileWriting() throws InterruptedException {
     StampedWriterPreferredLock lock = new StampedWriterPreferredLock();
     lock.writeLock();
     AtomicInteger counter = new AtomicInteger();
-    new Thread(
+    CountDownLatch readerFinished = new CountDownLatch(1);
+    Thread readerThread =
+        new Thread(
             () -> {
               lock.threadReadLock();
-              counter.incrementAndGet();
-              lock.threadReadUnlock();
-            })
-        .start();
-    // block reader util writer release write lock
-    Assert.assertEquals(0, counter.get());
-    lock.writeUnlock();
-    Awaitility.await().atMost(2, TimeUnit.SECONDS).until(() -> counter.get() == 1);
-    Assert.assertEquals(1, counter.get());
+              try {
+                counter.incrementAndGet();
+              } finally {
+                lock.threadReadUnlock();
+                readerFinished.countDown();
+              }
+            });
+    readerThread.setDaemon(true);
+    readerThread.start();
+
+    boolean writeLocked = true;
+    try {
+      // block reader until writer release write lock
+      Awaitility.await()
+          .atMost(10, TimeUnit.SECONDS)
+          .until(() -> readerThread.getState() == Thread.State.WAITING);
+      Assert.assertEquals(0, counter.get());
+      lock.writeUnlock();
+      writeLocked = false;
+      Assert.assertTrue(readerFinished.await(10, TimeUnit.SECONDS));
+      Assert.assertEquals(1, counter.get());
+    } finally {
+      if (writeLocked) {
+        lock.writeUnlock();
+      }
+    }
   }
 
   @Test