From aea48a9d7bb24df94d5de61a5dd39eda3dd028bd Mon Sep 17 00:00:00 2001
From: David Hoover <dhoover@google.com>
Date: Thu, 16 Apr 2015 14:04:24 -0700
Subject: [PATCH] Synchronize all appropriate methods in ConsistentHash.

Only add(T, int) was synchronized (and its javadoc claimed it was the
only one modifying .items), but all the bulk add methods were going
straight to addInternal() and then refreshTable() on their own. Then
countAllPoints() was iterating over items.values() which has undefined
behavior if the underlying Map changed.

Technically I didn't need to synchronize the current implementations of
remove(T) and add(T) since they delegated fully to add(T, int), but
that's leaking implementation details and inviting trouble down the road
when someone doesn't realize that.

Likewise, addInternal() and refreshTable() need not be synchronized
since they're only ever called from already synchronized methods, but
I'm inclined to go for belt & suspenders in the wake of these other
issues.
---
 .../main/java/hudson/util/ConsistentHash.java | 33 +++++++++----------
 .../java/hudson/util/ConsistentHashTest.java  | 13 ++++++++
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/core/src/main/java/hudson/util/ConsistentHash.java b/core/src/main/java/hudson/util/ConsistentHash.java
index 83bdb3f53d..0919975ac5 100644
--- a/core/src/main/java/hudson/util/ConsistentHash.java
+++ b/core/src/main/java/hudson/util/ConsistentHash.java
@@ -64,6 +64,7 @@ public class ConsistentHash<T> {
      * All the items in the hash, to their replication factors.
      */
     private final Map<T,Point[]> items = new HashMap<T,Point[]>();
+    private int numPoints;
 
     private final int defaultReplication;
     private final Hash<T> hash;
@@ -100,8 +101,13 @@ public class ConsistentHash<T> {
         private final Object[] owner; // really T[]
 
         private Table() {
+            int r=0;
+            for (Point[] v : items.values())
+                r+=v.length;
+            numPoints = r;
+
             // merge all points from all nodes and sort them into a single array
-            Point[] allPoints = new Point[countAllPoints()];
+            Point[] allPoints = new Point[numPoints];
             int p=0;
             for (Point[] v : items.values()) {
                 System.arraycopy(v,0,allPoints,p,v.length);
@@ -211,23 +217,20 @@ public class ConsistentHash<T> {
     }
 
     public int countAllPoints() {
-        int r=0;
-        for (Point[] v : items.values())
-            r+=v.length;
-        return r;
+        return numPoints;
     }
 
     /**
      * Adds a new node with the default number of replica.
      */
-    public void add(T node) {
+    public synchronized void add(T node) {
         add(node,defaultReplication);
     }
 
     /**
      * Calls {@link #add(Object)} with all the arguments.
      */
-    public void addAll(T... nodes) {
+    public synchronized void addAll(T... nodes) {
         for (T node : nodes)
             addInternal(node,defaultReplication);
         refreshTable();
@@ -236,7 +239,7 @@ public class ConsistentHash<T> {
     /**
      * Calls {@link #add(Object)} with all the arguments.
      */
-    public void addAll(Collection<? extends T> nodes) {
+    public synchronized void addAll(Collection<? extends T> nodes) {
         for (T node : nodes)
             addInternal(node,defaultReplication);
         refreshTable();
@@ -245,7 +248,7 @@ public class ConsistentHash<T> {
     /**
      * Calls {@link #add(Object,int)} with all the arguments.
      */
-    public void addAll(Map<? extends T,Integer> nodes) {
+    public synchronized void addAll(Map<? extends T,Integer> nodes) {
         for (Map.Entry<? extends T,Integer> node : nodes.entrySet())
             addInternal(node.getKey(),node.getValue());
         refreshTable();
@@ -254,23 +257,20 @@ public class ConsistentHash<T> {
     /**
      * Removes the node entirely. This is the same as {@code add(node,0)}
      */
-    public void remove(T node) {
+    public synchronized void remove(T node) {
         add(node, 0);
     }
 
     /**
      * Adds a new node with the given number of replica.
-     *
-     * <p>
-     * This is the only function that manipulates {@link #items}.
      */
     public synchronized void add(T node, int replica) {
         addInternal(node, replica);
         refreshTable();
     }
 
-    private void addInternal(T node, int replica) {
-        if(replica==0) {
+    private synchronized void addInternal(T node, int replica) {
+        if (replica==0) {
             items.remove(node);
         } else {
             Point[] points = new Point[replica];
@@ -281,11 +281,10 @@ public class ConsistentHash<T> {
         }
     }
 
-    private void refreshTable() {
+    private synchronized void refreshTable() {
         table = new Table();
     }
 
-
     /**
      * Compresses a string into an integer with MD5.
      */
diff --git a/core/src/test/java/hudson/util/ConsistentHashTest.java b/core/src/test/java/hudson/util/ConsistentHashTest.java
index ed0704a49c..0d10fff42a 100644
--- a/core/src/test/java/hudson/util/ConsistentHashTest.java
+++ b/core/src/test/java/hudson/util/ConsistentHashTest.java
@@ -122,11 +122,24 @@ public class ConsistentHashTest {
     @Test
     public void emptyBehavior() {
         ConsistentHash<String> hash = new ConsistentHash<String>();
+        assertEquals(0, hash.countAllPoints());
         assertFalse(hash.list(0).iterator().hasNext());
         assertNull(hash.lookup(0));
         assertNull(hash.lookup(999));
     }
 
+    @Test
+    public void countAllPoints() {
+        ConsistentHash<String> hash = new ConsistentHash<String>();
+        assertEquals(0, hash.countAllPoints());
+        hash.add("foo", 10);
+        assertEquals(10, hash.countAllPoints());
+        hash.add("bar", 5);
+        assertEquals(15, hash.countAllPoints());
+        hash.remove("foo");
+        assertEquals(5, hash.countAllPoints());
+    }
+
     /**
      * This test doesn't fail but it's written to measure the performance of the consistent hash function with large data set.
      */