聊天室模块源码
消息队列
(这里用的是自己写的消息队列,也可以换成别的RabbitMQ或RocketMQ之类的)
消息推送模块结构图
基础消息模板
type Message struct {
Id int64
GroupId int64
UserId int64
MessageData string
CreateTime string
UpdateTime string
Status int64
}
import (
"encoding/json"
"github.com/dpwgc/kapokmq-go-client/conn"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"strconv"
"sync"
)
var clients = make(map[models.WsKey]*websocket.Conn)
var broadcast = make(chan models.UserMessage, 10000)
var lock sync.RWMutex
InitChat 聊天室模块初始化
这里开启两个协程。
pushMessages广播协程:持续接收用户的消息并广播推送给其他用户。
insertMessage数据库连接协程:持续接收消息队列中的消息并将消息插入到数据库
func InitChat() {
go pushMessages()
go insertMessage()
}
ChatLink 监听模块
建立WebSocket连接并监听用户发送的消息,将用户发送的消息插入广播通道与消息队列。
func ChatLink(c *gin.Context) {
userId := c.Param("userId")
groupId := c.Param("groupId")
gid, _ := strconv.ParseInt(groupId, 10, 64)
uid, _ := strconv.ParseInt(userId, 10, 64)
ws, err := config.UpGrader.Upgrade(c.Writer, c.Request, nil)
wsKey := models.WsKey{
GroupId: gid,
UserId: uid,
}
lock.RLock()
clients[wsKey] = ws
lock.RUnlock()
if err != nil {
fmt.Println(err)
delete(clients, wsKey)
return
}
defer ws.Close()
for {
_, message, err := ws.ReadMessage()
if err != nil {
fmt.Println(err)
delete(clients, wsKey)
break
}
broadcast <- userMessage
byteMsg, err := json.Marshal(userMessage)
if err == nil {
conn.ProducerSend(string(byteMsg), 0)
}
}
}
pushMessages 广播推送模块
广播推送消息到其他在线客户端,将用户消息推送给该群组(同一GroupId)里所有在线的用户。
func pushMessages() {
for {
msg := <-broadcast
for key, client := range clients {
gid := key.GroupId
if msg.GroupId == gid {
for i := 0; i < 3; i++ {
err := client.WriteJSON(msg)
if err == nil {
break
}
if i == 2 && err != nil {
client.Close()
delete(clients, key)
}
}
}
}
}
}
数据库插入模块
从消息队列中取出消息,将消息插入数据库。
func insertMessage() {
for {
message, _ := utils.GetMsg()
tx := mysql.Db().Begin()
tx.Commit()
}
}
消息队列连接模块
import (
"encoding/json"
"github.com/dpwgc/kapokmq-go-client/conf"
"github.com/dpwgc/kapokmq-go-client/conn"
"github.com/spf13/viper"
"local/models"
)
func InitMQ() {
producer := conf.Producer{
MqAddr: viper.GetString("kapokmq.mqAddr"),
MqPort: viper.GetString("kapokmq.mqPort"),
MqProtocol: viper.GetString("kapokmq.mqProtocol"),
Topic: viper.GetString("kapokmq.topic"),
ProducerId: viper.GetString("kapokmq.producerId"),
SecretKey: viper.GetString("kapokmq.secretKey"),
}
err := conn.NewProducerConn(producer)
if err != nil {
panic(err)
}
consumer := conf.Consumer{
MqAddr: viper.GetString("kapokmq.mqAddr"),
MqPort: viper.GetString("kapokmq.mqPort"),
MqProtocol: viper.GetString("kapokmq.mqProtocol"),
Topic: viper.GetString("kapokmq.topic"),
ConsumerId: viper.GetString("kapokmq.consumerId"),
SecretKey: viper.GetString("kapokmq.secretKey"),
}
err = conn.NewConsumerConn(consumer)
if err != nil {
panic(err)
}
}
func GetMsg() (models.UserMessage, bool) {
mqMsg, isOk := conn.ConsumerReceive()
userMessage := models.UserMessage{}
if isOk {
err := json.Unmarshal([]byte(mqMsg.MessageData), &userMessage)
if err != nil {
return userMessage, false
}
return userMessage, true
}
return userMessage, false
}