Golang 高性能 Websocket 库 gws 使用与设计（一）-数字星空

前言

大家好这里是，白泽，这期分析一下 golang 开源高性能 websocket 库 gws。

视频讲解请关注📺B站：白泽talk

介绍

gws：https://github.com/lxzan/gws ｜GitHub 🌟 1.2k，高性能的 websocket 库，代码双语注释，适合有开发经验的同学进阶学习。
gws 的两个特性

High IOPS Low Latency（高I/O，低延迟）
Low Memory Usage（低内存占用）

可以从下图看到： payload 越高，性能相比其他 websocket 库越是优越，如何做到？

gws chatroom 架构图

这是 gws 的官方聊天室 demo 的架构图，绘制在这里帮助各位理解什么是全双工的通信模式。

WebSocket 与 HTTP 一样是应用层的协议，只需要 TCP 完成三次握手之后，Golang 的 net/http 库提供了 Hijack() 方法，将 TCP 套接字（活跃的一个会话），从 HTTP 劫持，此后 tcp 的连接将由 WebSocket 管理，脱离了 HTTP 协议的范畴。

而只要获取了 TCP 的套接字，何时发送和接受数据，都是由应用层决定的，传输层的 TCP 套接字只是被编排的对象（单工/双工），自然可以实现服务端主动发送数据。

缓冲池

为什么 payload 越高，性能相比其他 websocket 库越是优越？

原因：gws 中的读写操作，全部使用了缓冲池。

binaryPool    = internal.NewBufferPool(128, 256*1024) // 缓冲池

读缓冲：每次读取是一次系统调用，因此可以读取一段数据，且用一个 offset 定位消费的位置，减少读取次数。

写缓冲：每次写入是一次系统调用，因此可以多次写入 buffer，统一 flush。

缓冲池：为不同大小的 buffer 提供了缓冲池，大段 buffer 的创建次数减少，减少 GC 压力 & 创建对象和销毁对象时间。

// NewBufferPool Creating a memory pool
// Left, right indicate the interval range of the memory pool, they will be transformed into pow(2,n)。
// Below left, Get method will return at least left bytes; above right, Put method will not reclaim the buffer.
func NewBufferPool(left, right uint32) *BufferPool {
   var begin, end = int(binaryCeil(left)), int(binaryCeil(right))
   var p = &BufferPool{
      begin:  begin,
      end:    end,
      shards: map[int]*sync.Pool{},
   }
   for i := begin; i <= end; i *= 2 {
      capacity := i
      p.shards[i] = &sync.Pool{
         New: func() any { return bytes.NewBuffer(make([]byte, 0, capacity)) },
      }
   }
   return p
}

使用循环从 begin 到 end，每次容量翻倍（乘以2），为每个容量创建一个 sync.Pool 实例。sync.Pool 是Go语言标准库中的一个类型，用于存储和回收临时对象。

使用缓冲池中的 buffer 从 conn（网络连接）中读取和写入数据时，通常会执行以下步骤：

从缓冲池获取缓冲区：使用 Get 方法从缓冲池中获取一个 buffer。
读取数据：如果需要从 conn 读取数据，可以将 buffer 用作读取操作的目的地。
处理数据：根据需要处理读取到的数据。
写入数据：如果需要写入数据，可以将数据写入从缓冲池获取的 buffer，然后从 buffer 写入 conn。
释放缓冲区：使用完毕后，将 buffer 放回缓冲池，以便重用。

设计一个 WebScket 库

编写WebSocket库时，有几个关键点会影响其性能，尤其是在高并发场景下。

下面针对这些场景，部分给出一些 demo 写法（伪代码），可以从中提炼一些通用的项目设计方法：

事件驱动模型：使用非阻塞的事件驱动架构可以提高性能，因为它允许WebSocket库在单个线程内处理多个连接，而不会因等待I/O操作而阻塞。

package main

import (
	"fmt"
	"time"
)

