第三章节: WebSocketServer连接管理

前言

上一章节中我们在Server端通过HttpHeader中携带的Token识别了每个客户端连接，本章节我们将用户和WebSocket连接一一绑定并统一管理。

连接管理器设计

因此我们设计一个ConnManager，它需要具备以下功能：

需要维护一个Map，Key为UserID, Value为WebSocketConnection，表示用户与连接的对应关系
需要有方法AddConn(UserID, WebSocketConnection)来添加新的用户和连接
需要通过UserID来获取WebSocketConnection
需要有移除WebSocketConnection的方法
需要考虑并发下的线程安全

在go-im/src目录下新建文件夹conn, 在conn下新建conn.go

go 复制代码

package conn

import (
    "fmt"
    "sync"

    "aoki.com/go-im/src/model"
    "github.com/gorilla/websocket"
)

type ConnManager struct {
    // 考虑线程安全
    connections sync.Map
}

var m *ConnManager = &ConnManager{
    connections: sync.Map{},
}

// 保证单例
func GetConnManager() *ConnManager {
    return m
}

// 根据UserID获取连接
func (m *ConnManager) FindConn(userID int64) *websocket.Conn {
     conn, has := m.connections.Load(userID)
     if has {
        return conn.(*websocket.Conn)
     }
     return nil
}

// 添加新的连接
func (m *ConnManager) AddConn(userID int64, conn *websocket.Conn) {
    m.connections.Store(userID, conn)
}

// 移除连接
func (m *ConnManager) DelConn(userID int64) {
    m.connections.Delete(userID)
}

处理连接逻辑

在有新连接时需要通过ConnManager来添加新绑定，当有连接断开时需要通过ConnManager来释放绑定

修改server.go

go 复制代码

package main

import (
    "log"
    "net/http"

    "aoki.com/go-im/src/conn"
    "aoki.com/go-im/src/model"
    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    // Allow Cross Origin
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

var ConnManager *conn.ConnManager = conn.GetConnManager()

func ws(c *gin.Context) {
    conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
    if err != nil {
        panic(err)
    }
    defer conn.Close()
    user := model.ResolveUser(c.Request)
    if user == nil {
        log.Println("ResolveUser failed...")
        return
    }
    // 添加新的连接绑定
    ConnManager.AddConn(user.ID, conn)
    log.Printf("%s connected... \n", user.Name)
    // 方法退出时移除连接绑定
    defer OnDisconnect(*user)
    for {
        mt, message, err := conn.ReadMessage()
        if err != nil {
            log.Printf("Read Message Failed... \n", err)
            break
        }
        // 读到什么往WS客户端发什么
        log.Printf("Received Message %s... \n", message)
        err = conn.WriteMessage(mt, message)
        if err != nil {
            log.Printf("Write Message Failed... \n", err)
            break
        }
    }

}

func OnDisconnect(user model.User) {
    log.Printf("%v disconnect...\n", user.Name)
    ConnManager.DelConn(user.ID)
}

func main() {
    server := gin.Default()
    server.GET("/ws", ws)
    server.Run("localhost:8848")
}

调试

分别在两个Terminal启动server和client，然后关闭client进程，可以在Server端的Terminal中看到以下日志输出，说明调试结果符合期望

小结

本章节中我们成功将每个WebSocket连接和用户进行绑定，并通过ConnManager统一管理，这样已经为客户端之间的相互通信做好了铺垫，下一章节我们将实现用户A和用户B之间进行一对一聊天