diff --git a/spring-cloud-alibaba-dubbo/src/main/java/com/alibaba/cloud/dubbo/autoconfigure/DubboLoadBalancedRestTemplateAutoConfiguration.java b/spring-cloud-alibaba-dubbo/src/main/java/com/alibaba/cloud/dubbo/autoconfigure/DubboLoadBalancedRestTemplateAutoConfiguration.java index 74bb1e38..b1a75923 100644 --- a/spring-cloud-alibaba-dubbo/src/main/java/com/alibaba/cloud/dubbo/autoconfigure/DubboLoadBalancedRestTemplateAutoConfiguration.java +++ b/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 DUBBO_TRANSPORTED_CLASS = DubboTransported.class; @@ -89,6 +92,9 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration @Autowired(required = false) private Map 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() { + @EventListener(ContextRefreshedEvent.class) + public void adaptRestTemplates(ContextRefreshedEvent event) { - DubboTransportedAttributesResolver attributesResolver = new DubboTransportedAttributesResolver( - environment); + if (event.getApplicationContext() == this.applicationContext) { - for (Map.Entry entry : restTemplates.entrySet()) { - String beanName = entry.getKey(); - Map dubboTranslatedAttributes = getDubboTranslatedAttributes( - beanName, attributesResolver); - if (!CollectionUtils.isEmpty(dubboTranslatedAttributes)) { - adaptRestTemplate(entry.getValue(), dubboTranslatedAttributes); + DubboTransportedAttributesResolver attributesResolver = new DubboTransportedAttributesResolver( + environment); + + for (Map.Entry entry : restTemplates.entrySet()) { + String beanName = entry.getKey(); + Map 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 dubboTranslatedAttributes) { @@ -188,4 +197,9 @@ public class DubboLoadBalancedRestTemplateAutoConfiguration this.classLoader = classLoader; } + @Override + public void setApplicationContext(ApplicationContext applicationContext) { + this.applicationContext = applicationContext; + } + }