golang + vue + websocket 实现的简单在线消息转发例子

实现的功能是简单的websocket消息在线转发

应该算是学习产物吧，现在的版本应该没bug了（之前版本会挂）
之前的bug原因，close管道之后继续发送会panic
另外，在后期建立大量连接测试的时候，发送会产生broken pipe的错误，并导致发送服务阻塞
解决办法：1、如果其他goroutine中有对关闭管道的写操作，尽量不要用close
2、原因是关闭信号没有正确传达，设置超时时间，规定时间未接收到数据，关闭；发送超时，关闭

• 服务端只做数据转发，格式由客户端定义和解析

go代码

gin 启动一个web服务，在对应路由处理函数里升级为 websocket。要在https的站点使用的话，TLS相关的配置是必须的（主要是下载证书文件，之后只能以 wss://域名/* 的方式使用）。如果实在没条件，把对应的中间件注释掉,RunTLS改为Run,但只能在http站点下使用，wss://变为ws://

package main

import (
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/unrolled/secure"
	"log"
	"net/http"
	"sync"
	"time"
)

// 自定义客户端
type Client struct {
	// 引入锁
	sync.Mutex
	// websocket连接
	conn *websocket.Conn
	// 读到的信息在这里
	readChan chan []byte
	// 需要发送的信息在这里
	writeChan chan []byte
	// 协程之间通信，连接是否已关闭
	closeChan chan struct{}
	// close closeChan 使用，避免panic
	closed bool
}

// 读取一条信息
// false 代表连接关闭失效/等待超时
func (c *Client) ReadMessage() ([]byte, bool) {
	select {
	// 超时，关闭连接
	// todo 心跳信息没有处理，在外部处理(自定义)
	case <-time.After(time.Minute*5 + time.Second*10):
		c.Close()
		return nil, false
	case data := <-c.readChan:
		return data, true
	case <-c.closeChan:
		return nil, false
	}
}

// 写入一条信息
// 不能判断成功/失败
func (c *Client) WriteMessage(data []byte) {
	select {
	case c.writeChan <- data:
	case <-c.closeChan:
	case <-time.After(time.Second * 3):
		// 为了防止阻塞，如果处理过慢，这里可以增加服务的稳定，不会导致一直阻塞影响其他接收者接收
		c.Close()
	}
}

// 关闭连接
func (c *Client) Close() {
	c.Lock()
	if !c.closed {
		close(c.closeChan)
		c.closed = true
		// ===
		// conn.Close 可以不加锁,多次使用,线程安全
		c.conn.Close()
		// ===
	}
	c.Unlock()
}

func (c *Client) readLoop() {
	for {
		_, msg, err := c.conn.ReadMessage()
		if err != nil {
			c.Close()
			return
		}
		select {
		case c.readChan <- msg:
		case <-c.closeChan:
			return
		}
	}
}

func (c *Client) writeLoop() {
	for {
		select {
		case msg := <-c.writeChan:
			// 因为缓冲区的存在所以，不会立即发送，加上异常关闭的连接也被视为存在
			// 可能会出现 broken pipe
			err := c.conn.WriteMessage(websocket.TextMessage, msg)
			if err != nil {
				c.Close()
				return
			}
		case <-c.closeChan:
			return
		}
	}
}

// 升级成websocket，开启读（生产者）、写（消费者）的协程，返回自定义*Client
func NewClient(c *gin.Context) (*Client, error) {
	//升级get请求为webSocket协议
	ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println("NewClient err:", err)
		return nil, err
	}

	client := &Client{
		conn: ws,
		// 没必要设置太大，1也没问题，一般情况消费的足够快，除非是程序在发消息
		readChan:  make(chan []byte, 5),
		writeChan: make(chan []byte, 5),
		closeChan: make(chan struct{}),
	}

	go client.readLoop()
	go client.writeLoop()

	return client, nil
}

