From cf3fb987067e504e3c9105e95161008d496f5f55 Mon Sep 17 00:00:00 2001
From: SNORRIS721 <stn721@gmail.com>
Date: Fri, 3 May 2024 13:50:51 -0500
Subject: [PATCH] Unregistering sessions on acceptor stop

---
 acceptor.go  | 7 +++++++
 initiator.go | 4 ++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/acceptor.go b/acceptor.go
index 968be9efb..da4cd9bb0 100644
--- a/acceptor.go
+++ b/acceptor.go
@@ -144,6 +144,13 @@ func (a *Acceptor) Stop() {
 		session.stop()
 	}
 	a.sessionGroup.Wait()
+
+	for sessionID := range a.sessions {
+		err := UnregisterSession(sessionID)
+		if err != nil {
+			return
+		}
+	}
 }
 
 // RemoteAddr gets remote IP address for a given session.
diff --git a/initiator.go b/initiator.go
index b8735e388..eeff30fc0 100644
--- a/initiator.go
+++ b/initiator.go
@@ -74,14 +74,14 @@ func (i *Initiator) Stop() {
 	}
 	close(i.stopChan)
 
+	i.wg.Wait()
+
 	for sessionID := range i.sessionSettings {
 		err := UnregisterSession(sessionID)
 		if err != nil {
 			return
 		}
 	}
-
-	i.wg.Wait()
 }
 
 // NewInitiator creates and initializes a new Initiator.