Allow only specific actors defined as a sticky leader to become a leader

CHANGELOG.md

@@ -26,6 +26,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   [PR178](https://github.com/lerna-stack/akka-entity-replication/pull/178)
 - Add diagnostic info to logs of sending replication results
   [PR#179](https://github.com/lerna-stack/akka-entity-replication/pull/179)
+- Allow only specific actors defined as a sticky leader to become a leader
+  [PR#186](https://github.com/lerna-stack/akka-entity-replication/pull/186)
 
 ### Fixed
 - RaftActor might delete committed entries

docs/implementation_guide.md

@@ -164,6 +164,31 @@ AtLeastOnceComplete.askTo(
 )
 ```
 
+### Fixing leader actor per shard
+
+By default, all raft actors can become a Leader.
+You can fix a Leader actor per shard by setting a sticky leader.
+Fixing leader actor prevents committed event loss during recovery.
+
+```scala
+import akka.actor.typed.ActorSystem
+import lerna.akka.entityreplication.typed._
+
+val system: ActorSystem[_] = ???
+val clusterReplication = ClusterReplication(system)
+
+// Settings for sticky leader. Only a raft actor which have role "replica-group-1" can become Leader in shard given id "1"
+val settings =
+  ClusterReplicationSettings(system)
+    .withStickyLeaders(Map("1" -> "replica-group-1"))
+
+val entity = 
+  ReplicatedEntity(BankAccountBehavior.TypeKey)(entityContext => BankAccountBehavior(entityContext))
+    .withSettings(settings)
+
+clusterReplication.init(entity)
+```
+
 ### Configuration
 
 On the command side, the related settings are defined at `lerna.akka.entityreplication`(except `lerna.akka.entityreplication.raft.eventsourced`) in [reference.conf](/src/main/resources/reference.conf).

src/main/mima-filters/2.1.0.backwards.excludes/pr-186-allow-only-specific-actors-to-become-leader.excludes

@@ -0,0 +1,2 @@
+# ClusterReplicationSettings should not be extended by users
+ProblemFilters.exclude[ReversedMissingMethodProblem]("lerna.akka.entityreplication.ClusterReplicationSettings.withStickyLeaders")

src/main/scala/lerna/akka/entityreplication/ClusterReplicationSettings.scala

@@ -41,6 +41,8 @@ trait ClusterReplicationSettings {
 
   def withDisabledShards(disabledShards: Set[String]): ClusterReplicationSettings
 
+  def withStickyLeaders(stickyLeaders: Map[String, String]): ClusterReplicationSettings
+
   def withRaftJournalPluginId(pluginId: String): ClusterReplicationSettings
 
   def withRaftSnapshotPluginId(pluginId: String): ClusterReplicationSettings

src/main/scala/lerna/akka/entityreplication/internal/ClusterReplicationSettingsImpl.scala

@@ -21,6 +21,9 @@ private[entityreplication] final case class ClusterReplicationSettingsImpl(
   override def withDisabledShards(disabledShards: Set[String]): ClusterReplicationSettings =
     copy(raftSettings = raftSettings.withDisabledShards(disabledShards))
 
+  override def withStickyLeaders(stickyLeaders: Map[String, String]): ClusterReplicationSettingsImpl =
+    copy(raftSettings = raftSettings.withStickyLeaders(stickyLeaders))
+
   override def withRaftJournalPluginId(pluginId: String): ClusterReplicationSettingsImpl =
     copy(raftSettings = raftSettings.withJournalPluginId(pluginId))
 

src/main/scala/lerna/akka/entityreplication/raft/Candidate.scala

@@ -11,25 +11,7 @@ private[raft] trait Candidate { this: RaftActor =>
   import RaftActor._
 
   def candidateBehavior: Receive = {
-
-    case ElectionTimeout =>
-      if (log.isInfoEnabled)
-        log.info("[Candidate] Election timeout at {}. Retrying leader election.", currentData.currentTerm)
-      val newTerm = currentData.currentTerm.next()
-      cancelElectionTimeoutTimer()
-      broadcast(
-        RequestVote(
-          shardId,
-          newTerm,
-          selfMemberIndex,
-          currentData.replicatedLog.lastLogIndex,
-          currentData.replicatedLog.lastLogTerm,
-        ),
-      ) // TODO: 永続化前に broadcast して問題ないか調べる
-      applyDomainEvent(BegunNewTerm(newTerm)) { _ =>
-        become(Candidate)
-      }
-
+    case ElectionTimeout                                 => receiveElectionTimeout()
     case request: RequestVote                            => receiveRequestVote(request)
     case response: RequestVoteResponse                   => receiveRequestVoteResponse(response)
     case request: AppendEntries                          => receiveAppendEntries(request)
@@ -55,6 +37,27 @@ private[raft] trait Candidate { this: RaftActor =>
 
   }
 
+  private def receiveElectionTimeout(): Unit = {
+    if (log.isInfoEnabled)
+      log.info("[Candidate] Election timeout at {}. Retrying leader election.", currentData.currentTerm)
+    val newTerm = currentData.currentTerm.next()
+    cancelElectionTimeoutTimer()
+    if (canBecomeCandidate) {
+      broadcast(
+        RequestVote(
+          shardId,
+          newTerm,
+          selfMemberIndex,
+          currentData.replicatedLog.lastLogIndex,
+          currentData.replicatedLog.lastLogTerm,
+        ),
+      ) // TODO: 永続化前に broadcast して問題ないか調べる
+      applyDomainEvent(BegunNewTerm(newTerm)) { _ =>
+        become(Candidate)
+      }
+    }
+  }
+
   private[this] def receiveRequestVote(request: RequestVote): Unit =
     request match {
 

src/main/scala/lerna/akka/entityreplication/raft/Follower.scala

@@ -12,15 +12,7 @@ private[raft] trait Follower { this: RaftActor =>
   import RaftActor._
 
   def followerBehavior: Receive = {
-
-    case ElectionTimeout =>
-      if (currentData.leaderMember.isEmpty) {
-        if (log.isDebugEnabled) log.debug("=== [Follower] election timeout ===")
-      } else {
-        if (log.isWarningEnabled) log.warning("[{}] election timeout. Leader will be changed", currentState)
-      }
-      requestVote(currentData)
-
+    case ElectionTimeout                                 => receiveElectionTimeout()
     case request: RequestVote                            => receiveRequestVote(request)
     case request: AppendEntries                          => receiveAppendEntries(request)
     case request: InstallSnapshot                        => receiveInstallSnapshot(request)
@@ -45,6 +37,18 @@ private[raft] trait Follower { this: RaftActor =>
 
   }
 
+  private[this] def receiveElectionTimeout(): Unit = {
+    if (currentData.leaderMember.isEmpty) {
+      if (log.isDebugEnabled) log.debug("=== [Follower] election timeout ===")
+    } else {
+      if (log.isWarningEnabled) log.warning("[{}] election timeout. Leader will be changed", currentState)
+    }
+    cancelElectionTimeoutTimer()
+    if (canBecomeCandidate) {
+      requestVote(currentData)
+    }
+  }
+
   private[this] def receiveRequestVote(request: RequestVote): Unit =
     request match {
 
@@ -150,7 +154,6 @@ private[raft] trait Follower { this: RaftActor =>
 
   private[this] def requestVote(data: RaftMemberData): Unit = {
     val newTerm = data.currentTerm.next()
-    cancelElectionTimeoutTimer()
     broadcast(
       RequestVote(shardId, newTerm, selfMemberIndex, data.replicatedLog.lastLogIndex, data.replicatedLog.lastLogTerm),
     ) // TODO: 永続化前に broadcast して問題ないか調べる

src/main/scala/lerna/akka/entityreplication/raft/RaftActor.scala

@@ -157,6 +157,21 @@ private[raft] class RaftActor(
   import RaftActor._
   import context.dispatcher
 
+  if (isStickyLeader && log.isWarningEnabled) {
+    log.warning(
+      "This raft actor is defined as sticky leader. shard id = {} , role = {}, stickyLeaders = {}. sticky leaders are not able to failover to other multi-raft-roles. This means that some entities can be unavailable by some failures. Please unset sticky-leaders settings if there is no need.",
+      shardId.raw,
+      selfMemberIndex.role,
+      settings.stickyLeaders,
+    )
+  }
+
+  private def isStickyLeader: Boolean =
+    settings.stickyLeaders.get(this.shardId.raw).fold(false)(_ == this.selfMemberIndex.role)
+
+  protected def canBecomeCandidate: Boolean =
+    settings.stickyLeaders.get(this.shardId.raw).fold(true)(_ == this.selfMemberIndex.role)
+
   protected[this] def shardId: NormalizedShardId = NormalizedShardId.from(self.path)
 
   protected[this] def region: ActorRef = _region

src/main/scala/lerna/akka/entityreplication/raft/RaftSettings.scala

@@ -14,6 +14,8 @@ trait RaftSettings {
 
   def electionTimeout: FiniteDuration
 
+  def stickyLeaders: Map[String, String]
+
   private[raft] def randomizedElectionTimeout(): FiniteDuration
 
   def heartbeatInterval: FiniteDuration
@@ -85,6 +87,8 @@ trait RaftSettings {
 
   private[entityreplication] def withDisabledShards(disabledShards: Set[String]): RaftSettings
 
+  private[entityreplication] def withStickyLeaders(stickyLeaders: Map[String, String]): RaftSettings
+
   private[entityreplication] def withJournalPluginId(pluginId: String): RaftSettings
 
   private[entityreplication] def withSnapshotPluginId(pluginId: String): RaftSettings

src/main/scala/lerna/akka/entityreplication/raft/RaftSettingsImpl.scala

@@ -11,6 +11,7 @@ import scala.util.Random
 private[entityreplication] final case class RaftSettingsImpl(
     config: Config,
     electionTimeout: FiniteDuration,
+    stickyLeaders: Map[String, String],
     heartbeatInterval: FiniteDuration,
     electionTimeoutMin: Duration,
     multiRaftRoles: Set[String],
@@ -58,6 +59,9 @@ private[entityreplication] final case class RaftSettingsImpl(
   override private[entityreplication] def withDisabledShards(disabledShards: Set[String]): RaftSettings =
     copy(disabledShards = disabledShards)
 
+  override private[entityreplication] def withStickyLeaders(stickyLeaders: Map[String, String]) =
+    copy(stickyLeaders = stickyLeaders)
+
   override private[entityreplication] def withJournalPluginId(pluginId: String): RaftSettings =
     copy(journalPluginId = pluginId)
 
@@ -83,6 +87,8 @@ private[entityreplication] object RaftSettingsImpl {
 
     val electionTimeout: FiniteDuration = config.getDuration("election-timeout").toScala
 
+    val stickyLeaders: Map[String, String] = Map.empty
+
     val heartbeatInterval: FiniteDuration = config.getDuration("heartbeat-interval").toScala
 
     val electionTimeoutMin: Duration = electionTimeout * randomizedMinFactor
@@ -224,6 +230,7 @@ private[entityreplication] object RaftSettingsImpl {
     RaftSettingsImpl(
       config = config,
       electionTimeout = electionTimeout,
+      stickyLeaders = stickyLeaders,
       heartbeatInterval = heartbeatInterval,
       electionTimeoutMin = electionTimeoutMin,
       multiRaftRoles = multiRaftRoles,

src/test/scala/lerna/akka/entityreplication/ClusterReplicationSettingsSpec.scala

@@ -61,6 +61,12 @@ class ClusterReplicationSettingsSpec extends WordSpec with Matchers {
       newSettings.raftSettings.disabledShards should be(Set("1", "3"))
     }
 
+    "change value of raftSettings.stickyLeaders by withStickyLeaders" in {
+      val settings    = ClusterReplicationSettingsImpl(config, correctClusterRoles.headOption.toSet)
+      val newSettings = settings.withStickyLeaders(Map("1" -> "replica-group-1"))
+      newSettings.raftSettings.stickyLeaders should be(Map("1" -> "replica-group-1"))
+    }
+
     "change value of raftSettings.journalPluginId by withRaftJournalPluginId" in {
       val settings         = ClusterReplicationSettingsImpl(config, correctClusterRoles.headOption.toSet)
       val expectedPluginId = "new-raft-journal-plugin-id"

src/test/scala/lerna/akka/entityreplication/raft/RaftActorCandidateSpec.scala

@@ -7,13 +7,7 @@ import akka.persistence.testkit.scaladsl.PersistenceTestKit
 import akka.testkit.{ TestKit, TestProbe }
 import lerna.akka.entityreplication.ReplicationRegion
 import lerna.akka.entityreplication.model.{ EntityInstanceId, NormalizedEntityId, NormalizedShardId }
-import lerna.akka.entityreplication.raft.RaftProtocol.{
-  Command,
-  ForwardedCommand,
-  ProcessCommand,
-  Replicate,
-  ReplicationFailed,
-}
+import lerna.akka.entityreplication.raft.RaftProtocol._
 import lerna.akka.entityreplication.raft.model._
 import lerna.akka.entityreplication.raft.protocol.RaftCommands._
 import lerna.akka.entityreplication.raft.routing.MemberIndex
@@ -212,6 +206,105 @@ class RaftActorCandidateSpec
       }
     }
 
+    "send RequestVote on ElectionTimeout if a shard which it's belongs to doesn't have a sticky leader" in {
+      val shardId              = createUniqueShardId()
+      val candidateMemberIndex = createUniqueMemberIndex()
+      val regionProbe          = TestProbe()
+      val customSettings       = RaftSettings(defaultRaftConfig).withStickyLeaders(Map.empty)
+      val candidate = createRaftActor(
+        shardId = shardId,
+        selfMemberIndex = candidateMemberIndex,
+        region = regionProbe.ref,
+        settings = customSettings,
+      )
+      val currentTerm   = Term(2)
+      val candidateData = createCandidateData(currentTerm, ReplicatedLog())
+      setState(candidate, Candidate, candidateData)
+
+      // ElectionTimeout triggers that this candidate sends a RequestVote.
+      candidate ! ElectionTimeout
+      val expectedRequestVote =
+        RequestVote(
+          shardId,
+          term = Term(3),
+          candidate = candidateMemberIndex,
+          lastLogIndex = LogEntryIndex(0),
+          lastLogTerm = Term(0),
+        )
+      regionProbe.expectMsg(ReplicationRegion.Broadcast(expectedRequestVote))
+      awaitAssert {
+        assert(getState(candidate).stateName == Candidate)
+        val stateData = getState(candidate).stateData
+        assert(stateData.currentTerm == Term(3))
+        assert(stateData.votedFor.isEmpty)
+        assert(stateData.acceptedMembers.isEmpty)
+      }
+    }
+
+    "send RequestVote on ElectionTimeout if it is defined as sticky leader" in {
+      val shardId              = createUniqueShardId()
+      val candidateMemberIndex = createUniqueMemberIndex()
+      val regionProbe          = TestProbe()
+      val customSettings =
+        RaftSettings(defaultRaftConfig).withStickyLeaders(Map(shardId.raw -> candidateMemberIndex.role))
+      val candidate = createRaftActor(
+        shardId = shardId,
+        selfMemberIndex = candidateMemberIndex,
+        region = regionProbe.ref,
+        settings = customSettings,
+      )
+      val currentTerm   = Term(2)
+      val candidateData = createCandidateData(currentTerm, ReplicatedLog())
+      setState(candidate, Candidate, candidateData)
+
+      // ElectionTimeout triggers that this candidate sends a RequestVote.
+      candidate ! ElectionTimeout
+      val expectedRequestVote =
+        RequestVote(
+          shardId,
+          term = Term(3),
+          candidate = candidateMemberIndex,
+          lastLogIndex = LogEntryIndex(0),
+          lastLogTerm = Term(0),
+        )
+      regionProbe.expectMsg(ReplicationRegion.Broadcast(expectedRequestVote))
+      awaitAssert {
+        assert(getState(candidate).stateName == Candidate)
+        val stateData = getState(candidate).stateData
+        assert(stateData.currentTerm == Term(3))
+        assert(stateData.votedFor.isEmpty)
+        assert(stateData.acceptedMembers.isEmpty)
+      }
+    }
+
+    "not send RequestVote on ElectionTimeout if another raft actor which belongs same shard is defined as a sticky leader" in {
+      val shardId              = createUniqueShardId()
+      val candidateMemberIndex = createUniqueMemberIndex()
+      val regionProbe          = TestProbe()
+      val customSettings =
+        RaftSettings(defaultRaftConfig).withStickyLeaders(Map(s"${shardId.raw}" -> "other-actors-role"))
+      val candidate = createRaftActor(
+        shardId = shardId,
+        selfMemberIndex = candidateMemberIndex,
+        region = regionProbe.ref,
+        settings = customSettings,
+      )
+      val currentTerm   = Term(2)
+      val candidateData = createCandidateData(currentTerm, ReplicatedLog())
+      setState(candidate, Candidate, candidateData)
+
+      // ElectionTimeout should not triggers that this candidate sends a RequestVote.
+      candidate ! ElectionTimeout
+      regionProbe.expectNoMessage()
+      awaitAssert {
+        assert(getState(candidate).stateName == Candidate)
+        val stateData = getState(candidate).stateData
+        assert(stateData.currentTerm == Term(2))
+        assert(stateData.votedFor.isEmpty)
+        assert(stateData.acceptedMembers.isEmpty)
+      }
+    }
+
     "メンバーの過半数に Accept されると Leader になる" in {
       val selfMemberIndex      = createUniqueMemberIndex()
       val follower1MemberIndex = createUniqueMemberIndex()