优化建立连接

abstract/IWebsocket.go

@@ -22,7 +22,7 @@ type IWebsocket interface {
 	// Author : go_developer@163.com<张德满>
 	//
 	// Date : 6:39 下午 2021/3/27
-	Connect(ctx *context.WSContext)
+	Connect(ctx *context.WSContext) error
 	// Disconnect 断开连接时处理的方法
 	//
 	// Author : go_developer@163.com<张德满>

construct.go

@@ -89,10 +89,11 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
 				ctx.Session = session
-				if nil != storage.Connection {
-					storage.Connection.Store(ctx)
+				if err := wsInstance.Connect(ctx); nil == err {
+					if nil != storage.Connection {
+						storage.Connection.Store(ctx)
+					}
 				}
-				wsInstance.Connect(ctx)
 			})
 			// 2. 指令处理的函数
 			ginRouterTable[wsInstance.GetServerPort()].wsServer.HandleMessage(func(session *melody.Session, bytes []byte) {
@@ -106,6 +107,9 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
 				defer func() {
+					if nil == storage.Connection {
+						return
+					}
 					storage.Connection.Del(ctx)
 				}()
 				return wsInstance.Close(ctx, i, s)
@@ -114,6 +118,12 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 			ginRouterTable[wsInstance.GetServerPort()].wsServer.HandleDisconnect(func(session *melody.Session) {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
+				defer func() {
+					if nil == storage.Connection {
+						return
+					}
+					storage.Connection.Del(ctx)
+				}()
 				wsInstance.Disconnect(ctx)
 			})
 			// 注册指令

example/server.go

@@ -24,9 +24,10 @@ func main() {
 type Example struct {
 }
 
-func (e Example) Connect(ctx *context.WSContext) {
+func (e Example) Connect(ctx *context.WSContext) error {
 	fmt.Println("建立连接成功")
 	message.Response(ctx, map[string]interface{}{"say": "hello world!", "cid": ctx.ConnectionID})
+	return nil
 }
 
 func (e Example) Disconnect(ctx *context.WSContext) {