我的代码:imooc-go-websocket
参考代码:go-websocket
项目笔记:GO实现千万级WebSocket消息推送服务-项目笔记
package main
import "net/http"
// 请求处理器
func wsHandler(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("hello"))
}
func main() {
// 处理HTTP请求/ws
http.HandleFunc("/ws", wsHandler)
// 监听本地7777端口,并处理所有HTTP请求
http.ListenAndServe("0.0.0.0:7777", nil)
}
注意学习这里goto语法的用法及WebSocket协议的通信流程
package main
import (
"github.com/gorilla/websocket"
"net/http"
)
var (
upgrader = websocket.Upgrader{
// 允许跨域
CheckOrigin: func(r *http.Request) bool {
return true
},
}
)
// 请求处理器
func wsHandler(w http.ResponseWriter, r *http.Request) {
var (
conn *websocket.Conn
err error
data []byte
)
// 完成WebSocket握手,建立长连接
// Upgrade: websocket
if conn, err = upgrader.Upgrade(w, r, nil); err != nil {
return
}
// 读取通过WebSocket协议接收到的数据并写回浏览器
// websocket.Conn
for {
// Text, Binary
if _, data, err = conn.ReadMessage(); err != nil {
// 注意学习这里goto语法的用法
goto ERR
}
// 写回WebSocket协议接收到的数据
if err = conn.WriteMessage(websocket.TextMessage, data); err != nil {
goto ERR
}
}
ERR:
conn.Close()
}
func main() {
// 处理HTTP请求/ws
http.HandleFunc("/ws", wsHandler)
// 监听本地7777端口,并处理所有HTTP请求
http.ListenAndServe("0.0.0.0:7777", nil)
}
client.html
<html>
<head>
<meta charset="utf-8">
<script>
window.addEventListener("load", function(evt) {
var output = document.getElementById("output");
var input = document.getElementById("input");
var ws;
var print = function(message) {
var d = document.createElement("div");
d.innerHTML = message;
output.appendChild(d);
};
document.getElementById("open").onclick = function(evt) {
if (ws) {
return false;
}
ws = new WebSocket("ws://localhost:7777/ws");
ws.onopen = function(evt) {
print("OPEN");
}
ws.onclose = function(evt) {
print("CLOSE");
ws = null;
}
ws.onmessage = function(evt) {
print("RESPONSE: " + evt.data);
}
ws.onerror = function(evt) {
print("ERROR: " + evt.data);
}
return false;
};
document.getElementById("send").onclick = function(evt) {
if (!ws) {
return false;
}
print("SEND: " + input.value);
ws.send(input.value);
return false;
};
document.getElementById("close").onclick = function(evt) {
if (!ws) {
return false;
}
ws.close();
return false;
};
});
script>
head>
<body>
<table>
<tr><td valign="top" width="50%">
<p>Click "Open" to create a connection to the server,
"Send" to send a message to the server and "Close" to close the connection.
You can change the message and send multiple times.
p>
<form>
<button id="open">Openbutton>
<button id="close">Closebutton>
<input id="input" type="text" value="Hello world!">
<button id="send">Sendbutton>
form>
td><td valign="top" width="50%">
<div id="output">div>
td>tr>table>
body>
html>
connection.go
package impl
import (
"errors"
"github.com/gorilla/websocket"
"sync"
)
type Connection struct {
wsConn *websocket.Conn
inChan chan []byte
outChan chan []byte
closeChan chan byte
mutex sync.Mutex
isClosed bool
}
func InitConnection(wsConn *websocket.Conn) (conn *Connection, err error) {
conn = &Connection{
wsConn: wsConn,
inChan: make(chan []byte, 1000),
outChan: make(chan []byte, 1000),
closeChan: make(chan byte, 1),
}
// 启动读协程
go conn.readLoop()
// 启动写协程
go conn.writeLoop()
return
}
///////////////////////////////////////////////// API /////////////////////////////////////////////////
func (conn *Connection) ReadMessage() (data []byte, err error) {
select {
case data = <- conn.inChan:
case <- conn.closeChan:
err = errors.New("connection is closed")
}
return
}
func (conn *Connection) WriteMessage(data []byte) (err error) {
select {
case conn.outChan <- data:
case <- conn.closeChan:
err = errors.New("connection is closed")
}
return
}
func (conn *Connection) Close() {
// 线程安全,可重入的Close
conn.wsConn.Close()
// 这一行代码只执行一次
conn.mutex.Lock()
if !conn.isClosed {
close(conn.closeChan)
conn.isClosed = true
}
conn.mutex.Unlock()
}
///////////////////////////////////////////////// 内部实现 ////////////////////////////////////////////
func (conn *Connection) readLoop() {
var (
data []byte
err error
)
for {
if _, data, err = conn.wsConn.ReadMessage(); err != nil {
goto ERR
}
// 阻塞在这里,等待inChan有空闲的位置
select {
case conn.inChan <- data:
case <- conn.closeChan:
// closeChan关闭的时候
goto ERR
}
}
ERR:
conn.Close()
}
func (conn *Connection) writeLoop() {
var (
data []byte
err error
)
for {
select {
case data = <- conn.outChan:
case <- conn.closeChan:
goto ERR
}
data = <- conn.outChan
if err = conn.wsConn.WriteMessage(websocket.TextMessage, data); err != nil {
goto ERR
}
}
ERR:
conn.Close()
}
server.go
package main
import (
"github.com/gorilla/websocket"
"imooc-go-websocket/impl"
"net/http"
"time"
)
var (
upgrader = websocket.Upgrader{
// 允许跨域
CheckOrigin: func(r *http.Request) bool {
return true
},
}
)
// 请求处理器
func wsHandler(w http.ResponseWriter, r *http.Request) {
var (
wsConn *websocket.Conn
err error
data []byte
conn *impl.Connection
)
// 完成WebSocket握手,建立长连接
// Upgrade: websocket
if wsConn, err = upgrader.Upgrade(w, r, nil); err != nil {
return
}
if conn, err = impl.InitConnection(wsConn); err != nil {
goto ERR
}
// 每隔1s给客户端发送一个心跳信息
go func() {
var (
err error
)
for {
if err = conn.WriteMessage([]byte("heartbeat")); err != nil {
return
}
time.Sleep(1 * time.Second)
}
}()
for {
if data, err = conn.ReadMessage(); err != nil {
goto ERR
}
if err = conn.WriteMessage(data); err != nil {
goto ERR
}
}
ERR:
conn.Close()
}
func main() {
// 处理HTTP请求/ws
http.HandleFunc("/ws", wsHandler)
// 监听本地7777端口,并处理所有HTTP请求
http.ListenAndServe("0.0.0.0:7777", nil)
}