[improve][broker] Cancel the loadShedding task when closing pulsar service

pulsar-broker/src/main/java/org/apache/pulsar/broker/PulsarService.java

@@ -489,7 +489,10 @@ public CompletableFuture<Void> closeAsync() {
                 this.leaderElectionService = null;
             }
 
-            // shutdown loadmanager before shutting down the broker
+            // cancel loadShedding task and shutdown the loadManager executor before shutting down the broker
+            if (this.loadSheddingTask != null) {
+                this.loadSheddingTask.cancel();
+            }
             executorServicesShutdown.shutdown(loadManagerExecutor);
 
             if (adminClient != null) {

pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/LoadSheddingTask.java

@@ -19,6 +19,7 @@
 package org.apache.pulsar.broker.loadbalance;
 
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.pulsar.broker.ServiceConfiguration;
@@ -37,6 +38,8 @@ public class LoadSheddingTask implements Runnable {
 
     private volatile boolean isCancel = false;
 
+    private volatile ScheduledFuture<?> future;
+
     public LoadSheddingTask(AtomicReference<LoadManager> loadManager,
                             ScheduledExecutorService loadManagerExecutor,
                             ServiceConfiguration config) {
@@ -47,30 +50,32 @@ public LoadSheddingTask(AtomicReference<LoadManager> loadManager,
 
     @Override
     public void run() {
+        if (isCancel) {
+            return;
+        }
         try {
             loadManager.get().doLoadShedding();
         } catch (Exception e) {
             LOG.warn("Error during the load shedding", e);
         } finally {
-            if (!isCancel && loadManagerExecutor != null && config != null) {
-                loadManagerExecutor.schedule(
-                        new LoadSheddingTask(loadManager, loadManagerExecutor, config),
-                        config.getLoadBalancerSheddingIntervalMinutes(),
-                        TimeUnit.MINUTES);
-            }
+            start();
         }
     }
 
     public void start() {
-        if (loadManagerExecutor != null && config != null) {
-            loadManagerExecutor.schedule(
-                    new LoadSheddingTask(loadManager, loadManagerExecutor, config),
+        if (!isCancel && loadManagerExecutor != null && config != null) {
+            future = loadManagerExecutor.schedule(
+                    this,
                     config.getLoadBalancerSheddingIntervalMinutes(),
                     TimeUnit.MINUTES);
         }
     }
 
     public void cancel() {
         isCancel = true;
+        if (future != null) {
+            future.cancel(false);
+        }
     }
+
 }

pulsar-broker/src/test/java/org/apache/pulsar/broker/PulsarServiceCloseTest.java

@@ -0,0 +1,78 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.broker;
+
+import static org.mockito.Mockito.spy;
+import static org.testng.AssertJUnit.assertFalse;
+import static org.testng.AssertJUnit.assertTrue;
+
+import java.util.concurrent.ScheduledFuture;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.broker.auth.MockedPulsarServiceBaseTest;
+import org.apache.pulsar.broker.loadbalance.LoadSheddingTask;
+import org.awaitility.reflect.WhiteboxImpl;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+@Slf4j
+public class PulsarServiceCloseTest extends MockedPulsarServiceBaseTest {
+
+    @BeforeMethod
+    @Override
+    protected void setup() throws Exception {
+        super.internalSetup();
+    }
+
+    @AfterMethod(alwaysRun = true)
+    @Override
+    protected void cleanup() throws Exception {
+        super.internalCleanup();
+    }
+
+    protected PulsarService startBrokerWithoutAuthorization(ServiceConfiguration conf) throws Exception {
+        conf.setBrokerShutdownTimeoutMs(1000 * 60 * 5);
+        conf.setLoadBalancerSheddingIntervalMinutes(30);
+        PulsarService pulsar = spy(newPulsarService(conf));
+        setupBrokerMocks(pulsar);
+        beforePulsarStartMocks(pulsar);
+        pulsar.start();
+        log.info("Pulsar started. brokerServiceUrl: {} webServiceAddress: {}", pulsar.getBrokerServiceUrl(),
+                pulsar.getWebServiceAddress());
+        return pulsar;
+    }
+
+    @Test(timeOut = 30_000)
+    public void closeInTimeTest() throws Exception {
+        LoadSheddingTask task = pulsar.getLoadSheddingTask();
+        boolean isCancel = WhiteboxImpl.getInternalState(task, "isCancel");
+        assertFalse(isCancel);
+        ScheduledFuture<?> loadSheddingFuture = WhiteboxImpl.getInternalState(task, "future");
+        assertFalse(loadSheddingFuture.isCancelled());
+
+        // The pulsar service is not used, so it should be closed gracefully in short time.
+        pulsar.close();
+
+        isCancel = WhiteboxImpl.getInternalState(task, "isCancel");
+        assertTrue(isCancel);
+        loadSheddingFuture = WhiteboxImpl.getInternalState(task, "future");
+        assertTrue(loadSheddingFuture.isCancelled());
+    }
+
+}