(pulsar) 12/15: [fix][test] Cover scheduler shutdown membership-lock deadlock (#26557)

[lh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lhotari pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/pulsar.git

commit 2a05e1dcdba0f84e7d45a25aa9147f9032b940d5
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Mon Sep 14 01:00:46 2026 +0300

    [fix][test] Cover scheduler shutdown membership-lock deadlock (#26557)
    
    (cherry picked from commit 6075c2b4ddf7fd89a2e08cec4b4b8fe6286ee0af)
---
 pulsar-functions/worker/pom.xml                    |  6 +++
 .../functions/worker/SchedulerManagerTest.java     | 43 ++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/pulsar-functions/worker/pom.xml b/pulsar-functions/worker/pom.xml
index 9daabb673f87..84da1acdb232 100644
--- a/pulsar-functions/worker/pom.xml
+++ b/pulsar-functions/worker/pom.xml
@@ -209,6 +209,12 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <scope>test</scope>
+    </dependency>
+
   </dependencies>
 
   <build>
diff --git a/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/SchedulerManagerTest.java b/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/SchedulerManagerTest.java
index 87d4e82450d3..076cae8d25f9 100644
--- a/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/SchedulerManagerTest.java
+++ b/pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/SchedulerManagerTest.java
@@ -26,11 +26,13 @@ import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
 import static org.testng.Assert.assertTrue;
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.netty.util.concurrent.DefaultThreadFactory;
@@ -44,12 +46,15 @@ import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
 import org.apache.pulsar.client.admin.LongRunningProcessStatus;
@@ -67,6 +72,7 @@ import org.apache.pulsar.functions.runtime.thread.ThreadRuntimeFactory;
 import org.apache.pulsar.functions.runtime.thread.ThreadRuntimeFactoryConfig;
 import org.apache.pulsar.functions.utils.FunctionCommon;
 import org.apache.pulsar.functions.worker.scheduler.RoundRobinScheduler;
+import org.awaitility.Awaitility;
 import org.mockito.Mockito;
 import org.mockito.invocation.Invocation;
 import org.testng.Assert;
@@ -155,6 +161,43 @@ public class SchedulerManagerTest {
         this.executor.shutdownNow();
     }
 
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testCloseDoesNotBlockMembershipChecksWhileWaitingForScheduler() throws Exception {
+        doReturn(List.of(WorkerInfo.of("worker-1", "localhost", 5000)))
+                .when(membershipManager).getCurrentMembership();
+        schedulerManager.initialize(producer);
+        ReentrantLock schedulerLock = (ReentrantLock) schedulerManager.getSchedulerLock();
+        AtomicReference<Thread> closingThread = new AtomicReference<>();
+        ExecutorService closeExecutor = Executors.newSingleThreadExecutor(new DefaultThreadFactory("scheduler-close"));
+        Future<?> closeFuture = null;
+        schedulerLock.lock();
+        try {
+            closeFuture = closeExecutor.submit(() -> {
+                closingThread.set(Thread.currentThread());
+                schedulerManager.close();
+            });
+            Awaitility.await().until(() -> closingThread.get() != null
+                    && schedulerLock.hasQueuedThread(closingThread.get()));
+
+            // A scheduler holding schedulerLock must still be able to read membership while close waits.
+            // Use a separate task with a bounded wait so a regression can release the lock during cleanup.
+            executor.submit(() -> assertThrows(SchedulerManager.TooFewWorkersException.class,
+                    () -> schedulerManager.rebalanceIfNotInprogress())).get(5, TimeUnit.SECONDS);
+            verify(producer, never()).close();
+        } finally {
+            schedulerLock.unlock();
+            try {
+                if (closeFuture != null) {
+                    closeFuture.get(5, TimeUnit.SECONDS);
+                }
+            } finally {
+                closeExecutor.shutdownNow();
+            }
+        }
+        verify(producer).close();
+    }
+
     @Test
     public void testSchedule() throws Exception {