com.sandinh.seed.ClusterLeaveAwait.scala Maven / Gradle / Ivy

Go to download
package com.sandinh.seed

import scala.concurrent.{Await, CanAwait, Awaitable}
import java.util.concurrent.{TimeoutException, CountDownLatch}
import scala.concurrent.duration.Duration
import akka.actor.{Props, ActorSystem, ActorLogging, Actor}
import akka.cluster.ClusterEvent.{InitialStateAsEvents, MemberRemoved}
import akka.cluster.Cluster

object ClusterLeaveAwait extends Awaitable[Unit] {
  private val latch = new CountDownLatch(1)

  private[seed] def countDown(): Unit = latch.countDown()

  def ready(atMost: Duration)(implicit permit: CanAwait): this.type = {
    if (atMost.isFinite()) {
      if (!latch.await(atMost.length, atMost.unit))
        throw new TimeoutException("Await termination timed out after [%s]" format atMost.toString)
    } else latch.await()

    this
  }

  def result(atMost: Duration)(implicit permit: CanAwait): Unit = ready(atMost)

  def awaitLeave(system: ActorSystem, atMost: Duration = Duration.Inf): Unit = {
    val cluster = Cluster(system)
    val trackingActor = system.actorOf(Props[ClusterTracking])
    cluster.subscribe(trackingActor, InitialStateAsEvents, classOf[MemberRemoved])
    cluster.leave(cluster.selfAddress)
    Await.ready(this, atMost)
  }
}

private class ClusterTracking extends Actor with ActorLogging {
  override def receive = {
    case MemberRemoved(member, _) =>
      log.debug("{} is leaving", member.address)
      if (member.address == Cluster(context.system).selfAddress)
        ClusterLeaveAwait.countDown()
  }
}