Include workers in state data to thread safe

common/scala/src/main/scala/org/apache/openwhisk/core/service/LeaseKeepAliveService.scala

@@ -40,6 +40,7 @@ case object Active extends KeepAliveServiceState
 sealed trait KeepAliveServiceData
 case object NoData extends KeepAliveServiceData
 case class Lease(id: Long, ttl: Long) extends KeepAliveServiceData
+case class ActiveStates(worker: Cancellable, lease: Lease) extends KeepAliveServiceData
 
 // Events received by the actor
 case object RegrantLease
@@ -48,7 +49,7 @@ case object GrantLease
 
 // Events internally used
 case class SetLease(lease: Lease)
-case object SetWatcher
+case class SetWatcher(worker: Cancellable)
 
 class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watcherService: ActorRef)(
  implicit logging: Logging,
@@ -59,15 +60,14 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
  implicit val ec: ExecutionContextExecutor = context.dispatcher
 
  private val leaseTimeout = loadConfigOrThrow[Int](ConfigKeys.etcdLeaseTimeout).seconds
- private var worker: Option[Cancellable] = None
  private val key = instanceLease(instanceId)
  private val watcherName = "lease-service"
 
  self ! GrantLease
  startWith(Ready, NoData)
 
  when(Ready) {
- case Event(GrantLease, _) =>
+ case Event(GrantLease, NoData) =>
  etcdClient
  .grant(leaseTimeout.toSeconds)
  .map { res =>
@@ -76,14 +76,14 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
  .pipeTo(self)
  stay
 
- case Event(SetLease(lease), _) =>
+ case Event(SetLease(lease), NoData) =>
  startKeepAliveService(lease)
  .pipeTo(self)
  logging.info(this, s"Granted a new lease $lease")
  stay using lease
 
- case Event(SetWatcher, lease: Lease) =>
- goto(Active) using lease
+ case Event(SetWatcher(w), l: Lease) =>
+ goto(Active) using ActiveStates(w, l)
 
  case Event(t: FailureMessage, _) =>
  logging.warn(this, s"Failed to grant new lease caused by: $t")
@@ -94,15 +94,15 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
  }
 
  when(Active) {
- case Event(WatchEndpointRemoved(`key`, `key`, _, false), lease: Lease) =>
+ case Event(WatchEndpointRemoved(`key`, `key`, _, false), ActiveStates(worker, lease)) =>
  logging.info(this, s"endpoint ie removed so recreate a lease")
- recreateLease(lease)
+ recreateLease(worker, lease)
 
- case Event(RegrantLease, lease: Lease) =>
+ case Event(RegrantLease, ActiveStates(worker, lease)) =>
  logging.info(this, s"ReGrant a lease, old lease:${lease}")
- recreateLease(lease)
+ recreateLease(worker, lease)
 
- case Event(GetLease, lease: Lease) =>
+ case Event(GetLease, ActiveStates(_, lease)) =>
  logging.info(this, s"send the lease(${lease}) to ${sender()}")
  sender() ! lease
  stay()
@@ -112,18 +112,17 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
 
  initialize()
 
- private def startKeepAliveService(lease: Lease): Future[SetWatcher.type] = {
- worker = Some {
+ private def startKeepAliveService(lease: Lease): Future[SetWatcher] = {
+ val worker =
  actorSystem.scheduler.schedule(initialDelay = 0.second, interval = 500.milliseconds)(keepAliveOnce(lease))
- }
 
  /**
  * To verify that lease has been deleted since timeout,
  * create a key using lease, watch the key, and receive an event for deletion.
  */
  etcdClient.put(key, s"${lease.id}", lease.id).map { _ =>
  watcherService ! WatchEndpoint(key, s"${lease.id}", false, watcherName, Set(DeleteEvent))
- SetWatcher
+ SetWatcher(worker)
  }
  }
 
@@ -139,14 +138,14 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
  }
  }
 
- private def recreateLease(lease: Lease) = {
+ private def recreateLease(worker: Cancellable, lease: Lease) = {
  logging.info(this, s"recreate a lease, old lease: $lease")
- worker.foreach(_.cancel()) // stop scheduler
+ worker.cancel() // stop scheduler
  watcherService ! UnwatchEndpoint(key, false, watcherName) // stop watcher
  etcdClient
  .revoke(lease.id) // delete lease
  .onComplete(_ => self ! GrantLease) // create lease
- goto(Ready)
+ goto(Ready) using NoData
  }
 
  // Unstash all messages stashed while in intermediate state
@@ -162,7 +161,10 @@ class LeaseKeepAliveService(etcdClient: EtcdClient, instanceId: InstanceId, watc
  }
 
  override def postStop(): Unit = {
- worker.foreach(_.cancel()) // stop scheduler if that exist
+ stateData match {
+ case ActiveStates(w, _) => w.cancel() // stop scheduler if that exist
+ case _ => // do nothing
+ }
  watcherService ! UnwatchEndpoint(key, false, watcherName)
  }
 }

tests/src/test/scala/org/apache/openwhisk/core/service/LeaseKeepAliveServiceTests.scala

@@ -101,7 +101,11 @@ class LeaseKeepAliveServiceTests
 
  Thread.sleep(1000)
  service.stateName shouldBe Active
- service.stateData shouldBe testLease
+ service.stateData shouldBe a[ActiveStates]
+ service.stateData match {
+ case ActiveStates(_, lease) => lease shouldBe testLease
+ case _ => fail()
+ }
  watcher.expectMsg(WatchEndpoint(testKey, testLease.id.toString, false, watcherName, Set(DeleteEvent)))
 
  }
@@ -134,7 +138,11 @@ class LeaseKeepAliveServiceTests
  val service = TestFSMRef(new LeaseKeepAliveService(mockEtcd, testInstanceId, watcher.ref))
 
  service.stateName shouldBe Active
- service.stateData shouldBe testLease
+ service.stateData shouldBe a[ActiveStates]
+ service.stateData match {
+ case ActiveStates(_, lease) => lease shouldBe testLease
+ case _ => fail()
+ }
  watcher.expectMsg(WatchEndpoint(testKey, testLease.id.toString, false, watcherName, Set(DeleteEvent)))
 
  service ! WatchEndpointRemoved(testKey, testKey, testLease.id.toString, false)
@@ -143,7 +151,11 @@ class LeaseKeepAliveServiceTests
  Thread.sleep(500) //wait for the lease to be granted
 
  service.stateName shouldBe Active
- service.stateData shouldBe newTestLease
+ service.stateData shouldBe a[ActiveStates]
+ service.stateData match {
+ case ActiveStates(_, lease) => lease shouldBe newTestLease
+ case _ => fail()
+ }
  watcher.expectMsg(WatchEndpoint(newTestKey, newTestLease.id.toString, false, watcherName, Set(DeleteEvent)))
  }
 
@@ -198,14 +210,22 @@ class LeaseKeepAliveServiceTests
  val watcher = TestProbe()
  val service = TestFSMRef(new LeaseKeepAliveService(mockEtcd, testInstanceId, watcher.ref))
  service.stateName shouldBe Active
- service.stateData shouldBe testLease
+ service.stateData shouldBe a[ActiveStates]
+ service.stateData match {
+ case ActiveStates(_, lease) => lease shouldBe testLease
+ case _ => fail()
+ }
  watcher.expectMsg(WatchEndpoint(testKey, testLease.id.toString, false, watcherName, Set(DeleteEvent)))
 
  watcher.expectMsg(UnwatchEndpoint(testKey, false, watcherName))
  Thread.sleep(1500) //wait for the lease to be granted
 
  service.stateName shouldBe Active
- service.stateData shouldBe newTestLease
+ service.stateData shouldBe a[ActiveStates]
+ service.stateData match {
+ case ActiveStates(_, lease) => lease shouldBe newTestLease
+ case _ => fail()
+ }
  watcher.expectMsg(WatchEndpoint(newTestKey, newTestLease.id.toString, false, watcherName, Set(DeleteEvent)))
  }