Merge pull request Netflix#849 from qiangdavidliu/master

Add optional args to ApplicationInfoManager
xiaoshao · Sep 20, 2016 · 44c138e · 44c138e
2 parents 435f231 + 3f3fbb1
commit 44c138e
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 12 deletions.
diff --git a/eureka-client/src/main/java/com/netflix/appinfo/ApplicationInfoManager.java b/eureka-client/src/main/java/com/netflix/appinfo/ApplicationInfoManager.java
@@ -47,33 +47,68 @@
 @Singleton
 public class ApplicationInfoManager {
     private static final Logger logger = LoggerFactory.getLogger(ApplicationInfoManager.class);
-    private static ApplicationInfoManager instance = new ApplicationInfoManager();
+
+    private static final InstanceStatusMapper NO_OP_MAPPER = new InstanceStatusMapper() {
+        @Override
+        public InstanceStatus map(InstanceStatus prev) {
+            return prev;
+        }
+    };
+
+    private static ApplicationInfoManager instance = new ApplicationInfoManager(null, null, null);
+
+    protected final Map<String, StatusChangeListener> listeners;
+    private final InstanceStatusMapper instanceStatusMapper;
 
     private InstanceInfo instanceInfo;
     private EurekaInstanceConfig config;
 
-    protected Map<String, StatusChangeListener> listeners;
+    public static class OptionalArgs {
+        private InstanceStatusMapper instanceStatusMapper;
 
-    private ApplicationInfoManager() {
-        listeners = new ConcurrentHashMap<String, StatusChangeListener>();
+        @com.google.inject.Inject(optional = true)
+        public void setInstanceStatusMapper(InstanceStatusMapper instanceStatusMapper) {
+            this.instanceStatusMapper = instanceStatusMapper;
+        }
+
+        InstanceStatusMapper getInstanceStatusMapper() {
+            return instanceStatusMapper == null ? NO_OP_MAPPER : instanceStatusMapper;
+        }
     }
 
     /**
-     * public for spring DI use. This class should be in singleton scope so do not create explicitly.
-     * Either use DI or use getInstance().initComponent() if not using DI
+     * public for DI use. This class should be in singleton scope so do not create explicitly.
+     * Either use DI or create this explicitly using one of the other public constructors.
      */
     @Inject
-    public ApplicationInfoManager(EurekaInstanceConfig config, InstanceInfo instanceInfo) {
+    public ApplicationInfoManager(EurekaInstanceConfig config, InstanceInfo instanceInfo, OptionalArgs optionalArgs) {
         this.config = config;
         this.instanceInfo = instanceInfo;
         this.listeners = new ConcurrentHashMap<String, StatusChangeListener>();
+        if (optionalArgs != null) {
+            this.instanceStatusMapper = optionalArgs.getInstanceStatusMapper();
+        } else {
+            this.instanceStatusMapper = NO_OP_MAPPER;
+        }
 
         // Hack to allow for getInstance() to use the DI'd ApplicationInfoManager
         instance = this;
     }
 
+    public ApplicationInfoManager(EurekaInstanceConfig config, /* nullable */ OptionalArgs optionalArgs) {
+        this(config, new EurekaConfigBasedInstanceInfoProvider(config).get(), optionalArgs);
+    }
+
+    public ApplicationInfoManager(EurekaInstanceConfig config, InstanceInfo instanceInfo) {
+        this(config, instanceInfo, null);
+    }
+
+    /**
+     * @deprecated 2016-09-19 prefer {@link #ApplicationInfoManager(EurekaInstanceConfig, com.netflix.appinfo.ApplicationInfoManager.OptionalArgs)}
+     */
+    @Deprecated
     public ApplicationInfoManager(EurekaInstanceConfig config) {
-        this(config, new EurekaConfigBasedInstanceInfoProvider(config).get());
+        this(config, (OptionalArgs) null);
     }
 
     /**
@@ -130,11 +165,16 @@ public void registerAppMetadata(Map<String, String> appMetadata) {
      * @param status Status of the instance
      */
     public synchronized void setInstanceStatus(InstanceStatus status) {
-        InstanceStatus prev = instanceInfo.setStatus(status);
+        InstanceStatus next = instanceStatusMapper.map(status);
+        if (next == null) {
+            return;
+        }
+
+        InstanceStatus prev = instanceInfo.setStatus(next);
         if (prev != null) {
             for (StatusChangeListener listener : listeners.values()) {
                 try {
-                    listener.notify(new StatusChangeEvent(prev, status));
+                    listener.notify(new StatusChangeEvent(prev, next));
                 } catch (Exception e) {
                     logger.warn("failed to notify listener: {}", listener.getId(), e);
                 }
@@ -203,4 +243,16 @@ public static interface StatusChangeListener {
 
         void notify(StatusChangeEvent statusChangeEvent);
     }
+
+    public static interface InstanceStatusMapper {
+
+        /**
+         * given a starting {@link com.netflix.appinfo.InstanceInfo.InstanceStatus}, apply a mapping to return
+         * the follow up status, if applicable.
+         *
+         * @return the mapped instance status, or null if the mapping is not applicable.
+         */
+        InstanceStatus map(InstanceStatus prev);
+    }
+
 }
diff --git a/eureka-client/src/test/java/com/netflix/appinfo/ApplicationInfoManagerTest.java b/eureka-client/src/test/java/com/netflix/appinfo/ApplicationInfoManagerTest.java
@@ -8,6 +8,8 @@
 import static com.netflix.appinfo.AmazonInfo.MetaDataKey.publicHostname;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertThat;
 import static org.mockito.Matchers.anyBoolean;
 import static org.mockito.Mockito.spy;
@@ -29,7 +31,7 @@ public void setUp() {
 
         config = spy(new CloudInstanceConfig(initialAmazonInfo));
         instanceInfo = InstanceInfoGenerator.takeOne();
-        this.applicationInfoManager = new ApplicationInfoManager(config, instanceInfo);
+        this.applicationInfoManager = new ApplicationInfoManager(config, instanceInfo, null);
         when(config.getDefaultAddressResolutionOrder()).thenReturn(new String[]{
                 publicHostname.name(),
                 localIpv4.name()
@@ -47,4 +49,42 @@ public void testRefreshDataCenterInfoWithAmazonInfo() {
 
         assertThat(instanceInfo.getHostName(), is(newPublicHostname));
     }
+
+    @Test
+    public void testCustomInstanceStatusMapper() {
+        ApplicationInfoManager.OptionalArgs optionalArgs = new ApplicationInfoManager.OptionalArgs();
+        optionalArgs.setInstanceStatusMapper(new ApplicationInfoManager.InstanceStatusMapper() {
+            @Override
+            public InstanceInfo.InstanceStatus map(InstanceInfo.InstanceStatus prev) {
+                return InstanceInfo.InstanceStatus.UNKNOWN;
+            }
+        });
+
+        applicationInfoManager = new ApplicationInfoManager(config, instanceInfo, optionalArgs);
+        InstanceInfo.InstanceStatus existingStatus = applicationInfoManager.getInfo().getStatus();
+        assertNotEquals(existingStatus, InstanceInfo.InstanceStatus.UNKNOWN);
+
+        applicationInfoManager.setInstanceStatus(InstanceInfo.InstanceStatus.UNKNOWN);
+        existingStatus = applicationInfoManager.getInfo().getStatus();
+        assertEquals(existingStatus, InstanceInfo.InstanceStatus.UNKNOWN);
+    }
+
+    @Test
+    public void testNullResultInstanceStatusMapper() {
+        ApplicationInfoManager.OptionalArgs optionalArgs = new ApplicationInfoManager.OptionalArgs();
+        optionalArgs.setInstanceStatusMapper(new ApplicationInfoManager.InstanceStatusMapper() {
+            @Override
+            public InstanceInfo.InstanceStatus map(InstanceInfo.InstanceStatus prev) {
+                return null;
+            }
+        });
+
+        applicationInfoManager = new ApplicationInfoManager(config, instanceInfo, optionalArgs);
+        InstanceInfo.InstanceStatus existingStatus1 = applicationInfoManager.getInfo().getStatus();
+        assertNotEquals(existingStatus1, InstanceInfo.InstanceStatus.UNKNOWN);
+
+        applicationInfoManager.setInstanceStatus(InstanceInfo.InstanceStatus.UNKNOWN);
+        InstanceInfo.InstanceStatus existingStatus2 = applicationInfoManager.getInfo().getStatus();
+        assertEquals(existingStatus2, existingStatus1);
+    }
 }
diff --git a/eureka-client/src/test/java/com/netflix/discovery/DiscoveryClientRegisterUpdateTest.java b/eureka-client/src/test/java/com/netflix/discovery/DiscoveryClientRegisterUpdateTest.java
@@ -102,7 +102,7 @@ public void testRegistrationDisabled() throws Exception {
 
     public class TestApplicationInfoManager extends ApplicationInfoManager {
         TestApplicationInfoManager() {
-            super(null, null);
+            super(null, null, null);
         }
 
         public Map<String, StatusChangeListener> getStatusChangeListeners() {