Support Libmemcached ketama weighted

- Add support for alternative Ketama Node key format

  Libmemcached uses the format for `[hostname or ip][port unless
  default]-[repetition]` while spymemcached has been using
  `[hostname]/[ip]:[port]-[repetition]` the added `KetamaNodeKeyFormat`
  allows to choose either format while retaining the caching
  optimization.

- Add support for weighted ketama
  Straight port of the weighting based on Libmemcached, configured via
  passing a map of node socketaddress to weight to the configuration.
  This code is only active if the weight is actually configured
  otherwise the old ketama code is used.

- Split testLibKetamaCompatTwo into 2 pieces as the length breaks
  compiling on Eclipse and Java 8 "In java a methods can't have more than
  65535 bytes." http://stackoverflow.com/questions/12257398/how-to-fix-the-code-of-method-is-exceeding-the-65535-bytes-limit

Change-Id: I0263b9afc513f9a135d5d17318b3fe6bd4593437
Reviewed-on: http://review.couchbase.org/47624
Reviewed-by: Michael Nitschinger <[email protected]>
Tested-by: Michael Nitschinger <[email protected]>

Author: sideshowcoder

src/main/java/net/spy/memcached/KetamaConnectionFactory.java

@@ -33,7 +33,10 @@
 
 package net.spy.memcached;
 
+import java.net.InetSocketAddress;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * ConnectionFactory instance that sets up a ketama compatible connection.
@@ -51,6 +54,9 @@
  */
 public class KetamaConnectionFactory extends DefaultConnectionFactory {
 
+    private final KetamaNodeKeyFormatter.Format ketamaNodeKeyFormat;
+    private Map<InetSocketAddress, Integer> weights;
+
   /**
    * Create a KetamaConnectionFactory with the given maximum operation
    * queue length, and the given read buffer size.
@@ -62,6 +68,32 @@ public KetamaConnectionFactory(ClientMode clientMode, int qLen, int bufSize,
       long opQueueMaxBlockTime) {
     super(clientMode, qLen, bufSize, DefaultHashAlgorithm.KETAMA_HASH);
   }
+    
+  public KetamaConnectionFactory(int qLen, int bufSize,
+      long opQueueMaxBlockTime) {
+      this(qLen, bufSize, opQueueMaxBlockTime,
+              DefaultHashAlgorithm.KETAMA_HASH,
+              KetamaNodeKeyFormatter.Format.SPYMEMCACHED,
+              new HashMap<InetSocketAddress, Integer>());
+  }
+
+  /** Create a KetamaConnectionFactory with the maximum operation queue length,
+   * the given read buffer size, the maximum time to block waiting operations,
+   * a specific hash algorithm, a set ring key format, and a given set of
+   * weights.
+   *
+   */
+  public KetamaConnectionFactory(
+          int qLen,
+          int bufSize,
+          long opQueueMaxBlockTime,
+          HashAlgorithm hash,
+          KetamaNodeKeyFormatter.Format nodeKeyFormat,
+          Map<InetSocketAddress, Integer> weights) {
+      super(qLen, bufSize, hash);
+      this.ketamaNodeKeyFormat = nodeKeyFormat;
+      this.weights = weights;
+  }
 
   /**
    * Create a KetamaConnectionFactory with the default parameters.
@@ -78,6 +110,21 @@ public KetamaConnectionFactory() {
    */
   @Override
   public NodeLocator createLocator(List<MemcachedNode> nodes) {
-    return new KetamaNodeLocator(nodes, getHashAlg());
+        return new KetamaNodeLocator(nodes, getHashAlg(),
+                getKetamaNodeKeyFormat(), getWeights());
+    }
+
+  /**
+   * @return the ketamaNodeKeyFormat
+   */
+  public KetamaNodeKeyFormatter.Format getKetamaNodeKeyFormat() {
+      return ketamaNodeKeyFormat;
+  }
+
+  /**
+   * @return the ketama node weights
+   */
+  public Map<InetSocketAddress, Integer> getWeights() {
+      return weights;
   }
 }

src/main/java/net/spy/memcached/KetamaNodeKeyFormatter.java

