From c8eded5c122d5f4b8b1150ca2dda1680c65c1508 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Tue, 18 Aug 2026 12:53:05 +0800 Subject: [PATCH] [fix][broker] Ensure split and unload waits complete after in-flight cleanup Assisted-by: Codex --- .../extensions/manager/SplitManager.java | 61 +++--- .../extensions/manager/UnloadManager.java | 47 ++-- .../extensions/manager/SplitManagerTest.java | 192 ++++++++++++++--- .../extensions/manager/UnloadManagerTest.java | 204 ++++++++++++++---- 4 files changed, 384 insertions(+), 120 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManager.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManager.java index ab4cb5ab668e5..fd295403447f0 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManager.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManager.java @@ -20,10 +20,12 @@ import static org.apache.pulsar.broker.loadbalance.extensions.models.SplitDecision.Label.Failure; import static org.apache.pulsar.broker.loadbalance.extensions.models.SplitDecision.Reason.Unknown; +import com.google.common.annotations.VisibleForTesting; import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import lombok.CustomLog; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitState; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitStateData; @@ -47,16 +49,15 @@ public SplitManager(SplitCounter splitCounter) { } private void complete(String serviceUnit, Throwable ex) { - inFlightSplitRequests.computeIfPresent(serviceUnit, (__, future) -> { - if (!future.isDone()) { - if (ex != null) { - future.completeExceptionally(ex); - } else { - future.complete(null); - } - } - return null; - }); + CompletableFuture future = inFlightSplitRequests.remove(serviceUnit); + if (future == null || future.isDone()) { + return; + } + if (ex != null) { + future.completeExceptionally(ex); + } else { + future.complete(null); + } } public CompletableFuture waitAsync(CompletableFuture eventPubFuture, @@ -64,22 +65,22 @@ public CompletableFuture waitAsync(CompletableFuture eventPubFuture, SplitDecision decision, long timeout, TimeUnit timeoutUnit) { - return eventPubFuture - .thenCompose(__ -> inFlightSplitRequests.computeIfAbsent(bundle, ignore -> { - log.info().attr("bundle", bundle).attr("timeout", timeout) - .attr("timeoutUnit", timeoutUnit) - .log("Published the bundle split event for bundle: . " - + "Waiting the split event to complete. Timeout"); - CompletableFuture future = new CompletableFuture<>(); - future.orTimeout(timeout, timeoutUnit).whenComplete((v, ex) -> { - if (ex != null) { - inFlightSplitRequests.remove(bundle); - log.warn().attr("bundle", bundle).exception(ex) - .log("Timed out while waiting for the bundle split event"); - } - }); - return future; - })) + return eventPubFuture.thenCompose(__ -> { + CompletableFuture future = inFlightSplitRequests.computeIfAbsent(bundle, ignore -> { + log.info().attr("bundle", bundle).attr("timeout", timeout) + .attr("timeoutUnit", timeoutUnit) + .log("Published the bundle split event for bundle: . " + + "Waiting the split event to complete. Timeout"); + return new CompletableFuture().orTimeout(timeout, timeoutUnit); + }); + // Return the dependent stage so callers cannot observe completion before timeout cleanup finishes. + return future.whenComplete((v, ex) -> { + if (ex instanceof TimeoutException && inFlightSplitRequests.remove(bundle, future)) { + log.warn().attr("bundle", bundle).exception(ex) + .log("Timed out while waiting for the bundle split event"); + } + }); + }) .whenComplete((__, ex) -> { if (ex != null) { log.error().attr("bundle", bundle).exception(ex) @@ -109,12 +110,16 @@ public void handleEvent(String serviceUnit, ServiceUnitStateData data, Throwable public void close() { inFlightSplitRequests.forEach((bundle, future) -> { - if (!future.isDone()) { + if (inFlightSplitRequests.remove(bundle, future) && !future.isDone()) { String msg = String.format("Splitting bundle: %s, but the manager already closed.", bundle); log.warn(msg); future.completeExceptionally(new IllegalStateException(msg)); } }); - inFlightSplitRequests.clear(); + } + + @VisibleForTesting + int getInFlightSplitRequestCount() { + return inFlightSplitRequests.size(); } } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManager.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManager.java index 13e3579475811..5f4a21e12ce6b 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManager.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManager.java @@ -30,6 +30,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import lombok.CustomLog; import org.apache.commons.lang3.StringUtils; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitState; @@ -117,16 +118,15 @@ private void complete(String serviceUnit, Throwable ex) { LatencyMetric.ASSIGN.endMeasurement(serviceUnit); } - inFlightUnloadRequest.computeIfPresent(serviceUnit, (__, future) -> { - if (!future.isDone()) { - if (ex != null) { - future.completeExceptionally(ex); - } else { - future.complete(null); - } - } - return null; - }); + CompletableFuture future = inFlightUnloadRequest.remove(serviceUnit); + if (future == null || future.isDone()) { + return; + } + if (ex != null) { + future.completeExceptionally(ex); + } else { + future.complete(null); + } } public CompletableFuture waitAsync(CompletableFuture eventPubFuture, @@ -134,18 +134,19 @@ public CompletableFuture waitAsync(CompletableFuture eventPubFuture, UnloadDecision decision, long timeout, TimeUnit timeoutUnit) { - return eventPubFuture.thenCompose(__ -> inFlightUnloadRequest.computeIfAbsent(bundle, ignore -> { - log.debug().attr("bundle", bundle).attr("timeout", timeout).attr("timeoutUnit", timeoutUnit) - .log("Handle unload bundle: , timeout"); - CompletableFuture future = new CompletableFuture<>(); - future.orTimeout(timeout, timeoutUnit).whenComplete((v, ex) -> { - if (ex != null) { - inFlightUnloadRequest.remove(bundle); + return eventPubFuture.thenCompose(__ -> { + CompletableFuture future = inFlightUnloadRequest.computeIfAbsent(bundle, ignore -> { + log.debug().attr("bundle", bundle).attr("timeout", timeout).attr("timeoutUnit", timeoutUnit) + .log("Handle unload bundle: , timeout"); + return new CompletableFuture().orTimeout(timeout, timeoutUnit); + }); + // Return the dependent stage so callers cannot observe completion before timeout cleanup finishes. + return future.whenComplete((v, ex) -> { + if (ex instanceof TimeoutException && inFlightUnloadRequest.remove(bundle, future)) { log.warn().attr("bundle", bundle).exception(ex).log("Failed to wait unload for serviceUnit"); } }); - return future; - })).whenComplete((__, ex) -> { + }).whenComplete((__, ex) -> { if (ex != null) { counter.update(Failure, Unknown); log.warn().attr("bundle", bundle).exception(ex).log("Failed to unload bundle"); @@ -207,12 +208,16 @@ public void handleEvent(String serviceUnit, ServiceUnitStateData data, Throwable public void close() { inFlightUnloadRequest.forEach((bundle, future) -> { - if (!future.isDone()) { + if (inFlightUnloadRequest.remove(bundle, future) && !future.isDone()) { String msg = String.format("Unloading bundle: %s, but the unload manager already closed.", bundle); log.warn(msg); future.completeExceptionally(new IllegalStateException(msg)); } }); - inFlightUnloadRequest.clear(); + } + + @VisibleForTesting + int getInFlightUnloadRequestCount() { + return inFlightUnloadRequest.size(); } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManagerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManagerTest.java index bebe842f2425a..5c35f63a17835 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManagerTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/SplitManagerTest.java @@ -22,15 +22,19 @@ import static org.apache.pulsar.broker.loadbalance.extensions.models.SplitDecision.Reason.Sessions; import static org.apache.pulsar.broker.loadbalance.extensions.models.SplitDecision.Reason.Unknown; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; -import java.util.Map; +import java.util.ArrayList; +import java.util.List; 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.TimeUnit; import java.util.concurrent.TimeoutException; import lombok.CustomLog; -import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitState; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitStateData; import org.apache.pulsar.broker.loadbalance.extensions.models.SplitCounter; @@ -42,6 +46,8 @@ @Test(groups = "broker") public class SplitManagerTest { + private static final int CONCURRENT_REQUEST_COUNT = 32; + String bundle = "bundle-1"; String dstBroker = "broker-1"; @@ -69,16 +75,17 @@ public void testEventPubFutureHasException() { } @Test - public void testTimeout() throws IllegalAccessException { + public void testTimeout() { var counter = new SplitCounter(); SplitManager manager = new SplitManager(counter); var decision = new SplitDecision(); CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, decision, 3, TimeUnit.SECONDS); - var inFlightUnloadRequests = getinFlightUnloadRequests(manager); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); try { future.get(); @@ -87,15 +94,22 @@ public void testTimeout() throws IllegalAccessException { assertTrue(ex.getCause() instanceof TimeoutException); } - assertEquals(inFlightUnloadRequests.size(), 0); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); + assertEquals(manager.getInFlightSplitRequestCount(), 0); var counterExpected = new SplitCounter(); counterExpected.update(SplitDecision.Label.Failure, Unknown); assertEquals(counter.toMetrics(null).toString(), counterExpected.toMetrics(null).toString()); + + CompletableFuture nextFuture = manager.waitAsync(CompletableFuture.completedFuture(null), + bundle, decision, 5, TimeUnit.SECONDS); + assertFalse(nextFuture.isDone()); + assertEquals(manager.getInFlightSplitRequestCount(), 1); + manager.close(); } @Test - public void testSuccess() throws IllegalAccessException, ExecutionException, InterruptedException { + public void testSuccess() throws ExecutionException, InterruptedException { var counter = new SplitCounter(); SplitManager manager = new SplitManager(counter); var counterExpected = new SplitCounter(); @@ -104,56 +118,57 @@ public void testSuccess() throws IllegalAccessException, ExecutionException, Int CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, decision, 5, TimeUnit.SECONDS); - var inFlightUnloadRequests = getinFlightUnloadRequests(manager); - - assertEquals(inFlightUnloadRequests.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Assigning, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Splitting, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Releasing, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Free, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Deleted, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Owned, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 1); + assertEquals(manager.getInFlightSplitRequestCount(), 1); // Success with Init state. manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Init, dstBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequests.size(), 0); + assertEquals(manager.getInFlightSplitRequestCount(), 0); counterExpected.update(SplitDecision.Label.Success, Sessions); assertEquals(counter.toMetrics(null).toString(), counterExpected.toMetrics(null).toString()); future.get(); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); } @Test - public void testFailedStage() throws IllegalAccessException { + public void testFailedStage() { var counter = new SplitCounter(); SplitManager manager = new SplitManager(counter); var decision = new SplitDecision(); CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, decision, 5, TimeUnit.SECONDS); - var inFlightUnloadRequests = getinFlightUnloadRequests(manager); - - assertEquals(inFlightUnloadRequests.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Owned, dstBroker, VERSION_ID_INIT), @@ -167,7 +182,8 @@ public void testFailedStage() throws IllegalAccessException { assertEquals(ex.getCause().getMessage(), "Failed stage."); } - assertEquals(inFlightUnloadRequests.size(), 0); + assertEquals(manager.getInFlightSplitRequestCount(), 0); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); var counterExpected = new SplitCounter(); counterExpected.update(SplitDecision.Label.Failure, Unknown); assertEquals(counter.toMetrics(null).toString(), @@ -175,16 +191,17 @@ public void testFailedStage() throws IllegalAccessException { } @Test - public void testClose() throws IllegalAccessException { + public void testClose() { SplitManager manager = new SplitManager(new SplitCounter()); var decision = new SplitDecision(); CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, decision, 5, TimeUnit.SECONDS); - var inFlightUnloadRequests = getinFlightUnloadRequests(manager); - assertEquals(inFlightUnloadRequests.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightSplitRequestCount(), 1); manager.close(); - assertEquals(inFlightUnloadRequests.size(), 0); + assertEquals(manager.getInFlightSplitRequestCount(), 0); try { future.get(); @@ -192,15 +209,126 @@ public void testClose() throws IllegalAccessException { } catch (Exception ex) { assertTrue(ex.getCause() instanceof IllegalStateException); } + assertEquals(inFlightRequestCountOnCompletion.join(), 0); + } + + @Test(timeOut = 10_000) + public void testConcurrentWaitersShareOneInFlightRequest() { + SplitManager manager = new SplitManager(new SplitCounter()); + var decision = new SplitDecision(); + decision.succeed(Sessions); + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + CountDownLatch start = new CountDownLatch(1); + List>> registrations = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + registrations.add(CompletableFuture.supplyAsync(() -> { + await(start); + return manager.waitAsync(CompletableFuture.completedFuture(null), + bundle, decision, 30, TimeUnit.SECONDS); + }, executor)); + } + + start.countDown(); + List> waiters = registrations.stream() + .map(CompletableFuture::join) + .toList(); + assertEquals(manager.getInFlightSplitRequestCount(), 1); + List> countsOnCompletion = waiters.stream() + .map(future -> captureInFlightRequestCountOnCompletion(future, manager)) + .toList(); + + manager.handleEvent(bundle, + new ServiceUnitStateData(ServiceUnitState.Init, dstBroker, VERSION_ID_INIT), null); + + waiters.forEach(CompletableFuture::join); + countsOnCompletion.forEach(count -> assertEquals(count.join().intValue(), 0)); + assertEquals(manager.getInFlightSplitRequestCount(), 0); + + String timeoutBundle = "concurrent-timeout"; + List> timeoutWaiters = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + timeoutWaiters.add(manager.waitAsync(CompletableFuture.completedFuture(null), + timeoutBundle, decision, 1, TimeUnit.SECONDS)); + } + assertEquals(manager.getInFlightSplitRequestCount(), 1); + List> timeoutCountsOnCompletion = timeoutWaiters.stream() + .map(future -> captureInFlightRequestCountOnCompletion(future, manager)) + .toList(); + List> timeoutFailures = timeoutWaiters.stream() + .map(future -> future.handle((__, ex) -> ex)) + .toList(); + + CompletableFuture.allOf(timeoutWaiters.stream() + .map(SplitManagerTest::ignoreFailure) + .toArray(CompletableFuture[]::new)).join(); + timeoutCountsOnCompletion.forEach(count -> assertEquals(count.join().intValue(), 0)); + timeoutFailures.forEach(failure -> + assertTrue(FutureUtil.unwrapCompletionException(failure.join()) instanceof TimeoutException)); + assertEquals(manager.getInFlightSplitRequestCount(), 0); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeOut = 10_000) + public void testCloseRacingWithConcurrentCompletions() { + SplitManager manager = new SplitManager(new SplitCounter()); + var decision = new SplitDecision(); + decision.succeed(Sessions); + List bundles = new ArrayList<>(); + List> waiters = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + String requestBundle = "close-race-" + i; + bundles.add(requestBundle); + waiters.add(manager.waitAsync(CompletableFuture.completedFuture(null), + requestBundle, decision, 30, TimeUnit.SECONDS)); + } + assertEquals(manager.getInFlightSplitRequestCount(), CONCURRENT_REQUEST_COUNT); + + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + CountDownLatch start = new CountDownLatch(1); + List> terminalSignals = new ArrayList<>(); + terminalSignals.add(CompletableFuture.runAsync(() -> { + await(start); + manager.close(); + }, executor)); + for (String requestBundle : bundles) { + terminalSignals.add(CompletableFuture.runAsync(() -> { + await(start); + manager.handleEvent(requestBundle, + new ServiceUnitStateData(ServiceUnitState.Init, dstBroker, VERSION_ID_INIT), null); + }, executor)); + } + + start.countDown(); + CompletableFuture.allOf(terminalSignals.toArray(CompletableFuture[]::new)).join(); + CompletableFuture.allOf(waiters.stream() + .map(SplitManagerTest::ignoreFailure) + .toArray(CompletableFuture[]::new)).join(); + assertEquals(manager.getInFlightSplitRequestCount(), 0); + } finally { + executor.shutdownNow(); + } + } + + private static CompletableFuture captureInFlightRequestCountOnCompletion( + CompletableFuture future, SplitManager manager) { + return future.handle((__, ex) -> manager.getInFlightSplitRequestCount()); } - @SuppressWarnings("unchecked") - private Map> getinFlightUnloadRequests(SplitManager manager) - throws IllegalAccessException { - var inFlightUnloadRequest = - (Map>) FieldUtils.readField(manager, "inFlightSplitRequests", true); + private static CompletableFuture ignoreFailure(CompletableFuture future) { + return future.handle((__, ex) -> null); + } - return inFlightUnloadRequest; + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManagerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManagerTest.java index 8d593f3541f2b..eca1d7d59888e 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManagerTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/manager/UnloadManagerTest.java @@ -24,15 +24,19 @@ import static org.apache.pulsar.broker.loadbalance.extensions.models.UnloadDecision.Reason.Admin; import static org.apache.pulsar.broker.loadbalance.extensions.models.UnloadDecision.Reason.Unknown; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; -import java.util.Map; +import java.util.ArrayList; +import java.util.List; 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.TimeUnit; import java.util.concurrent.TimeoutException; import lombok.CustomLog; -import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitState; import org.apache.pulsar.broker.loadbalance.extensions.channel.ServiceUnitStateData; import org.apache.pulsar.broker.loadbalance.extensions.models.Unload; @@ -45,6 +49,8 @@ @Test(groups = "broker") public class UnloadManagerTest { + private static final int CONCURRENT_REQUEST_COUNT = 32; + @Test public void testEventPubFutureHasException() { UnloadCounter counter = new UnloadCounter(); @@ -66,7 +72,7 @@ public void testEventPubFutureHasException() { } @Test - public void testTimeout() throws IllegalAccessException { + public void testTimeout() { UnloadCounter counter = new UnloadCounter(); UnloadManager manager = new UnloadManager(counter, "mockBrokerId"); var unloadDecision = @@ -74,9 +80,10 @@ public void testTimeout() throws IllegalAccessException { CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), "bundle-1", unloadDecision, 3, TimeUnit.SECONDS); - Map> inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); try { future.get(); @@ -85,12 +92,19 @@ public void testTimeout() throws IllegalAccessException { assertTrue(ex.getCause() instanceof TimeoutException); } - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); assertEquals(counter.getBreakdownCounters().get(Failure).get(Unknown).get(), 1); + + CompletableFuture nextFuture = manager.waitAsync(CompletableFuture.completedFuture(null), + "bundle-1", unloadDecision, 5, TimeUnit.SECONDS); + assertFalse(nextFuture.isDone()); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); + manager.close(); } @Test - public void testSuccess() throws IllegalAccessException, ExecutionException, InterruptedException { + public void testSuccess() throws ExecutionException, InterruptedException { UnloadCounter counter = new UnloadCounter(); UnloadManager manager = new UnloadManager(counter, "mockBrokerId"); String dstBroker = "broker-2"; @@ -101,61 +115,60 @@ public void testSuccess() throws IllegalAccessException, ExecutionException, Int CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, unloadDecision, 5, TimeUnit.SECONDS); - Map> inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); - - assertEquals(inFlightUnloadRequestMap.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Assigning, null, srcBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Deleted, null, srcBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Splitting, null, srcBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Releasing, null, srcBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Free, null, srcBroker, true, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); // Success with Init state. manager.handleEvent(bundle, null, null); - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); future.get(); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); assertEquals(counter.getBreakdownCounters().get(Success).get(Admin).get(), 1); // Success with Owned state. future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, unloadDecision, 5, TimeUnit.SECONDS); - inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Owned, dstBroker, null, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Owned, dstBroker, srcBroker, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); future.get(); assertEquals(counter.getBreakdownCounters().get(Success).get(Admin).get(), 2); // Success with Free state. future = manager.waitAsync(CompletableFuture.completedFuture(null), bundle, unloadDecision, 5, TimeUnit.SECONDS); - inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Free, dstBroker, srcBroker, true, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 1); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent(bundle, new ServiceUnitStateData(ServiceUnitState.Free, dstBroker, srcBroker, false, VERSION_ID_INIT), null); - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); future.get(); assertEquals(counter.getBreakdownCounters().get(Success).get(Admin).get(), 3); @@ -163,7 +176,7 @@ public void testSuccess() throws IllegalAccessException, ExecutionException, Int } @Test - public void testFailedStage() throws IllegalAccessException { + public void testFailedStage() { UnloadCounter counter = new UnloadCounter(); UnloadManager manager = new UnloadManager(counter, "mockBrokerId"); var unloadDecision = @@ -171,9 +184,9 @@ public void testFailedStage() throws IllegalAccessException { CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), "bundle-1", unloadDecision, 5, TimeUnit.SECONDS); - Map> inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); - - assertEquals(inFlightUnloadRequestMap.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.handleEvent("bundle-1", new ServiceUnitStateData(ServiceUnitState.Owned, null, "broker-1", VERSION_ID_INIT), @@ -187,12 +200,13 @@ public void testFailedStage() throws IllegalAccessException { assertEquals(ex.getCause().getMessage(), "Failed stage."); } - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); + assertEquals(inFlightRequestCountOnCompletion.join(), 0); assertEquals(counter.getBreakdownCounters().get(Failure).get(Unknown).get(), 1); } @Test - public void testClose() throws IllegalAccessException { + public void testClose() { UnloadCounter counter = new UnloadCounter(); UnloadManager manager = new UnloadManager(counter, "mockBrokerId"); var unloadDecision = @@ -200,10 +214,11 @@ public void testClose() throws IllegalAccessException { CompletableFuture future = manager.waitAsync(CompletableFuture.completedFuture(null), "bundle-1", unloadDecision, 5, TimeUnit.SECONDS); - Map> inFlightUnloadRequestMap = getInFlightUnloadRequestMap(manager); - assertEquals(inFlightUnloadRequestMap.size(), 1); + CompletableFuture inFlightRequestCountOnCompletion = + captureInFlightRequestCountOnCompletion(future, manager); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); manager.close(); - assertEquals(inFlightUnloadRequestMap.size(), 0); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); try { future.get(); @@ -211,16 +226,127 @@ public void testClose() throws IllegalAccessException { } catch (Exception ex) { assertTrue(ex.getCause() instanceof IllegalStateException); } + assertEquals(inFlightRequestCountOnCompletion.join(), 0); assertEquals(counter.getBreakdownCounters().get(Failure).get(Unknown).get(), 1); } - @SuppressWarnings("unchecked") - private Map> getInFlightUnloadRequestMap(UnloadManager manager) - throws IllegalAccessException { - Map> inFlightUnloadRequest = - (Map>) FieldUtils.readField(manager, "inFlightUnloadRequest", true); + @Test(timeOut = 10_000) + public void testConcurrentWaitersShareOneInFlightRequest() { + UnloadManager manager = new UnloadManager(new UnloadCounter(), "mockBrokerId"); + var decision = newUnloadDecision("bundle-1"); + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + CountDownLatch start = new CountDownLatch(1); + List>> registrations = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + registrations.add(CompletableFuture.supplyAsync(() -> { + await(start); + return manager.waitAsync(CompletableFuture.completedFuture(null), + "bundle-1", decision, 30, TimeUnit.SECONDS); + }, executor)); + } + + start.countDown(); + List> waiters = registrations.stream() + .map(CompletableFuture::join) + .toList(); + assertEquals(manager.getInFlightUnloadRequestCount(), 1); + List> countsOnCompletion = waiters.stream() + .map(future -> captureInFlightRequestCountOnCompletion(future, manager)) + .toList(); + + manager.handleEvent("bundle-1", null, null); + + waiters.forEach(CompletableFuture::join); + countsOnCompletion.forEach(count -> assertEquals(count.join().intValue(), 0)); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); + + String timeoutBundle = "concurrent-timeout"; + var timeoutDecision = newUnloadDecision(timeoutBundle); + List> timeoutWaiters = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + timeoutWaiters.add(manager.waitAsync(CompletableFuture.completedFuture(null), + timeoutBundle, timeoutDecision, 1, TimeUnit.SECONDS)); + } + assertEquals(manager.getInFlightUnloadRequestCount(), 1); + List> timeoutCountsOnCompletion = timeoutWaiters.stream() + .map(future -> captureInFlightRequestCountOnCompletion(future, manager)) + .toList(); + List> timeoutFailures = timeoutWaiters.stream() + .map(future -> future.handle((__, ex) -> ex)) + .toList(); + + CompletableFuture.allOf(timeoutWaiters.stream() + .map(UnloadManagerTest::ignoreFailure) + .toArray(CompletableFuture[]::new)).join(); + timeoutCountsOnCompletion.forEach(count -> assertEquals(count.join().intValue(), 0)); + timeoutFailures.forEach(failure -> + assertTrue(FutureUtil.unwrapCompletionException(failure.join()) instanceof TimeoutException)); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeOut = 10_000) + public void testCloseRacingWithConcurrentCompletions() { + UnloadManager manager = new UnloadManager(new UnloadCounter(), "mockBrokerId"); + List bundles = new ArrayList<>(); + List> waiters = new ArrayList<>(); + for (int i = 0; i < CONCURRENT_REQUEST_COUNT; i++) { + String bundle = "close-race-" + i; + bundles.add(bundle); + waiters.add(manager.waitAsync(CompletableFuture.completedFuture(null), + bundle, newUnloadDecision(bundle), 30, TimeUnit.SECONDS)); + } + assertEquals(manager.getInFlightUnloadRequestCount(), CONCURRENT_REQUEST_COUNT); + + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + CountDownLatch start = new CountDownLatch(1); + List> terminalSignals = new ArrayList<>(); + terminalSignals.add(CompletableFuture.runAsync(() -> { + await(start); + manager.close(); + }, executor)); + for (String bundle : bundles) { + terminalSignals.add(CompletableFuture.runAsync(() -> { + await(start); + manager.handleEvent(bundle, null, null); + }, executor)); + } + + start.countDown(); + CompletableFuture.allOf(terminalSignals.toArray(CompletableFuture[]::new)).join(); + CompletableFuture.allOf(waiters.stream() + .map(UnloadManagerTest::ignoreFailure) + .toArray(CompletableFuture[]::new)).join(); + assertEquals(manager.getInFlightUnloadRequestCount(), 0); + } finally { + executor.shutdownNow(); + } + } + + private static CompletableFuture captureInFlightRequestCountOnCompletion( + CompletableFuture future, UnloadManager manager) { + return future.handle((__, ex) -> manager.getInFlightUnloadRequestCount()); + } + + private static UnloadDecision newUnloadDecision(String bundle) { + return new UnloadDecision(new Unload("broker-1", bundle), Success, Admin); + } - return inFlightUnloadRequest; + private static CompletableFuture ignoreFailure(CompletableFuture future) { + return future.handle((__, ex) -> null); + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } } }