先创建一个 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)
}