从0到1的弹幕系统--弹幕功能实现

今天我们开始来实现弹幕功能，虽然我们实现了一个超简单、粗暴的websocket协议，但前面已经说了，那个只是学习用，实际我们使用第三方的包。

弹幕功能实现

依赖安装

我们使用go mod来安装依赖，关于go的包管理，可以看官方文档。

安装websocket实现，我们使用github.com/gorilla/websocket包：

$ go get github.com/gorilla/websocket
复制代码

我们使用到了Redis，安装Redis包，我们使用github.com/go-redis/redis：

$ go get github.com/go-redis/redis/v8
复制代码

golang.org/x包失败

因为众所周知墙的问题，会导致安装失败，可以看看这篇文章

设置环境变量：

$ export GOPROXY=https://goproxy.io
复制代码

这样，就可以成功安装了。

使用go mod进行依赖管理，实际的包并没有下载到src目录下，而是在go mod仓库里，可以通过go env命令查看go mod仓库目录。另外，本人使用的Goland编辑器，由于依赖包不在src目录下，import时，编辑器一直报红，提示找不到包，找不到包，相应的也没有代码提示了，这就很烦了。还好，Goland支持go mod，我们需要简单的设置一下：

好了，不会提示找不到包了，代码提示也正常了。

连接Redis

var ctx = context.Background()

var rdb = NewRedisConn()
func NewRedisConn() (client *redis.Client) {
	client = redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // no password set
		DB:       0,  // use default DB
	})

	_, err := client.Ping(ctx).Result()
	if err != nil {
		log.Fatal(err)
	}
	return
}
复制代码

照着官方的demo来就行了，没什么好说的。

URI定义

RFC6455中定义了websocket的路由规则，我们定义URI为/chat?room=xxx

其中，room表示房间号。众所周知，直播中肯定有房间号，视频网站中，将视频的id或者当前页码id，作为房间号。

http.HandleFunc("/chat", func(writer http.ResponseWriter, request *http.Request) {
    serveWs(writer, request)
})
err := http.ListenAndServe("localhost:9527", nil)
if err != nil {
    log.Fatal(err)
}
复制代码

因为websocket是基于HTTP协议的，我们可以按照HTTP的方式获取路由参数：

queries := request.URL.Query()
roomId := queries.Get("room")
复制代码

参数获取到后，需要进行一些校验：

if roomId == "" {
    writer.WriteHeader(400)
    log.Println("room必须")
    return
}
_, err := rdb.Get(ctx, fmt.Sprintf("room:%s", roomId)).Result()
if err == redis.Nil {
    writer.WriteHeader(400)
    log.Println("房间不存在")
    return
} else if err != nil {
    writer.WriteHeader(500)
    log.Println(err)
    return
}
复制代码

一些定义

// 消息类型
const (
	danmuMsg = 1 // 弹幕
	bannedMsg = 2 // 禁言
	joinMsg = 3 // 进入房间
	metaMsg = 4 // 房间信息
)

const (
	deadline = 5 * time.Second // 超时时间
	pongWait = 60 * time.Second // pong等待时间
	pingPeriod = (pongWait * 9)/10 // ping频率
	maxMessageSize = 512 // 最大读取数据大小
)

type Client struct {
	conn *websocket.Conn
	msg RsvData

	room *Room

	// 客户端弹幕缓存数据量
	bufChan chan []byte
}

// 消息格式
type Data struct {
	MsgType int
	Msg string
}

// 结束消息的格式
type RsvData struct {
	Data
	Token []byte // 预留
}

// 写消息格式
type WriteData struct {
	Data
}
复制代码

房间定义：

type Room struct {
	id string
	clients map[*Client]bool

	broadcastMsg chan []byte

	leaving chan *Client
	entering chan *Client
}

type Rooms struct {
	rooms map[string]*Room

	mux sync.Mutex
}
复制代码

房间处理

用户进入某个房间，需要获取这个房间的信息，如果房间不存在就创建：

room, ok := rooms.rooms[roomId]
if !ok {
    room = &Room{
        id: roomId,
        clients: make(map[*Client]bool),
        leaving: make(chan *Client),
        broadcastMsg: make(chan []byte),
        entering: make(chan *Client),
    }
    go room.broadcast()
    rooms.mux.Lock()
    rooms.rooms[roomId] = room
    rooms.mux.Unlock()
}
复制代码

我们以房间为单位，一个房间一个goroutine，一个用户来了，需要告诉把他加入到该房间里：

room.entering <-client
复制代码

我们在房间的goroutine中监听entering：

cli := <-r.entering
r.clients[cli] = true
复制代码

一般我们在看直播时，会经常看到”欢迎XXX进入房间“

msg := WriteData{
    Data{MsgType:joinMsg, Msg: cli.conn.RemoteAddr().String()},
}
byteMsg, _ := json.Marshal(msg)
r.broadcastMsg <- byteMsg
复制代码

我们还没添加用户相关功能，这里先用用户的IP代替。

用户看了会直播，不想看了或者想看其他直播了，这时抽象来看，用户离开了当前房间，

cli := <-r.leaving
delete(r.clients, cli)
close(cli.bufChan)
复制代码

