Kitex 框架入门系列（2）Middleware

持续创作，加速成长！这是我参与「掘金日新计划 · 6 月更文挑战」的第4天，点击查看活动详情

开篇

在第一期文章中，我们了解了 Kitex 的一些基本背景以及上手用法。其实作为开发者，很多时候我们更在意的是用法，比如业务有多种多样的需求，框架能否灵活支持。

从一个 RPC 框架的角度看，对于服务发现，负载均衡，编解码，元信息传递等能力，都可能根据场景有所区分，Kitex 作为一款在字节内部充分验证的优质框架，在扩展性上也非常优秀（这部分如果感兴趣，可以看看官方此前对 Kitex 扩展性设计方面做的分享，对开发者来说有很多借鉴意义 视频地址。

作为框架的使用方，普通的业务开发者们，最经常使用的其实是 Kitex 提供的中间件扩展。今天我们来结合官方文档以及源码，看看这一点是怎样支持的，以及给出几个常见的案例。

Endpoint

Kitex 对于 Middleware 的定义在 pkg/endpoint/endpoint.go 目录下

// Endpoint represent one method for calling from remote.
type Endpoint func(ctx context.Context, req, resp interface{}) (err error)

// Middleware deal with input Endpoint and output Endpoint.
type Middleware func(Endpoint) Endpoint
复制代码

Endpoint 和 Middleware 都是函数定义。我们可以看到，Endpoint 包含了RPC调用的上下文 context.Context，以及请求 req，响应 resp，返回一个 error 标识调用的状态。在调用发起前，resp 还是一个空对象，直到 invoke 远端的接口才会给 resp 赋值。

从便于理解的角度看，可以理解为，业务 server 所实现的类似下面这样函数签名的接口，本质上就是一个 Endpoint。

func(ctx context.Context, req interface{}) (resp interface{}, err error)
复制代码

Middleware

从上一节的函数签名就可以看出，Middleware 就是基于 Endpoint 实现的扩展能力。

作为一个业务开发者，我们可能有各种各样针对「请求」，「响应」，「context.Context」，乃至耗时的观察处理，这样的业务诉求可能在请求前，也可能在拿到响应后。

Middleware 提供的能力就是对于 Endpoint 的嵌套，Context + req + resp + error 都给你，要做什么处理你直接自己看就 ok。开发者无需，也不能感知外层还有什么 Endpoint 包裹着自己。你只需要知道，在给你的这个 Endpoint 最内层，包含了实际执行 RPC 接口逻辑的那个 Endpoint 就好，你可以基于这个假设来开发自己的中间件逻辑。至于中间是不是还有别的 Middleware，乃至自己这个 Middleware 执行之后会不会还有别的更上一层的 Middleware，这些无需关心。

中间件是串连使用的，通过调用传入的 next，可以得到后一个中间件返回的 response（如果有）和 err，据此作出相应处理后，向前一个中间件返回 err（务必判断 next err 返回，勿吞了 err）或者设置 response。

扫描二维码关注公众号，回复： 14228397 查看本文章

通过源码里组合多个 Middleware 的函数也可以发现，这里就是链式调用，在你的 Middleware 中 return 的 Endpoint 就会作为外层 Middleware 的参数。

// Chain connect middlewares into one middleware.
func Chain(mws ...Middleware) Middleware {
	return func(next Endpoint) Endpoint {
		for i := len(mws) - 1; i >= 0; i-- {
			next = mws[i](next)
		}
		return next
	}
}

// Build builds the given middlewares into one middleware.
func Build(mws []Middleware) Middleware {
	if len(mws) == 0 {
		return DummyMiddleware
	}
	return func(next Endpoint) Endpoint {
		return mws[0](Build(mws[1:])(next))
	}
}

// DummyMiddleware is a dummy middleware.
func DummyMiddleware(next Endpoint) Endpoint {
	return next
}

// DummyEndpoint is a dummy endpoint.
func DummyEndpoint(ctx context.Context, req, resp interface{}) (err error) {
	return nil
}
复制代码

生效逻辑

参考官方示例，我们先写一个 Middleware，假设我们需要打印请求和响应出来：

func PrintRequestResponseMW(next endpoint.Endpoint) endpoint.Endpoint {
    return func(ctx context.Context, request, response interface{}) error {
        fmt.Printf("request: %v\n", request)
        err := next(ctx, request, response)
        fmt.Printf("response: %v", response)
        return err
    }
}
复制代码

上一节的 DummyMiddleware 大家可以注意一下，如果什么都不做，只是把原有的入参 Endpoint 返回，就相当于一个空实现，只是包了一层，没有任何逻辑。

那么在 PrintRequestResponseMW 里面，我们其实只是在调用 next 这个 Endpoint 前后打印了 request 和 response。（还记得么？这里的 err := next(ctx, request, response) 可以简化理解为就是实际 RPC 调用的接口逻辑）

下面我们来看看应该如何让中间件生效。

参考官方文档说明：

在扩展过程中，要记得两点原则：

• 中间件和套件都只允许在初始化 Server、Client 的时候设置，不允许动态修改。
• 后设置的会覆盖先设置的。

简单说，我们只能设置 server 或 client 级别的 middleware，一个 server/client 对应一组中间件，不能更细粒度，比如到方法级。下面我们分别看看怎样在双端设置中间件。

Server 中间件

server/option.go 里面提供了 WithMiddleware 函数，我们可以据此来添加 server middleware：

