Gracefully leave the cluster and terminate the container instance when no running user actors left

Author: izeigerman

akkeeper/src/main/scala/akkeeper/container/service/ContainerInstanceService.scala

@@ -48,15 +48,16 @@ class ContainerInstanceService(userActors: Seq[ActorLaunchContext],
     instanceStorage.stop()
   }
 
-  private def launchUserActors: Unit = {
+  private def launchUserActors(): Unit = {
     userActors.foreach(actor => {
       log.debug(s"Deploying actor ${actor.name} (${actor.fqn})")
       val clazz = Class.forName(actor.fqn)
-      context.actorOf(Props(clazz), actor.name)
+      val userActorRef = context.actorOf(Props(clazz), actor.name)
+      context.watch(userActorRef)
     })
   }
 
-  private def notifyMonitoringService: Unit = {
+  private def notifyMonitoringService(): Unit = {
     try {
       thisInstance.foreach(info => {
         val monitoringService = MonitoringService.createRemote(context.system)
@@ -69,7 +70,7 @@ class ContainerInstanceService(userActors: Seq[ActorLaunchContext],
     }
   }
 
-  private def registerThisInstance: Unit = {
+  private def registerThisInstance(): Unit = {
     if (!thisInstance.isDefined) {
       val actors = context.children.map(r => r.path.toStringWithoutAddress)
       val info = InstanceInfo(
@@ -91,11 +92,16 @@ class ContainerInstanceService(userActors: Seq[ActorLaunchContext],
     })
   }
 
+  private def terminateThisInstance(): Unit = {
+    cluster.leave(cluster.selfAddress)
+    cluster.registerOnMemberRemoved(context.system.terminate())
+  }
+
   private def initializedReceive: Receive = {
     case _: InstanceId =>
       // The record was successfully saved to a storage.
       log.debug("Successfully registered this instance")
-      notifyMonitoringService
+      notifyMonitoringService()
     case OperationFailed(_, e) =>
       // Failed to save the record to a storage.
       log.error(e, "Failed to store this instance information. " +
@@ -105,24 +111,25 @@ class ContainerInstanceService(userActors: Seq[ActorLaunchContext],
         self, RetryRegistration)
     case RetryRegistration =>
       log.info("Retrying instance registration process")
-      registerThisInstance
+      registerThisInstance()
     case StopInstance =>
       log.info("Termination command received. Stopping this instance")
-      cluster.leave(cluster.selfAddress)
-      context.system.terminate()
+      terminateThisInstance()
+    case Terminated(_) =>
+      if (context.children.isEmpty) {
+        log.info("No running user actors left. Terminating this instance")
+        terminateThisInstance()
+      }
     case JoinClusterTimeout =>
       // Safely ignore the timeout command.
   }
 
   private def joiningTheClusterReceive: Receive = {
-    case MemberUp(member) =>
-      if (member.address == cluster.selfAddress) {
-        log.debug("Successfully joined the cluster")
-        launchUserActors
-        cluster.unsubscribe(self)
-        context.become(initializedReceive)
-        registerThisInstance
-      }
+    case InstanceJoinedCluster =>
+      log.debug("Successfully joined the cluster")
+      launchUserActors()
+      context.become(initializedReceive)
+      registerThisInstance()
     case JoinClusterTimeout =>
       log.error(s"Couldn't join the cluster during ${joinClusterTimeout.toSeconds} seconds. " +
         "Terminating this instance...")
@@ -134,7 +141,7 @@ class ContainerInstanceService(userActors: Seq[ActorLaunchContext],
       context.become(joiningTheClusterReceive)
       log.debug(s"Joining the cluster (master: $masterAddress)")
       cluster.join(masterAddress)
-      cluster.subscribe(self, initialStateMode = InitialStateAsEvents, classOf[MemberUp])
+      cluster.registerOnMemberUp(self ! InstanceJoinedCluster)
       // Scheduling a timeout command.
       context.system.scheduler.scheduleOnce(joinClusterTimeout, self, JoinClusterTimeout)
   }
@@ -146,6 +153,7 @@ object ContainerInstanceService {
   private case object JoinCluster
   private case object RetryRegistration
   private case object JoinClusterTimeout
+  private case object InstanceJoinedCluster
   private[akkeeper] val DefaultRegistrationRetryInterval = 30 seconds
   private[akkeeper] val DefaultJoinClusterTimeout = 120 seconds
 

akkeeper/src/main/scala/akkeeper/master/service/MonitoringService.scala

@@ -60,7 +60,7 @@ private[akkeeper] class MonitoringService(instanceStorage: InstanceStorage.Async
   }
 
   private def memberAutoDown(instanceId: InstanceId, instanceAddr: UniqueAddress): Unit = {
-    val autoDownService = MemberAutoDownService.createLocal(context.system,
+    val autoDownService = MemberAutoDownService.createLocal(context,
       instanceAddr, instanceId, instanceStorage)
     autoDownService ! MemberAutoDownService.PollInstanceStatus
   }

akkeeper/src/test/scala/akkeeper/container/service/ContainerInstanceServiceSpec.scala

@@ -17,7 +17,7 @@ package akkeeper.container.service
 
 import akka.actor._
 import akka.cluster.{Cluster, UniqueAddress}
-import akka.testkit.{ImplicitSender, TestKit}
+import akka.testkit.{ImplicitSender, TestKit, TestProbe}
 import akkeeper.{ActorTestUtils, AkkeeperException}
 import akkeeper.common._
 import akkeeper.master.service._
@@ -27,7 +27,7 @@ import com.typesafe.config.ConfigFactory
 import org.scalamock.scalatest.MockFactory
 import org.scalatest.{BeforeAndAfterAll, FlatSpecLike, Matchers}
 
-import scala.concurrent.Future
+import scala.concurrent.{Future, Promise}
 import scala.concurrent.duration._
 import ContainerInstanceService._
 import ContainerInstanceServiceSpec._
@@ -50,14 +50,16 @@ class ContainerInstanceServiceSpec(system: ActorSystem) extends TestKit(system)
   }
 
   private def createExpectedInstanceInfo(instanceId: InstanceId,
-                                         addr: UniqueAddress): InstanceInfo = {
+                                         addr: UniqueAddress,
+                                         actorPath: String = "/system/testActor-1/akkeeperInstance/testActor"
+                                        ): InstanceInfo = {
     InstanceInfo(
       instanceId = instanceId,
       status = InstanceUp,
       containerName = instanceId.containerName,
       roles = Set("akkeeperMaster", "dc-default"),
       address = Some(addr),
-      actors = Set("/system/testActor-1/akkeeperInstance/testActor")
+      actors = Set(actorPath)
     )
   }
 
@@ -98,7 +100,6 @@ class ContainerInstanceServiceSpec(system: ActorSystem) extends TestKit(system)
     expectMsg(TestPong)
 
     gracefulActorStop(service)
-    gracefulActorStop(masterServiceMock)
   }
 
   it should "retry if the registration failed" in {
@@ -124,7 +125,6 @@ class ContainerInstanceServiceSpec(system: ActorSystem) extends TestKit(system)
     val maxWaitForNoMsg = numberOfAttempts.seconds
     expectNoMessage(maxWaitForNoMsg)
 
-    gracefulActorStop(masterServiceMock)
     gracefulActorStop(service)
   }
 
@@ -145,6 +145,48 @@ class ContainerInstanceServiceSpec(system: ActorSystem) extends TestKit(system)
 
     await(newSystem.whenTerminated)
   }
+
+  it should "terminate this instance if all user actors have been terminated" in {
+    val newSystem = ActorSystem("ContainerInstanceServiceSpecTemp",
+      ConfigFactory.load().withMasterPort.withMasterRole)
+    val newCluster = Cluster(newSystem)
+    val instanceId = InstanceId("container")
+    val expectedInstanceInfo = createExpectedInstanceInfo(instanceId,
+      newCluster.selfUniqueAddress, actorPath = "/user/akkeeperInstance/testActor")
+
+    val storage = mock[InstanceStorage.Async]
+    (storage.start _).expects()
+    (storage.stop _).expects()
+    (storage.registerInstance _)
+      .expects(expectedInstanceInfo)
+      .returns(Future successful instanceId)
+
+    val actors = Seq(ActorLaunchContext("testActor", classOf[TestUserActor].getName))
+
+    ContainerInstanceService.createLocal(newSystem, actors, storage, instanceId, newCluster.selfAddress)
+
+    // Verify that the user actor was actually launched.
+    val testProbe = TestProbe()(newSystem)
+    val userActor = newSystem.actorSelection("/user/akkeeperInstance/testActor")
+
+    val joinPromise = Promise[Unit]
+    newCluster.registerOnMemberUp(joinPromise.success(()))
+    await(joinPromise.future)
+    val waitTimeoutMs = 1000
+    Thread.sleep(waitTimeoutMs)
+
+    userActor.tell(TestPing, testProbe.ref)
+    testProbe.expectMsg(TestPong)
+
+    // Terminate the user actor.
+    userActor ! TestTerminate
+
+    val leavePromise = Promise[Unit]
+    newCluster.registerOnMemberRemoved(leavePromise.success(()))
+    await(leavePromise.future)
+
+    await(newSystem.whenTerminated)
+  }
 }
 
 object ContainerInstanceServiceSpec {

akkeeper/src/test/scala/akkeeper/container/service/TestUserActor.scala

@@ -21,10 +21,12 @@ import TestUserActor._
 class TestUserActor extends Actor {
   override def receive: Receive = {
     case TestPing => sender() ! TestPong
+    case TestTerminate => context.stop(self)
   }
 }
 
 object TestUserActor {
   case object TestPing
   case object TestPong
+  case object TestTerminate
 }