解决websocket官方没有检测客户端是否断开链接的问题(不用心跳检测,简单写法)

130 阅读2分钟

问题还原:

在使用websocket做消息推送的过程中遇到一个问题,如何检测websocket链接池中的链接是否已经断开,断开了就要移除这个链接,不然的话链接池里的链接越来越多,导致内存泄漏,迟早有一天内存溢出,服务就会挂掉,下面是我之前的写法

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}
var conns []*websocket.Conn
//客户端出发WS方法,则注册一个链接,加入到链接池conns中
func (this UserController) WS(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		println("upgrade错误:", err)
		return
	}
	defer c.Close()
	conns = append(conns, c)
	for {
		_, _, err := c.ReadMessage()
		if err != nil {
			println("read:", err)
			break
		}
	}
}
...
...
...
//下面的方法是循环给链接池里的链接推送消息
for i := range conns {
    if index := strings.Index(scene, "jianli"); index != -1 {
        conns[i].WriteMessage(websocket.TextMessage, []byte("jianli"))
    } else if index := strings.Index(scene, "music"); index != -1 {
        conns[i].WriteMessage(websocket.TextMessage, []byte("music"))
    }
}

以上代码有个很严重的漏洞,conns链接池中的元素一直在增加,而没有主动移除的,会导致内存溢出,我在网上搜索相关问题,大多数人只管链接,没有说移除链接的事,有一部分人使用心跳检测,客户端定时给服务端发送心跳数据,表示自己还活着,也是一种思路

我的解决方案

监听推送消息的方法,WriteMessage,该方法推送失败会返回error,出现error时,就移除这个链接,修改后的推送方法如下

func removeIndex(slice []*websocket.Conn, index int) []*websocket.Conn {
    return append(slice[:index], slice[index+1:]...)
}
...
...
...
for i := 0; i < len(conns); i++ {
    if index := strings.Index(scene, "jianli"); index != -1 {
        e := conns[i].WriteMessage(websocket.TextMessage, []byte("jianli"))
        if e != nil {
                conns = removeIndex(conns, i)
                i--
        }
    } else if index := strings.Index(scene, "music"); index != -1 {
        e := conns[i].WriteMessage(websocket.TextMessage, []byte("music"))
        if e != nil {
                conns = removeIndex(conns, i)
                i--
        }
    }
}

这样这次推送失败就移除了链接池中的该链接,下次循环推送就不再给它推送,不用心跳检测,简单解决问题,ok