Golang实现完整聊天室（内附源码）

项目github地址：
由于我们项目的需要，我就研究了一下关于websocket的相关内容，去实现一个聊天室的功能。
经过几天的探索，现在使用Gin框架实现了一个完整的聊天室+消息实时通知系统。有什么不完善的地方还请大佬指正。

用到的技术

websocket、gin、mysql、redis、协程、通道

实现思路

说到聊天室可以有多种方法实现，例如：使用单纯的MySQL也可以实现，但是为什么要选择使用websocket去实现呢？有什么优势呢？
websocket是基于TCP/IP，独立的HTTP协议的双向通信协议，这就使实时的消息通知成为可能，同时又符合Go高效处理高并发的语言特点，结合聊天室又是高并发的，所以采取的室websocket进行消息的转接，MySQL持久化聊天消息，redis用于做一些判断。
首先用户在进入App时，客户端和服务端建立一个websocket连接，并开启一个通道。
当服务端收到客户端的消息后，将消息写入通道里，服务端监听通道的消息，并将消息取出，使用接收人的websocket连接将消息广播到接收人那里。

实现代码

下面开始实现：
创建模型，用于关系的确立及数据的传输

//数据库存储消息结构体，用于持久化历史记录
type ChatMessage struct {gorm.ModelDirection   string //这条消息是从谁发给谁的SendID      int    //发送者idRecipientID int    //接受者idGroupID     string //群id，该消息要发到哪个群里面去Content     string //内容Read        bool   //是否读了这条消息
}//群聊结构体
type Group struct {ID           string ` gorm:"primaryKey"` //群idCreatedAt    time.TimeUpdatedAt    time.TimeDeletedAt    gorm.DeletedAt `gorm:"index"`GroupName    string         `json:"group_name"`    //群名GroupContent string         `json:"group_content"` //群签名GroupIcon    string         `json:"group_icon"`    //群头像GroupNum     int            //群人数GroupOwnerId int            //群主idUsers        []User         `gorm:"many2many:users_groups;"` //群成员
}type UsersGroup struct {GroupId string `json:"group_id"`UserId  int    `json:"user_id"`
}// 用于处理请求后返回一些数据
type ReplyMsg struct {From    string `json:"from"`Code    int    `json:"code"`Content string `json:"content"`
}// 发送消息的类型
type SendMsg struct {Type        int    `json:"type"`RecipientID int    `json:"recipient_id"` //接受者idContent     string `json:"content"`
}// 用户类
type Client struct {ID          string          //消息的去向RecipientID int             //接受者idSendID      int             //发送人的idGroupID     string          //群聊idSocket      *websocket.Conn //websocket连接对象Send        chan []byte     //发送消息用的管道
}// 广播类，包括广播内容和源用户
type Broadcast struct {Client  *ClientMessage []byteType    int
}// 用户管理,用于管理用户的连接及断开连接
type ClientManager struct {Clients    map[string]*ClientBroadcast  chan *BroadcastReply      chan *ClientRegister   chan *ClientUnregister chan *Client
}//创建一个用户管理对象
var Manager = ClientManager{Clients:    make(map[string]*Client), // 参与连接的用户，出于性能的考虑，需要设置最大连接数Broadcast:  make(chan *Broadcast),Register:   make(chan *Client), //新建立的连接访放入这里面Reply:      make(chan *Client),Unregister: make(chan *Client), //新断开的连接放入这里面
}

创建连接

