From 10f91288bca51cbbcaffb2cdde9e21d8f605d41a Mon Sep 17 00:00:00 2001
From: NoahLan <6995syu@163.com>
Date: Thu, 14 Sep 2023 15:16:45 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E6=99=AE=E9=80=9Ahear?=
 =?UTF-8?q?tbeat=E4=BB=A5=E5=8F=8Aws-heartbeat=E7=9A=84return=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 middleware/heartbeat.go    | 50 +++++++++++++++++++-------------------
 middleware/heartbeat_ws.go |  2 +-
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/middleware/heartbeat.go b/middleware/heartbeat.go
index 0a8a65d..596c000 100644
--- a/middleware/heartbeat.go
+++ b/middleware/heartbeat.go
@@ -28,8 +28,32 @@ func WithHeartbeat(interval time.Duration, hbdFn func(conn *conn.Connection) []b
 		panic("dataFn must not be nil")
 	}
 
+	var startFn event.OnConnectedFn = func(nc *conn.Connection) {
+		ticker := time.NewTicker(m.interval)
+
+		defer func() {
+			ticker.Stop()
+		}()
+
+		for {
+			select {
+			case <-ticker.C:
+				deadline := time.Now().Add(-2 * m.interval).Unix()
+				if atomic.LoadInt64(&m.lastAt) < deadline {
+					nlog.Errorf("Heartbeat timeout, LastTime=%d, Deadline=%d", atomic.LoadInt64(&m.lastAt), deadline)
+					break
+				}
+				err := nc.SendBytes(m.hbdFn(nc))
+				if err != nil {
+					nlog.Errorf("Heartbeat err: %v", err)
+					return
+				}
+			}
+		}
+	}
+
 	return func(ngin *nnet.Engine) {
-		ngin.EventManager().RegisterEvent(event.EvtOnConnected, m.start)
+		ngin.EventManager().RegisterEvent(event.EvtOnConnected, startFn)
 
 		ngin.Use(func(next rt.HandlerFunc) rt.HandlerFunc {
 			return func(conn *conn.Connection, pkg packet.IPacket) {
@@ -41,30 +65,6 @@ func WithHeartbeat(interval time.Duration, hbdFn func(conn *conn.Connection) []b
 	}
 }
 
-func (m *HeartbeatMiddleware) start(conn *conn.Connection) {
-	ticker := time.NewTicker(m.interval)
-
-	defer func() {
-		ticker.Stop()
-	}()
-
-	for {
-		select {
-		case <-ticker.C:
-			deadline := time.Now().Add(-2 * m.interval).Unix()
-			if atomic.LoadInt64(&m.lastAt) < deadline {
-				nlog.Errorf("Heartbeat timeout, LastTime=%d, Deadline=%d", atomic.LoadInt64(&m.lastAt), deadline)
-				return
-			}
-			err := conn.SendBytes(m.hbdFn(conn))
-			if err != nil {
-				nlog.Errorf("Heartbeat err: %v", err)
-				return
-			}
-		}
-	}
-}
-
 func (m *HeartbeatMiddleware) handle(_ *conn.Connection, _ packet.IPacket) {
 	atomic.StoreInt64(&m.lastAt, time.Now().Unix())
 }
diff --git a/middleware/heartbeat_ws.go b/middleware/heartbeat_ws.go
index 4fcdb29..fa385ea 100644
--- a/middleware/heartbeat_ws.go
+++ b/middleware/heartbeat_ws.go
@@ -48,7 +48,7 @@ func WithHeartbeatWS(interval time.Duration, hbdFn WsHeartbeatFn) nnet.RunOption
 				deadline := time.Now().Add(-2 * m.interval).Unix()
 				if atomic.LoadInt64(&m.lastAt) < deadline {
 					nlog.Errorf("Heartbeat timeout, LastTime=%d, Deadline=%d", atomic.LoadInt64(&m.lastAt), deadline)
-					return
+					break
 				}
 
 				err := m.hbdFn(nc.WsConn())