From 73cc5ee7a8e0d765909e5bbdcd4d8241fd685a83 Mon Sep 17 00:00:00 2001 From: Francisco Rojas Date: Tue, 15 Sep 2026 17:21:39 -0700 Subject: [PATCH] Handle poller resubmission during shutdown --- .../internal/worker/MultiThreadedPoller.java | 11 +- .../worker/MultiThreadedPollerTest.java | 121 ++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/worker/MultiThreadedPollerTest.java diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/MultiThreadedPoller.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/MultiThreadedPoller.java index e82d162665..5d6744b07d 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/MultiThreadedPoller.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/MultiThreadedPoller.java @@ -190,8 +190,15 @@ public void run() { uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e); } finally { if (!shouldTerminate()) { - // Resubmit itself back to pollExecutor - pollExecutor.execute(this); + try { + // Resubmit itself back to pollExecutor. + pollExecutor.execute(this); + } catch (RejectedExecutionException e) { + // Shutdown can race with the termination check above. + if (!pollExecutor.isShutdown()) { + throw e; + } + } } else { log.debug( "poll loop is terminated: {}", diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/MultiThreadedPollerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/MultiThreadedPollerTest.java new file mode 100644 index 0000000000..b2768244f6 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/MultiThreadedPollerTest.java @@ -0,0 +1,121 @@ +package io.temporal.internal.worker; + +import static org.junit.Assert.*; +import static org.mockito.AdditionalAnswers.delegatesTo; +import static org.mockito.Mockito.*; + +import com.uber.m3.tally.NoopScope; +import io.temporal.worker.tuning.PollerBehaviorSimpleMaximum; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; + +public class MultiThreadedPollerTest { + + @Test(timeout = 10000) + public void shutdownDuringResubmissionPreservesTaskHandoff() throws Exception { + assertRejectedResubmission(true, false, new Object()); + } + + @Test(timeout = 10000) + public void shutdownDuringResubmissionAfterEmptyPoll() throws Exception { + assertRejectedResubmission(true, false, null); + } + + @Test(timeout = 10000) + public void rejectionWhileRunningIsReported() throws Exception { + assertRejectedResubmission(false, false, new Object()); + } + + @Test(timeout = 10000) + public void interruptionDoesNotSuppressRejectionWhileRunning() throws Exception { + assertRejectedResubmission(false, true, new Object()); + } + + private void assertRejectedResubmission(boolean shutdown, boolean interrupt, Object task) + throws Exception { + CountDownLatch pollStarted = new CountDownLatch(1); + CountDownLatch finishPoll = new CountDownLatch(1); + AtomicReference pollThread = new AtomicReference<>(); + AtomicReference uncaught = new AtomicReference<>(); + AtomicInteger polls = new AtomicInteger(); + AtomicInteger resubmissions = new AtomicInteger(); + ShutdownableTaskExecutor taskExecutor = mock(ShutdownableTaskExecutor.class); + PollerOptions options = + PollerOptions.newBuilder() + .setPollThreadNamePrefix("test-poller") + .setPollerBehavior(new PollerBehaviorSimpleMaximum(1)) + .setUncaughtExceptionHandler((thread, error) -> uncaught.set(error)) + .build(); + MultiThreadedPoller poller = + new MultiThreadedPoller<>( + "test", + () -> { + polls.incrementAndGet(); + pollThread.set(Thread.currentThread()); + pollStarted.countDown(); + try { + assertTrue(finishPoll.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return task; + }, + taskExecutor, + options, + new NoopScope(), + new NamespaceCapabilities()); + + assertTrue(poller.start()); + ExecutorService executor = poller.pollExecutor; + try { + assertTrue(pollStarted.await(5, TimeUnit.SECONDS)); + ExecutorService resubmissionExecutor = mock(ExecutorService.class, delegatesTo(executor)); + RejectedExecutionException rejection = new RejectedExecutionException("Executor is running"); + doAnswer( + invocation -> { + resubmissions.incrementAndGet(); + // Force shutdown after shouldTerminate() but before the executor accepts the task. + if (shutdown) { + executor.shutdown(); + executor.execute(invocation.getArgument(0)); + } else { + if (interrupt) { + Thread.currentThread().interrupt(); + } + throw rejection; + } + return null; + }) + .when(resubmissionExecutor) + .execute(any(Runnable.class)); + poller.pollExecutor = resubmissionExecutor; + finishPoll.countDown(); + + // Join the thread so its uncaught exception handler has also finished. + pollThread.get().join(5000); + assertFalse(pollThread.get().isAlive()); + assertEquals(1, polls.get()); + assertEquals(1, resubmissions.get()); + if (shutdown) { + assertNull(uncaught.get()); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } else { + assertSame(rejection, uncaught.get()); + } + if (task != null) { + verify(taskExecutor).process(task); + } + verifyNoMoreInteractions(taskExecutor); + } finally { + finishPoll.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } +}