From 25b7d72fc9e121a0a7e86e5c8bea11cafb41ddff Mon Sep 17 00:00:00 2001
From: Daniele Alessandri <suppakilla@gmail.com>
Date: Sun, 11 Dec 2011 14:03:02 +0100
Subject: [PATCH] Make it possible to remove connections from a cluster.

---
 lib/Predis/Network/IConnectionCluster.php  |  8 +++++
 lib/Predis/Network/PredisCluster.php       | 40 ++++++++++++++++++++-
 tests/Predis/Network/PredisClusterTest.php | 41 ++++++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/lib/Predis/Network/IConnectionCluster.php b/lib/Predis/Network/IConnectionCluster.php
index 655b99a9..ea48621d 100644
--- a/lib/Predis/Network/IConnectionCluster.php
+++ b/lib/Predis/Network/IConnectionCluster.php
@@ -28,6 +28,14 @@ interface IConnectionCluster extends IConnection
      */
     public function add(IConnectionSingle $connection);
 
+    /**
+     * Removes the specified connection instance from the cluster.
+     *
+     * @param IConnectionSingle $connection Instance of a connection.
+     * @return Boolean Returns true if the connection was in the pool.
+     */
+    public function remove(IConnectionSingle $connection);
+
     /**
      * Gets the actual connection instance in charge of the specified command.
      *
diff --git a/lib/Predis/Network/PredisCluster.php b/lib/Predis/Network/PredisCluster.php
index b8c517a4..7d8e8eec 100644
--- a/lib/Predis/Network/PredisCluster.php
+++ b/lib/Predis/Network/PredisCluster.php
@@ -25,7 +25,7 @@ use Predis\Distribution\HashRing;
  * @author Daniele Alessandri <suppakilla@gmail.com>
  * @todo Add the ability to remove connections from pool.
  */
-class PredisCluster implements IConnectionCluster, \IteratorAggregate
+class PredisCluster implements IConnectionCluster, \IteratorAggregate, \Countable
 {
     private $pool;
     private $distributor;
@@ -91,6 +91,36 @@ class PredisCluster implements IConnectionCluster, \IteratorAggregate
         $this->distributor->add($connection, $weight);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function remove(IConnectionSingle $connection)
+    {
+        if (($id = array_search($connection, $this->pool, true)) !== false) {
+            unset($this->pool[$id]);
+            $this->distributor->remove($connection);
+
+            return true;
+        }
+
+        return false;
+    }
+
+    /**
+     * Removes a connection instance using its alias or index.
+     *
+     * @param string $connectionId Alias or index of a connection.
+     * @return Boolean Returns true if the connection was in the pool.
+     */
+    public function removeById($connectionId)
+    {
+        if ($connection = $this->getConnectionById($connectionId)) {
+            return $this->remove($connection);
+        }
+
+        return false;
+    }
+
     /**
      * {@inheritdoc}
      */
@@ -131,6 +161,14 @@ class PredisCluster implements IConnectionCluster, \IteratorAggregate
         return $this->distributor->get($keyHash);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function count()
+    {
+        return count($this->pool);
+    }
+
     /**
      * {@inheritdoc}
      */
diff --git a/tests/Predis/Network/PredisClusterTest.php b/tests/Predis/Network/PredisClusterTest.php
index 2fefa12b..f452112a 100644
--- a/tests/Predis/Network/PredisClusterTest.php
+++ b/tests/Predis/Network/PredisClusterTest.php
@@ -33,6 +33,7 @@ class PredisClusterTest extends StandardTestCase
         $cluster->add($connection1);
         $cluster->add($connection2);
 
+        $this->assertSame(2, count($cluster));
         $this->assertSame($connection1, $cluster->getConnectionById(0));
         $this->assertSame($connection2, $cluster->getConnectionById(1));
     }
@@ -49,10 +50,50 @@ class PredisClusterTest extends StandardTestCase
         $cluster->add($connection1);
         $cluster->add($connection2);
 
+        $this->assertSame(2, count($cluster));
         $this->assertSame($connection1, $cluster->getConnectionById('node1'));
         $this->assertSame($connection2, $cluster->getConnectionById('node2'));
     }
 
+    /**
+     * @group disconnected
+     */
+    public function testRemovingConnectionsFromCluster()
+    {
+        $connection1 = $this->getMockConnection();
+        $connection2 = $this->getMockConnection();
+        $connection3 = $this->getMockConnection();
+
+        $cluster = new PredisCluster();
+        $cluster->add($connection1);
+        $cluster->add($connection2);
+
+        $this->assertTrue($cluster->remove($connection1));
+        $this->assertFalse($cluster->remove($connection3));
+        $this->assertSame(1, count($cluster));
+    }
+
+    /**
+     * @group disconnected
+     */
+    public function testRemovingConnectionsFromClusterByAlias()
+    {
+        $connection1 = $this->getMockConnection();
+        $connection2 = $this->getMockConnection('tcp://host1:7001?alias=node2');
+        $connection3 = $this->getMockConnection('tcp://host1:7002?alias=node3');
+        $connection4 = $this->getMockConnection('tcp://host1:7003?alias=node4');
+
+        $cluster = new PredisCluster();
+        $cluster->add($connection1);
+        $cluster->add($connection2);
+        $cluster->add($connection3);
+
+        $this->assertTrue($cluster->removeById(0));
+        $this->assertTrue($cluster->removeById('node2'));
+        $this->assertFalse($cluster->removeById('node4'));
+        $this->assertSame(1, count($cluster));
+    }
+
     /**
      * @group disconnected
      */