Add host tracker in onUp/onAdd

olim7t · olim7t · commit 4281d0710666 · 2018-01-11T17:33:09.000-08:00
This avoids a race when onDown/onRemove removes the tracker, but pending
requests complete after that and re-add it.
diff --git a/driver-core/src/main/java/com/datastax/driver/core/policies/LatencyAwarePolicy.java b/driver-core/src/main/java/com/datastax/driver/core/policies/LatencyAwarePolicy.java
@@ -178,6 +178,9 @@ private static ThreadFactory threadFactory(String nameFormat) {
     @Override
     public void init(Cluster cluster, Collection<Host> hosts) {
         childPolicy.init(cluster, hosts);
+        for (Host host : hosts) {
+            latencyTracker.addHost(host);
+        }
         cluster.register(latencyTracker);
         metrics = cluster.getMetrics();
         if (metrics != null) {
@@ -303,6 +306,7 @@ public Snapshot getScoresSnapshot() {
     @Override
     public void onUp(Host host) {
         childPolicy.onUp(host);
+        latencyTracker.addHost(host);
     }
 
     @Override
@@ -314,6 +318,7 @@ public void onDown(Host host) {
     @Override
     public void onAdd(Host host) {
         childPolicy.onAdd(host);
+        latencyTracker.addHost(host);
     }
 
     @Override
@@ -423,34 +428,15 @@ private class Tracker implements LatencyTracker {
 
         @Override
         public void update(final Host host, Statement statement, Exception exception, long newLatencyNanos) {
-            if (shouldConsiderNewLatency(statement, exception)) {
-                HostLatencyTracker hostTracker = latencies.get(host);
-                if (hostTracker == null) {
-                    hostTracker = new HostLatencyTracker(scale, (30L * minMeasure) / 100L);
-                    HostLatencyTracker old = latencies.putIfAbsent(host, hostTracker);
-                    if (old != null) {
-                        hostTracker = old;
-                    } else if (hostMetricsEnabled()) {
-                        String metricName = MetricsUtil.hostMetricName("LatencyAwarePolicy.latencies.", host);
-                        if (!metrics.getRegistry().getNames().contains(metricName)) {
-                            logger.info("Adding gauge " + metricName);
-                            metrics.getRegistry().register(
-                                    metricName,
-                                    new Gauge<Long>() {
-                                        @Override
-                                        public Long getValue() {
-                                            TimestampedAverage latency = latencyTracker.latencyOf(host);
-                                            return (latency == null) ? -1 : latency.average;
-                                        }
-                                    });
-                        }
-                    }
+            HostLatencyTracker hostTracker = latencies.get(host);
+            if (hostTracker != null) {
+                if (shouldConsiderNewLatency(statement, exception)) {
+                    hostTracker.add(newLatencyNanos);
+                } else if (hostMetricsEnabled()) {
+                    metrics.getRegistry()
+                            .counter(MetricsUtil.hostMetricName("LatencyAwarePolicy.ignored-latencies.", host))
+                            .inc();
                 }
-                hostTracker.add(newLatencyNanos);
-            } else if (hostMetricsEnabled()) {
-                metrics.getRegistry()
-                        .counter(MetricsUtil.hostMetricName("LatencyAwarePolicy.ignored-latencies.", host))
-                        .inc();
             }
         }
 
@@ -490,7 +476,29 @@ public Map<Host, TimestampedAverage> currentLatencies() {
             return map;
         }
 
+        public void addHost(final Host host) {
+            logger.debug("Adding tracker for {}", host);
+            HostLatencyTracker old = latencies.putIfAbsent(host,
+                    new HostLatencyTracker(scale, (30L * minMeasure) / 100L));
+            if (old == null && hostMetricsEnabled()) {
+                String metricName = MetricsUtil.hostMetricName("LatencyAwarePolicy.latencies.", host);
+                if (!metrics.getRegistry().getNames().contains(metricName)) {
+                    logger.debug("Adding gauge " + metricName);
+                    metrics.getRegistry().register(
+                            metricName,
+                            new Gauge<Long>() {
+                                @Override
+                                public Long getValue() {
+                                    TimestampedAverage latency = latencyTracker.latencyOf(host);
+                                    return (latency == null) ? -1 : latency.average;
+                                }
+                            });
+                }
+            }
+        }
+
         public void resetHost(Host host) {
+            logger.debug("Removing tracker for {}", host);
             latencies.remove(host);
         }
 
