使用Gin框架搭配WebSocket完成实时聊天

文章目录

- 前言
- 实时聊天
- - 聊天功能
  - 测试
  - 发送信息

前言

在写项目的时候，需要完成实时聊天的功能，于是简单的学习下WebSocket，想知道WebSocket是什么的小伙伴可以去网上别的地方学习一下。

要实现实时聊天，网上的大部分内容都是SpringBoot和WebSocket完成的，但是我使用Go写的，所以让我们来学习一下Gin框架搭配WebSocket完成实时聊天功能。

实时聊天

实时聊天的数据可以存到数据库，redis中，在这个小demo中没写相应的代码

聊天功能

package mainimport ("encoding/json""log""net/http""sync""github.com/gorilla/websocket"
)var (// 使用 sync.Map 来保存 WebSocket 连接的映射// key 是 userID，value 是 WebSocket 连接webSocketMap sync.Map// 记录在线人数onlineCount int
)type Message struct {FromUserID string `json:"fromUserId"`ToUserID   string `json:"toUserId"`// 其他字段...
}func handleWebSocket(w http.ResponseWriter, r *http.Request) {upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true },}// 从 URL 中获取 userId// userId 为发送者// toUserId 为接收者userId := r.URL.Query().Get("userId")toUserId := r.URL.Query().Get("toUserId")// 升级 HTTP 连接为 WebSocket 连接conn, err := upgrader.Upgrade(w, r, nil)if err != nil {log.Println("WebSocket upgrade error:", err)return}// 注册 WebSocket 连接webSocketMap.Store(userId, conn)addOnlineCount()log.Println("用户连接:", userId, ",当前在线人数为:", getOnlineCount())// 发送消息sendMessage(userId, "连接成功")// 接收消息并处理for {_, message, err := conn.ReadMessage()if err != nil {// 关闭连接并从映射中删除conn.Close()webSocketMap.Delete(userId)subOnlineCount()log.Println("用户退出:", userId, ",当前在线人数为:", getOnlineCount())break}log.Println("用户消息:", userId, ",报文:", string(message))// 解析 JSONvar msg Messageif err := json.Unmarshal(message, &msg); err != nil {log.Println("Error parsing JSON:", err)continue}// 追加发送人(防止串改)msg.FromUserID = userIdmsg.ToUserID = toUserId// 传送给对应 ToUserID 用户的 WebSocketif msg.ToUserID != "" {sendMessage(msg.ToUserID, "收到来自用户 "+userId+" 的消息："+string(message))} else {sendMessage(userId, "消息发送失败，目标用户不存在")}}
}func sendMessage(userID, message string) {if conn, ok := webSocketMap.Load(userID); ok {conn.(*websocket.Conn).WriteMessage(websocket.TextMessage, []byte(message))}
}func getOnlineCount() int {return onlineCount
}func addOnlineCount() {onlineCount++
}func subOnlineCount() {onlineCount--
}func main() {http.HandleFunc("/ws", handleWebSocket)http.ListenAndServe(":8080", nil)
}