From 2fb136569e9eebe384c6bb4a0417b8076143981a Mon Sep 17 00:00:00 2001
From: neokami <ryan@megidov.com>
Date: Fri, 17 Nov 2023 12:48:03 +0100
Subject: [PATCH] Allow poisoning router actors

---
 router/process.go | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/router/process.go b/router/process.go
index 46ca7d0f2..560f4473c 100644
--- a/router/process.go
+++ b/router/process.go
@@ -23,6 +23,12 @@ var _ actor.Process = &process{}
 
 func (ref *process) SendUserMessage(pid *actor.PID, message interface{}) {
 	_, msg, _ := actor.UnwrapEnvelope(message)
+
+	// Add support for PoisonPill. Originally only Stop is supported.
+	if _, ok := msg.(*actor.PoisonPill); ok {
+		ref.Poison(pid)
+		return
+	}
 	if _, ok := msg.(ManagementMessage); !ok {
 		ref.state.RouteMessage(message)
 	} else {
@@ -85,3 +91,13 @@ func (ref *process) Stop(pid *actor.PID) {
 	ref.actorSystem.ProcessRegistry.Remove(pid)
 	ref.SendSystemMessage(pid, &actor.Stop{})
 }
+
+func (ref *process) Poison(pid *actor.PID) {
+	if atomic.SwapInt32(&ref.stopping, 1) == 1 {
+		return
+	}
+
+	_ = ref.actorSystem.Root.PoisonFuture(ref.router).Wait()
+	ref.actorSystem.ProcessRegistry.Remove(pid)
+	ref.SendSystemMessage(pid, &actor.Stop{})
+}