Merge pull request #905 from echooymxq/master

Advance adapt Dubbo restTemplates

spring-cloud-alibaba-dubbo/src/main/java/com/alibaba/cloud/dubbo/autoconfigure/DubboLoadBalancedRestTemplateAutoConfiguration.java

@@ -25,16 +25,19 @@ import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
-import org.springframework.boot.context.event.ApplicationStartedEvent;
 import org.springframework.cloud.client.loadbalancer.LoadBalanced;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerInterceptor;
 import org.springframework.cloud.client.loadbalancer.RestTemplateCustomizer;
 import org.springframework.cloud.client.loadbalancer.RetryLoadBalancerInterceptor;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.context.event.EventListener;
 import org.springframework.core.env.Environment;
 import org.springframework.core.type.MethodMetadata;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
+import org.springframework.lang.Nullable;
 import org.springframework.util.CollectionUtils;
 import org.springframework.web.client.RestTemplate;
 
@@ -56,8 +59,8 @@ import com.alibaba.cloud.dubbo.service.DubboGenericServiceFactory;
 @ConditionalOnClass(name = { "org.springframework.web.client.RestTemplate" })
 @AutoConfigureAfter(name = {
 		"org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration" })
-public class DubboLoadBalancedRestTemplateAutoConfiguration
-		implements BeanClassLoaderAware, SmartInitializingSingleton {
+public class DubboLoadBalancedRestTemplateAutoConfiguration implements
+		BeanClassLoaderAware, ApplicationContextAware, SmartInitializingSingleton {
 
 	private static final Class<DubboTransported> DUBBO_TRANSPORTED_CLASS = DubboTransported.class;
 
@@ -89,6 +92,9 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 	@Autowired(required = false)
 	private Map<String, RestTemplate> restTemplates = Collections.emptyMap();
 
+	@Nullable
+	private ApplicationContext applicationContext;
+
 	private ClassLoader classLoader;
 
 	/**
@@ -111,18 +117,21 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 	 * {@link SmartInitializingSingleton} beans or
 	 * {@link RestTemplateCustomizer#customize(RestTemplate) customization})
 	 */
-	@EventListener(ApplicationStartedEvent.class)
-	public void adaptRestTemplates() {
-
-		DubboTransportedAttributesResolver attributesResolver = new DubboTransportedAttributesResolver(
-				environment);
-
-		for (Map.Entry<String, RestTemplate> entry : restTemplates.entrySet()) {
-			String beanName = entry.getKey();
-			Map<String, Object> dubboTranslatedAttributes = getDubboTranslatedAttributes(
-					beanName, attributesResolver);
-			if (!CollectionUtils.isEmpty(dubboTranslatedAttributes)) {
-				adaptRestTemplate(entry.getValue(), dubboTranslatedAttributes);
+	@EventListener(ContextRefreshedEvent.class)
+	public void adaptRestTemplates(ContextRefreshedEvent event) {
+
+		if (event.getApplicationContext() == this.applicationContext) {
+
+			DubboTransportedAttributesResolver attributesResolver = new DubboTransportedAttributesResolver(
+					environment);
+
+			for (Map.Entry<String, RestTemplate> entry : restTemplates.entrySet()) {
+				String beanName = entry.getKey();
+				Map<String, Object> dubboTranslatedAttributes = getDubboTranslatedAttributes(
+						beanName, attributesResolver);
+				if (!CollectionUtils.isEmpty(dubboTranslatedAttributes)) {
+					adaptRestTemplate(entry.getValue(), dubboTranslatedAttributes);
+				}
 			}
 		}
 	}
@@ -132,7 +141,7 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 	 * {@link DubboTransported @DubboTransported}
 	 *
 	 * @param beanName the bean name of {@link LoadBalanced @LoadBalanced}
-	 * {@link RestTemplate}
+	 *     {@link RestTemplate}
 	 * @param attributesResolver {@link DubboTransportedAttributesResolver}
 	 * @return non-null {@link Map}
 	 */
@@ -144,10 +153,10 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 			AnnotatedBeanDefinition annotatedBeanDefinition = (AnnotatedBeanDefinition) beanDefinition;
 			MethodMetadata factoryMethodMetadata = annotatedBeanDefinition
 					.getFactoryMethodMetadata();
-			attributes = factoryMethodMetadata != null
-					? Optional.ofNullable(factoryMethodMetadata
-							.getAnnotationAttributes(DUBBO_TRANSPORTED_CLASS_NAME)).orElse(attributes)
-					: Collections.emptyMap();
+			attributes = factoryMethodMetadata != null ? Optional
+					.ofNullable(factoryMethodMetadata
+							.getAnnotationAttributes(DUBBO_TRANSPORTED_CLASS_NAME))
+					.orElse(attributes) : Collections.emptyMap();
 		}
 		return attributesResolver.resolve(attributes);
 	}
@@ -158,8 +167,8 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 	 *
 	 * @param restTemplate {@link LoadBalanced @LoadBalanced} {@link RestTemplate} Bean
 	 * @param dubboTranslatedAttributes the annotation dubboTranslatedAttributes
-	 * {@link RestTemplate} bean being annotated
-	 * {@link DubboTransported @DubboTransported}
+	 *     {@link RestTemplate} bean being annotated
+	 *     {@link DubboTransported @DubboTransported}
 	 */
 	private void adaptRestTemplate(RestTemplate restTemplate,
 			Map<String, Object> dubboTranslatedAttributes) {
@@ -188,4 +197,9 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration
 		this.classLoader = classLoader;
 	}
 
+	@Override
+	public void setApplicationContext(ApplicationContext applicationContext) {
+		this.applicationContext = applicationContext;
+	}
+
 }