optimize: enhance close() logic of discovery module (#7375)

Author: YoWuwuuuw

changes/en-us/2.x.md

@@ -48,6 +48,7 @@ Add changes here for all PR submitted to the 2.x branch.
 - [[#7360](https://github.com/apache/incubator-seata/pull/7360)] Update resource cleanup logic for channel disconnection
 - [[#7363](https://github.com/apache/incubator-seata/pull/7363)] Upgrade npmjs dependencies
 - [[#7372](https://github.com/apache/incubator-seata/pull/7372)] optimize license ignore
+- [[#7375](https://github.com/apache/incubator-seata/pull/7375)] optimize close() logic of discovery module
 - [[#7388](https://github.com/apache/incubator-seata/pull/7388)] optimize binary packaging directory structure
 - [[#7412](https://github.com/apache/incubator-seata/pull/7412)] Helm template adapted to the new version of seata
 - [[#7414](https://github.com/apache/incubator-seata/pull/7414)] Remove the unused defaultEventExecutorGroup from the NettyClientBootstrap
@@ -88,7 +89,6 @@ Add changes here for all PR submitted to the 2.x branch.
 - [[#7379](https://github.com/apache/incubator-seata/issues/7379)] add UT for TccAnnotationProcessor class
 - [[#7422](https://github.com/apache/incubator-seata/pull/7422)] add UT for seata-spring-boot-starter module
 
-
 ### refactor:
 
 - [[#7315](https://github.com/apache/incubator-seata/pull/7315)] Refactor log testing to use ListAppender for more accurate and efficient log capture

changes/zh-cn/2.x.md

@@ -48,6 +48,7 @@
 - [[#7360](https://github.com/apache/incubator-seata/pull/7360)] 更新通道断开连接时的资源清理逻辑
 - [[#7363](https://github.com/apache/incubator-seata/pull/7363)] 升级 npmjs 依赖项
 - [[#7372](https://github.com/apache/incubator-seata/pull/7372)] 改进忽略许可证标头检查
+- [[#7375](https://github.com/apache/incubator-seata/pull/7375)] 优化 discovery 模块的 close 方法
 - [[#7388](https://github.com/apache/incubator-seata/pull/7388)] 优化二进制打包目录结构
 - [[#7412](https://github.com/apache/incubator-seata/pull/7412)] 适配新版本 Seata 的 Helm 模板
 - [[#7414](https://github.com/apache/incubator-seata/pull/7414)] 移除 NettyClientBootstrap 中 defaultEventExecutorGroup

discovery/seata-discovery-consul/src/main/java/org/apache/seata/discovery/registry/consul/ConsulRegistryServiceImpl.java

@@ -364,7 +364,23 @@ void stop() {
 
     @Override
     public void close() throws Exception {
-        client = null;
-    }
+        notifiers.values().forEach(ConsulNotifier::stop);
+        notifiers.clear();
+
+        // Shut down the ThreadPoolExecutor
+        if (notifierExecutor != null && !notifierExecutor.isShutdown()) {
+            notifierExecutor.shutdown();
+            try {
+                if (!notifierExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
+                    notifierExecutor.shutdownNow();
+                }
+            } catch (InterruptedException e) {
+                notifierExecutor.shutdownNow();
+            } finally {
+                notifierExecutor = null;
+            }
+        }
 
+        RegistryHeartBeats.close(REGISTRY_TYPE);
+    }
 }

discovery/seata-discovery-consul/src/test/java/org/apache/seata/discovery/registry/consul/ConsulRegistryServiceImplTest.java renamed to discovery/seata-discovery-consul/src/test/java/org/apache/seata/discovery/registry/consul/ConsulRegistryServiceImplMockTest.java

@@ -16,33 +16,40 @@
  */
 package org.apache.seata.discovery.registry.consul;
 
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import com.ecwid.consul.transport.RawResponse;
 import com.ecwid.consul.v1.ConsulClient;
 import com.ecwid.consul.v1.Response;
 import com.ecwid.consul.v1.health.model.HealthService;
-import org.apache.seata.config.Configuration;
-import org.apache.seata.config.ConfigurationFactory;
-import org.apache.seata.config.exception.ConfigNotFoundException;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.mockito.MockedStatic;
-import org.mockito.Mockito;
 
 import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
-import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
+import org.apache.seata.config.Configuration;
+import org.apache.seata.config.ConfigurationFactory;
+import org.apache.seata.config.exception.ConfigNotFoundException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
-public class ConsulRegistryServiceImplTest {
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+public class ConsulRegistryServiceImplMockTest {
 
     final String TEST_CLUSTER_NAME = "testCluster";
 
@@ -52,18 +59,22 @@ public class ConsulRegistryServiceImplTest {
 
     @BeforeEach
     public void init() throws Exception {
+        configuration = mock(Configuration.class);
         service = (ConsulRegistryServiceImpl) new ConsulRegistryProvider().provide();
         client = mock(ConsulClient.class);
-        this.setClient(service, client);
 
-        configuration = mock(Configuration.class);
+        Field clientField = ConsulRegistryServiceImpl.class.getDeclaredField("client");
+        clientField.setAccessible(true);
+        clientField.set(service, client);
     }
 
+    @Order(1)
     @Test
     public void testGetInstance() {
         Assertions.assertEquals(ConsulRegistryServiceImpl.getInstance(), service);
     }
 
+    @Order(2)
     @Test
     public void testRegister() throws Exception {
         InetSocketAddress inetSocketAddress = new InetSocketAddress("127.0.0.1", 8080);
@@ -77,6 +88,7 @@ public void testRegister() throws Exception {
         verify(client).agentServiceDeregister(any(), any());
     }
 
+    @Order(3)
     @Test
     public void testSubscribeAndLookup() throws Exception {
         ConsulListener consulListener = mock(ConsulListener.class);
@@ -110,14 +122,54 @@ public void testSubscribeAndLookup() throws Exception {
         }
 
         service.unsubscribe(TEST_CLUSTER_NAME, consulListener);
-        Assertions.assertNull(getMap("notifiers").get(TEST_CLUSTER_NAME));
+        assertNull(getMap("notifiers").get(TEST_CLUSTER_NAME));
     }
 
+    @Order(4)
+    @Test
+    public void testClose() throws Exception {
+        ExecutorService executorService1 = mockExecutorService(false, new InterruptedException("Test interruption"));
+        service.close();
+        verifyCloseResults(executorService1, true);
 
-    private void setClient(ConsulRegistryServiceImpl service, ConsulClient client) throws Exception {
-        Field clientField = ConsulRegistryServiceImpl.class.getDeclaredField("client");
+        ExecutorService executorService = mockExecutorService(false, null);
+        service.close();
+
+        verifyCloseResults(executorService, true);
+    }
+
+    private ExecutorService mockExecutorService(boolean awaitTerminationResult, InterruptedException exception) throws Exception {
+        ExecutorService executorService = mock(ExecutorService.class);
+        when(executorService.isShutdown()).thenReturn(false);
+
+        if (exception != null) {
+            when(executorService.awaitTermination(5, TimeUnit.SECONDS)).thenThrow(exception);
+        } else {
+            when(executorService.awaitTermination(5, TimeUnit.SECONDS)).thenReturn(awaitTerminationResult);
+        }
+
+        setExecutorService(executorService);
+        return executorService;
+    }
+
+    /**
+     * Verify the results of the closure method
+     */
+    private void verifyCloseResults(ExecutorService executorService, boolean expectShutdownNow) throws Exception {
+        verify(executorService).shutdown();
+        verify(executorService).awaitTermination(5, TimeUnit.SECONDS);
+        if (expectShutdownNow) {
+            verify(executorService).shutdownNow();
+        }
+
+        Field clientField = ConsulRegistryServiceImpl.class.getDeclaredField("notifiers");
         clientField.setAccessible(true);
-        clientField.set(service, client);
+        ConcurrentMap notifiers = (ConcurrentMap)clientField.get(service);
+        assertTrue(notifiers.isEmpty());
+
+        Field executorServiceField = ConsulRegistryServiceImpl.class.getDeclaredField("notifierExecutor");
+        executorServiceField.setAccessible(true);
+        assertNull(executorServiceField.get(service));
     }
 
     private void setExecutorService(ExecutorService executorService) throws Exception {

discovery/seata-discovery-core/src/main/java/org/apache/seata/discovery/registry/RegistryHeartBeats.java

@@ -75,8 +75,10 @@ public static void addHeartBeat(String registryType, InetSocketAddress serverAdd
         }, period, period, TimeUnit.MILLISECONDS);
     }
 
-    public static void close() {
-        HEARTBEAT_SCHEDULED.shutdown();
+    public static void close(String registryType) {
+        if (getHeartbeatEnabled(registryType)) {
+            HEARTBEAT_SCHEDULED.shutdown();
+        }
     }
 
     private static long getHeartbeatPeriod(String registryType) {

discovery/seata-discovery-etcd3/src/main/java/org/apache/seata/discovery/registry/etcd3/EtcdRegistryServiceImpl.java

@@ -225,12 +225,34 @@ public void onCompleted() {}
 
     @Override
     public void close() throws Exception {
-        if (lifeKeeper != null) {
-            lifeKeeper.stop();
-            if (lifeKeeperFuture != null) {
-                lifeKeeperFuture.get(3, TimeUnit.SECONDS);
+        // Shut down the ThreadPoolExecutor
+        if (executorService != null && !executorService.isShutdown()) {
+            executorService.shutdown();
+
+            try {
+                if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
+                    executorService.shutdownNow();
+                }
+            } catch (InterruptedException e) {
+                LOGGER.warn("ExecutorService shutdown interrupted. Forcing shutdown.");
+                executorService.shutdownNow();
+            } finally {
+                executorService = null;
+            }
+        }
+
+        // Close the Etcd client and release the underlying connection
+        if (client != null) {
+            try {
+                client.close();
+            } catch (Exception e) {
+                LOGGER.warn("Failed to close Etcd client: {}", e.getMessage());
+            } finally {
+                client = null;
             }
         }
+
+        RegistryHeartBeats.close(REGISTRY_TYPE);
     }
 
     /**

discovery/seata-discovery-etcd3/src/test/java/org/apache/seata/discovery/registry/etcd/EtcdRegistryProviderTest.java renamed to discovery/seata-discovery-etcd3/src/test/java/org/apache/seata/discovery/registry/etcd3/EtcdRegistryProviderTest.java

@@ -14,10 +14,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.discovery.registry.etcd;
+package org.apache.seata.discovery.registry.etcd3;
 
-import org.apache.seata.discovery.registry.etcd3.EtcdRegistryProvider;
-import org.apache.seata.discovery.registry.etcd3.EtcdRegistryServiceImpl;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;