From b9d953338f7efdac1d9ca97c7494a3ff0718adcd Mon Sep 17 00:00:00 2001
From: kercylan98 <kercylan@gmail.com>
Date: Tue, 1 Aug 2023 15:07:33 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20server.Server=20=E6=96=B0=E5=A2=9E?=
 =?UTF-8?q?=E5=88=86=E6=B5=81=E9=80=9A=E9=81=93=E5=88=9B=E5=BB=BA=E5=92=8C?=
 =?UTF-8?q?=E5=85=B3=E9=97=AD=E4=BA=8B=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/event.go   | 32 ++++++++++++++++++++++++++++++++
 server/options.go |  3 +++
 server/server.go  |  2 ++
 3 files changed, 37 insertions(+)

diff --git a/server/event.go b/server/event.go
index a492e87..c201cc5 100644
--- a/server/event.go
+++ b/server/event.go
@@ -22,6 +22,8 @@ type MessageLowExecEventHandle func(srv *Server, message *Message, cost time.Dur
 type ConsoleCommandEventHandle func(srv *Server)
 type ConnectionOpenedAfterEventHandle func(srv *Server, conn *Conn)
 type ConnectionWritePacketBeforeEventHandle func(srv *Server, conn *Conn, packet Packet) Packet
+type ShuntChannelCreatedEventHandle func(srv *Server, guid int64)
+type ShuntChannelClosedEventHandle func(srv *Server, guid int64)
 
 type event struct {
 	*Server
@@ -36,6 +38,8 @@ type event struct {
 	messageLowExecEventHandles          []MessageLowExecEventHandle
 	connectionOpenedAfterEventHandles   []ConnectionOpenedAfterEventHandle
 	connectionWritePacketBeforeHandles  []ConnectionWritePacketBeforeEventHandle
+	shuntChannelCreatedEventHandles     []ShuntChannelCreatedEventHandle
+	shuntChannelClosedEventHandles      []ShuntChannelClosedEventHandle
 
 	consoleCommandEventHandles        map[string][]ConsoleCommandEventHandle
 	consoleCommandEventHandleInitOnce sync.Once
@@ -251,6 +255,34 @@ func (slf *event) OnConnectionWritePacketBeforeEvent(conn *Conn, packet Packet)
 	return newPacket
 }
 
+// RegShuntChannelCreatedEvent 在分流通道创建时将立刻执行被注册的事件处理函数
+func (slf *event) RegShuntChannelCreatedEvent(handle ShuntChannelCreatedEventHandle) {
+	slf.shuntChannelCreatedEventHandles = append(slf.shuntChannelCreatedEventHandles, handle)
+	log.Info("Server", log.String("RegEvent", runtimes.CurrentRunningFuncName()), log.String("handle", reflect.TypeOf(handle).String()))
+}
+
+func (slf *event) OnShuntChannelCreatedEvent(guid int64) {
+	PushSystemMessage(slf.Server, func() {
+		for _, handle := range slf.shuntChannelCreatedEventHandles {
+			handle(slf.Server, guid)
+		}
+	}, "ShuntChannelCreatedEvent")
+}
+
+// RegShuntChannelCloseEvent 在分流通道关闭时将立刻执行被注册的事件处理函数
+func (slf *event) RegShuntChannelCloseEvent(handle ShuntChannelClosedEventHandle) {
+	slf.shuntChannelClosedEventHandles = append(slf.shuntChannelClosedEventHandles, handle)
+	log.Info("Server", log.String("RegEvent", runtimes.CurrentRunningFuncName()), log.String("handle", reflect.TypeOf(handle).String()))
+}
+
+func (slf *event) OnShuntChannelClosedEvent(guid int64) {
+	PushSystemMessage(slf.Server, func() {
+		for _, handle := range slf.shuntChannelClosedEventHandles {
+			handle(slf.Server, guid)
+		}
+	}, "ShuntChannelCloseEvent")
+}
+
 func (slf *event) check() {
 	switch slf.network {
 	case NetworkHttp, NetworkGRPC:
diff --git a/server/options.go b/server/options.go
index 1aa4cf9..209e0cb 100644
--- a/server/options.go
+++ b/server/options.go
@@ -238,6 +238,9 @@ func WithPProf(pattern ...string) Option {
 //
 // 将被分流的消息类型（更多类型有待斟酌）：
 //   - MessageTypePacket
+//
+// 注意事项：
+//   - 需要在分流通道使用完成后主动调用 Server.ShuntChannelFreed 函数释放分流通道，避免内存泄漏
 func WithShunt(channelGenerator func(guid int64) chan *Message, shuntMatcher func(conn *Conn) (guid int64, allowToCreate bool)) Option {
 	return func(srv *Server) {
 		if channelGenerator == nil || shuntMatcher == nil {
diff --git a/server/server.go b/server/server.go
index d957cf6..1223cbd 100644
--- a/server/server.go
+++ b/server/server.go
@@ -489,6 +489,7 @@ func (slf *Server) ShuntChannelFreed(channelGuid int64) {
 	if exist {
 		close(channel)
 		slf.shuntChannels.Delete(channelGuid)
+		slf.OnShuntChannelClosedEvent(channelGuid)
 	}
 }
 
@@ -510,6 +511,7 @@ func (slf *Server) pushMessage(message *Message) {
 					slf.dispatchMessage(message)
 				}
 			}(channel)
+			defer slf.OnShuntChannelCreatedEvent(channelGuid)
 		}
 		if channel != nil {
 			channel <- message