Added awaitTermination method + Redis now waits for all actors to be terminated before shutting down ActorSystem

Author: Alexandre Curreli

build.sbt

@@ -6,7 +6,7 @@ organization := "com.livestream"
 
 name := "scredis"
 
-version := "2.0.2"
+version := "2.0.3"
 
 scalaVersion := "2.11.2"
 

src/main/scala/scredis/BlockingClient.scala

@@ -163,6 +163,8 @@ final class BlockingClient(
    */
   def quit()(implicit timeout: Duration): Try[Unit] = sendBlocking(Quit())
 
+  watchTermination()
+  
 }
 
 /**

src/main/scala/scredis/Client.scala

@@ -137,6 +137,8 @@ final class Client(
     RedisConfig(configName, path)
   )
 
+  watchTermination()
+  
 }
 
 /**

src/main/scala/scredis/Redis.scala

@@ -255,12 +255,15 @@ final class Redis(
     if (shouldShutdownBlockingClient) {
       try {
         blocking.quit()(5 seconds)
+        blocking.awaitTermination(3 seconds)
       } catch {
         case e: Throwable => logger.error("Could not shutdown blocking client", e)
       }
     }
     val future = if (shouldShutdownSubscriberClient) {
-      subscriber.quit()
+      subscriber.quit().map { _ =>
+        subscriber.awaitTermination(3 seconds)
+      }
     } else {
       Future.successful(())
     }
@@ -269,6 +272,7 @@ final class Redis(
     }.flatMap { _ =>
       super.quit()
     }.map { _ =>
+      awaitTermination(3 seconds)
       system.shutdown()
     }
   }
@@ -292,6 +296,8 @@ final class Redis(
     super.select(database)
   }
 
+  watchTermination()
+  
 }
 
 /**

src/main/scala/scredis/SubscriberClient.scala

@@ -148,6 +148,8 @@ final class SubscriberClient(
    */
   def quit(): Future[Unit] = shutdown()
 
+  watchTermination()
+  
 }
 
 /**

src/main/scala/scredis/io/AbstractAkkaConnection.scala

@@ -11,6 +11,8 @@ import scredis.protocol.requests.ServerRequests.{ ClientSetName, Shutdown }
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.concurrent.duration._
 
+import java.util.concurrent.{ CountDownLatch, TimeUnit }
+
 abstract class AbstractAkkaConnection(
   protected val system: ActorSystem,
   protected val host: String,
@@ -29,10 +31,14 @@ abstract class AbstractAkkaConnection(
   protected val akkaDecoderDispatcherPath: String
 ) extends Connection with LazyLogging {
 
+  private val shutdownLatch = new CountDownLatch(1)
+  
   @volatile protected var isShuttingDown = false
 
   override implicit val dispatcher = system.dispatcher
 
+  protected val listenerActor: ActorRef
+  
   protected def updateState(request: Request[_]): Unit = request match {
     case Auth(password) => if (password.isEmpty) {
       passwordOpt = None
@@ -53,4 +59,37 @@ abstract class AbstractAkkaConnection(
   protected def getDatabase: Int = database
   protected def getNameOpt: Option[String] = nameOpt
 
+  protected def watchTermination(): Unit = system.actorOf(
+    Props(
+      classOf[WatchActor],
+      listenerActor,
+      shutdownLatch
+    )
+  )
+  
+  /**
+   * Waits for all the internal actors to be shutdown.
+   * 
+   * @note This method is usually called after issuing a QUIT command
+   * 
+   * @param timeout amount of time to wait
+   */
+  def awaitTermination(timeout: Duration = Duration.Inf): Unit = {
+    if (timeout.isFinite) {
+      shutdownLatch.await(timeout.toMillis, TimeUnit.MILLISECONDS)
+    } else {
+      shutdownLatch.await()
+    }
+  }
+  
+}
+
+class WatchActor(actor: ActorRef, shutdownLatch: CountDownLatch) extends Actor {
+  def receive: Receive = {
+    case Terminated(_) => {
+      shutdownLatch.countDown()
+      context.stop(self)
+    }
+  }
+  context.watch(actor)
 }

src/main/scala/scredis/io/IOActor.scala

@@ -74,12 +74,6 @@ class IOActor(
     become(awaitingAbort)
   }
 
