From 32be79557a3e0a9da162f388cd2d7e0c03e03f36 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Tue, 7 May 2024 20:45:02 +0800
Subject: [PATCH] :zap: Use map to improve message delivery time

---
 pkg/server/messages_ws.go   |  7 ++-----
 pkg/services/connections.go | 18 ++++++++++++++----
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/pkg/server/messages_ws.go b/pkg/server/messages_ws.go
index a51475b..5fdcb58 100644
--- a/pkg/server/messages_ws.go
+++ b/pkg/server/messages_ws.go
@@ -5,14 +5,13 @@ import (
 	"git.solsynth.dev/hydrogen/messaging/pkg/services"
 	"github.com/gofiber/contrib/websocket"
 	jsoniter "github.com/json-iterator/go"
-	"github.com/samber/lo"
 )
 
 func messageGateway(c *websocket.Conn) {
 	user := c.Locals("principal").(models.Account)
 
 	// Push connection
-	services.WsConn[user.ID] = append(services.WsConn[user.ID], c)
+	services.ClientRegister(user, c)
 
 	// Event loop
 	var task models.UnifiedCommand
@@ -42,7 +41,5 @@ func messageGateway(c *websocket.Conn) {
 	}
 
 	// Pop connection
-	services.WsConn[user.ID] = lo.Filter(services.WsConn[user.ID], func(item *websocket.Conn, idx int) bool {
-		return item != c
-	})
+	services.ClientUnregister(user, c)
 }
diff --git a/pkg/services/connections.go b/pkg/services/connections.go
index 53117a1..3c413d0 100644
--- a/pkg/services/connections.go
+++ b/pkg/services/connections.go
@@ -5,14 +5,24 @@ import (
 	"github.com/gofiber/contrib/websocket"
 )
 
-var WsConn = make(map[uint][]*websocket.Conn)
+var wsConn = make(map[uint]map[*websocket.Conn]bool)
 
-func CheckOnline(user models.Account) bool {
-	return len(WsConn[user.ID]) > 0
+func ClientRegister(user models.Account, conn *websocket.Conn) {
+	if wsConn[user.ID] == nil {
+		wsConn[user.ID] = make(map[*websocket.Conn]bool)
+	}
+	wsConn[user.ID][conn] = true
+}
+
+func ClientUnregister(user models.Account, conn *websocket.Conn) {
+	if wsConn[user.ID] == nil {
+		wsConn[user.ID] = make(map[*websocket.Conn]bool)
+	}
+	delete(wsConn[user.ID], conn)
 }
 
 func PushCommand(userId uint, task models.UnifiedCommand) {
-	for _, conn := range WsConn[userId] {
+	for conn := range wsConn[userId] {
 		_ = conn.WriteMessage(1, task.Marshal())
 	}
 }