[Pitaya Demo解读笔记]6.rate_limiting demo

项目目录：examples/demo/rate_limiting

pitaya 前端服务器提供了消息限频率，本例就是通过内置的 ratelimiting 系列配置，展示了如何限制消息频率，避免一段时间处理过多请求。

限频是一个很重要的功能，部署到公网上的服务，如果会被恶意刷爆（如DDOS攻击），就会影响到正常用户的访问，所以一般来说，我们会在服务端加上访问频率限制。不多说，看一下这个简单的小 demo。

运行

开启服务器

go run main.go

服务器默认监听 3250 端口，且配置为每分钟最多处理5个请求

// main.go:24
vConfig.Set("pitaya.conn.ratelimiting.limit", 5)
vConfig.Set("pitaya.conn.ratelimiting.interval", time.Minute)

为了触发限频率，我们快速发送几个消息

> pitaya-cli
Pitaya REPL Client
>>> connect localhost:3250
Using json client
connected!
>>> request room.room.ping
>>> sv->pong
>>> request room.room.ping
>>> sv->pong
>>> request room.room.ping
>>> sv->pong
>>> request room.room.ping
>>> request room.room.ping
>>> sv->{"Code":"PIT-504","Message":"request timeout","Metadata":null}

疯狂发送 ping 消息，达到限频后，服务端不再处理请求，即客户端没有收到 pong 回复，而是收到了请求超时，服务端输出了限频日志：

time="2023-12-04T17:38:37+08:00" level=error msg="Data=\x04\x00\x00\x11\x00\v\x0eroom.room.ping, Error=rate limit exceeded" source=pitaya

代码分析

限频功能不是直接在默认config里去支持配置的，所以并不能像 chat demo 那样，通过conf.Pitaya.Buffer.Handler.LocalProcess = 15 来直接设置。限频功能类似于一种"插件"，通过 wrapper（包装）的方式来实现功能注入。

我们先看调用代码：

func createAcceptor(port int, reporters []metrics.Reporter) acceptor.Acceptor {

	// 5 requests in 1 minute. Doesn't make sense, just to test
	// rate limiting
	vConfig := viper.New()
	vConfig.Set("pitaya.conn.ratelimiting.limit", 5)
	vConfig.Set("pitaya.conn.ratelimiting.interval", time.Minute)
	pConfig := config.NewConfig(vConfig)

	rateLimitConfig := config.NewRateLimitingConfig(pConfig)

	tcp := acceptor.NewTCPAcceptor(fmt.Sprintf(":%d", port))
	return acceptorwrapper.WithWrappers(
		tcp,
		acceptorwrapper.NewRateLimitingWrapper(reporters, *rateLimitConfig))
}

重点在最后一句，对真正的acceptor tcp 进行了包装，返回一个包装过的变量，也是一个 acceptor。跟进去看 WithWrappers 的实现：

// Wrapper has a method that receives an acceptor and the struct
// that implements must encapsulate it. The main goal is to create
// a middleware for packets of net.Conn from acceptor.GetConnChan before
// giving it to serviceHandler.
type Wrapper interface {
	Wrap(acceptor.Acceptor) acceptor.Acceptor
}

// WithWrappers walks through wrappers calling Wrapper
func WithWrappers(
	a acceptor.Acceptor,
	wrappers ...Wrapper,
) acceptor.Acceptor {
	for _, w := range wrappers {
		a = w.Wrap(a)
	}
	return a
}

这种形式感觉跟函数式编程模式 Functional Options 有点类似啊？让我们暂时放下这些 Wrappers，先了解一下更简单一点的函数式选项编程模式。

Functional Options 函数式选项编程模式

这个编程模式在 go 语言里非常常见，一般也是以 WithXXX 来命名函数。比如 pitaya 在 component 包里就用到了函数式选项编程模式：

// options.go
package component

type (
	options struct {
		name     string              // component name
		nameFunc func(string) string // rename handler name
	}

	// Option used to customize handler
	Option func(options *options)
)

// WithName used to rename component name
func WithName(name string) Option {
	return func(opt *options) {
		opt.name = name
	}
}

// WithNameFunc override handler name by specific function
// such as: strings.ToUpper/strings.ToLower
func WithNameFunc(fn func(string) string) Option {
	return func(opt *options) {
		opt.nameFunc = fn
	}
}

这些 options 最终在 NewService 中应用了：

// service.go
func NewService(comp Component, opts []Option) *Service {
    ...
    // apply options
	for i := range opts {
		opt := opts[i]
		opt(&s.Options)
	}
    ...
}

但是这个示例的包裹层次比较深，不是很方便学习，建议移步耗子叔的文章，他将演进过程讲解得非常通透（R.I.P.）：

Go 编程模式：Functional Options | 酷 壳 - CoolShell

