增加建立以及断开连接时,移除连接

construct.go

@@ -12,6 +12,8 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/go-developer/websocket/storage"
+
 	"github.com/go-developer/websocket/context"
 
 	"gopkg.in/olahol/melody.v1"
@@ -87,6 +89,9 @@ 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)
+				}
 				wsInstance.Connect(ctx)
 			})
 			// 2. 指令处理的函数
@@ -100,6 +105,9 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 			ginRouterTable[wsInstance.GetServerPort()].wsServer.HandleClose(func(session *melody.Session, i int, s string) error {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
+				defer func() {
+					storage.Connection.Del(ctx)
+				}()
 				return wsInstance.Close(ctx, i, s)
 			})
 			// 4. 断开连接的处理函数