go 如何优雅的实现函数级超时控制

2024-09-15 ⏳2.7分钟(1.1千字)

一般情况下后端的代码都是树状的，随着业务的堆积，支线可能越来越多，拿漫画的 ComicDetail 举例，这个接口的目的在于返回漫画的章节和详情信息，但是随着业务的增长可能它还返回了特典、个人追漫、打折等信息，一般情况下 golang 开发 err 都是往上抛的，有时候也会有 err = nil 的方案，进行业务降级。这里就会有一个问题 A 在调用边缘业务 B 时，你虽然忽略 B 的 err 但是 B 的耗时过长，仍然会影响主流程的正常返回，在一些 http/grpc/db/redis ..etc 的请求下我们会给请求上一个 timeout，下一个操作进行超时检测，当然，如果业务比较简单，封装比较少，这样还行，当领域划分多的情况下，很明显我们需要能给某个函数同样的加超时时间。

题外话，当然我们可以从软件工程和代码重构的角度来规避这些问题，比如 ComicDetail 就是仅返回漫画详情，其他的都拆出来，这会衍生另外的问题:

  1. 拆分粒度的粗细该当如何?
  2. 还有调用方的感受，过多的接口可能让调用方,比如 web 和 客户端的同学因此红温。

期望的方案

既然牵扯到降级，那么对业务代码的入侵是无法避免的，只有多少的问题。

泛型

假设我们有个函数如下:

func Add(a, b int) (r int) {
    return a+b
}

如果给它加超时控制的话, 那我可能期望是 Hook(Add, time.Second) Hook 函数有 Add 函数一样的返回值, 只是多了一个入参 time.Duration. 这样对业务代码的改造比较小, 入侵少负担就小.

起初的想法就是泛型, 我们能不能对原函数简单的包装达到目的？很遗憾泛型有局限：

func Hook[T any](ori T, t time.Duration) T {
        var x T
        return x
}

虽然可以声明变量 x，但是没办法实现 T, 只用泛型行不通。

不过我们还有旁门左道能用, 反射。

func Hook[T any](ori T, t time.Duration) T {
     v := reflect.MakeFunc(reflect.TypeOf(ori), func(in []reflect.Value) []reflect.Value {
             // 超时控制
             // ...
             f := reflect.ValueOf(ori)
             return f.Call(in)
     })
     return v.Interface().(T)
}

对性能要求不高的朋友到这里就可以不用往下看了, 实测下来反射的性能比直接调用差了几百倍, 酌情使用。

装饰器

计算机遇到不能解决的问题就加一层, 如果不要求跟原函数一样的函数签名(意味着更多的改动量), 装饰器模式也能实现。

func SyncWithTimeout(ctx context.Context, timeout time.Duration, task func(ctx context.Context) error) Resp {
	nctx := context.Background()
	// 创建一个有超时的上下文
	nctx, cancel := context.WithTimeout(nctx, timeout)
	defer cancel()

	// 创建一个通道，来接收任务执行的结果
	resultChan := make(chan error, 1)

	// 启动一个 goroutine 来执行任务
	go func() {
		resultChan <- task(nctx)
	}()

	// 选择监听任务的结果或上下文的超时
	select {
	case <-nctx.Done(): // 超时或取消
		return Resp{err: ErrTimeout}
	case err := <-resultChan: // 正常完成
		return Resp{err: err}
	}
}

这个时候我们如果对 Add 进行超时控制就要这么改:

var res int
chronos.SyncWithTimeout(ctx, time.Millisecond*20, func(ctx context.Context) error {
    res = Add(a, b)
	return nil
}).Do(func() {
        // not timeout
})

本来简单的一行 Add(a, b) 现在要写很多行, 而且需要处理变量竞争问题, 很明显这个方案不是很优雅, 问题变复杂了。

代码生成

要说优雅, 我们能不能给 Add(a, b) 加行注释就搞定呢？比如:

// @timeout
func Add(a, b int) int

看过我之前的给 gopls 贡献源码的同学可能就立马想到了, 我们可以解析 ast 来构造一个Hook(Add, time.Second), 既然选择代码生成, 这个形式也没有那么符合直觉, 我们可以期望 AddWithTimeout(a, b, time.Duration) 的写法, 且返回值保持一致。

思路:

  1. 扫描项目所有包含关键字 @timeout  的 package

  2. 加载 package 构造抽象语法树

cfg := &packages.Config{
	Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax | packages.NeedImports,
	Tests:      false,
	Logf:       g.logf,
}

pkgs, err := packages.Load(cfg, patterns...)

这里详细的内容可以参考 golang.org/x/tools/go/packages, 会按照 package 加载抽象语法树. 注意具体加载哪些数据是要进行 mode 字段调整的, 不指定就不会去解析.

  1. 通过扫描 comment 来锚定函数
// 遍历所有包
for _, pkg := range pkgs {
	// 遍历包中的每个 Go 文件
	for _, syntax := range pkg.Syntax {
		// 遍历文件中的声明
		for _, decl := range syntax.Decls {
			// 只处理函数声明
			funcDecl, ok := decl.(*ast.FuncDecl)
			if !ok {
				continue
			}
			// 检查函数声明中的注释
			if funcDecl.Doc == nil {
				continue

			}
			for _, comment := range funcDecl.Doc.List {
				// 查找 go:generate timeout 指令
				if strings.HasPrefix(comment.Text, "// @timeout") {
					continue
				}
				pos := pkg.Fset.Position(funcDecl.Pos())
				fmt.Printf("Found @timeout in %s at line %d for function %s: %s, type:%+v\n",
					pos.Filename, pos.Line, funcDecl.Name.Name, comment.Text, funcDecl.Type)
				g.timeFuncs = append(g.timeFuncs, funcDecl)
			}
		}
	}
}

锚定函数也很简单, 只要根据注释反查函数即可.

  1. 使用装饰器模式操作函数签名来构造对应的超时函数

获取到函数签名后这里其实就很简单了, 只要按照自己的逻辑去拼接代码然后写入文件即可。

测试下效果:

// @timeout
func GetByID(id int) (string, error) {
	return "foo", nil
}

执行 timeout -r ., 得到以下结果:

func GetByIDWithTimeout(id int, timeout time.Duration) (string, error) {
	if timeout == 0 {
		return GetByID(id)
	}

	nctx := context.Background()
	nctx, cancel := context.WithTimeout(nctx, timeout)
	defer cancel()

	resultChan := make(chan error, 1)
	var r0 string
	var r1 error
	go func() {
		r0, r1 = GetByID(id)
		resultChan <- nil
	}()

	select {
	case <-nctx.Done():
		return r0, ErrTimeout
	case _ = <-resultChan:
		return r0, r1
	}

	return r0, r1
}

调整业务代码的时候也只需要把 GetByID(1) ==> GetByIDWithTimeout(1, time.Second). 跟其他方案比可以说改动是相当小了。

小结

源码可以参考 timeout.