究其根本，就是通过 WithXXX 函数闭包了参数信息，在 for...range 遍历赋值给了真正的接收方：

for _, option := range options {
    option(&srv)
}

我们看这段代码，好像是跟前面的 WithWrappers 有点像：

那么找不同，不同在于什么？Wrap 完了又赋值回去了，这就是包装的意义，一层一层包裹起来，即修饰器模式。

修饰器编程模式

我看这段代码的时候也是云里雾里，平时没打好设计模式的基础，到这里就有点蒙圈了，幸得朋友推荐了耗子叔的文章，才总算是解了惑。依然建议移步耗子叔的文章（最末的泛型部分可以不看，go 1.18 已经正式支持泛型了）：

Go编程模式：修饰器 | 酷 壳 - CoolShell

(耗子叔千古，多说一句，希望我也能做到持续学习、持续分享，即使达不到陈皓的高度，也努力去继承这样的精神。缅怀，感谢......)

回到代码里来，如果你看完以上文章，应该对修饰器有了一定的理解，修饰器模式是对原型的一种扩展和包装，所以它需要将 Wrap 后的变量再赋值回来，实现一层一层的包裹。

AcceptorWrapper

pitaya的 Wrapper 本质上就是对 Acceptor 的修饰，核心就是耗子叔给出的示例的修饰器编程模式，然而，它，更复杂......也不知道是什么样的脑子能写出这样的代码，我看都看不懂（菜呀= =），让我们结合代码、画图，一步步探究下去。

先从调用部分入手：

// part1
tcp := acceptor.NewTCPAcceptor(fmt.Sprintf(":%d", port))
return acceptorwrapper.WithWrappers(
    tcp,
    acceptorwrapper.NewRateLimitingWrapper(reporters, *rateLimitConfig))

// part2
func WithWrappers(a acceptor.Acceptor, wrappers ...Wrapper) acceptor.Acceptor {
	for _, w := range wrappers {
		a = w.Wrap(a)
	}
	return a
}

// part3
func (r *RateLimitingWrapper) Wrap(a acceptor.Acceptor) acceptor.Acceptor {
	r.Acceptor = a
	return r
}

真正的 acceptor 即 tcp，通过 Wrap 函数，将其包装了一层，赋值给了RateLimitingWrapper.BaseWrapper.Acceptor，如图示：

再看WithWrappers的第二个实参，就是构建的一个RateLimitingWrapper，我们看看它是如何构建的：

// part1
// NewRateLimitingWrapper returns an instance of *RateLimitingWrapper
func NewRateLimitingWrapper(reporters []metrics.Reporter, c config.RateLimitingConfig) *RateLimitingWrapper {
	r := &RateLimitingWrapper{}

	r.BaseWrapper = NewBaseWrapper(func(conn acceptor.PlayerConn) acceptor.PlayerConn {
		return NewRateLimiter(reporters, conn, c.Limit, c.Interval, c.ForceDisable)
	})

	return r
}

// part2
// NewBaseWrapper returns an instance of BaseWrapper.
func NewBaseWrapper(wrapConn func(acceptor.PlayerConn) acceptor.PlayerConn) BaseWrapper {
	return BaseWrapper{
		connChan: make(chan acceptor.PlayerConn),
		wrapConn: wrapConn,
	}
}

核心逻辑就是将一个 func 赋值给了 BaseWrapper.wrapConn，画图帮助理解：

现在我们在 createAcceptor 返回得到的 Acceptor 其实是包装后的，即上图中 RateLimitingWrapper 包裹的这部分，Acceptor 是它的第二层嵌入结构，所以该结构体也实现了 Acceptor interface，也是一种 Acceptor。

那么，当 main 函数中 app.Start() 内部调用了 Acceptor.ListenAndServe() 开启服务时，实际调用的是 RateLimitingWrapper.ListenAndServe，由于它没有直接实现，那就看它的第一层嵌入结构，这里是有实现的：

// acceptorwrapper/base.go:47
// ListenAndServe starts a goroutine that wraps acceptor's conn
// and calls acceptor's listenAndServe
func (b *BaseWrapper) ListenAndServe() {
	go b.pipe()
	b.Acceptor.ListenAndServe()
}

先看第二行 b.Acceptor.ListenAndServe()，我们看图可以清楚的知道，这里的b.Accetpor 即最底层的真实 Acceptor -> tcp，这一行代码在本例中就是调用了

// tcp_acceptor.g:129
// ListenAndServe using tcp acceptor
func (a *TCPAcceptor) ListenAndServe() {
    ...
}

即真正的开启 tcp 服务。

回到第一行 go b.pipe()，看看这行代码开了一个 goroutine 做了些啥：