// WithMiddleware adds middleware for server to handle request.
func WithMiddleware(mw endpoint.Middleware) Option {
	mwb := func(ctx context.Context) endpoint.Middleware {
		return mw
	}
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMiddleware(%+v)", utils.GetFuncName(mw)))
		o.MWBs = append(o.MWBs, mwb)
	}}
}
复制代码

使用起来其实很简单，import "github.com/cloudwego/kitex/server"，然后直接 server.WithMiddleware(PrintRequestResponseMW) ，随后在 NewServer 时传入即可，当然可以把 PrintRequestResponseMW 替换成任何符合业务场景的中间件。

（kstudy 就是第一篇文章中我们创建的 demo 项目）

这里的 Option 其实我们在系列第一篇文章也提到过，经典的设计模式，Kitex 提供了相关服务治理的能力，如果你需要，比如这个创建 server Middleware 的诉求，那么就直接一个 WithMiddleware 传入进来就ok。框架会感知这个 Opiton，应用到 server 的配置中。

// Option is the only way to config a server.
type Option struct {
	F func(o *Options, di *utils.Slice)
}
复制代码

Client 中间件

跟 server 中间件类似，client middleware 也是通过 option 的形式配置的。

你需要 import "github.com/cloudwego/kitex/client"，然后直接用 client.WithMiddleware(XXX) 即可。

// WithMiddleware adds middleware for client to handle request.
func WithMiddleware(mw endpoint.Middleware) Option {
	mwb := func(ctx context.Context) endpoint.Middleware {
		return mw
	}
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMiddleware(%+v)", utils.GetFuncName(mw)))
		o.MWBs = append(o.MWBs, mwb)
	}}
}
复制代码

二者实现唯一的区别在于 F 依赖的 Options 结构体不同，一个是 *internal_server.Options，一个是 *client.Options。

实战场景

打印请求响应

func LogParam(next endpoint.Endpoint) endpoint.Endpoint {
	return func(ctx context.Context, req, resp interface{}) (err error) {
		fmt.Printf("request:\n%s", req)

		err = next(ctx, req, resp)
		if err != nil {
			fmt.Printf("request faild, error: %+v", err)
		}

		fmt.Printf("response:\n%s", resp)
		return err
	}
}
复制代码

线上直接打印确实有风险，因为你不知道，或者知道了也可能不注意，请求体或者响应体可能会很大，每次全都打印其实是一个很耗性能的操作。另一点在于安全，直接在 server/client 层面全部打印，会导致你的日志里出现很多敏感信息，尤其是B端业务。数据的安全性如果不能得到保障，客户的信息随随便便一搜日志就全都明文出现，是有很大的隐患的。如果必须要做，请一定要评估业务场景，做好加密。

recover panic

func RecoverPanic(next endpoint.Endpoint) endpoint.Endpoint {
	return func(ctx context.Context, req, resp interface{}) (err error) {
		defer func() {
			if e := recover(); e != nil {
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				fmt.Printf("KITEX: panic in handler: %s: %s", e, buf)
			}
		}()
		return next(ctx, req, resp)
	}
}
复制代码

打印慢请求

func SlowReqLogMW(next endpoint.Endpoint) endpoint.Endpoint {
        // 默认慢请求阈值为500ms
        var threshold int64 = 500
	return func(ctx context.Context, req, resp interface{}) (err error) {
		begin := time.Now()
		err = next(ctx, req, resp)
		took := time.Since(begin).Nanoseconds() / 1e6 // ms
		if took > threshold {
			fmt.Printf("slow request, cost=%v", took)
		}
		return err
	}
}
复制代码

Suite 扩展

Suite 是一种更高层次的组合和封装，更加推荐第三方开发者能够基于 Suite 对外提供 Kitex 的扩展，Suite 可以允许在创建的时候，动态地去注入一些值，或者在运行时动态地根据自身的某些值去指定自己的 middleware 中的值，这使得用户的使用以及第三方开发者的开发都更加地方便，无需再依赖全局变量，也使得每个 client 使用不同的配置成为可能。

Suite 是对于 Option 和 Middleware（通过 Option 设置）的组合和封装。使得我们可以直接定制自己的业务插件，只要实现 Options() []Option 方法即可。

// A Suite is a collection of Options. It is useful to assemble multiple associated
// Options as a single one to keep the order or presence in a desired manner.
type Suite interface {
	Options() []Option
}

// WithSuite adds an option suite for server.
func WithSuite(suite Suite) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		var nested struct {
			Suite   string
			Options utils.Slice
		}
		nested.Suite = fmt.Sprintf("%T(%+v)", suite, suite)

		for _, op := range suite.Options() {
			op.F(o, &nested.Options)
		}
		di.Push(nested)
	}}
}
复制代码

Server 端和 Client 端都是通过 WithSuite 这个方法来启用新的套件。用法跟前一节提到的 Middleware 基本一样，在 client 以及 server 两个包下都提供了 WithSuite 的函数，直接用即可。

事实上字节内部也是依赖的开源 Kitex 实现，对于一些内部的配置，比如 mesh，限流，ACL 等，都是通过 Suite 实现的。只是调整了代码生成工具的逻辑，补充上 byted suite，这一点也体现了强大的扩展能力。生成的代码类似这样：

// NewServer creates a server.Server with the given handler and options.
func NewServer(handler XXXXService, opts ...server.Option) server.Server {
	var options []server.Option
	options = append(options, byted.ServerSuite(serviceInfo()))
	options = append(options, opts...)

	svr := server.NewServer(options...)
	if err := svr.RegisterService(serviceInfo(), handler); err != nil {
		panic(err)
	}
	return svr
}
复制代码

附录

• 官方文档 Middleware 扩展
• Kitex 框架入门系列（1）