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 协议》，转载必须注明作者和本文链接