func (b *BaseWrapper) pipe() {
	for conn := range b.Acceptor.GetConnChan() {
		b.connChan <- b.wrapConn(conn)
	}
}

遍历 ConnChan，对 Acceptor 接收到的新连接[1]做一层 wrapConn(conn) 处理[2]后再扔到自己的 connChan 里[3]，在上图的基础上再画一下这个数据流向：

(PS: 这不是什么标准的UML啥的类图、流程图画法，只是为了帮助理解自创瞎画的)

再看第[2]步，wrapConn 在图中我们已经标识出了，就是 NewRateLimiter：

type RateLimiter struct {
	acceptor.PlayerConn
	reporters    []metrics.Reporter
	limit        int
	interval     time.Duration
	times        list.List
	forceDisable bool
}

// NewRateLimiter returns an initialized *RateLimiting
func NewRateLimiter(
	reporters []metrics.Reporter,
	conn acceptor.PlayerConn,
	limit int,
	interval time.Duration,
	forceDisable bool,
) *RateLimiter {
	r := &RateLimiter{
		PlayerConn:   conn,
		reporters:    reporters,
		limit:        limit,
		interval:     interval,
		forceDisable: forceDisable,
	}

	r.times.Init()

	return r
}

这个 RateLimiter 又嵌入了一个 PlayerConn，当我们通过 wrapConn(conn) 处理完连接、加入到 connChan 时，这个 conn 已经变成了被 RateLimiter 修饰过的连接（这一大段实现，感觉实质上就是各种的嵌入结构，各种的修饰器模式）。

所以，当底层 HandlerService 调用 conn.GetNextMessage 时：

// service/handler.go
func (h *HandlerService) Handle(conn acceptor.PlayerConn) {
    ...
    msg, err := conn.GetNextMessage()
    ...
}

实际上，是对已经包装过的conn来调用，也就会进入到 RateLimiter.GetNextMessage 函数：

// GetNextMessage gets the next message in the connection
func (r *RateLimiter) GetNextMessage() (msg []byte, err error) {
	if r.forceDisable {
		return r.PlayerConn.GetNextMessage()
	}

	for {
		msg, err := r.PlayerConn.GetNextMessage()
		if err != nil {
			return nil, err
		}

		now := time.Now()
		if r.shouldRateLimit(now) {
			logger.Log.Errorf("Data=%s, Error=%s", msg, constants.ErrRateLimitExceeded)
			metrics.ReportExceededRateLimiting(r.reporters)
			continue
		}

		return msg, err
	}
}

在这里，真正实现了对每个连接的请求限频。

我们不具体看限频逻辑了，来看看行8：msg, err := r.PlayerConn.GetNextMessage()，既然我们劫持了真正的消息处理逻辑，自然要把这段逻辑给续上，所以还得要调用真正的连接的 GetNextMessage。

扩展

pitaya目前只提供了请求限频这一种 wrapper 插件，那如果我们想要再写一个 wrapper，再包裹一层，那整个代码逻辑会是什么样子的呢？

让我们照抄 RateLimitingWrapper 写两个简单的 wrapper，并输出一些 log 翻遍查看逻辑流转

MsgCounterWrapper

// msg_counter_wrapper.go
type MsgCounterWrapper struct {
	BaseWrapper
}

func NewMsgCounterWrapper() *MsgCounterWrapper {
	r := &MsgCounterWrapper{}
	r.BaseWrapper = NewBaseWrapper(func(conn acceptor.PlayerConn) acceptor.PlayerConn {
		return NewMsgCounter(conn)
	})

	return r
}

func (r *MsgCounterWrapper) Wrap(a acceptor.Acceptor) acceptor.Acceptor {
	r.Acceptor = a
	return r
}

// msg_counter.go
type MsgCounter struct {
	acceptor.PlayerConn
	msgCount int
}

func NewMsgCounter(conn acceptor.PlayerConn) *MsgCounter {
	logrus.Infof("===================[MsgCounter] New")
	return &MsgCounter{
		PlayerConn: conn,
	}
}

func (r *MsgCounter) GetNextMessage() (msg []byte, err error) {
	msg, err = r.PlayerConn.GetNextMessage()
	r.msgCount++
	logrus.Infof("===================[MsgCounter] msgCount: %d", r.msgCount)
	return msg, err
}

WriteCounterWrapper

RateLimitingWrapper 和 MsgCounterWrapper 可以"劫持" GetNextMessage，因为它嵌入了 PlayerConn，就"继承"到了这个函数。如果我们再看 PlayerConn，发现它又嵌入了 net.Conn，那说明我们还可以"继承"、"劫持"到 Conn 的方法。

// write_counter_wrapper.go
type WriteCounterWrapper struct {
	BaseWrapper
}

