websocket消息服务
目的：搭建websocket服务，用浏览器与服务进行消息交互（写的第一个Go程序）
代码目录结构：
前端html页面：
1 <!DOCTYPE html> 2 <html> 3 <head> 4 <meta charset="utf-8"> 5 <script> 6 window.addEventListener("load", function(evt) { 7 var output = document.getElementById("output"); 8 var input = document.getElementById("input"); 9 var ws; 10 var print = function(message) { 11 var d = document.createElement("div"); 12 d.innerHTML = message; 13 output.appendChild(d); 14 }; 15 document.getElementById("open").onclick = function(evt) { 16 if (ws) { 17 return false; 18 } 19 ws = new WebSocket("ws://127.0.0.1:7777/ws"); 20 ws.onopen = function(evt) { 21 print("OPEN"); 22 } 23 ws.onclose = function(evt) { 24 print("CLOSE"); 25 ws = null; 26 } 27 ws.onmessage = function(evt) { 28 print("RESPONSE: " + evt.data); 29 } 30 ws.onerror = function(evt) { 31 print("ERROR: " + evt.data); 32 } 33 return false; 34 }; 35 document.getElementById("send").onclick = function(evt) { 36 if (!ws) { 37 return false; 38 } 39 print("SEND: " + input.value); 40 ws.send(input.value); 41 return false; 42 }; 43 document.getElementById("close").onclick = function(evt) { 44 if (!ws) { 45 return false; 46 } 47 ws.close(); 48 return false; 49 }; 50 }); 51 </script> 52 </head> 53 <body> 54 <table> 55 <tr><td valign="top" width="50%"> 56 <p>Click "Open" to create a connection to the server, 57 "Send" to send a message to the server and "Close" to close the connection. 58 You can change the message and send multiple times. 59 </p> 60 <form> 61 <button id="open">Open</button> 62 <button id="close">Close</button> 63 <input id="input" type="text" value="Hello world!"> 64 <button id="send">Send</button> 65 </form> 66 </td><td valign="top" width="50%"> 67 <div id="output"></div> 68 </td></tr></table> 69 </body> 70 </html>
server.go代码：
package main
import (
"fmt"
"github.com/gorilla/websocket"
"go_websocket"
"net/http"
)
// http升级websocket协议的配置
var wsUpgrader = websocket.Upgrader{
// 允许跨域CORS
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// 消息处理
func wsHandler(resp http.ResponseWriter, req *http.Request) {
wsSocket, err := wsUpgrader.Upgrade(resp, req, nil)
if err != nil {
return
}
wsConn := go_websocket.WsConnectionInit(wsSocket)
wsConn.Run()
for {
wsmsg, err := wsConn.ReadMessage()
if err != nil {
goto error
}
err = wsConn.WriteMessage(wsmsg)
if err != nil {
goto error
}
}
error:
fmt.Println("websocket is closed")
return
}
func main() {
fmt.Println("websocket start")
http.HandleFunc("/ws", wsHandler)
http.ListenAndServe("0.0.0.0:7777", nil)
}
connection.go代码：
package go_websocket
import (
"errors"
"fmt"
"github.com/gorilla/websocket"
"sync"
"time"
)
// 客户端读写消息
type WsMessage struct {
msgType int
data []byte
}
// 客户端连接
type wsConnection struct {
wsSocket *websocket.Conn
inChan chan *WsMessage
outChan chan *WsMessage
isClosed bool
closeChan chan []byte
mutex sync.Mutex
}
// 连接初始化
func WsConnectionInit(wsSocket *websocket.Conn) (wsConn *wsConnection) {
wsConn = &wsConnection{
wsSocket: wsSocket,
inChan: make(chan *WsMessage, 1000),
outChan: make(chan *WsMessage, 1000),
closeChan: make(chan []byte, 1),
}
return wsConn
}
// 启动
func (wsConn *wsConnection) Run() {
go wsConn.readLoop()
go wsConn.writeLoop()
go wsConn.heartbeat()
}
// 心跳检测
func (wsConn *wsConnection) heartbeat() {
for {
time.Sleep(2 * time.Second)
wsmsg := &WsMessage{msgType: websocket.TextMessage, data: []byte("heartbeat")}
err := wsConn.WriteMessage(wsmsg)
if err != nil {
fmt.Println("send heartbeat stop")
return
}
}
}
// 循环接收
func (wsConn *wsConnection) readLoop() {
var ()
for {
msgType, data, err := wsConn.wsSocket.ReadMessage()
if err != nil {
goto error
}
select {
case wsConn.inChan <- &WsMessage{msgType: msgType, data: data}:
case <-wsConn.closeChan:
goto closed
}
}
error:
wsConn.Close()
closed:
fmt.Println("readLoop closed")
}
// 循环发送
func (wsConn *wsConnection) writeLoop() {
for {
select {
case wsmsg := <-wsConn.outChan:
if err := wsConn.wsSocket.WriteMessage(wsmsg.msgType, wsmsg.data); err != nil {
goto error
}
case <-wsConn.closeChan:
goto closed
}
}
error:
wsConn.Close()
closed:
fmt.Println("writeLoop close")
}
// 取消息，外部可调用
func (wsConn *wsConnection) ReadMessage() (wsmsg *WsMessage, err error) {
select {
case wsmsg = <-wsConn.inChan:
return wsmsg, nil
case <-wsConn.closeChan:
return nil, errors.New("websocket is closed")
}
}
// 写消息，外部可调用
func (wsConn *wsConnection) WriteMessage(wsmsg *WsMessage) (err error) {
select {
case wsConn.outChan <- wsmsg:
case <-wsConn.closeChan:
return errors.New("websocket is closed")
}
return nil
}
// 关闭wsSocket
func (wsConn *wsConnection) Close() {
wsConn.wsSocket.Close()
// 加锁
wsConn.mutex.Lock()
if !wsConn.isClosed {
wsConn.isClosed = true
close(wsConn.closeChan)
}
wsConn.mutex.Unlock()
}
效果展示：
