优化websocket

network/ws_server.go

@@ -21,7 +21,7 @@ type WSServer struct {
 	NewAgent        func(*WSConn) Agent
 	ln              net.Listener
 	handler         *WSHandler
-	messageType 	int
+	messageType     int
 }
 
 type WSHandler struct {
@@ -33,10 +33,10 @@ type WSHandler struct {
 	conns           WebsocketConnSet
 	mutexConns      sync.Mutex
 	wg              sync.WaitGroup
-	messageType 	int
+	messageType     int
 }
 
-func (handler *WSHandler) SetMessageType(messageType int){
+func (handler *WSHandler) SetMessageType(messageType int) {
 	handler.messageType = messageType
 }
 
@@ -51,7 +51,9 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	conn.SetReadLimit(int64(handler.maxMsgLen))
-	handler.messageType = websocket.TextMessage
+	if handler.messageType == 0 {
+		handler.messageType = websocket.TextMessage
+	}
 
 	handler.wg.Add(1)
 	defer handler.wg.Done()
@@ -71,7 +73,7 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	handler.conns[conn] = struct{}{}
 	handler.mutexConns.Unlock()
 
-	wsConn := newWSConn(conn, handler.pendingWriteNum, handler.maxMsgLen,handler.messageType)
+	wsConn := newWSConn(conn, handler.pendingWriteNum, handler.maxMsgLen, handler.messageType)
 	agent := handler.newAgent(wsConn)
 	agent.Run()
 
@@ -83,9 +85,9 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	agent.OnClose()
 }
 
-func (server *WSServer) SetMessageType(messageType int){
+func (server *WSServer) SetMessageType(messageType int) {
 	server.messageType = messageType
-	if server.handler!= nil {
+	if server.handler != nil {
 		server.handler.SetMessageType(messageType)
 	}
 }