⚡ Use map to improve message delivery time

pkg/server/notify_ws.go

@@ -4,14 +4,13 @@ import (
 	"git.solsynth.dev/hydrogen/passport/pkg/models"
 	"git.solsynth.dev/hydrogen/passport/pkg/services"
 	"github.com/gofiber/contrib/websocket"
-	"github.com/samber/lo"
 )
 
 func listenNotifications(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 err error
@@ -22,7 +21,5 @@ func listenNotifications(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)
 }

pkg/services/connections.go

@@ -5,13 +5,18 @@ import (
 	"github.com/gofiber/contrib/websocket"
 )
 
-type WsPushRequest struct {
-	Payload     []byte
-	RecipientID uint
+var wsConn = make(map[uint]map[*websocket.Conn]bool)
+
+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
 }
 
-var WsConn = make(map[uint][]*websocket.Conn)
-
-func CheckOnline(user models.Account) bool {
-	return len(WsConn[user.ID]) > 0
+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)
 }

pkg/services/notifications.go

@@ -41,7 +41,7 @@ func NewNotification(notification models.Notification) error {
 
 func PushNotification(notification models.Notification) error {
 	raw, _ := jsoniter.Marshal(notification)
-	for _, conn := range WsConn[notification.RecipientID] {
+	for conn := range wsConn[notification.RecipientID] {
 		_ = conn.WriteMessage(1, raw)
 	}