Change BaseLoadBalancer to choose the LoadBalancerStats class based o…

…n a new ClientConfigKey so that a different impl can be used
Netflix · Mar 30, 2018 · 72df8b1 · 72df8b1
1 parent ffd1fdd
commit 72df8b1
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/ribbon-core/src/main/java/com/netflix/client/config/CommonClientConfigKey.java b/ribbon-core/src/main/java/com/netflix/client/config/CommonClientConfigKey.java
@@ -157,6 +157,8 @@ public abstract class CommonClientConfigKey<T> implements IClientConfigKey<T> {
     public static final IClientConfigKey<Integer> NFLoadBalancerPingInterval = new CommonClientConfigKey<Integer>("NFLoadBalancerPingInterval"){};
 
     public static final IClientConfigKey<Integer> NFLoadBalancerMaxTotalPingTime = new CommonClientConfigKey<Integer>("NFLoadBalancerMaxTotalPingTime"){};
+
+    public static final IClientConfigKey<String> NFLoadBalancerStatsClassName = new CommonClientConfigKey<String>("NFLoadBalancerStatsClassName"){};
 
     public static final IClientConfigKey<String> NIWSServerListClassName = new CommonClientConfigKey<String>("NIWSServerListClassName"){};
 

diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/BaseLoadBalancer.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/BaseLoadBalancer.java
@@ -160,10 +160,10 @@ public BaseLoadBalancer(IClientConfig config) {
     }
 
     public BaseLoadBalancer(IClientConfig config, IRule rule, IPing ping) {
-        initWithConfig(config, rule, ping);
+        initWithConfig(config, rule, ping, new LoadBalancerStats((config.getClientName())));
     }
 
-    void initWithConfig(IClientConfig clientConfig, IRule rule, IPing ping) {
+    void initWithConfig(IClientConfig clientConfig, IRule rule, IPing ping, LoadBalancerStats stats) {
         this.config = clientConfig;
         String clientName = clientConfig.getClientName();
         this.name = clientName;
@@ -184,7 +184,8 @@ void initWithConfig(IClientConfig clientConfig, IRule rule, IPing ping) {
         // LB, these are your Ping and Rule guys ...
         setRule(rule);
         setPing(ping);
-        setLoadBalancerStats(new LoadBalancerStats(clientName));
+
+        setLoadBalancerStats(stats);
         rule.setLoadBalancer(this);
         if (ping instanceof AbstractLoadBalancerPing) {
             ((AbstractLoadBalancerPing) ping).setLoadBalancer(this);
@@ -209,18 +210,23 @@ public void initWithNiwsConfig(IClientConfig clientConfig) {
                 .getProperty(CommonClientConfigKey.NFLoadBalancerRuleClassName);
         String pingClassName = (String) clientConfig
                 .getProperty(CommonClientConfigKey.NFLoadBalancerPingClassName);
+        String loadBalancerStatsClassName = (String) clientConfig
+                .getProperty(CommonClientConfigKey.NFLoadBalancerStatsClassName, LoadBalancerStats.class.getName());
 
         IRule rule;
         IPing ping;
+        LoadBalancerStats stats;
         try {
             rule = (IRule) ClientFactory.instantiateInstanceWithClientConfig(
                     ruleClassName, clientConfig);
             ping = (IPing) ClientFactory.instantiateInstanceWithClientConfig(
                     pingClassName, clientConfig);
+            stats = (LoadBalancerStats) ClientFactory.instantiateInstanceWithClientConfig(
+                    loadBalancerStatsClassName, clientConfig);
         } catch (Exception e) {
             throw new RuntimeException("Error initializing load balancer", e);
         }
-        initWithConfig(clientConfig, rule, ping);
+        initWithConfig(clientConfig, rule, ping, stats);
     }
 
     public void addServerListChangeListener(ServerListChangeListener listener) {

diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/LoadBalancerStats.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/LoadBalancerStats.java
@@ -33,6 +33,8 @@
 import com.google.common.cache.LoadingCache;
 import com.google.common.cache.RemovalListener;
 import com.google.common.cache.RemovalNotification;
+import com.netflix.client.IClientConfigAware;
+import com.netflix.client.config.IClientConfig;
 import com.netflix.config.CachedDynamicIntProperty;
 import com.netflix.config.DynamicIntProperty;
 import com.netflix.config.DynamicPropertyFactory;
@@ -51,7 +53,7 @@
  * @author stonse
  * 
  */
-public class LoadBalancerStats {
+public class LoadBalancerStats implements IClientConfigAware {
 
     private static final String PREFIX = "LBStats_";
 
@@ -95,7 +97,7 @@ protected ServerStats createServerStats(Server server) {
         return ss;        
     }
 
-    private LoadBalancerStats(){
+    public LoadBalancerStats(){
         zoneStatsMap = new ConcurrentHashMap<String, ZoneStats>();  
         upServerListZoneMap = new ConcurrentHashMap<String, List<? extends Server>>();        
     }
@@ -105,7 +107,14 @@ public LoadBalancerStats(String name){
         this.name = name;
         Monitors.registerObject(name, this); 
     }
-
+
+    @Override
+    public void initWithNiwsConfig(IClientConfig clientConfig)
+    {
+        this.name = clientConfig.getClientName();
+        Monitors.registerObject(name, this);
+    }
+
     public String getName() {
         return name;
     }