diff --git a/abstract/IWebsocket.go b/abstract/IWebsocket.go index a1a4640..053c6a2 100644 --- a/abstract/IWebsocket.go +++ b/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<张德满> diff --git a/construct.go b/construct.go index fdae93d..c7f3880 100644 --- a/construct.go +++ b/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) }) // 注册指令 diff --git a/example/server.go b/example/server.go index 4e0869e..dc31a50 100644 --- a/example/server.go +++ b/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) {