优化类结构

discovery-plugin-framework-eureka/src/main/java/com/nepxion/discovery/plugin/framework/context/EurekaApplicationContextInitializer.java

@@ -10,22 +10,16 @@ package com.nepxion.discovery.plugin.framework.context;
  */
 
 import org.springframework.beans.BeansException;
-import org.springframework.cloud.client.discovery.DiscoveryClient;
 import org.springframework.cloud.netflix.eureka.EurekaInstanceConfigBean;
 import org.springframework.cloud.netflix.eureka.serviceregistry.EurekaServiceRegistry;
 import org.springframework.context.ConfigurableApplicationContext;
 
-import com.nepxion.discovery.plugin.framework.decorator.DiscoveryClientDecorator;
 import com.nepxion.discovery.plugin.framework.decorator.EurekaServiceRegistryDecorator;
 
 public class EurekaApplicationContextInitializer extends PluginApplicationContextInitializer {
     @Override
     protected Object afterInitialization(ConfigurableApplicationContext applicationContext, Object bean, String beanName) throws BeansException {
-        if (bean instanceof DiscoveryClient) {
-            DiscoveryClient discoveryClient = (DiscoveryClient) bean;
-
-            return new DiscoveryClientDecorator(discoveryClient, applicationContext);
-        } else if (bean instanceof EurekaServiceRegistry) {
+        if (bean instanceof EurekaServiceRegistry) {
             EurekaServiceRegistry eurekaServiceRegistry = (EurekaServiceRegistry) bean;
 
             return new EurekaServiceRegistryDecorator(eurekaServiceRegistry, applicationContext);

discovery-plugin-framework-zookeeper/src/main/java/com/nepxion/discovery/plugin/framework/context/ZookeeperApplicationContextInitializer.java

@@ -10,21 +10,15 @@ package com.nepxion.discovery.plugin.framework.context;
  */
 
 import org.springframework.beans.BeansException;
-import org.springframework.cloud.client.discovery.DiscoveryClient;
 import org.springframework.cloud.zookeeper.serviceregistry.ZookeeperServiceRegistry;
 import org.springframework.context.ConfigurableApplicationContext;
 
-import com.nepxion.discovery.plugin.framework.decorator.DiscoveryClientDecorator;
 import com.nepxion.discovery.plugin.framework.decorator.ZookeeperServiceRegistryDecorator;
 
 public class ZookeeperApplicationContextInitializer extends PluginApplicationContextInitializer {
     @Override
     protected Object afterInitialization(ConfigurableApplicationContext applicationContext, Object bean, String beanName) throws BeansException {
-        if (bean instanceof DiscoveryClient) {
-            DiscoveryClient discoveryClient = (DiscoveryClient) bean;
-
-            return new DiscoveryClientDecorator(discoveryClient, applicationContext);
-        } else if (bean instanceof ZookeeperServiceRegistry) {
+        if (bean instanceof ZookeeperServiceRegistry) {
             ZookeeperServiceRegistry zookeeperServiceRegistry = (ZookeeperServiceRegistry) bean;
 
             return new ZookeeperServiceRegistryDecorator(zookeeperServiceRegistry, applicationContext);

discovery-plugin-framework/src/main/java/com/nepxion/discovery/plugin/framework/context/PluginApplicationContextInitializer.java

@@ -11,17 +11,26 @@ package com.nepxion.discovery.plugin.framework.context;
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.InstantiationAwareBeanPostProcessorAdapter;
+import org.springframework.cloud.client.discovery.DiscoveryClient;
 import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.ConfigurableApplicationContext;
 
+import com.nepxion.discovery.plugin.framework.decorator.DiscoveryClientDecorator;
+
 public abstract class PluginApplicationContextInitializer implements ApplicationContextInitializer<ConfigurableApplicationContext> {
     @Override
     public void initialize(ConfigurableApplicationContext applicationContext) {
         applicationContext.getBeanFactory().addBeanPostProcessor(new InstantiationAwareBeanPostProcessorAdapter() {
             @Override
             public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
+                if (bean instanceof DiscoveryClient) {
+                    DiscoveryClient discoveryClient = (DiscoveryClient) bean;
+
+                    return new DiscoveryClientDecorator(discoveryClient, applicationContext);
+                } else {
                     return afterInitialization(applicationContext, bean, beanName);
                 }
+            }
         });
     }