用户离开房间了，将用户从当前房间中删除，也就不再收发信息了，关闭相应的channel

下面最主要的来了，弹幕呢，看直播或视频时，那些弹幕应该怎么处理呢？我们还是抽象成房间，以房间为单位，将这个房间里的弹幕发送给在这个房间里的所有用户：

msg := <-r.broadcastMsg
for cli := range r.clients {
    cli.bufChan <- msg
}
复制代码

以上就是房间需要做的事情，处理进入、离开、消息：

func (r *Room) broadcast() {
	for {
		select {
		case cli := <-r.entering:
			msg := WriteData{
				Data{MsgType:joinMsg, Msg: cli.conn.RemoteAddr().String()},
			}
			byteMsg, _ := json.Marshal(msg)
			r.broadcastMsg <- byteMsg
			r.clients[cli] = true
		case cli := <-r.leaving:
			delete(r.clients, cli)
			close(cli.bufChan)
		case msg := <-r.broadcastMsg:
			for cli := range r.clients {
				cli.bufChan <- msg
			}
		}
	}
}
复制代码

收发弹幕

一个用户一个websocket连接，抽象来说，一个连接上只处理2种事情，收发信息：

go client.readMsg()
client.writeMsg()
复制代码

go的HTTP服务中，当有一个连接过来，会开启一个goroutine处理这个连接，所以一个websocket一个goroutine，我们再开一个goroutine处理读信息。

写消息

我们对每个客户端连接定义了一个bufChan的channel，这个channel可能已经被关闭（如离开房间），所以我们需要考虑channel被关闭的情况：

msg, ok := <-c.bufChan
// c.bufChan通道已关闭
if !ok {
    c.conn.WriteMessage(websocket.CloseMessage, []byte{})
    return
}
复制代码

如果没有关闭，正常的向客户端写消息就行了：

// 暂时只支持文本信息
w, err := c.conn.NextWriter(websocket.TextMessage)
if err != nil {
    return
}
w.Write(msg)

// 发送缓存中的信息
n := len(c.bufChan)
for i := 0; i < n; i++ {
    w.Write(<-c.bufChan)
}

// github.com/gorilla/websocket包的通用做法，照做就行
// 冲洗信息写入流
// 并不是关闭TCP连接
// 这里应该w.Flush()更贴切
if err := w.Close(); err != nil {
    return
}
复制代码

心跳处理

为什么要加心跳处理呢？？防止一端意外断开了TCP连接，而另一端还不知道，还保持着连接句柄，随着时间的推移，导致内存耗尽。TCP有个keep-alive选项，实际中很少使用，一般都是由应用自己做心跳处理。

websocket协议中定义了ping、pong控制帧，用于心跳包，关于websocket的控制帧可以去看RFC64555.5节

心跳包可以由客户端或者服务器发送，我们没办法要求客户端做这做那，所以一般都是由服务端自己发送心跳包。

创建定时器：

ticker := time.NewTicker(pingPeriod)
<-ticker.C
c.conn.SetWriteDeadline(time.Now().Add(deadline))
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
    return
}
复制代码

写消息已经处理完了，下面是完整的代码：

func (c *Client) writeMsg() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case msg, ok := <-c.bufChan:
			c.conn.SetWriteDeadline(time.Now().Add(deadline))
			// c.bufChan通道已关闭
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			// todo 消息类型可配置
			// 暂时只支持文本信息
			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(msg)

			// 发送缓存中的信息
			n := len(c.bufChan)
			for i := 0; i < n; i++ {
				w.Write(<-c.bufChan)
			}

			// 冲洗信息写入流
			// 并不是关闭TCP连接
			// 这里应该w.Flush()更贴切
			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(deadline))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}
复制代码

读信息

读取信息就很简单了，直接读就行了，然后将读取到的信息发送给客户端：

c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Now().Add(pongWait))
_, msg, err:= c.conn.ReadMessage()
if err != nil {
    if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
        log.Printf("error: %v", err)
    }
    break
}

// todo json解析 其他处理
writeMsg := WriteData{
    Data{MsgType:danmuMsg, Msg: string(msg)},
}
byteMsg, _ := json.Marshal(writeMsg)
c.room.broadcastMsg <- byteMsg
复制代码

这里读取的信息，我们先不做任何处理，直接转发给客户端。之后我们再对读取的信息做处理，如敏感词过滤、禁言处理等。

测试一下： 消息格式先不处理，现在只要保证数据能正常收发就行，后期再处理消息。

看似都正常了，但是过了大概一分钟后，连接突然断了： 这是怎么回事呢？？

因为我们上面设置了ReadDeadline的值为1分钟。一分钟后，服务端自己会断开链接。1分钟为pong等待时间，所以当我们向客户端发送ping，客户端相应时，我们都需要重新设置ReadDeadline时间。

c.conn.SetPongHandler(func(string) error {
    c.conn.SetReadDeadline(time.Now().Add(pongWait))
    return nil
})
复制代码

收到pong时重新设置deadline，这样就可以了。现在基本的弹幕发送功能已经实现了。

关于deadline，大家可以查看官方的文档说明：

$ go doc net.Conn.SetDeadline
复制代码

代码已提交到gitee仓库。