diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelImpl.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelImpl.java index 3f5bd2569faf9..972ed10689356 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelImpl.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelImpl.java @@ -232,6 +232,16 @@ public ServiceUnitStateChannelImpl(PulsarService pulsar) { this.channelState = Constructed; } + @VisibleForTesting + Map> getOwnerRequests() { + return getOwnerRequests; + } + + @VisibleForTesting + Map> getCleanupJobs() { + return cleanupJobs; + } + @Override public void scheduleOwnershipMonitor() { if (monitorTask == null) { @@ -854,13 +864,15 @@ private void log(Throwable e, String serviceUnit, ServiceUnitStateData data, Ser } } - private void handleSkippedEvent(String serviceUnit) { + @VisibleForTesting + void handleSkippedEvent(String serviceUnit) { var getOwnerRequest = getOwnerRequests.get(serviceUnit); if (getOwnerRequest != null) { var data = tableview.get(serviceUnit); if (data != null && data.state() == Owned) { getOwnerRequest.complete(data.dstBroker()); - getOwnerRequests.remove(serviceUnit); + // Completing the request can run callbacks that install a newer request for the same service unit. + getOwnerRequests.remove(serviceUnit, getOwnerRequest); stateChangeListeners.notify(serviceUnit, data, null); } } @@ -1024,7 +1036,8 @@ private CompletableFuture deferGetOwner(String serviceUnit) { return future; } - private CompletableFuture dedupeGetOwnerRequest(String serviceUnit) { + @VisibleForTesting + CompletableFuture dedupeGetOwnerRequest(String serviceUnit) { var requested = new MutableObject>(); try { @@ -1057,7 +1070,8 @@ private CompletableFuture dedupeGetOwnerRequest(String serviceUnit) { var future = requested.get(); if (future != null) { future.whenComplete((__, e) -> { - getOwnerRequests.remove(serviceUnit); + // The request may have been replaced by a later lookup before this callback runs. + getOwnerRequests.remove(serviceUnit, future); if (e != null) { log.warn().attr("broker", brokerId).attr("serviceUnit", serviceUnit).exception(e) .log("failed to getOwner for serviceUnit"); @@ -1340,14 +1354,15 @@ private MetadataState getMetadataState() { private void handleBrokerCreationEvent(String broker) { - if (!cleanupJobs.isEmpty() && cleanupJobs.containsKey(broker)) { + CompletableFuture cleanupJob = cleanupJobs.get(broker); + if (cleanupJob != null) { healthCheckBrokerAsync(broker) .orTimeout(MAX_BROKER_HEALTH_CHECK_DELAY_IN_MILLIS * (MAX_BROKER_HEALTH_CHECK_RETRY + 1) , MILLISECONDS) .thenAccept(__ -> { - CompletableFuture future = cleanupJobs.remove(broker); - if (future != null) { - future.cancel(false); + // The health check is async; only cancel the cleanup job observed before the check. + if (cleanupJobs.remove(broker, cleanupJob)) { + cleanupJob.cancel(false); totalInactiveBrokerCleanupCancelledCnt++; log.info().attr("broker", broker).attr("count", cleanupJobs.size()) .log("Successfully cancelled the ownership cleanup for broker: ." @@ -1413,7 +1428,8 @@ private boolean channelDisabled() { return false; } - private void scheduleCleanup(String broker, long delayInSecs) { + @VisibleForTesting + void scheduleCleanup(String broker, long delayInSecs) { var scheduled = new MutableObject>(); try { if (channelDisabled()) { @@ -1443,7 +1459,8 @@ private void scheduleCleanup(String broker, long delayInSecs) { var future = scheduled.get(); if (future != null) { future.whenComplete((v, ex) -> { - cleanupJobs.remove(broker); + // The job may have been replaced by a later cleanup schedule before this callback runs. + cleanupJobs.remove(broker, future); }); } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelTest.java index 4b279d5f1e424..373068c38d746 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/loadbalance/extensions/channel/ServiceUnitStateChannelTest.java @@ -280,6 +280,104 @@ public void channelOwnerTest() throws Exception { } } + @Test(priority = 1) + public void testCompletedGetOwnerRequestDoesNotRemoveNewRequest() { + ServiceUnitStateChannelImpl channel = (ServiceUnitStateChannelImpl) channel1; + String serviceUnit = namespaceName + "/0x10000000_0x10000001"; + var getOwnerRequests = channel.getOwnerRequests(); + getOwnerRequests.remove(serviceUnit); + CompletableFuture oldRequest = channel.dedupeGetOwnerRequest(serviceUnit); + assertEquals(getOwnerRequests.get(serviceUnit), oldRequest); + + CompletableFuture newRequest = null; + try { + // State-event handlers remove the current request before completing it, allowing a later lookup + // to install a new request generation before the old request's completion cleanup runs. + assertTrue(getOwnerRequests.remove(serviceUnit, oldRequest)); + newRequest = channel.dedupeGetOwnerRequest(serviceUnit); + assertTrue(newRequest != oldRequest); + assertTrue(getOwnerRequests.get(serviceUnit) == newRequest); + + // The previous unconditional removal would remove newRequest here. + assertTrue(oldRequest.complete(brokerId1)); + + assertTrue(getOwnerRequests.get(serviceUnit) == newRequest, + "A stale get-owner cleanup must not remove a newer request future"); + + assertTrue(newRequest.complete(brokerId2)); + assertFalse(getOwnerRequests.containsKey(serviceUnit), + "The newer request must remove itself after completion"); + } finally { + getOwnerRequests.remove(serviceUnit); + oldRequest.cancel(false); + if (newRequest != null) { + newRequest.cancel(false); + } + } + } + + @Test(priority = 1) + public void testSkippedEventDoesNotRemoveNewGetOwnerRequest() throws Exception { + ServiceUnitStateChannelImpl channel = (ServiceUnitStateChannelImpl) channel1; + String serviceUnit = namespaceName + "/0x10000002_0x10000003"; + var getOwnerRequests = channel.getOwnerRequests(); + CompletableFuture oldRequest = new CompletableFuture<>(); + CompletableFuture newRequest = new CompletableFuture<>(); + try { + overrideTableView(channel, serviceUnit, new ServiceUnitStateData(Owned, brokerId1, 1)); + getOwnerRequests.put(serviceUnit, oldRequest); + oldRequest.whenComplete((__, ___) -> getOwnerRequests.put(serviceUnit, newRequest)); + + channel.handleSkippedEvent(serviceUnit); + + assertEquals(oldRequest.getNow(null), brokerId1); + assertTrue(getOwnerRequests.get(serviceUnit) == newRequest, + "A stale skipped-event cleanup must not remove a newer request future"); + } finally { + getOwnerRequests.remove(serviceUnit); + oldRequest.cancel(false); + overrideTableView(channel, serviceUnit, null); + } + } + + @Test(priority = 1) + public void testCompletedCleanupJobDoesNotRemoveNewCleanupJob() { + ServiceUnitStateChannelImpl channel = (ServiceUnitStateChannelImpl) channel1; + String broker = brokerId3; + var cleanupJobs = channel.getCleanupJobs(); + cleanupJobs.remove(broker); + channel.scheduleCleanup(broker, 60L); + CompletableFuture oldJob = cleanupJobs.get(broker); + assertNotNull(oldJob); + + CompletableFuture newJob = null; + try { + // Broker-creation handling removes a cleanup job before cancelling it. A later broker-deletion + // event can therefore schedule a new job before the old job's completion cleanup runs. + assertTrue(cleanupJobs.remove(broker, oldJob)); + channel.scheduleCleanup(broker, 60L); + newJob = cleanupJobs.get(broker); + assertNotNull(newJob); + assertTrue(newJob != oldJob); + + // The previous unconditional removal would remove newJob here. + assertTrue(oldJob.cancel(false)); + + assertTrue(cleanupJobs.get(broker) == newJob, + "A stale cleanup job completion must not remove a newer cleanup job future"); + + assertTrue(newJob.cancel(false)); + assertFalse(cleanupJobs.containsKey(broker), + "The newer cleanup job must remove itself after completion"); + } finally { + cleanupJobs.remove(broker); + oldJob.cancel(false); + if (newJob != null) { + newJob.cancel(false); + } + } + } + @Test(priority = 100) public void channelValidationTest() throws ExecutionException, InterruptedException, IllegalAccessException, PulsarServerException, @@ -822,6 +920,59 @@ public void handleBrokerCreationEventTest() throws IllegalAccessException { } + @Test(priority = 8) + public void handleBrokerCreationEventDoesNotCancelNewCleanupJobTest() { + ServiceUnitStateChannelImpl channel = (ServiceUnitStateChannelImpl) channel1; + var cleanupJobs = channel.getCleanupJobs(); + String broker = brokerId2; + CompletableFuture healthCheck = new CompletableFuture<>(); + cleanupJobs.remove(broker); + channel.scheduleCleanup(broker, 60L); + CompletableFuture oldJob = cleanupJobs.get(broker); + assertNotNull(oldJob); + + reset(brokers); + doReturn(healthCheck).when(brokers).healthcheckAsync(any()); + doReturn(brokers).when(pulsarAdmin).brokers(); + CompletableFuture newJob = null; + try { + channel.handleBrokerRegistrationEvent(broker, NotificationType.Created); + verify(brokers, times(1)).healthcheckAsync(any()); + + // The old cleanup can finish while the asynchronous health check is still pending. A later + // broker-deletion event can then schedule a new cleanup job for the same broker. + assertTrue(oldJob.complete(null)); + assertFalse(cleanupJobs.containsKey(broker)); + channel.scheduleCleanup(broker, 60L); + newJob = cleanupJobs.get(broker); + assertNotNull(newJob); + assertTrue(newJob != oldJob); + + healthCheck.complete(null); + + CompletableFuture expectedNewJob = newJob; + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + assertTrue(cleanupJobs.get(broker) == expectedNewJob, + "A stale broker-creation callback must not remove a newer cleanup job"); + assertFalse(expectedNewJob.isCancelled()); + }); + + assertTrue(newJob.cancel(false)); + assertFalse(cleanupJobs.containsKey(broker), + "The newer cleanup job must remove itself after cancellation"); + } finally { + cleanupJobs.remove(broker); + oldJob.cancel(false); + if (newJob != null) { + newJob.cancel(false); + } + reset(brokers); + doReturn(CompletableFuture.failedFuture(new RuntimeException("failed"))).when(brokers) + .healthcheckAsync(any()); + reset(pulsarAdmin); + } + } + @Test(priority = 9) public void handleBrokerDeletionEventTest() throws Exception { @@ -2075,9 +2226,9 @@ public void testHandleExistingResolvesAssigningStateOnChannelRestart() } } - private static ConcurrentHashMap>> getOwnerRequests( + private static ConcurrentHashMap> getOwnerRequests( ServiceUnitStateChannel channel) throws IllegalAccessException { - return (ConcurrentHashMap>>) + return (ConcurrentHashMap>) FieldUtils.readDeclaredField(channel, "getOwnerRequests", true); }