go 系列实现websocket

一、简介

websocket是个二进制协议，需要先通过Http协议进行握手，从而协商完成从Http协议向websocket协议的转换。一旦握手结束，当前的TCP连接后续将采用二进制websocket协议进行双向双工交互，自此与Http协议无关。

二、websocket demo

ws.go

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"time"
)

type wsMessage struct {
	mType int
	data  []byte
}

type wsConnection struct {
	wsSocket *websocket.Conn // 底层websocket
	inChan   chan *wsMessage // 读队列
	outChan  chan *wsMessage // 写队列
	state    int             // 连接状态
}

func (wsCon *wsConnection) wsReadLoop() {
	for {
		if wsCon.state == 1 {
			return
		}
		mType, data, err := wsCon.wsSocket.ReadMessage() // 阻塞
		if err != nil {
			fmt.Println("read err: %v", err)
			continue
		}
		msg := &wsMessage{
			mType,
			data,
		}
		wsCon.inChan <- msg
	}

}

func (wsCon *wsConnection) wsWriteLoop() {
	for {
		if wsCon.state == 1 {
			return
		}
		msg := <-wsCon.outChan // 阻塞
		err := wsCon.wsSocket.WriteMessage(msg.mType, msg.data)
		if err != nil {
			fmt.Printf("write err: %v", err)
			continue
		}
	}
}

func (wsCon *wsConnection) bizLoop() {
	for {
		if wsCon.state == 1 {
			return
		}
		// 处理消息
		msg := <-wsCon.inChan
		fmt.Println("process msg", string(msg.data))
		// 返回响应
		wsCon.outChan <- &wsMessage{
			mType: 2,
			data:  []byte(fmt.Sprintf("%s done", string(msg.data))),
		}
	}
}

func (wsCon *wsConnection) healthLoop() {
	defer wsCon.wsSocket.Close()
	for {
		time.Sleep(3 * time.Second)
		// 返回响应
		if err := wsCon.wsSocket.WriteMessage(2,
			[]byte("心跳检查")); err != nil {
			fmt.Println("write err: %v", err)
			wsCon.state = 1
			return
		}

	}
}

func wsHandler(resp http.ResponseWriter, req *http.Request) {
	// 升级websocket协议
	wsSocket, err := websocket.Upgrade(resp, req, nil, 1024, 1024)
	if err != nil {
		fmt.Println("upgrade err: %v", err)
	}
	// 一个请求，一个websocket 连接
	wsCon := &wsConnection{
		wsSocket: wsSocket,
		inChan:   make(chan *wsMessage, 1000),
		outChan:  make(chan *wsMessage, 1000),
	}

	go wsCon.wsReadLoop()
	go wsCon.wsWriteLoop()
	go wsCon.healthLoop()
	go wsCon.bizLoop()

}

func main() {
	http.HandleFunc("/ws", wsHandler)
	http.ListenAndServe(":8080", nil)
}

ws.html

<!DOCTYPE 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:8080/ws");
                ws.onopen = function (evt) {
                    print("OPEN");
                }
                ws.onclose = function (evt) {
                    print("CLOSE");
                    ws = null;
                }
                ws.onmessage = function (evt) {
                    const reader = new FileReader()
                    reader.onload = function (event) {
                        print("RESPONSE: " + event.target.result)
                    }
                    reader.readAsText(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">Open</button>
                <button id="close">Close</button>
                <input id="input" type="text" value="Hello!">
                <button id="send">Send</button>
            </form>
        </td>
        <td valign="top" width="50%">
            <div id="output"></div>
        </td>
    </tr>
</table>
</body>
</html>

实现效果