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

客户端切换网络下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 options

func 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
}