Go使用websocket實現彈幕功能

使用websocket協議，客戶端傳送一個訊息，服務端廣播到所有有效連線中。
主要思路：
1.封裝*websocket.conn，用client結構表示一個客戶端。
2.維持一個map[client]bool，表示有效的客戶端對映，用於廣播訊息
3.除了處理websocket連線外，還要開啟一個廣播協程，監聽客戶端連線，斷開，發彈幕事件。

主要的結構：

type Client struct{
    wsConnect *websocket.Conn
    inChan chan []byte
    outChan chan []byte
    closeChan chan byte
    Name string //客戶的名稱
    Id string //客戶id，唯一
    mutex sync.Mutex  // 對closeChan關閉上鎖
    IsClosed bool  // 防止closeChan被關閉多次
}
type Message struct {
    EventType byte  `json:"type"`       // 0表示使用者釋出訊息；1表示使用者進入；2表示使用者退出
    Name string     `json:"name"`       // 使用者名稱稱
    Message string  `json:"message"`    // 訊息內容
}
    clients = make(map [*util.Client] bool)      // 使用者組對映
    join = make(chan *util.Client, 10)        // 使用者加入通道
    leave = make(chan *util.Client, 10)       // 使用者退出通道
    message = make(chan Message, 10)    // 訊息通道

server端程式碼

package main

import (
    "encoding/json"
    "fmt"
    "github.com/gorilla/websocket"
    "goGin/server/util"
    "net/http"
)

var(
    upgrader = websocket.Upgrader{
        // 允許跨域
        CheckOrigin:func(r *http.Request) bool{
            return true
        },
    }
    clients = make(map [*util.Client] bool)      // 使用者組對映
    join = make(chan *util.Client, 10)        // 使用者加入通道
    leave = make(chan *util.Client, 10)       // 使用者退出通道
    message = make(chan Message, 10)    // 訊息通道
)
type Message struct {
    EventType byte  `json:"type"`       // 0表示使用者釋出訊息；1表示使用者進入；2表示使用者退出
    Name string     `json:"name"`       // 使用者名稱稱
    Message string  `json:"message"`    // 訊息內容
}

//處理每個連線，每個連線，go都會重新起一個協程來處理
func wsHandler(w http.ResponseWriter , r *http.Request){
    var(
        wsConn *websocket.Conn
        err error
        client *util.Client
        data []byte
    )
    r.ParseForm() //返回一個map,並且賦值給r.Form
    name := r.Form["name"][0]
    id := r.Form["id"][0]

    if wsConn , err = upgrader.Upgrade(w,r,nil); err != nil{
        return
    }

    if client , err = util.InitConnection(wsConn); err != nil{
        goto ERR
    }
    client.Id = id
    client.Name = name

    // 如果使用者列表中沒有該使用者
    if !clients[client] {
        join <- client
    }

    for {
        if data , err = client.ReadMessage();err != nil{ //一直讀訊息，沒有訊息就阻塞
            goto ERR
        }
        var msg Message
        msg.EventType = 0
        msg.Name = client.Name
        msg.Message = string(data)
        message <- msg
    }

ERR:
    leave<-client//這個客戶斷開
    client.Close()

}

func broadcaster() {
    for {
        select {
        // 訊息通道中有訊息則執行，否則堵塞
        case msg := <-message:
            // 將資料編碼成json形式，data是[]byte型別
            // json.Marshal()只會編碼結構體中公開的屬性(即大寫字母開頭的屬性)
            data, err := json.Marshal(msg)
            if err != nil {
                return
            }
            for client := range clients {
                if client.IsClosed == true {
                    leave<-client//這個客戶斷開
                    continue
                }
                // fmt.Println("=======the json message is", string(data))  // 轉換成字串型別便於檢視
                if client.WriteMessage(data) != nil {
                    continue //傳送失敗就跳過
                }
            }

        // 有使用者加入
        case client := <-join:
            clients[client] = true  // 將使用者加入對映
            // 將使用者加入訊息放入訊息通道
            var msg Message
            msg.Name = client.Name
            msg.EventType = 1
            msg.Message = fmt.Sprintf("%s join in, there are %d preson in room", client.Name, len(clients))
            message <- msg

        // 有使用者退出
        case client := <-leave:
            // 如果該使用者已經被刪除
            if !clients[client] {
                break
            }
            delete(clients, client) // 將使用者從對映中刪除
            // 將使用者退出訊息放入訊息通道
            var msg Message
            msg.Name = client.Name
            msg.EventType = 2
            msg.Message = fmt.Sprintf("%s leave, there are %d preson in room", client.Name, len(clients))
            message <- msg
        }
    }
}