// 消息广播时使用
// 要广播的信息，源客户端信息和要发送的信息
type Message struct {
	client *Client
	msg    []byte
}

var (
	// key:client, value:struct{}{}
	clients sync.Map
	// 广播消息队列
	messages = make(chan Message, 100)

	// 广播消息发送 worker
	sendWorkers = make(chan struct{}, 100)

	// 测试时使用，客户端连接数
	//clientNum int32

	// 升级为 websocket
	upGrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func Chat(c *gin.Context) {
	// 获取封装好的自定义client
	conn, err := NewClient(c)
	if err != nil {
		return
	}

	// 存储至连接队列
	clients.Store(conn, struct{}{})

	//atomic.AddInt32(&clientNum, 1)
	//log.Println("now clients:", clientNum)

	go func() {
		// 收尾的函数
		defer func() {
			clients.Delete(conn)

			//atomic.AddInt32(&clientNum, -1)
			//log.Println("now clients:", clientNum)
		}()
		for {
			i, ok := conn.ReadMessage()
			if !ok {
				return
			}
			// 处理自定义心跳信息
			if string(i) != "{\"msg\":\"beat\"}" {
				tmp := Message{client: conn, msg: i}
				messages <- tmp
			}
		}
	}()
}

func sendWorker(c *Client, msg []byte) {
	sendWorkers <- struct{}{}
	go func() {
		c.WriteMessage(msg)
		<-sendWorkers
	}()
}

// 广播函数
func BroadCast() {
	// 取消息
	for msg := range messages {
		c := msg.client
		// 遍历客户端队列
		clients.Range(func(key, value interface{}) bool {
			// 不是本身就发送
			if key != c {
				// 因为 key 的类型确定，所以可以这么用
				sendWorker(key.(*Client), msg.msg)
			}
			// 继续迭代
			return true
		})
	}
}

// https 处理, 即变为 wss://*
func TlsHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		secureMiddleware := secure.New(secure.Options{
			SSLRedirect: true,
			SSLHost:     ":8849",
		})
		err := secureMiddleware.Process(c.Writer, c.Request)

		if err != nil {
			return
		}
	}
}

func main() {
	// goroutine 广播消息
	go BroadCast()

	// 发布模式
	gin.SetMode(gin.ReleaseMode)
	r := gin.Default()
	// 路由设定
	r.GET("/chat", Chat)
	// TLS 中间件设置
	r.Use(TlsHandler())
	// 启动服务
	// TLS 需要对应文件
	err := r.RunTLS(":8849", "*.pem", "*.key")
	if err != nil {
		panic(err)
	}
}

/*
没有条件升级https或不会升级的
注释掉
r.Use(TlsHandler())
改
err := r.RunTLS(":8849", "*.pem", "*.key")
为
err := r.Run(":8849")

使用时前端改wss:// 为 ws://
此时只能在极少部分的网站可以使用，现在大部分是https站点
*/

前端部分

• 如果有需要，拿去用，自行修改

<!DOCTYPE html>
<html lang="ch">

<head>
    <meta charset="UTF-8">
    <title>Title</title>
    <script src="https://cdn.jsdelivr.net/npm/vue"></script>
    <style>
        #chat-container {
            top: 100px;
            left: 100px;
            margin: 0 auto;
            border-radius: 8px;
            box-shadow: 0 0 0 -20px rgba(0, 0, 0, .2), 0 24px 38px 3px rgba(0, 0, 0, .14), 0 9px 46px 8px rgba(0, 0, 0, .12);
            position: fixed;
            background: rgba(255, 255, 255, 0.95);
            width: 150px;
            max-height: 100px;
            overflow: scroll;
            opacity: 40%;
            z-index: 999999999;
        }

        #chat-container div li p,
        #chat-container div input {
            font-weight: 100;
            padding: 0;
            margin: 0;
        }

        #chat-container div input {
            border-width: 0;
        }

        #chat-container div input:focus {
            outline: none;
        }

        #chat-container::-webkit-scrollbar {
            display: none;
        }

        #chat-container div li {
            list-style: none
        }

        .chat-right p {
            text-align: right;
        }

        .chat-left p {
            text-align: left;
        }

        [v-cloak] {
            display: none
        }
    </style>
