我們在設計產品的時候通常都會遇到消息通知的時候，比如用戶下單支付成功，比如用戶有站內信來可以實時通知。而http是單向的，客戶端請求，服務端返回，這次請求就已經結束。而websocket可以保持連接實現長連接，在遇到通知時往往使用websocket來達到服務端主動向客戶端發送消息的目的。
我們的目標是實現服務端主動向某個用戶發消息。所以要做到一下四步。
- 建立連接(保持連接)
- 斷開連接(刪除連接)
- 維護連接(心跳檢測)
- 接收消息
- 發送消息
我們這里使用到 github.com/gorilla/websocket 包。
准備工作
首先是定義一個客戶端連接的結構，先有了連接的結構才能保存連接，ID是一個客戶端連接的id，而Socket是真正的客戶端連接
// 客戶端連接信息
type Client struct {
ID string // 連接ID
AccountId string // 賬號id, 一個賬號可能有多個連接
Socket *websocket.Conn // 連接
HeartbeatTime int64 // 前一次心跳時間
}
然后定義一個客戶端管理，來管理所有的客戶端連接,並且實例化為一個全局的變量。
// 消息類型
const (
MessageTypeHeartbeat = "heartbeat" // 心跳
MessageTypeRegister = "register" // 注冊
HeartbeatCheckTime = 9 // 心跳檢測幾秒檢測一次
HeartbeatTime = 20 // 心跳距離上一次的最大時間
ChanBufferRegister = 100 // 注冊chan緩沖
ChanBufferUnregister = 100 // 注銷chan大小
)
// 客戶端管理
type ClientManager struct {
Clients map[string]*Client // 保存連接
Accounts map[string][]string // 賬號和連接關系,map的key是賬號id即：AccountId，這里主要考慮到一個賬號多個連接
mu *sync.Mutex
}
// 定義一個管理Manager
var Manager = ClientManager{
Clients: make(map[string]*Client), // 參與連接的用戶，出於性能的考慮，需要設置最大連接數
Accounts: make(map[string][]string), // 賬號和連接關系
mu: new(sync.Mutex),
}
var (
RegisterChan = make(chan *Client, ChanBufferRegister) // 注冊
unregisterChan = make(chan *Client, ChanBufferUnregister) // 注銷
)
這里還要封裝一下服務器給客戶端發消息的格式，以便客戶端連接成功后服務端給客戶端回復消息
// 封裝回復消息
type ServiceMessage struct {
Type string `json:"type"` // 類型
Content ServiceMessageContent `json:"content"`
}
type ServiceMessageContent struct {
Body string `json:"body"` // 主要數據
MetaData string `json:"meta_data"` // 擴展數據
}
func CreateReplyMsg(t string, content ServiceMessageContent) []byte {
replyMsg := ServiceMessage{
Type: t,
Content: content,
}
msg, _ := json.Marshal(replyMsg)
return msg
}
管理連接
建立連接和斷開連接
連接保持在Manager里的Clients,和Accounts。Clients用於保存每個與客戶端通信的連接。而Account保持者連接id與連接分類(Category)的綁定關系。
// 注冊注銷
func register() {
for {
select {
case conn := <-RegisterChan: // 新注冊，新連接
// 加入連接,進行管理
accountBind(conn)
// 回復消息
content := CreateReplyMsg(MessageTypeRegister, ServiceMessageContent{})
_ = conn.Socket.WriteMessage(websocket.TextMessage, content)
case conn := <-unregisterChan: // 注銷，或者沒有心跳
// 關閉連接
_ = conn.Socket.Close()
// 刪除Client
unAccountBind(conn)
}
}
}
// 綁定賬號
func accountBind(c *Client) {
Manager.mu.Lock()
defer Manager.mu.Unlock()
// 加入到連接
Manager.Clients[c.ID] = c
// 加入到綁定
if _, ok := Manager.Accounts[c.AccountId]; ok { // 該賬號已經有綁定，就追加一個綁定
Manager.Accounts[c.AccountId] = append(Manager.Accounts[c.AccountId], c.ID)
} else { // 沒有就新增一個賬號的綁定切片
Manager.Accounts[c.AccountId] = []string{c.ID}
}
}
// 解綁賬號
func unAccountBind(c *Client) {
Manager.mu.Lock()
defer Manager.mu.Unlock()
// 取消連接
delete(Manager.Clients, c.ID)
// 取消綁定
if len(Manager.Accounts[c.AccountId]) > 0 {
for k, clientId := range Manager.Accounts[c.AccountId] {
if clientId == c.ID { // 找到綁定客戶端Id
Manager.Accounts[c.AccountId] = append(Manager.Accounts[c.AccountId][:k], Manager.Accounts[c.AccountId][k+1:]...)
}
}
}
}
維護連接(心跳檢測)
每隔一段時間，就檢測一次心跳，如果上次心跳時間超過了HeartbeatTime時間視為已經斷開連接。
// 維持心跳
func heartbeat() {
for {
// 獲取所有的Clients
Manager.mu.Lock()
clients := make([]*Client, len(Manager.Clients))
for _, c := range Manager.Clients {
clients = append(clients, c)
}
Manager.mu.Unlock()
for _, c := range clients {
if time.Now().Unix()-c.HeartbeatTime > HeartbeatTime {
unAccountBind(c)
}
}
time.Sleep(time.Second * HeartbeatCheckTime)
}
}
管理連接
// 管理連接
func Start() {
// 檢查心跳
go func() {
defer func() {
if r := recover(); r != nil {
log.Println(r)
}
}()
heartbeat()
}()
// 注冊注銷
go func() {
defer func() {
if r := recover(); r != nil {
log.Println(r)
}
}()
register()
}()
}
收發消息
根據賬號獲取連接
// 根據賬號獲取連接
func GetClient (accountId string) []*Client{
clients := make([]*Client,0)
Manager.mu.Lock()
defer Manager.mu.Unlock()
if len(Manager.Accounts[accountId]) > 0 {
for _,clientId := range Manager.Accounts[accountId] {
if c,ok := Manager.Clients[clientId]; ok {
clients = append(clients,c)
}
}
}
return clients
}
讀取客戶端的消息
我們這是只是心跳用到了，所以只要判斷客戶端是心跳消息，然后回復即可。
// 讀取信息，即收到消息
func (c *Client) Read() {
defer func() {
_ = c.Socket.Close()
}()
for {
// 讀取消息
_, body, err := c.Socket.ReadMessage()
if err != nil {
break
}
var msg struct {
Type string `json:"type"`
}
err = json.Unmarshal(body, &msg)
if err != nil {
log.Println(err)
continue
}
if msg.Type == MessageTypeHeartbeat { // 維持心跳消息
// 刷新連接時間
c.HeartbeatTime = time.Now().Unix()
// 回復心跳
replyMsg := CreateReplyMsg(MessageTypeHeartbeat, ServiceMessageContent{})
err = c.Socket.WriteMessage(websocket.TextMessage, replyMsg)
if err != nil {
log.Println(err)
}
continue
}
}
}
對客戶端發送消息
只要找到連接，對連接發送消息即可。
// 發送消息
func Send(accounts []string,message ServiceMessage) error{
msg,err := json.Marshal(message)
if err != nil {
return err
}
for _,accountId := range accounts{
// 獲取連接id
clients := GetClient(accountId)
// 發送消息
for _,c := range clients {
_ = c.Socket.WriteMessage(websocket.TextMessage, msg)
}
}
return nil
}
請求調用
這里對http請求升級為websocket,然后單獨建立一個goroutine去維持連接。下面類似這樣調用，但是鑒權，日志等很多細節不完善，只是提供一個思路。
package wesocket
import (
websocket2 "demo/websocket"
"fmt"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"github.com/rs/xid"
"log"
"net/http"
"time"
)
type MessageNotifyRequest struct {
UserId string `form:"user_id"`
}
func MessageNotify(ctx *gin.Context) {
// 獲取參數
var params MessageNotifyRequest
if err := ctx.ShouldBindQuery(¶ms); err != nil {
log.Println(err)
return
}
// TODO: 鑒權
// 將http升級為websocket
conn, err := (&websocket.Upgrader{
// 1. 解決跨域問題
CheckOrigin: func(r *http.Request) bool {
return true
},
}).Upgrade(ctx.Writer, ctx.Request, nil) // 升級
if err != nil {
log.Println(err)
http.NotFound(ctx.Writer, ctx.Request)
return
}
// 創建一個實例連接
ConnId := xid.New().String()
client := &websocket2.Client{
ID: ConnId, // 連接id
AccountId: fmt.Sprintf("%s", params.UserId),
HeartbeatTime: time.Now().Unix(),
Socket: conn,
}
// 用戶注冊到用戶連接管理
websocket2.RegisterChan <- client
// 讀取信息
go func() {
defer func() {
if r := recover(); r != nil {
log.Printf("MessageNotify read panic: %+v\n",r)
}
}()
client.Read()
}()
}
總結
用websocket做消息通知，對於后端來說，主要是綁定連接和管理連接，綁定連接就是用戶id和websocket連接建立一個綁定關系，而管理連接就是存儲連接，刪除連接，維護連接的健康(心跳檢測)，其次就是定義服務端接收和發送數據的格式。總體大概就是這樣一個思路。