Merge pull request mDialog#28 from damienlevin/master

Listeners for state changes are now cleaned up

build.sbt

@@ -2,7 +2,7 @@ name := "brando"
 
 organization := "com.digital-achiever"
 
-version := "1.0.0"
+version := "1.0.1"
 
 scalaVersion := "2.10.3"
 

src/main/scala/Brando.scala

@@ -1,6 +1,6 @@
 package brando
 
-import akka.actor.{ Actor, ActorContext, ActorRef, Props, Status, Stash }
+import akka.actor.{ Actor, ActorContext, ActorRef, Props, Status, Stash, Terminated }
 import akka.io.{ IO, Tcp }
 import akka.pattern.ask
 import akka.util.{ ByteString, Timeout }
@@ -119,7 +119,7 @@ class Brando(
     port: Int,
     database: Option[Int],
     auth: Option[String],
-    listeners: Set[ActorRef]) extends Actor with Stash {
+    private[brando] var listeners: Set[ActorRef]) extends Actor with Stash {
   import context.dispatcher
 
   val config = context.system.settings.config
@@ -136,21 +136,23 @@ class Brando(
   val connection = context.actorOf(Props(classOf[Connection],
     self, address, connectionRetry, maxConnectionAttempts))
 
-  def receive = disconnected
+  listeners.map(context.watch(_))
+
+  def receive = disconnected orElse cleanListeners
 
   def authenticated: Receive = {
     case request: Request ⇒ connection forward request
     case x: Tcp.ConnectionClosed ⇒
       notifyStateChange(Disconnected)
-      context.become(disconnected)
+      context.become(disconnected orElse cleanListeners)
   }
 
   def disconnected: Receive = {
     case request: Request ⇒ stash()
 
     case x: Tcp.Connected ⇒
 
-      context.become(authenticating)
+      context.become(authenticating orElse cleanListeners)
 
       (for {
         auth ← if (auth.isDefined) connection ? Request(ByteString("AUTH"), ByteString(auth.get)) else Future.successful()
@@ -171,18 +173,23 @@ class Brando(
 
     case x: Tcp.ConnectionClosed ⇒
       notifyStateChange(Disconnected)
-      context.become(disconnected)
+      context.become(disconnected orElse cleanListeners)
 
     case Connected ⇒
       unstashAll()
       notifyStateChange(Connected)
-      context.become(authenticated)
+      context.become(authenticated orElse cleanListeners)
 
     case AuthenticationFailed ⇒
       notifyStateChange(AuthenticationFailed)
 
   }
 
+  def cleanListeners: Receive = {
+    case Terminated(l) ⇒
+      listeners = listeners - l
+  }
+
   private def notifyStateChange(newState: BrandoStateChange) {
     listeners foreach { _ ! newState }
   }

src/main/scala/ShardManager.scala

@@ -1,6 +1,6 @@
 package brando
 
-import akka.actor.{ Actor, ActorRef, Props }
+import akka.actor.{ Actor, ActorRef, Props, Terminated }
 import akka.util.ByteString
 import collection.mutable
 import java.util.zip.CRC32
@@ -34,12 +34,13 @@ object ShardManager {
 class ShardManager(
     shards: Seq[Shard],
     hashFunction: (Array[Byte] ⇒ Long),
-    val listeners: Set[ActorRef] = Set()) extends Actor {
+    private[brando] var listeners: Set[ActorRef] = Set()) extends Actor {
 
   val pool = mutable.Map.empty[String, ActorRef]
   val shardLookup = mutable.Map.empty[ActorRef, Shard]
 
   shards.map(create(_))
+  listeners.map(context.watch(_))
 
   def receive = {
 
@@ -64,6 +65,9 @@ class ShardManager(
         case None ⇒ println("Update received for unknown shard actorRef " + sender + "\r\n")
       }
 
+    case Terminated(l) ⇒
+      listeners = listeners - l
+
     case x ⇒ println("ShardManager received unexpected " + x + "\r\n")
   }
 

src/test/scala/HealthMonitorTest.scala

@@ -36,6 +36,21 @@ class HealthMonitorTest extends TestKit(ActorSystem("HealthMonitorTest"))
       manager ! PoisonPill
     }
 
+    it("should cleaned up any dead listeners") {
+
+      val probe1 = TestProbe()
+      val probe2 = TestProbe()
+
+      val brando = TestActorRef(new Brando("localhost", 6379, None, None, listeners = Set(probe1.ref, probe2.ref))).underlyingActor
+      assertResult(2)(brando.listeners.size)
+
+      probe1.ref ! PoisonPill
+      probe2.expectMsg(Connected)
+
+      assertResult(1)(brando.listeners.size)
+
+    }
+
     it("should restart the shard, and notify, when healthcheck fails") {
       val probe = TestProbe()
       val listener = TestProbe()

src/test/scala/ResponseTest.scala

@@ -7,51 +7,51 @@ class ResponseTest extends FunSpec {
   describe("Utf8String") {
     it("should extract an utf8 String from a ByteString") {
       val ok = Response.AsString.unapply(Some(ByteString("ok")))
-      expectResult(Some("ok"))(ok)
+      assertResult(Some("ok"))(ok)
     }
   }
 
   describe("Strings") {
     it("should extract a list of string from a option bytestring list ") {
       val resp = Some(List(Some(ByteString("l1")), Some(ByteString("l2")), Some(ByteString("l3"))))
       val seq = Response.AsStrings.unapply(resp)
-      expectResult(Some(Seq("l1", "l2", "l3")))(seq)
+      assertResult(Some(Seq("l1", "l2", "l3")))(seq)
     }
 
     it("shouldn't extract if it is another type") {
       val seq = Response.AsStrings.unapply(Some(List(Some(12L))))
-      expectResult(None)(seq)
+      assertResult(None)(seq)
     }
   }
 
   describe("Bytes Sequences") {
     it("should extract a list of string from a option bytestring list ") {
       val resp = Some(List(Some(ByteString(0, 1)), Some(ByteString(2, 3))))
       val seq = Response.AsByteSeqs.unapply(resp)
-      expectResult(Some(Seq(Seq(0, 1), Seq(2, 3))))(seq)
+      assertResult(Some(Seq(Seq(0, 1), Seq(2, 3))))(seq)
     }
 
     it("shouldn't extract if it is another type") {
       val seq = Response.AsByteSeqs.unapply(Some(List(Some(12L))))
-      expectResult(None)(seq)
+      assertResult(None)(seq)
     }
   }
 
   describe("Strings Hashes") {
     it("should extract a map when the result list has an heaven size") {
       val resp = Some(List(Some(ByteString("k1")), Some(ByteString("v1")), Some(ByteString("k2")), Some(ByteString("v2"))))
       val map = Response.AsStringsHash.unapply(resp)
-      expectResult(Some(Map("k1" -> "v1", "k2" -> "v2")))(map)
+      assertResult(Some(Map("k1" -> "v1", "k2" -> "v2")))(map)
     }
 
     it("should extract an empty map when the result list is empty") {
       val map = Response.AsStringsHash.unapply(Some(List.empty))
-      expectResult(Some(Map.empty))(map)
+      assertResult(Some(Map.empty))(map)
     }
 
     it("should fails when the result list has an odd size") {
       val map = Response.AsStringsHash.unapply(Some(List(Some(ByteString("k1")))))
-      expectResult(None)(map)
+      assertResult(None)(map)
     }
   }
 

src/test/scala/ShardManagerTest.scala

@@ -87,6 +87,25 @@ class ShardManagerTest extends TestKit(ActorSystem("ShardManagerTest"))
       probe.expectMsg(ShardStateChange(shards(1), ConnectionFailed))
     }
 
+    it("should cleaned up any dead listeners") {
+
+      val shards = Seq(
+        Shard("server1", "localhost", 6379, Some(0)),
+        Shard("server2", "localhost", 13579, Some(1)))
+
+      val probe1 = TestProbe()
+      val probe2 = TestProbe()
+
+      val shardManager = TestActorRef(new ShardManager(shards, ShardManager.defaultHashFunction, Set(probe1.ref, probe2.ref))).underlyingActor
+      assertResult(2)(shardManager.listeners.size)
+
+      probe1.ref ! PoisonPill
+      probe2.expectMsg(ShardStateChange(shards(0), Connected))
+
+      assertResult(1)(shardManager.listeners.size)
+
+    }
+
     it("should notify listeners when a shard fails to authenticate") {
       val shards = Seq(
         Shard("server1", "localhost", 6379, Some(0)),