@@ -0,0 +1,138 @@
+/**
+ * Copyright (C) 2009-2015 Couchbase, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALING
+ * IN THE SOFTWARE.
+ */
+
+package net.spy.memcached;
+
+import java.net.InetSocketAddress;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Known key formats used in Ketama for assigning nodes around the ring
+ */
+
+public class KetamaNodeKeyFormatter {
+
+    public enum Format {
+        /**
+         * SPYMEMCACHED uses the format traditionally used by spymemcached to map
+         * nodes to names. The format is HOSTNAME/IP:PORT-ITERATION
+         *
+         * <p>
+         * This default implementation uses the socket-address of the MemcachedNode
+         * and concatenates it with a hyphen directly against the repetition number
+         * for example a key for a particular server's first repetition may look like:
+         * <p>
+         *
+         * <p>
+         * <code>myhost/10.0.2.1-0</code>
+         * </p>
+         *
+         * <p>
+         * for the second repetition
+         * </p>
+         *
+         * <p>
+         * <code>myhost/10.0.2.1-1</code>
+         * </p>
+         *
+         * <p>
+         * for a server where reverse lookups are failing the returned keys may look
+         * like
+         * </p>
+         *
+         * <p>
+         * <code>/10.0.2.1-0</code> and <code>/10.0.2.1-1</code>
+         * </p>
+         */
+        SPYMEMCACHED,
+
+        /**
+         * LIBMEMCACHED uses the format traditionally used by libmemcached to map
+         * nodes to names. The format is HOSTNAME:[PORT]-ITERATION the PORT is not
+         * part of the node identifier if it is the default memecached port (11211)
+         */
+        LIBMEMCACHED
+    }
+
+    private final Format format;
+
+    public Format getFormat() {
+        return format;
+    }
+
+    // Carrried over from the DefaultKetamaNodeLocatorConfiguration:
+    // Internal lookup map to try to carry forward the optimisation that was
+    // previously in KetamaNodeLocator
+    private Map<MemcachedNode, String> nodeKeys = new HashMap<MemcachedNode, String>();
+
+    public KetamaNodeKeyFormatter() {
+        this(Format.SPYMEMCACHED);
+    }
+
+    public KetamaNodeKeyFormatter(Format format) {
+        this.format = format;
+    }
+
+    /**
+     * Returns a uniquely identifying key, suitable for hashing by the
+     * KetamaNodeLocator algorithm.
+     *
+     * @param node The MemcachedNode to use to form the unique identifier
+     * @param repetition The repetition number for the particular node in question
+     *          (0 is the first repetition)
+     * @return The key that represents the specific repetition of the node
+     */
+    public String getKeyForNode(MemcachedNode node, int repetition) {
+        // Carrried over from the DefaultKetamaNodeLocatorConfiguration:
+        // Internal Using the internal map retrieve the socket addresses
+        // for given nodes.
+        // I'm aware that this code is inherently thread-unsafe as
+        // I'm using a HashMap implementation of the map, but the worst
+        // case ( I believe) is we're slightly in-efficient when
+        // a node has never been seen before concurrently on two different
+        // threads, so it the socketaddress will be requested multiple times!
+        // all other cases should be as fast as possible.
+        String nodeKey = nodeKeys.get(node);
+        if (nodeKey == null) {
+            switch(this.format) {
+                case LIBMEMCACHED:
+                    InetSocketAddress address = (InetSocketAddress)node.getSocketAddress();
+                    nodeKey = address.getHostName();
+                    if (address.getPort() != 11211) {
+                        nodeKey += ":" + address.getPort();
+                    }
+                    break;
+                case SPYMEMCACHED:
+                    nodeKey = String.valueOf(node.getSocketAddress());
+                    if (nodeKey.startsWith("/")) {
+                        nodeKey = nodeKey.substring(1);
+                    }
+                    break;
+                default:
+                    assert false;
+            }
+            nodeKeys.put(node, nodeKey);
+        }
+        return nodeKey + "-" + repetition;
+    }
+}

src/main/java/net/spy/memcached/KetamaNodeLocator.java

