From 889e6b0d84f75eafd9b53d600047f3d951cc2c89 Mon Sep 17 00:00:00 2001
From: Gordon <1432970085@qq.com>
Date: Thu, 23 Mar 2023 14:18:33 +0800
Subject: [PATCH] ws update

---
 internal/msggateway/hub_server.go  |  6 +++---
 internal/msggateway/init.go        | 13 ++-----------
 internal/msggateway/n_ws_server.go |  1 +
 3 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/internal/msggateway/hub_server.go b/internal/msggateway/hub_server.go
index ff22567da..f9f7bfe8d 100644
--- a/internal/msggateway/hub_server.go
+++ b/internal/msggateway/hub_server.go
@@ -16,7 +16,7 @@ import (
 )
 
 func (s *Server) InitServer(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) error {
-	s.notification = notification.NewCheck(client)
+	s.LongConnServer.SetMessageHandler(notification.NewCheck(client))
 	msggateway.RegisterMsgGatewayServer(server, s)
 	return nil
 }
@@ -42,8 +42,8 @@ func (s *Server) Notification() *notification.Check {
 	return s.notification
 }
 
-func NewServer(rpcPort int) *Server {
-	return &Server{rpcPort: rpcPort, pushTerminal: []int{constant.IOSPlatformID, constant.AndroidPlatformID}}
+func NewServer(rpcPort int, longConnServer LongConnServer) *Server {
+	return &Server{rpcPort: rpcPort, LongConnServer: longConnServer, pushTerminal: []int{constant.IOSPlatformID, constant.AndroidPlatformID}}
 }
 
 func (s *Server) OnlinePushMsg(context context.Context, req *msggateway.OnlinePushMsgReq) (*msggateway.OnlinePushMsgResp, error) {
diff --git a/internal/msggateway/init.go b/internal/msggateway/init.go
index 808abf817..9d0eee451 100644
--- a/internal/msggateway/init.go
+++ b/internal/msggateway/init.go
@@ -5,13 +5,10 @@ import (
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/config"
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/constant"
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/log"
-	"sync"
 	"time"
 )
 
 func RunWsAndServer(rpcPort, wsPort, prometheusPort int) error {
-	var wg sync.WaitGroup
-	wg.Add(1)
 	log.NewPrivateLog(constant.LogFileName)
 	fmt.Println("start rpc/msg_gateway server, port: ", rpcPort, wsPort, prometheusPort, ", OpenIM version: ", config.Version)
 	longServer, err := NewWsServer(
@@ -22,14 +19,8 @@ func RunWsAndServer(rpcPort, wsPort, prometheusPort int) error {
 	if err != nil {
 		return err
 	}
-	hubServer := NewServer(rpcPort)
+	hubServer := NewServer(rpcPort, longServer)
 	go hubServer.Start()
-	if hubServer.Notification() == nil {
-		panic("notification is nil")
-	}
-	longServer.SetMessageHandler(hubServer.Notification())
-	hubServer.SetLongConnServer(longServer)
-	go hubServer.LongConnServer.Run()
-	wg.Wait()
+	hubServer.LongConnServer.Run()
 	return nil
 }
diff --git a/internal/msggateway/n_ws_server.go b/internal/msggateway/n_ws_server.go
index 24bfa213b..412818bd1 100644
--- a/internal/msggateway/n_ws_server.go
+++ b/internal/msggateway/n_ws_server.go
@@ -21,6 +21,7 @@ type LongConnServer interface {
 	GetUserAllCons(userID string) ([]*Client, bool)
 	GetUserPlatformCons(userID string, platform int) ([]*Client, bool, bool)
 	Validate(s interface{}) error
+	SetMessageHandler(rpcClient *notification.Check)
 	UnRegister(c *Client)
 	Compressor
 	Encoder