From 7d47091e59c0ddce3f2aacd5b14c0d5de370069b Mon Sep 17 00:00:00 2001
From: Patrik Nordwall <patrik.nordwall@gmail.com>
Date: Thu, 27 Sep 2012 08:57:38 +0200
Subject: [PATCH] Add some extra safety to Guardian termination hook, see #2552

---
 .../scala/akka/actor/ActorRefProvider.scala   | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/akka-actor/src/main/scala/akka/actor/ActorRefProvider.scala b/akka-actor/src/main/scala/akka/actor/ActorRefProvider.scala
index 5aa6049c74..d0cfb91a04 100644
--- a/akka-actor/src/main/scala/akka/actor/ActorRefProvider.scala
+++ b/akka-actor/src/main/scala/akka/actor/ActorRefProvider.scala
@@ -396,11 +396,27 @@ class LocalActorRefProvider(
     var terminationHooks = Set.empty[ActorRef]
 
     def receive = {
-      case Terminated(_)           ⇒ terminationHooks foreach { _ ! TerminationHook }; stopWhenAllTerminationHooksDone()
-      case StopChild(child)        ⇒ context.stop(child)
-      case RegisterTerminationHook ⇒ terminationHooks += sender
-      case TerminationHookDone     ⇒ terminationHooks -= sender; stopWhenAllTerminationHooksDone()
-      case m                       ⇒ deadLetters ! DeadLetter(m, sender, self)
+      case Terminated(a) if terminationHooks.contains(a) ⇒ terminationHooks -= a
+      case Terminated(_) ⇒
+        context.become(terminating)
+        terminationHooks foreach { _ ! TerminationHook }
+        stopWhenAllTerminationHooksDone()
+      case StopChild(child) ⇒ context.stop(child)
+      case RegisterTerminationHook if sender != context.system.deadLetters ⇒
+        terminationHooks += sender
+        context watch sender
+      case m ⇒ deadLetters ! DeadLetter(m, sender, self)
+    }
+
+    def terminating: Receive = {
+      case Terminated(a) if terminationHooks.contains(a) ⇒ stopWhenAllTerminationHooksDone(a)
+      case TerminationHookDone ⇒ stopWhenAllTerminationHooksDone(sender)
+      case m ⇒ deadLetters ! DeadLetter(m, sender, self)
+    }
+
+    def stopWhenAllTerminationHooksDone(remove: ActorRef): Unit = {
+      terminationHooks -= sender
+      stopWhenAllTerminationHooksDone()
     }
 
     def stopWhenAllTerminationHooksDone(): Unit = if (terminationHooks.isEmpty) {