func NewWriteCounterWrapper() *WriteCounterWrapper {
	r := &WriteCounterWrapper{}
	r.BaseWrapper = NewBaseWrapper(func(conn acceptor.PlayerConn) acceptor.PlayerConn {
		return NewWriteCounter(conn)
	})

	return r
}

func (r *WriteCounterWrapper) Wrap(a acceptor.Acceptor) acceptor.Acceptor {
	r.Acceptor = a
	return r
}

// write_counter.go
type WriteCounter struct {
	acceptor.PlayerConn
	writeCount int
}

func NewWriteCounter(conn acceptor.PlayerConn) *WriteCounter {
	logrus.Infof("===================[WriteCounter] New")
	return &WriteCounter{
		PlayerConn: conn,
	}
}

func (c *WriteCounter) Write(b []byte) (int, error) {
	n, err := c.PlayerConn.Write(b)
	c.writeCount++
	logrus.Infof("===================[WriteCounter] writeCount: %d", c.writeCount)
	return n, err
}

测试

懒得再写测试代码了，把这个 demo 的 createAcceptor 改吧改吧：

func createAcceptor(port int, reporters []metrics.Reporter) acceptor.Acceptor {
    tcp := acceptor.NewTCPAcceptor(fmt.Sprintf(":%d", port))
    return acceptorwrapper.WithWrappers(
        tcp,
        acceptorwrapper.NewWriteCounterWrapper(),
        acceptorwrapper.NewMsgCounterWrapper())
}

所以我们是 tcp 包一层 writeCounter 再包一层 msgCounter，开启服务和 pitaya-cli 看看输出：

pitaya-cli
Pitaya REPL Client
>>> connect localhost:3250
Using json client
connected!

客户端只要连上就可以了，不需要消息，底层会发送握手消息和心跳包。

看看服务端的日志：

time="2023-12-05T16:16:13+08:00" level=info msg="===================[WriteCounter] New"
time="2023-12-05T16:16:13+08:00" level=info msg="===================[MsgCounter] New"
time="2023-12-05T16:16:13+08:00" level=debug msg="New session established: Remote=[::1]:59560, LastTime=1701764173" source=pitaya  
time="2023-12-05T16:16:13+08:00" level=info msg="===================[MsgCounter] msgCount: 1"
time="2023-12-05T16:16:13+08:00" level=debug msg="Received handshake packet" source=pitaya
time="2023-12-05T16:16:13+08:00" level=info msg="===================[WriteCounter] writeCount: 1"
time="2023-12-05T16:16:13+08:00" level=debug msg="Session handshake Id=4, Remote=[::1]:59560" source=pitaya
time="2023-12-05T16:16:13+08:00" level=debug msg="Successfully saved handshake data" source=pitaya
time="2023-12-05T16:16:13+08:00" level=info msg="===================[MsgCounter] msgCount: 2"
time="2023-12-05T16:16:13+08:00" level=debug msg="Receive handshake ACK Id=4, Remote=[::1]:59560" source=pitaya
time="2023-12-05T16:16:13+08:00" level=info msg="===================[MsgCounter] msgCount: 3"

包装顺序是，由内到外依次是：Acceptor -> WriteCounterWrapper -> MsgCounterWrapper。

当一个新连接进来时，包装结构也是由内而外创建的：PlayerConn -> WriteCounter -> MsgCounter。

如果你还想把最底层的 PlayerConn 创建时也输出日志，也可以，我们改一下源码：

// tcp_acceptor.go 新增一个函数
func NewTcpPlayerConn(conn net.Conn, remoteAddr net.Addr) PlayerConn {
	logrus.Infof("===================[PlayerConn] New")
	return &tcpPlayerConn{
		Conn:       conn,
		remoteAddr: remoteAddr,
	}
}

func (a *TCPAcceptor) serve() {
    ...
    // 方法末尾，注释掉之前的写channel代码，改成调用NewTcpPlayerConn
    //a.connChan <- &tcpPlayerConn{
    //	Conn:       conn,
    //	remoteAddr: remoteAddr,
    //}
    a.connChan <- NewTcpPlayerConn(conn, remoteAddr)
}

再次开启服务，使用pitaya-cli连接上，看看日志，与预期的一致，最先输出的是 PlayerConn 部分的创建，由内而外依次构建：

time="2023-12-05T16:28:47+08:00" level=info msg="===================[PlayerConn] New"
time="2023-12-05T16:28:47+08:00" level=info msg="===================[WriteCounter] New"
time="2023-12-05T16:28:47+08:00" level=info msg="===================[MsgCounter] New"
time="2023-12-05T16:28:47+08:00" level=debug msg="New session established: Remote=[::1]:62000, LastTime=1701764927" source=pitaya

好了，这个 demo 总算是告一段落，真的有点子复杂啊，不知道我有没有把这部分内容给讲明白，太难了......