服务端建立websocket监听端口代码分析

272 阅读4分钟

先创建一个 Manager结构体,用来管理所有websocket连接的信息

type Manager struct {
	Group    map[string]map[string]*Client
	groupCount, clientCount uint
	Lock  sync.Mutex
	Register, UnRegister  chan *Client
	Message          chan *MessageData
	GroupMessage     chan *GroupMessageData
	BroadCastMessage  chan *BroadCastMessageData
}

Group: 用于存放分组的所有连接的map 其中key是分组name value是一个( key为连接id v是连接对象Client指针的map)

groupCount: 分组总数

clientCount: 连接总数

Register,UnRegister:两个传递内容为连接对象Client指针的chan 用于建立连接和断开连接的时候

Message: 传递单个信息的chan 对象为MessageData指针

GroupMessage: 传递同一个分组内信息的chan 对象为GruopMessage指针

BroadCastMessage: 广播信息的chan 对象为BroadCastMessage指针

连接对象结构体:
type Client struct {
	Id,Group string
	Socket *websockt.Conn
	Message chan []byte
}

Socket:具体的连接对象指针 这里引入了"github.com/gorilla/websocket"包 Message: 传递消息的channel

单个数据信息:
type MessageData struct {
	Id, Group string
	Message []byte
}
组广播信息:
type GroupMessageData struct {
	Group string
	Message []byte
}
广播发送数据信息
type BroadCastMessageData struct {
	Message []byte
}

启动监听的服务流程:

注册新的连接: for循环监听Register 通道 从注册的chan中取一个Client 加入到map中 因为map不是并发安全的,所以要加锁

注销也是一样的:从注销的chan中取出Client 先判断map中有没有这个client 再关闭client中的message Chan Manager的ClientCount-1 如果需要的话清除掉空的分组,GroupCount-1 同样要加锁 具体代码如下:

func (manager *Manager) Start() {
	log.Printf("websocket manage start")
	for {
		select {
		// 注册
		case client := <-manager.Register:
			log.Printf("client [%s] connect", client.Id)
			log.Printf("register client [%s] to group [%s]", client.Id, client.Group)

			manager.Lock.Lock()
			if manager.Group[client.Group] == nil {
				manager.Group[client.Group] = make(map[string]*Client)
				manager.groupCount += 1
			}
			manager.Group[client.Group][client.Id] = client
			manager.clientCount += 1
			manager.Lock.Unlock()

		// 注销
		case client := <-manager.UnRegister:
			log.Printf("unregister client [%s] from group [%s]", client.Id, client.Group)
			manager.Lock.Lock()
			if _, ok := manager.Group[client.Group]; ok {
				if _, ok := manager.Group[client.Group][client.Id]; ok {
					close(client.Message)
					delete(manager.Group[client.Group], client.Id)
					manager.clientCount -= 1
					if len(manager.Group[client.Group]) == 0 {
						delete(manager.Group, client.Group)
						manager.groupCount -= 1
					}
				}
			}
			manager.Lock.Unlock()
		}
	}
}

manager 处理单个Client发送数据

func (manager *Manager) sendService(){
	for {
		select {
			case data := <- manager.Message:
				if groupMap, ok := manager.Groupp[data.Group];ok{
					if conn, ok := groupMap[data.Id];ok{
						conn.Message <- data.Message
					}
				}
		}
	}
}

for 循环从manager的message Chan取出一个data

根据data中的group和id找到map中的连接对象

向连接对象的message Chan中放入 data中的message 是一个[]byte类型的数据

相应的发送方法:

// Send 向指定的 client 发送数据
func (manager *Manager) Send(id string, group string, message []byte) {
	data := &MessageData{
		Id:      id,
		Group:   group,
		Message: message,
	}
	manager.Message <- data
}

manager 处理广播的数据

func (manager *Manager) sendGroupService(){
	for {
		select {
		case data := <- manager.GroupMessage:
			if groupMap, ok := manager.Group[data.Group];ok{
				for _, conn := range groupMap {
					conn.Message <- data.Message
				}
			}
		}
	}
}

for 循环从manager的GroupMessage Channel中取出一个data

