Go WebSocket 的使用

HTTP

先使用 net/http 标准库实现一个 http 服务：

func main() {
	http.HandleFunc("/", home)
	err := http.ListenAndServe("localhost:18000", nil)
	if err != nil {
		panic(err)
	}
}

func home(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("hello"))
}

执行 go run main.go ，在浏览器访问，就可以看到我们打印的 hello 辣。

WebSocket

gorilla/websocket 利用 Go 标准库实现了 ws，我们使用它来快速启用一个 ws 服务，修改我们的代码成这样：

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

func main() {
	http.HandleFunc("/", home)
	// 定义ws路径
	http.HandleFunc("/ws", ws)
	err := http.ListenAndServe("localhost:18000", nil)
	if err != nil {
		panic(err)
	}
}

func home(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("hello"))
}

func ws(w http.ResponseWriter, r *http.Request) {
	// 将http服务升级成ws服务
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		// 监听消息
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)
		// 监听到信息，向客户端响应
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
}

执行 go run main.go ，打开 ws://localhost:18000/ws，发送一条消息，就会原样返回消息。

广播模式

我们再启用一个 http 路由，当访问它时，像所有已经打开 ws 连接发送一条信息。这就需要使用一个 map 记录下所有打开的 ws 客户端连接，当 ws 连接关闭时在从其中去除。先定义一个全局变量，用来存放客户端连接：

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

var clients = make(map[*websocket.Conn]struct{})

...

改造一下 func ws，加入我们需要的业务逻辑：

func ws(w http.ResponseWriter, r *http.Request) {
	// 将http服务升级成ws服务
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	// 建立连接将客户端添加到map中
	clients[c] = struct{}{}
	defer c.Close()
	for {
		// 监听消息
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)
		// 监听到信息，向客户端响应
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
	// 关闭连接将客户端从map中剔除
	delete(clients, c)
}

新增一个路由，当访问它时，遍历 clients 进行广播：

func broadcast(w http.ResponseWriter, r *http.Request) {
	for client, _ := range clients {
		client.WriteMessage(websocket.TextMessage, []byte("能听到我的广播吗"))
	}
}

ok，最后我们的代码结构是这样的：

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

// 存放client
var clients = make(map[*websocket.Conn]struct{})

func main() {
	http.HandleFunc("/", home)
	// 定义ws路径
	http.HandleFunc("/ws", ws)
	// 广播
	http.HandleFunc("/broadcast", broadcast)
	err := http.ListenAndServe("localhost:18000", nil)
	if err != nil {
		panic(err)
	}
}

func home(w http.ResponseWriter, r *http.Request) {}

func broadcast(w http.ResponseWriter, r *http.Request) {}

func ws(w http.ResponseWriter, r *http.Request) {}

这里要特别注意，这里的 map 不是并发安全的，关于并发安全，可以查看《Go 并发编程 - 并发安全（二）》