func main(){
    go broadcaster()
    http.HandleFunc("/ws",wsHandler)
    http.ListenAndServe("0.0.0.0:7777",nil)
}

封裝client

package util
import (
    "github.com/gorilla/websocket"
    "sync"
    "errors"
)
type Client struct{
    wsConnect *websocket.Conn
    inChan chan []byte
    outChan chan []byte
    closeChan chan byte
    Name string //客戶的名稱
    Id string //客戶id，唯一

    mutex sync.Mutex  // 對closeChan關閉上鎖
    IsClosed bool  // 防止closeChan被關閉多次
}
func InitConnection(wsConn *websocket.Conn)(conn *Client ,err error){
    conn = &Client{
        wsConnect:wsConn,
        inChan: make(chan []byte,1000),
        outChan: make(chan []byte,1000),
        closeChan: make(chan byte,1),
        IsClosed:false,
    }
    // 啟動讀協程
    go conn.readLoop();
    // 啟動寫協程
    go conn.writeLoop();
    return
}
func (conn *Client)ReadMessage()(data []byte , err error){
    select{
    case data = <- conn.inChan:
    case <- conn.closeChan:
        err = errors.New("connection is closeed")
    }
    return
}
func (conn *Client)WriteMessage(data []byte)(err error){
    select{
    case conn.outChan <- data:
    case <- conn.closeChan:
        err = errors.New("connection is closeed")
    }
    return
}
func (conn *Client)Close(){
    // 執行緒安全，可多次呼叫
    conn.wsConnect.Close()
    // 利用標記，讓closeChan只關閉一次
    conn.mutex.Lock()
    if !conn.IsClosed {
        close(conn.closeChan)
        conn.IsClosed = true
    }
    conn.mutex.Unlock()
}

func (conn *Client)readLoop(){
    var(
        data []byte
        err error
    )
    for{
        if _, data , err = conn.wsConnect.ReadMessage(); err != nil{
            goto ERR
        }
        //阻塞在這裡，等待inChan有空閒位置
        select{
        case conn.inChan <- data:
        case <- conn.closeChan:        // closeChan 感知 conn斷開
            goto ERR
        }

    }

ERR:
    conn.Close()
}

func (conn *Client)writeLoop(){
    var(
        data []byte
        err error
    )

    for{
        select{
        case data= <- conn.outChan:
        case <- conn.closeChan:
            goto ERR
        }
        if err = conn.wsConnect.WriteMessage(websocket.TextMessage , data); err != nil{
            goto ERR
        }
    }

ERR:
    conn.Close()

}

客戶端程式碼

<!DOCTYPE html>
<html>
<head>
    <title>go websocket</title>
    <meta charset="utf-8" />
</head>
<body>
<script type="text/javascript">
    var wsUri ="ws://127.0.0.1:7777/ws?name=aaa&id=112";
    var output;

    function init() {
        output = document.getElementById("output");
        testWebSocket();
    }

    function testWebSocket() {
        websocket = new WebSocket(wsUri);
        websocket.onopen = function(evt) {
            onOpen(evt)
        };
        websocket.onclose = function(evt) {
            onClose(evt)
        };
        websocket.onmessage = function(evt) {
            onMessage(evt)
        };
        websocket.onerror = function(evt) {
            onError(evt)
        };
    }

    function onOpen(evt) {
        writeToScreen("CONNECTED");
        // doSend("WebSocket rocks");
    }

    function onClose(evt) {
        writeToScreen("DISCONNECTED");
    }

    function onMessage(evt) {
        writeToScreen('<span style="color: blue;">RESPONSE: '+ evt.data+'</span>');
        // websocket.close();
    }

    function onError(evt) {
        writeToScreen('<span style="color: red;">ERROR:</span> '+ evt.data);
    }

    function doSend(message) {
        // writeToScreen("SENT: " + message);
        websocket.send(message);
    }

    function writeToScreen(message) {
        var pre = document.createElement("p");
        pre.style.wordWrap = "break-word";
        pre.innerHTML = message;
        output.appendChild(pre);
    }

    window.addEventListener("load", init, false);
    function sendBtnClick(){
        var msg = document.getElementById("input").value;
        doSend(msg);
        document.getElementById("input").value = '';
    }
    function closeBtnClick(){
        websocket.close();
    }
</script>
<h2>WebSocket Test</h2>
<input type="text" id="input"></input>
<button onclick="sendBtnClick()" >send</button>
<button onclick="closeBtnClick()" >close</button>
<div id="output"></div>

</body>
</html>

本作品採用《CC 協議》，轉載必須註明作者和本文連結