@@ -27,8 +27,10 @@
 import net.spy.memcached.util.DefaultKetamaNodeLocatorConfiguration;
 import net.spy.memcached.util.KetamaNodeLocatorConfiguration;
 
+import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -51,6 +53,8 @@ public final class KetamaNodeLocator extends SpyObject implements NodeLocator {
   private volatile Collection<MemcachedNode> allNodes;
 
   private final HashAlgorithm hashAlg;
+  private final Map<InetSocketAddress, Integer> weights;
+  private final boolean isWeightedKetama;
   private final KetamaNodeLocatorConfiguration config;
 
   /**
@@ -63,7 +67,26 @@ public final class KetamaNodeLocator extends SpyObject implements NodeLocator {
    *          consistent hash continuum
    */
   public KetamaNodeLocator(List<MemcachedNode> nodes, HashAlgorithm alg) {
-    this(nodes, alg, new DefaultKetamaNodeLocatorConfiguration());
+    this(nodes, alg, KetamaNodeKeyFormatter.Format.SPYMEMCACHED, new HashMap<InetSocketAddress, Integer>());
+  }
+
+  /**
+   * Create a new KetamaNodeLocator with specific nodes, hash, node key format,
+   * and weight
+   *
+   * @param nodes The List of nodes to use in the Ketama consistent hash
+   *          continuum
+   * @param alg The hash algorithm to use when choosing a node in the Ketama
+   *          consistent hash continuum
+   * @param nodeKeyFormat the format used to name the nodes in Ketama, either
+   *          SPYMEMCACHED or LIBMEMCACHED
+   * @param weights node weights for ketama, a map from InetSocketAddress to
+   *          weight as Integer
+   */
+    public KetamaNodeLocator(List<MemcachedNode> nodes, HashAlgorithm alg,
+            KetamaNodeKeyFormatter.Format nodeKeyFormat,
+            Map<InetSocketAddress, Integer> weights) {
+    this(nodes, alg, weights, new DefaultKetamaNodeLocatorConfiguration(new KetamaNodeKeyFormatter(nodeKeyFormat)));
   }
 
   /**
@@ -78,21 +101,44 @@ public KetamaNodeLocator(List<MemcachedNode> nodes, HashAlgorithm alg) {
    */
   public KetamaNodeLocator(List<MemcachedNode> nodes, HashAlgorithm alg,
       KetamaNodeLocatorConfiguration conf) {
+    this(nodes, alg, new HashMap<InetSocketAddress, Integer>(), conf);
+  }
+
+  /**
+   * Create a new KetamaNodeLocator with specific nodes, hash, node key format,
+   * and weight
+   *
+   * @param nodes The List of nodes to use in the Ketama consistent hash
+   *          continuum
+   * @param alg The hash algorithm to use when choosing a node in the Ketama
+   *          consistent hash continuum
+   * @param weights node weights for ketama, a map from InetSocketAddress to
+   *          weight as Integer
+   * @param configuration node locator configuration
+   */
+    public KetamaNodeLocator(List<MemcachedNode> nodes, HashAlgorithm alg,
+            Map<InetSocketAddress, Integer> nodeWeights,
+            KetamaNodeLocatorConfiguration configuration) {
     super();
     allNodes = nodes;
     hashAlg = alg;
-    config = conf;
+    config = configuration;
+    weights = nodeWeights;
+    isWeightedKetama = !weights.isEmpty();
     setKetamaNodes(nodes);
   }
 
   private KetamaNodeLocator(TreeMap<Long, MemcachedNode> smn,
       Collection<MemcachedNode> an, HashAlgorithm alg,
+      Map<InetSocketAddress, Integer> nodeWeights,
       KetamaNodeLocatorConfiguration conf) {
     super();
     ketamaNodes = smn;
     allNodes = an;
     hashAlg = alg;
     config = conf;
+    weights = nodeWeights;
+    isWeightedKetama = !weights.isEmpty();
   }
 
   public Collection<MemcachedNode> getAll() {
@@ -147,7 +193,7 @@ public NodeLocator getReadonlyCopy() {
       an.add(new MemcachedNodeROImpl(n));
     }
 
-    return new KetamaNodeLocator(smn, an, hashAlg, config);
+    return new KetamaNodeLocator(smn, an, hashAlg, weights, config);
   }
 
   @Override
@@ -172,30 +218,62 @@ protected TreeMap<Long, MemcachedNode> getKetamaNodes() {
    */
   protected void setKetamaNodes(List<MemcachedNode> nodes) {
     TreeMap<Long, MemcachedNode> newNodeMap =
-        new TreeMap<Long, MemcachedNode>();
+            new TreeMap<Long, MemcachedNode>();
     int numReps = config.getNodeRepetitions();
+    int nodeCount = nodes.size();
+    int totalWeight = 0;
+
+    if (isWeightedKetama) {
+        for (MemcachedNode node : nodes) {
+            totalWeight += weights.get(node.getSocketAddress());
+        }
+    }
+
     for (MemcachedNode node : nodes) {
-      // Ketama does some special work with md5 where it reuses chunks.
-      if (hashAlg == DefaultHashAlgorithm.KETAMA_HASH) {
-        for (int i = 0; i < numReps / 4; i++) {
-          byte[] digest =
-              DefaultHashAlgorithm.computeMd5(config.getKeyForNode(node, i));
-          for (int h = 0; h < 4; h++) {
-            Long k = ((long) (digest[3 + h * 4] & 0xFF) << 24)
-                    | ((long) (digest[2 + h * 4] & 0xFF) << 16)
-                    | ((long) (digest[1 + h * 4] & 0xFF) << 8)
-                    | (digest[h * 4] & 0xFF);
-            newNodeMap.put(k, node);
-            getLogger().debug("Adding node %s in position %d", node, k);
+      if (isWeightedKetama) {
+
+          int thisWeight = weights.get(node.getSocketAddress());
+          float percent = (float)thisWeight / (float)totalWeight;
+          int pointerPerServer = (int)((Math.floor((float)(percent * (float)config.getNodeRepetitions() / 4 * (float)nodeCount + 0.0000000001))) * 4);
+          for (int i = 0; i < pointerPerServer / 4; i++) {
+              for(long position : ketamaNodePositionsAtIteration(node, i)) {
+                  newNodeMap.put(position, node);
+                  getLogger().debug("Adding node %s with weight %s in position %d", node, thisWeight, position);
+              }
           }
-        }
       } else {
-        for (int i = 0; i < numReps; i++) {
-          newNodeMap.put(hashAlg.hash(config.getKeyForNode(node, i)), node);
-        }
+          // Ketama does some special work with md5 where it reuses chunks.
+          // Check to be backwards compatible, the hash algorithm does not
+          // matter for Ketama, just the placement should always be done using
+          // MD5
+          if (hashAlg == DefaultHashAlgorithm.KETAMA_HASH) {
+              for (int i = 0; i < numReps / 4; i++) {
+                  for(long position : ketamaNodePositionsAtIteration(node, i)) {
+                    newNodeMap.put(position, node);
+                    getLogger().debug("Adding node %s in position %d", node, position);
+                  }
+              }
+          } else {
+              for (int i = 0; i < numReps; i++) {
+                  newNodeMap.put(hashAlg.hash(config.getKeyForNode(node, i)), node);
+              }
+          }
       }
     }
     assert newNodeMap.size() == numReps * nodes.size();
     ketamaNodes = newNodeMap;
   }
+
+  private List<Long> ketamaNodePositionsAtIteration(MemcachedNode node, int iteration) {
+      List<Long> positions = new ArrayList<Long>();
+      byte[] digest = DefaultHashAlgorithm.computeMd5(config.getKeyForNode(node, iteration));
+      for (int h = 0; h < 4; h++) {
+          Long k = ((long) (digest[3 + h * 4] & 0xFF) << 24)
+              | ((long) (digest[2 + h * 4] & 0xFF) << 16)
+              | ((long) (digest[1 + h * 4] & 0xFF) << 8)
+              | (digest[h * 4] & 0xFF);
+          positions.add(k);
+      }
+      return positions;
+  }
 }