From 307935a5d9638f81165c8ecae75e38ff56cf03b0 Mon Sep 17 00:00:00 2001
From: yuhuangbin <danielyu96@163.com>
Date: Mon, 7 Sep 2020 21:52:16 +0800
Subject: [PATCH] nacos re-register enhance

---
 .../cloud/nacos/NacosDiscoveryProperties.java | 26 ++++++++++++++++---
 .../nacos/registry/NacosServiceRegistry.java  |  8 +++++-
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/NacosDiscoveryProperties.java b/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/NacosDiscoveryProperties.java
index 52fedfe6..78096d2e 100644
--- a/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/NacosDiscoveryProperties.java
+++ b/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/NacosDiscoveryProperties.java
@@ -30,6 +30,7 @@ import java.util.regex.Pattern;
 import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
 
+import com.alibaba.cloud.nacos.registry.NacosAutoServiceRegistration;
 import com.alibaba.nacos.api.exception.NacosException;
 import com.alibaba.nacos.api.naming.NamingMaintainService;
 import com.alibaba.nacos.api.naming.NamingService;
@@ -210,12 +211,20 @@ public class NacosDiscoveryProperties {
 	 */
 	private boolean ephemeral = true;
 
+	/**
+	 * If the instance needs to be re-registered. The default value is false.
+	 */
+	private boolean needNewlyRegister = false;
+
 	@Autowired
 	private InetUtils inetUtils;
 
 	@Autowired
 	private Environment environment;
 
+	@Autowired
+	private NacosAutoServiceRegistration nacosAutoServiceRegistration;
+
 	private NamingService namingService;
 
 	private NamingMaintainService namingMaintainService;
@@ -272,12 +281,16 @@ public class NacosDiscoveryProperties {
 		Properties properties = getNacosProperties();
 		this.namingService = createNamingService(properties);
 		this.namingMaintainService = createMaintainService(properties);
+
+		if (needNewlyRegister) {
+			nacosAutoServiceRegistration.start();
+		}
 	}
 
-    @PreDestroy
-    public void destroy() throws NacosException {
-        namingService.shutDown();
-    }
+	@PreDestroy
+	public void destroy() throws NacosException {
+		shutdownNacosService();
+	}
 
 	public String getEndpoint() {
 		return endpoint;
@@ -598,4 +611,9 @@ public class NacosDiscoveryProperties {
 		return sb.toString();
 	}
 
+	private void shutdownNacosService() throws NacosException {
+		nacosAutoServiceRegistration.stop();
+		needNewlyRegister = true;
+	}
+
 }
diff --git a/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/registry/NacosServiceRegistry.java b/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/registry/NacosServiceRegistry.java
index 1877ceae..5e0abcb5 100644
--- a/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/registry/NacosServiceRegistry.java
+++ b/spring-cloud-alibaba-starters/spring-cloud-starter-alibaba-nacos-discovery/src/main/java/com/alibaba/cloud/nacos/registry/NacosServiceRegistry.java
@@ -19,6 +19,7 @@ package com.alibaba.cloud.nacos.registry;
 import java.util.List;
 
 import com.alibaba.cloud.nacos.NacosDiscoveryProperties;
+import com.alibaba.nacos.api.exception.NacosException;
 import com.alibaba.nacos.api.naming.NamingService;
 import com.alibaba.nacos.api.naming.pojo.Instance;
 import org.slf4j.Logger;
@@ -101,7 +102,12 @@ public class NacosServiceRegistry implements ServiceRegistry<Registration> {
 
 	@Override
 	public void close() {
-
+		try {
+			nacosDiscoveryProperties.namingServiceInstance().shutDown();
+		}
+		catch (NacosException e) {
+			log.error("Nacos namingService shutDown failed", e);
+		}
 	}
 
 	@Override