</head>

<body>
<div id="chat-container" @mousedown="move" v-cloak v-show="show">
    <div>
        <li v-for="item in dataList" :class="{'chat-right':item.send,  'chat-left':!item.send}">
            <p>{{item.msg}}</p>
        </li>
    </div>
    <div style="text-align: center;">
        <input type="text" v-model="msg" :placeholder="tips" @keypress.enter="send"
               style="width: 100%;text-align: center;">
    </div>
</div>
</body>
<script>
    // f12 测试使用
    // setInterval(()=>{app.msg="test";app.send()}, 3000)
    /*
    自定义格式方法：
        1、修改显示消息的 p 标签内容，<p>{ {item.msg} }</p>
        2、修改 send 方法,修改发送的 data 对象
        * 自定义的字段可以写在 vue 的数据段 data 中，如昵称name等
     */

    document.addEventListener("keypress", function (ev) {
        switch (ev.key) {
            case "~":
                if (!window.top.chatapp)
                    window.top.chatapp = new Vue({
                        el: "#chat-container",
                        mounted: function () {
                            this.ws = new WebSocket("wss://*")
                            this.ws.onopen = () => {
                                this.tips = "连接成功,在此输入"
                            }
                            this.ws.onerror = () => {
                                this.tips = "发送失败"
                            }
                            this.ws.onmessage = (ev) => {
                                let item = JSON.parse(ev.data);
                                this.dataList.push(item)
                            }
                            // 10s检查一次ws状态
                            setInterval(() => {
                                if (this.ws.readyState > 1) {
                                    this.ws = new WebSocket("wss://*")
                                    this.ws.onopen = () => {
                                        this.tips = "连接成功,在此输入"
                                    }
                                    this.ws.onerror = () => {
                                        this.tips = "发送失败"
                                    }
                                    this.ws.onmessage = (ev) => {
                                        let item = JSON.parse(ev.data);
                                        this.dataList.push(item)
                                    }
                                }
                            }, 10000);
                            // 5min发送一次心跳信息
                            setInterval(() => {
                                let data = {
                                    msg: "beat",
                                }
                                this.ws.send(JSON.stringify(data))
                            }, 1000 * 60 * 5)
                        },
                        // 新消息滚动值最下方
                        updated: function () {
                            let showContent = document.getElementById("chat-container");
                            showContent.scrollTop = showContent.scrollHeight;
                        },
                        data: {
                            ws: null,
                            dataList: [],
                            msg: "",
                            positionX: 0,
                            positionY: 0,
                            tips: "",
                            show: true,
                        },
                        methods: {
                            send() {
                                if (this.msg.trim() && this.ws) {
                                    let data = {
                                        msg: this.msg.trim(),
                                    }
                                    this.ws.send(JSON.stringify(data))
                                    data['send'] = true
                                    this.dataList.push(data)
                                    this.msg = ""
                                }
                            },
                            move(e) {
                                let odiv = document.getElementById("chat-container");
                                let disX = e.clientX - odiv.offsetLeft;
                                let disY = e.clientY - odiv.offsetTop;
                                document.onmousemove = (e) => {
                                    let left = e.clientX - disX;
                                    let top = e.clientY - disY;
                                    this.positionX = top;
                                    this.positionY = left;
                                    odiv.style.left = left + 'px';
                                    odiv.style.top = top + 'px';
                                };
                                document.onmouseup = () => {
                                    document.onmousemove = null;
                                    document.onmouseup = null;
                                };
                            },
                        }
                    })
                else
                    window.top.chatapp.show = !window.top.chatapp.show
                break
        }
    })
</script>
<html>