From 6246528e864f0d868c0cd04d29aca38d6ced29ed Mon Sep 17 00:00:00 2001
From: Patrik Nordwall <patrik.nordwall@gmail.com>
Date: Mon, 15 May 2017 15:27:53 +0200
Subject: [PATCH] cleanup the Shard idByRefs map when entity terminated, #22893

---
 .../src/main/scala/akka/cluster/sharding/Shard.scala  | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/akka-cluster-sharding/src/main/scala/akka/cluster/sharding/Shard.scala b/akka-cluster-sharding/src/main/scala/akka/cluster/sharding/Shard.scala
index 612ce25fec..e4b4e10d8a 100644
--- a/akka-cluster-sharding/src/main/scala/akka/cluster/sharding/Shard.scala
+++ b/akka-cluster-sharding/src/main/scala/akka/cluster/sharding/Shard.scala
@@ -224,6 +224,8 @@ private[akka] class Shard(
 
   def entityTerminated(ref: ActorRef): Unit = {
     val id = idByRef(ref)
+    idByRef -= ref
+    refById -= id
     if (messageBuffers.getOrEmpty(id).nonEmpty) {
       log.debug("Starting entity [{}] again, there are buffered messages for it", id)
       sendMsgBuffer(EntityStarted(id))
@@ -251,12 +253,7 @@ private[akka] class Shard(
 
   // EntityStopped handler
   def passivateCompleted(event: EntityStopped): Unit = {
-    log.debug("Entity stopped [{}]", event.entityId)
-
-    val ref = refById(event.entityId)
-    idByRef -= ref
-    refById -= event.entityId
-
+    log.debug("Entity stopped after passivation [{}]", event.entityId)
     state = state.copy(state.entities - event.entityId)
     messageBuffers.remove(event.entityId)
   }
@@ -352,6 +349,8 @@ private[akka] trait RememberingShard { selfType: Shard ⇒
   override def entityTerminated(ref: ActorRef): Unit = {
     import settings.tuningParameters._
     val id = idByRef(ref)
+    idByRef -= ref
+    refById -= id
     if (messageBuffers.getOrEmpty(id).nonEmpty) {
       //Note; because we're not persisting the EntityStopped, we don't need
       // to persist the EntityStarted either.