func WsHandle(c *gin.Context) {myid := c.Query("myid")userid, err := strconv.Atoi(myid)if err != nil {zap.L().Error("转换失败", zap.Error(err))ResponseError(c, CodeParamError)}//将http协议升级为ws协议conn, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {return true}}).Upgrade(c.Writer, c.Request, nil)if err != nil {http.NotFound(c.Writer, c.Request)return}//创建一个用户客户端实例，用于记录该用户的连接信息client := new(model.Client)client = &model.Client{ID:     myid + "->",SendID: userid,Socket: conn,Send:   make(chan []byte),}//使用管道将实例注册到用户管理上model.Manager.Register <- client//开启两个协程用于读写消息go Read(client)go Write(client)
}//用于读管道中的数据
func Read(c *model.Client) {//结束把通道关闭defer func() {model.Manager.Unregister <- c//关闭连接_ = c.Socket.Close()}()for {//先测试一下连接能不能连上c.Socket.PongHandler()sendMsg := new(model.SendMsg)err := c.Socket.ReadJSON(sendMsg)c.RecipientID = sendMsg.RecipientIDif err != nil {zap.L().Error("数据格式不正确", zap.Error(err))model.Manager.Unregister <- c_ = c.Socket.Close()return}//根据要发送的消息类型去判断怎么处理//消息类型的后端调度switch sendMsg.Type {case 1: //私信SingleChat(c, sendMsg)case 2: //获取未读消息UnreadMessages(c)case 3: //拉取历史消息记录HistoryMsg(c, sendMsg)case 4: //群聊消息广播GroupChat(c, sendMsg)}}
}//用于将数据写进管道中
func Write(c *model.Client) {defer func() {_ = c.Socket.Close()}()for {select {//读取管道里面的信息case message, ok := <-c.Send://连接不到就返回消息if !ok {_ = c.Socket.WriteMessage(websocket.CloseMessage, []byte{})return}fmt.Println(c.ID+"接收消息：", string(message))replyMsg := model.ReplyMsg{Code:    int(CodeConnectionSuccess),Content: fmt.Sprintf("%s", string(message)),}msg, _ := json.Marshal(replyMsg)//将接收的消息发送到对应的websocket连接里rwLocker.Lock()_ = c.Socket.WriteMessage(websocket.TextMessage, msg)rwLocker.Unlock()}}
}

后端调度

//聊天的后端调度逻辑
//单聊
func SingleChat(c *model.Client, sendMsg *model.SendMsg) {//获取当前用户发出到固定用户的消息r1, _ := redis.REDIS.Get(context.Background(), c.ID).Result()//从redis中取出固定用户发给当前用户的消息id := CreateId(strconv.Itoa(c.RecipientID), strconv.Itoa(c.SendID))r2, _ := redis.REDIS.Get(context.Background(), id).Result()//根据redis的结果去做未关注聊天次数限制if r2 >= "3" && r1 == "" {ResponseWebSocket(c.Socket, CodeLimiteTimes, "未相互关注，限制聊天次数")return} else {//将消息写入redisredis.REDIS.Incr(context.Background(), c.ID)//设置消息的过期时间_, _ = redis.REDIS.Expire(context.Background(), c.ID, time.Hour*24*30*3).Result()}fmt.Println(c.ID+"发送消息：", sendMsg.Content)//将消息广播出去model.Manager.Broadcast <- &model.Broadcast{Client:  c,Message: []byte(sendMsg.Content),}
}//查看未读消息
func UnreadMessages(c *model.Client) {//获取数据库中的未读消息msgs, err := mysql.GetMessageUnread(c.SendID)if err != nil {ResponseWebSocket(c.Socket, CodeServerBusy, "服务繁忙")}for i, msg := range msgs {replyMsg := model.ReplyMsg{From:    msg.Direction,Content: msg.Content,}message, _ := json.Marshal(replyMsg)_ = c.Socket.WriteMessage(websocket.TextMessage, message)//发送完后将消息设为已读msgs[i].Read = trueerr := mysql.UpdateMessage(&msgs[i])if err != nil {ResponseWebSocket(c.Socket, CodeServerBusy, "服务繁忙")}}
}//拉取历史消息记录
func HistoryMsg(c *model.Client, sendMsg *model.SendMsg) {//拿到传过来的时间timeT := TimeStringToGoTime(sendMsg.Content)//查找聊天记录//做一个分页处理，一次查询十条数据,根据时间去限制次数//别人发给当前用户的direction := CreateId(strconv.Itoa(c.RecipientID), strconv.Itoa(c.SendID))//当前用户发出的id := CreateId(strconv.Itoa(c.SendID), strconv.Itoa(c.RecipientID))msgs, err := mysql.GetHistoryMsg(direction, id, timeT, 10)if err != nil {ResponseWebSocket(c.Socket, CodeServerBusy, "服务繁忙")}//把消息写给用户for _, msg := range *msgs {replyMsg := model.ReplyMsg{From:    msg.Direction,Content: msg.Content,}message, _ := json.Marshal(replyMsg)_ = c.Socket.WriteMessage(websocket.TextMessage, message)//发送完后将消息设为已读if err != nil {ResponseWebSocket(c.Socket, CodeServerBusy, "服务繁忙")}}
}//群聊消息广播
func GroupChat(c *model.Client, sendMsg *model.SendMsg) {//根据消息类型判断是否为群聊消息//先去数据库查询该群下的所有用户users, err := mysql.GetAllGroupUser(strconv.Itoa(sendMsg.RecipientID))if err != nil {ResponseWebSocket(c.Socket, CodeServerBusy, "服务繁忙")}//向群里面的用户广播消息for _, user := range users {//获取群里每个用户的连接if int(user.ID) == c.SendID {continue}c.ID = strconv.Itoa(c.SendID) + "->"c.GroupID = strconv.Itoa(sendMsg.RecipientID)c.RecipientID = int(user.ID)model.Manager.Broadcast <- &model.Broadcast{Client:  c,Message: []byte(sendMsg.Content),}}
}

转发消息

//用于在启动时进行监听
func Start(manager *model.ClientManager) {for {fmt.Println("<-----监听通信管道----->")select {//监测model.Manager.Register这个的变化，有新的东西加入管道时会被监听到，从而建立连接case conn := <-model.Manager.Register:fmt.Println("建立新连接:", conn.ID)//将新建立的连接加入到用户管理的map中，用于记录连接对象,以连接人的id为键，以连接对象为值model.Manager.Clients[conn.ID] = conn//返回成功信息controller.ResponseWebSocket(conn.Socket, controller.CodeConnectionSuccess, "已连接至服务器")//断开连接,监测到变化，有用户断开连接case conn := <-model.Manager.Unregister:fmt.Println("连接失败:", conn.ID)if _, ok := model.Manager.Clients[conn.ID]; ok {controller.ResponseWebSocket(conn.Socket, controller.CodeConnectionBreak, "连接已断开")}//关闭当前用户使用的管道//close(conn.Send)//删除用户管理中的已连接的用户delete(model.Manager.Clients, conn.ID)case broadcast := <-model.Manager.Broadcast: //广播消息message := broadcast.MessagerecipientID := broadcast.Client.RecipientID//给一个变量用于确定状态flag := falsecontentid := createId(strconv.Itoa(broadcast.Client.SendID), strconv.Itoa(recipientID))rID := strconv.Itoa(recipientID) + "->"//遍历客户端连接map,查找该用户有没有在线,判断的是对方的连接例如:1要向2发消息,我现在是用户1,那么我需要判断2->1是否存在在用户管理中for id, conn := range model.Manager.Clients {//如果找不到就说明用户不在线,与接收人的id比较if id != rID {continue}//走到这一步,就说明用户在线,就把消息放入管道里面select {case conn.Send <- message:flag = truedefault: //否则就把该连接从用户管理中删除close(conn.Send)delete(model.Manager.Clients, conn.ID)}}//判断完之后就把将消息发给用户if flag {fmt.Println("用户在线应答")controller.ResponseWebSocket(model.Manager.Clients[rID].Socket, controller.CodeConnectionSuccess, string(message))//把消息插到数据库中msg := model.ChatMessage{Direction:   contentid,SendID:      broadcast.Client.SendID,RecipientID: recipientID,GroupID:     broadcast.Client.GroupID,Content:     string(message),Read:        true,}err := mysql.DB.Create(&msg).Errorif err != nil {zap.L().Error("在线发送消息出现了错误", zap.Error(err))}} else { //如果不在线controller.ResponseWebSocket(broadcast.Client.Socket, controller.CodeConnectionSuccess, "对方不在线")//把消息插到数据库中msg := model.ChatMessage{Direction:   contentid,SendID:      broadcast.Client.SendID,RecipientID: recipientID,GroupID:     broadcast.Client.GroupID,Content:     string(message),Read:        false,}err := mysql.DB.Create(&msg).Errorif err != nil {zap.L().Error("不在线发送消息出现了错误", zap.Error(err))}}}}}func createId(uid, toUid string) string {return uid + "->" + toUid
}