-  protected def stop(): Unit = {
-    logger.trace("Stopping Actor...")
-    //listenerActor ! PoisonPill
-    context.stop(self)
-  }
-  
   protected def encode(request: Request[_]): Int = {
     request.encode()
     request.encoded match {

src/main/scala/scredis/io/ListenerActor.scala

@@ -54,6 +54,7 @@ class ListenerActor(
   private var initializationRequestsCount = 0
   private var isConnecting = false
   private var isShuttingDown = false
+  private var isShuttingDownBeforeConnected = false
   private var isReceiveTimeout = false
   private var timeoutCancellableOpt: Option[Cancellable] = None
 
@@ -284,7 +285,7 @@ class ListenerActor(
     case request: Quit => {
       request.success(())
       failAllQueuedRequests(RedisIOException("Connection has been shutdown by QUIT command"))
-      shutdown()
+      isShuttingDownBeforeConnected = true
     }
     case request: Request[_] => {
       queuedRequests.addLast(request)
@@ -305,60 +306,64 @@ class ListenerActor(
     case Connected => {
       isConnecting = false
 
-      onConnect()
-      
-      val authRequestOpt = passwordOpt.map { password =>
-        Auth(password)
-      }
-      val selectRequestOpt = if (database > 0) {
-        Some(Select(database))
+      if (isShuttingDownBeforeConnected) {
+        shutdown()
       } else {
-        None
-      }
-      val setNameRequestOpt = nameOpt.map { name =>
-        ServerRequests.ClientSetName(name)
-      }
-      
-      val authFuture = authRequestOpt match {
-        case Some(request) => request.future
-        case None => Future.successful(())
-      }
-      val selectFuture = selectRequestOpt match {
-        case Some(request) => request.future
-        case None => Future.successful(())
-      }
-      val setNameFuture = setNameRequestOpt match {
-        case Some(request) => request.future
-        case None => Future.successful(())
-      }
-      
-      val requests = List[Option[Request[Unit]]](
-        authRequestOpt, selectRequestOpt, setNameRequestOpt
-      ).flatten
-      
-      initializationRequestsCount = requests.size
-      
-      if (initializationRequestsCount > 0) {
-        send(requests: _*)
-        become(initializing)
-      } else {
-        onInitialized()
-        sendAllQueuedRequests()
-        if (isShuttingDown) {
-          become(shuttingDown)
+        onConnect()
+        
+        val authRequestOpt = passwordOpt.map { password =>
+          Auth(password)
+        }
+        val selectRequestOpt = if (database > 0) {
+          Some(Select(database))
         } else {
-          become(initialized)
+          None
+        }
+        val setNameRequestOpt = nameOpt.map { name =>
+          ServerRequests.ClientSetName(name)
+        }
+        
+        val authFuture = authRequestOpt match {
+          case Some(request) => request.future
+          case None => Future.successful(())
+        }
+        val selectFuture = selectRequestOpt match {
+          case Some(request) => request.future
+          case None => Future.successful(())
+        }
+        val setNameFuture = setNameRequestOpt match {
+          case Some(request) => request.future
+          case None => Future.successful(())
+        }
+        
+        val requests = List[Option[Request[Unit]]](
+          authRequestOpt, selectRequestOpt, setNameRequestOpt
+        ).flatten
+        
+        initializationRequestsCount = requests.size
+        
+        if (initializationRequestsCount > 0) {
+          send(requests: _*)
+          become(initializing)
+        } else {
+          onInitialized()
+          sendAllQueuedRequests()
+          if (isShuttingDown) {
+            become(shuttingDown)
+          } else {
+            become(initialized)
+          }
+        }
+        
+        authFuture.recover {
+          case e: Throwable => logger.error(s"Could not authenticate to $remote", e)
+        }
+        selectFuture.recover {
+          case e: Throwable => logger.error(s"Could not select database '$database' in $remote", e)
+        }
+        setNameFuture.recover {
+          case e: Throwable => logger.error(s"Could not set client name in $remote", e)
         }
-      }
-      
-      authFuture.recover {
-        case e: Throwable => logger.error(s"Could not authenticate to $remote", e)
-      }
-      selectFuture.recover {
-        case e: Throwable => logger.error(s"Could not select database '$database' in $remote", e)
-      }
-      setNameFuture.recover {
-        case e: Throwable => logger.error(s"Could not set client name in $remote", e)
       }
     }
     case ReceiveTimeout =>