golang 从零实现一个中间件(middleware)

实现中间件的背景

先看如下代码

package main

func hello(wr http.ResponseWriter, r *http.Request) {
    wr.Write([]byte("hello"))
}

func main() {
    http.HandleFunc("/", hello)
    err := http.ListenAndServe(":8080", nil)
    ...
}

这是一个简单的HTTP接口服务

现在有了一个新需求，需要在接口上增加服务的耗时处理日志，我们对上面的程序进行少量修改

func hello(wr http.ResponseWriter, r *http.Request) {
    timeStart := time.Now()
    wr.Write([]byte("hello"))
    timeElapsed := time.Since(timeStart)
    fmt.Println(timeElapsed)
}

随着业务的逐渐增多后，路由可能变成了

package main

func helloHandler(wr http.ResponseWriter, r *http.Request) {
    // ...
}

func showInfoHandler(wr http.ResponseWriter, r *http.Request) {
    // ...
}

func showEmailHandler(wr http.ResponseWriter, r *http.Request) {
    // ...
}

func showFriendsHandler(wr http.ResponseWriter, r *http.Request) {
    timeStart := time.Now()
    wr.Write([]byte("your friends is tom and alex"))
    timeElapsed := time.Since(timeStart)
    logger.Println(timeElapsed)
}

func main() {
    http.HandleFunc("/", helloHandler)
    http.HandleFunc("/info/show", showInfoHandler)
    http.HandleFunc("/email/show", showEmailHandler)
    http.HandleFunc("/friends/show", showFriendsHandler)
    // ...
}

此时，如果需要在每个路由上加一个参数信息日志，我们会发现虽然代码很好编写，但是需要去所有的handle里面修改这段代码，并且以后新增的handle也需要了解并添加这段代码。导致代码维护起来越来越麻烦。因此我们需要使用中间件，将业务代码和非业务代码（例如上面的接口耗时统计，请求参数日志等）拆分开

实现代码

package main

import(
	"fmt"
	"net/http"
	"sync"
)


// HandlerFunc defines the handler used by  middleware as return value.
type HandlerFunc func(*Context)

// HandlersChain defines a HandlerFunc array.
type HandlersChain []HandlerFunc

//定义的上下文
type Context struct {
	Request   *http.Request
	Writer    http.ResponseWriter
	handlers HandlersChain
	index    int8

}

//模拟的调用堆栈
func (c *Context) Next() {
	c.index++
	for c.index < int8(len(c.handlers)) {
		//按顺序执行HandlersChain内的函数
		//如果函数内无c.Next()方法调用则函数顺序执行完
		//如果函数内有c.Next()方法调用则代码执行到c.Next()方法处压栈，等待后面的函数执行完在回来执行c.Next()后的命令
		c.handlers[c.index](c)
		c.index++
	}
}
func (c *Context) reset() {
	c.handlers = nil
	c.index = -1
}

//中间件组
type RouterGroup struct {
	//存储定义的中间件
	Handlers HandlersChain
	engine   *Engine
}

func (group *RouterGroup) Use(middleware ...HandlerFunc) {
	group.Handlers = append(group.Handlers, middleware...)
}

func (group *RouterGroup) AddRoute(absolutePath string, handlers ...HandlerFunc) {
	handlers = group.combineHandlers(handlers)
	//建立路由和相关中间件组的绑定
	group.engine.addRoute(absolutePath, handlers)
}

//将定义的公用中间件和路由相关的中间件合并
func (group *RouterGroup) combineHandlers(handlers HandlersChain) HandlersChain {
	finalSize := len(group.Handlers) + len(handlers)
	mergedHandlers := make(HandlersChain, finalSize)
	copy(mergedHandlers, group.Handlers)
	copy(mergedHandlers[len(group.Handlers):], handlers)
	return mergedHandlers
}

type Engine struct{

	tree map[string]HandlersChain	// tree为了简化做成了map路由路径完全匹配
	RouterGroup
	pool             sync.Pool		// 正常情况存在大量的上下文切换，所以使用一个临时对象存储
}

func NewEngine() *Engine {
	engine := &Engine{
		RouterGroup: RouterGroup{
			Handlers: nil,
		},
		tree: make(map[string]HandlersChain),
	}
	engine.RouterGroup.engine = engine
	engine.pool.New = func() interface{} {
		return engine.allocateContext()
	}
	return engine
}

func (engine *Engine) allocateContext() *Context {
	return &Context{}
}

//url请求时，默认执行入口
func (engine *Engine) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	c := engine.pool.Get().(*Context)
	c.Writer = w
	c.Request = req
	c.reset()

	engine.handleHTTPRequest(c)

	engine.pool.Put(c)
}

func (engine *Engine) handleHTTPRequest(c *Context) {
	rPath := c.Request.URL.Path

	handlers := engine.getValue(rPath)
	if handlers != nil {
		c.handlers = handlers
		//按顺序执行中间件
		c.Next()
		return
	}

}

//获取路由下的相关HandlersChain
func (engine *Engine)getValue(path string)(handlers HandlersChain){
	handlers,ok := engine.tree[path]
	if !ok {
		return nil
	}
	return
}

func (engine *Engine) addRoute(path string, handlers HandlersChain) {
	engine.tree[path]=handlers
}

func (engine *Engine) Use(middleware ...HandlerFunc)  {
	engine.RouterGroup.Use(middleware...)
}

func main(){
	engine := NewEngine()
	engine.Use(func(c * Context){
		fmt.Println("begin middle1")
		fmt.Println("end middle1")
	})
	engine.Use(func(c * Context){
		fmt.Println("begin middle2")
		c.Next()
		fmt.Println("end middle2")
	})
	engine.Use(func(c * Context){
		fmt.Println("begin middle3")
		c.Next()
		fmt.Println("end middle3")
	})
	engine.AddRoute("/path1",func( c *Context){
		fmt.Println("path1")
		c.Writer.Write([]byte("path1"))

	})
	engine.AddRoute("/path2",func( c *Context){
		fmt.Println("path2")
		c.Writer.Write([]byte("path2"))

	})
	http.ListenAndServe(":8080", engine)
}

浏览器输入：http://localhost:8080/path1

执行结果

执行结果的堆栈流程

应用场景

• 对http的响应体进行压缩处理

• 设置一个特殊的路由，例如/ping，/healthcheck，用来给负载均衡一类的前置服务进行探活

• 打印请求处理处理日志，例如请求处理时间，请求路由

• 从请求头中读取X-Forwarded-For和X-Real-IP，将http.Request中的RemoteAddr修改为得到的RealIP

• 为本次请求生成单独的requestid，可一路透传，用来生成分布式调用链路，也可用于在日志中串连单次请求的所有逻辑

• 用context.Timeout设置超时时间，并将其通过http.Request一路透传下去

• 通过定长大小的channel存储token，并通过这些token对接口进行限流

参考资料

gin框架源码中间价实现

GO高级编程