Go语言websocket

OStack程序员社区-中国程序员成长平台 › 门户 › 编程› Go›Go教程

原作者: [db:作者] 来自: [db:来源] 收藏邀请

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>

client.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()
}

效果展示：