func main() {
	eventChan := make(chan string)
	readyChan := make(chan bool)

	// 模拟WebSocket连接
	go func() {
		time.Sleep(2 * time.Second)
		eventChan <- "connected"
		readyChan <- true
	}()

	// 事件处理循环
	for {
		select {
		case event := <-eventChan:
			fmt.Println("Event received:", event)
		case <-readyChan:
			fmt.Println("WebSocket is ready to use")
			return
		}
	}
}

并发处理：库如何处理并发连接和消息是影响性能的重要因素。使用goroutines或线程池可以提高并发处理能力。
消息压缩：支持消息压缩（如permessage-deflate扩展）可以减少传输数据量，但同时也会增加CPU的使用率，需要找到合适的平衡点。
内存管理：优化内存使用，比如通过减少内存分配和重用缓冲区，可以提高性能并减少垃圾回收的压力。

var buffer = make([]byte, 0, 1024)

func readMessage(conn *websocket.Conn) {
	_, buffer, err := conn.ReadMessage()
	if err != nil {
		// 处理错误
	}
	// 使用buffer中的数据
}

连接池管理：有效的连接池管理可以减少连接建立和关闭的开销，特别是在长连接和频繁通信的场景下。

type WebSocketPool struct {
	pool map[*websocket.Conn]struct{}
}

func (p *WebSocketPool) Add(conn *websocket.Conn) {
	p.pool[conn] = struct{}{}
}

func (p *WebSocketPool) Remove(conn *websocket.Conn) {
	delete(p.pool, conn)
}

func (p *WebSocketPool) Broadcast(message []byte) {
	for conn := range p.pool {
		conn.WriteMessage(websocket.TextMessage, message)
	}
}

锁和同步机制：在多线程或goroutine环境中，合理的锁和同步机制是必要的，以避免竞态条件和死锁，但过多的锁竞争会降低性能。

import "sync"

var pool = &WebSocketPool{
	pool: make(map[*websocket.Conn]struct{}),
}
var mu sync.Mutex

func broadcast(message []byte) {
	mu.Lock()
	defer mu.Unlock()
	for conn := range pool.pool {
		conn.WriteMessage(websocket.TextMessage, message)
	}
}

I/O模型：使用非阻塞I/O或异步I/O模型可以提高性能，因为它们允许在等待网络数据时执行其他任务。

func handleConnection(conn *websocket.Conn) {
	go func() {
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				return // 处理错误
			}
			// 处理接收到的消息
		}
	}()
}

协议实现：精确且高效的WebSocket协议实现，包括帧的处理、掩码的添加和去除、以及控制帧的管理，都是影响性能的因素。

func (c *Conn) genFrame(opcode Opcode, payload internal.Payload, isBroadcast bool) (*bytes.Buffer, error) {
	if opcode == OpcodeText && !payload.CheckEncoding(c.config.CheckUtf8Enabled, uint8(opcode)) {
		return nil, internal.NewError(internal.CloseUnsupportedData, ErrTextEncoding)
	}

	var n = payload.Len()

	if n > c.config.WriteMaxPayloadSize {
		return nil, internal.CloseMessageTooLarge
	}

	var buf = binaryPool.Get(n + frameHeaderSize)
	buf.Write(framePadding[0:])

	if c.pd.Enabled && opcode.isDataFrame() && n >= c.pd.Threshold {
		return c.compressData(buf, opcode, payload, isBroadcast)
	}

	var header = frameHeader{}
	headerLength, maskBytes := header.GenerateHeader(c.isServer, true, false, opcode, n)
	_, _ = payload.WriteTo(buf)
	var contents = buf.Bytes()
	if !c.isServer {
		internal.MaskXOR(contents[frameHeaderSize:], maskBytes)
	}
	var m = frameHeaderSize - headerLength
	copy(contents[m:], header[:headerLength])
	buf.Next(m)
	return buf, nil
}