客户端切换网络下WebSocket重连的后台实现

来源:互联网 发布:kmp算法 july 编辑:程序博客网 时间:2024/05/22 12:49

客户端切换网络下WebSocket重连的后台实现

客户端断网重连(或者切换网络)会重新访问 CommandListenHandler接口,但服务端的WebSocket并未关闭,直接重启会导致数据表混乱,所以以下代码实现了控制原WebSocket的目的。

type connPoolItem struct {    Id       uint64    Ch       chan *apimodel.Command    Conn     *websocket.Conn    IsOk     bool //判断Conn是否被关闭,Conn关闭会触发前一个Listen关闭    IsClosed bool //判断前一个Listen是否被关闭}var connPool = map[uint64]*connPoolItem{}var poolMutex = &sync.Mutex{}var upgrader = websocket.Upgrader{} // use default optionsfunc CommandListenHandler(c *gin.Context) {    type param struct {        InstallationId uint64 `form:"installation_id" binding:"required"`    }    var p param    if err := c.Bind(&p); err != nil {        logger.Error("Invalid command listen param ", err)        c.AbortWithStatus(http.StatusBadRequest)        return    }    db := c.MustGet(constant.ContextDb).(*gorm.DB)    var device model.Installation // 获取设备信息    if err := db.Where("id = ?", p.InstallationId).First(&device).Error; err != nil {        logger.Error("Device not found", err)        c.JSON(http.StatusOK, gin.H{"err_code": constant.DeviceNotRegistered, "err_msg": constant.TranslateErrCode(constant.DeviceNotRegistered)})        return    }    logger.Debug("Command Listen begining...")    poolMutex.Lock()    item := connPool[device.Id]    poolMutex.Unlock()    if item != nil {        logger.Debug("find an old conn!will close it.")        item.IsOk = false        item.Conn.Close()        for {            if item.IsClosed {                break            }            time.Sleep(time.Millisecond * 20)        }    }    // 升级到WebSocket模式    conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)    if err != nil {        logger.Error("upgrade:", err)        c.AbortWithStatus(http.StatusInternalServerError)        return    }    logger.Debug("Websocket connected, device: ", device.DeviceId)    // 更新设备的连接服务器IP为本机IP    ip := util.GetLocalIp()    if ip == "" {        logger.Error("Server IP not found!")        c.AbortWithStatus(http.StatusInternalServerError)        return    }    device.ServerIp = ip    if err = db.Save(&device).Error; err != nil {        logger.Error(err)    }    // 定义命令数据通道    ch := make(chan *apimodel.Command)    addChannelPool(device.Id, ch, conn) // 将该连接专属命令通到放入通道池    defer closeChannel(device.Id, db)   // 关闭websocket连接并从通道池中删除    // 监听连接状态,如果连接失效则显式关闭ch    go func() {        for {            mt, message, err := conn.ReadMessage()            if err != nil {                close(ch)                logger.Error("Device websocket closed: ", err)                break            }            // 处理接收到的命令响应            logger.Debug("Websocket message received: ", string(message))            if mt == websocket.TextMessage {                var res model.CommandResponse                if err = json.Unmarshal(message, &res); err != nil {                    logger.Error(err)                    continue                }                // 修改数据库中指令的状态                if err := db.Exec("update command set status = ?, executed_at = ? where id = ?", res.Status, time.Unix(res.Timestamp, 0), res.Id).Error; err != nil {                    logger.Error(err)                }            }        }    }()    // 逐条读取channel里面的数据并发送指令到大屏端,直到ch被显式关闭    for cm := range ch {        if msg, err := json.Marshal(cm); err == nil && len(msg) > 0 {            logger.Debug("Prepare send command: ", string(msg))            if err = conn.WriteMessage(websocket.TextMessage, msg); err != nil {                logger.Error("Command write failed:", err)            }        }    }    logger.Debug("Already break the websocket!")}// 将设备指令通道放入通道池func addChannelPool(id uint64, ch chan *apimodel.Command, conn *websocket.Conn) {    c := connPoolItem{}    c.Id = id    c.Ch = ch    c.Conn = conn    c.IsOk = true    c.IsClosed = false    connPool[id] = &c    logger.Debug("Channel Pool size: ", len(connPool))}// 关闭设备指令通道并从通道池中删除func closeChannel(id uint64, db *gorm.DB) {    poolMutex.Lock()    item := connPool[id]    poolMutex.Unlock()    //Conn有没有被新Listen关闭    if item.IsOk {        item.Conn.Close()    }    delete(connPool, id) // 删除连接    // 置空设备连接服务器IP    if err := db.Exec("update installation set server_ip = '' where id = ?", id).Error; err != nil {        logger.Error(err)    }    //该Listen关闭后新的Listen才会启动    item.IsClosed = true}