From 09de7be7cb18c2639f53505366c6e749baa6f906 Mon Sep 17 00:00:00 2001
From: Daniele Alessandri <suppakilla@gmail.com>
Date: Sun, 13 Jan 2013 12:43:30 +0100
Subject: [PATCH] Add optional callable to drive extraction of node hash in
 distributor.

This is mainly in response to the longstanding issue #36 in which my
proposed solution was fine in terms of functionalities, but eventually
never made into the repository since it was far from being clean enough
for my taste.

Now developers can optionally pass a callable object when creating the
hashring instance to decide how the distributor should extract the hash
from a node (really a connection instance) to populate the ring:

  use Predis\Cluster\Distribution\HashRing;
  use Predis\Connection\PredisCluster;

  $servers = array(
    'tcp://10.0.0.1?alias=node01',
    'tcp://10.0.0.2?alias=node02',
  );

  $options = array(
    'nodehash' => function ($connection) {
      return $connection->getParameters()->alias;
    },
    'cluster' => function ($options) {
      $replicas = HashRing::DEFAULT_REPLICAS;
      $hashring = new HashRing($replicas, $options->nodehash);
      $cluster  = new PredisCluster($hashring);

      return $cluster;
    },
  );

  $client = new Predis\Client($servers, $options);

Both HashRing and KetamaPureRing in the Predis\Cluster\Distribution
namespace support this new approach.
---
 lib/Predis/Cluster/Distribution/HashRing.php  | 14 +++++++++----
 .../Cluster/Distribution/KetamaPureRing.php   |  4 ++--
 .../Cluster/Distribution/HashRingTest.php     | 21 +++++++++++++++++++
 .../Distribution/KetamaPureRingTest.php       | 20 ++++++++++++++++++
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/lib/Predis/Cluster/Distribution/HashRing.php b/lib/Predis/Cluster/Distribution/HashRing.php
index 7764d33d..182ca61f 100644
--- a/lib/Predis/Cluster/Distribution/HashRing.php
+++ b/lib/Predis/Cluster/Distribution/HashRing.php
@@ -26,19 +26,21 @@ class HashRing implements DistributionStrategyInterface, HashGeneratorInterface
     const DEFAULT_REPLICAS = 128;
     const DEFAULT_WEIGHT   = 100;
 
-    private $nodes;
     private $ring;
     private $ringKeys;
     private $ringKeysCount;
     private $replicas;
+    private $nodeHashCallback;
+    private $nodes = array();
 
     /**
      * @param int $replicas Number of replicas in the ring.
+     * @param mixed $nodeHashCallback Callback returning the string used to calculate the hash of a node.
      */
-    public function __construct($replicas = self::DEFAULT_REPLICAS)
+    public function __construct($replicas = self::DEFAULT_REPLICAS, $nodeHashCallback = null)
     {
         $this->replicas = $replicas;
-        $this->nodes = array();
+        $this->nodeHashCallback = $nodeHashCallback;
     }
 
     /**
@@ -164,7 +166,11 @@ class HashRing implements DistributionStrategyInterface, HashGeneratorInterface
      */
     protected function getNodeHash($nodeObject)
     {
-        return (string) $nodeObject;
+        if ($this->nodeHashCallback === null) {
+            return (string) $nodeObject;
+        }
+
+        return call_user_func($this->nodeHashCallback, $nodeObject);
     }
 
     /**
diff --git a/lib/Predis/Cluster/Distribution/KetamaPureRing.php b/lib/Predis/Cluster/Distribution/KetamaPureRing.php
index 9e3c5708..9f17ea0c 100644
--- a/lib/Predis/Cluster/Distribution/KetamaPureRing.php
+++ b/lib/Predis/Cluster/Distribution/KetamaPureRing.php
@@ -26,9 +26,9 @@ class KetamaPureRing extends HashRing
     /**
      *
      */
-    public function __construct()
+    public function __construct($nodeHashCallback = null)
     {
-        parent::__construct($this::DEFAULT_REPLICAS);
+        parent::__construct($this::DEFAULT_REPLICAS, $nodeHashCallback);
     }
 
     /**
diff --git a/tests/Predis/Cluster/Distribution/HashRingTest.php b/tests/Predis/Cluster/Distribution/HashRingTest.php
index ca94bcc7..e618ea00 100644
--- a/tests/Predis/Cluster/Distribution/HashRingTest.php
+++ b/tests/Predis/Cluster/Distribution/HashRingTest.php
@@ -129,4 +129,25 @@ class HashRingTest extends DistributionStrategyTestCase
         $this->assertSame($expected2, $actual2);
         $this->assertSame($expected3, $actual3);
     }
+
+    /**
+     * @todo This tests should be moved in Predis\Cluster\Distribution\DistributionStrategyTestCase
+     * @group disconnected
+     */
+    public function testCallbackToGetNodeHash()
+    {
+        $node = '127.0.0.1:7000';
+        $replicas = HashRing::DEFAULT_REPLICAS;
+        $callable = $this->getMock('stdClass', array('__invoke'));
+
+        $callable->expects($this->once())
+                 ->method('__invoke')
+                 ->with($node)
+                 ->will($this->returnValue($node));
+
+        $ring = new HashRing($replicas, $callable);
+        $ring->add($node);
+
+        $this->getNodes($ring);
+    }
 }
diff --git a/tests/Predis/Cluster/Distribution/KetamaPureRingTest.php b/tests/Predis/Cluster/Distribution/KetamaPureRingTest.php
index 47a0cca5..b5c3145e 100644
--- a/tests/Predis/Cluster/Distribution/KetamaPureRingTest.php
+++ b/tests/Predis/Cluster/Distribution/KetamaPureRingTest.php
@@ -130,4 +130,24 @@ class KetamaPureRingTest extends DistributionStrategyTestCase
         $this->assertSame($expected2, $actual2);
         $this->assertSame($expected3, $actual3);
     }
+
+    /**
+     * @todo This tests should be moved in Predis\Cluster\Distribution\DistributionStrategyTestCase
+     * @group disconnected
+     */
+    public function testCallbackToGetNodeHash()
+    {
+        $node = '127.0.0.1:7000';
+        $callable = $this->getMock('stdClass', array('__invoke'));
+
+        $callable->expects($this->once())
+                 ->method('__invoke')
+                 ->with($node)
+                 ->will($this->returnValue($node));
+
+        $ring = new KetamaPureRing($callable);
+        $ring->add($node);
+
+        $this->getNodes($ring);
+    }
 }