根据data的group从map中取到相应的map

for循环往map中的每个conn的message Channel放入data的message字段,是一个[]byte类型的数据

相应的发送方法:

// SendGroup 向指定的 Group 广播
func (manager *Manager) SendGroup(group string, message []byte) {
	data := &GroupMessageData{
		Group:   group,
		Message: message,
	}
	manager.GroupMessage <- data
}

处理广播数据

func (manager *Manager) SendAllService() {
	for {
		select {
		case data := <-manager.BroadCastMessage:
			for _, v := range manager.Group {
				for _, conn := range v {
					conn.Message <- data.Message
				}
			}
		}
	}
}

直接从manager的广播chan中取出数据

for循环manager的map中所有的conn 放入data的message

相应的发送方法:

// SendGroup 向指定的 Group 广播
func (manager *Manager) SendAll( message []byte) {
	data := &BroadCastMessageData{
		Message: message,
	}
	manager.BroadCastMessage <- data
}

注册函数
func (manager *Manager) RegisterClient(client *Client) {
	manager.Register <- client
}
注销函数
func (manager *Manager) UnRegisterClient(client *Client) {
	manager.UnRegister <- client
}

创建新的连接 WsClient gin 处理 websocket handler

func (manager *Manager) WsClient(ctx *gin.Context) {
	//将http请求升级处理
	upGrader := websocket.Upgrader{
		// cross origin domain
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		// 处理 Sec-WebSocket-Protocol Header
		Subprotocols: []string{ctx.GetHeader("Sec-WebSocket-Protocol")},
	}

	conn, err := upGrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		log.Printf("websocket connect error: %s", ctx.Param("channel"))
		return
	}

	//升级成功以后 创建一个客户端对象
	client := &Client{
		Id:      uuid.New().String(),
		Group:   ctx.Param("channel"),
		Socket:  conn,
		Message: make(chan []byte, 1024),
	}
	//将对象注册到manager中
	manager.RegisterClient(client)
	//开启协程负责读取客户端的消息
	go client.Read()
	//开启协程负责向客户端写消息
	go client.Write()
}

读Client的信息,从 websocket 连接直接读取数据

func (c *Client) Read() {
	defer func() {
		//协程退出之前,要将client对象放入manager的注销Channel
		//并且关闭client的socket连接
		WebsocketManager.UnRegister <- c
		log.Printf("client [%s] disconnect", c.Id)
		if err := c.Socket.Close(); err != nil {
			log.Printf("client [%s] disconnect err: %s", c.Id, err)
		}
	}()

	//for循环从client的socket中读取消息
	for {
		messageType, message, err := c.Socket.ReadMessage()
		if err != nil || messageType == websocket.CloseMessage {
			break
		}
		log.Printf("client [%s] receive message: %s", c.Id, string(message))
		msg := new(JsonMsg)
		_ = json.Unmarshal(message, &msg)
		log.Printf(msg.ID)
	}
}

向client写信息,从 channel 变量 Send 中读取数据写入 websocket 连接

func (c *Client) Write() {
	defer func() {
		//协程退出后关闭client的socket连接
		log.Printf("client [%s] disconnect", c.Id)
		if err := c.Socket.Close(); err != nil {
			log.Printf("client [%s] disconnect err: %s", c.Id, err)
		}
	}()
        
	//for循环阻塞client的message channel
	//取出[]byte结构,向client的socket里写入BinaryMessage类型消息
	//如果channel关闭
	//向client的socket写入CloseMessage类型消息
	for {
		select {
		case message, ok := <-c.Message:
			if !ok {
				_ = c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			log.Printf("client [%s] write message: %s", c.Id, string(message))
			err := c.Socket.WriteMessage(websocket.BinaryMessage, message)
			if err != nil {
				log.Printf("client [%s] writemessage err: %s", c.Id, err)
			}
		}
	}
}

利用gin接受clien的请求建立websocket连接


r := gin.Default()
wsGroup := r.Group("/ws"){
	wsGroup.GET("/:channel", ws.WebsocketManager.WsClient)
}
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
	log.Fatalf("Server Start Error: %s\n", err)
}