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

Go to download
package com.sandinh.seed

import java.util.concurrent.{CountDownLatch, TimeUnit, TimeoutException}
import akka.actor.{Actor, ActorLogging, ActorSystem, Props}
import akka.cluster.Cluster
import akka.cluster.ClusterEvent.{InitialStateAsEvents, MemberRemoved}
import org.slf4j.LoggerFactory
import scala.concurrent.duration._
import scala.concurrent.{Await, Awaitable, CanAwait}
import scala.util.Try

object ClusterLeaveAwait {
  private val logger = LoggerFactory.getLogger(this.getClass)

  def leaveAndAwait(system: ActorSystem): Unit = {
    val atMost =
      Try(system.settings.config.getDuration(s"akka.cluster-seed.await-leave-timeout", TimeUnit.SECONDS).seconds)
        .filter(_ > Duration.Zero)
        .getOrElse(Duration.Inf)
    leaveAndAwait(system, atMost)
  }

  def leaveAndAwait(system: ActorSystem, atMost: Duration): Unit = {
    if (atMost < 10.seconds) {
      logger.warn("cluster-seed.await-leave-timeout < 10s is known to have issue: 'Existing member is trying to join, ignoring..'")
    } else {
      logger.info("cluster leave & await for {} ...", atMost)
    }
    val await = new ClusterLeaveAwait
    await.leave(system)
    Await.ready(await, atMost)
  }

  @deprecated("config using akka.cluster-seed.await-leave-timeout then use leaveAndAwait instead", "1.5.0")
  def awaitLeave(system: ActorSystem, atMost: Duration) = leaveAndAwait(system, atMost)
}

class ClusterLeaveAwait extends Awaitable[Unit] {
  private[this] val latch = new CountDownLatch(1)
  def countDown() = latch.countDown()

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

  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)
}

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