除了捕捉 panic、延迟释放资源外，我们日常编码中还有哪些使用 defer 的小技巧呢？
使用 defer 可以跟踪函数的执行过程。
// trace.go
package main
func Trace(name string) func() {
println("enter:", name)
return func() {
println("exit:", name)
}
}
func foo() {
defer Trace("foo")()
bar()
}
func bar() {
defer Trace("bar")()
}
func main() {
defer Trace("main")()
foo()
}
在讲解这段代码的原理之前，我们先看一下这段代码的执行结果，直观感受一下什么是函数调用跟踪：
enter: main
enter: foo
enter: bar
exit: bar
exit: foo
exit: main
Go 会在 defer 设置 deferred 函数时对 defer 后面的表达式进行求值。
上面的 Trace 函数只接受一个参数，˙这个参数代表函数名，Trace 会首先打印进入某函数的日志，比如：“enter: foo”。然后返回一个闭包函数，这个闭包函数一旦被执行，就会输出离开某函数的日志。在 foo 函数中，这个由 Trace 函数返回的闭包函数就被设置为了 deferred 函数，于是当 foo 函数返回后，这个闭包函数就会被执行，输出“exit: foo”的日志。
搞清楚上面跟踪函数调用链的实现原理后，我们再来看看这个实现。我们会发现这里还是有一些“瑕疵”，也就是离我们期望的“跟踪函数调用链”的实现还有一些不足之处。这里我列举了几点：
- 调用 Trace 时需手动显式传入要跟踪的函数名；
- 如果是并发应用，不同 Goroutine 中函数链跟踪混在一起无法分辨；
- 输出的跟踪结果缺少层次感，调用关系不易识别；
- 对要跟踪的函数，需手动调用 Trace 函数。
那么，这一讲我们的任务就是逐一分析并解决上面提出的这几点问题进行，经过逐步地代码演进，最终实现一个自动注入跟踪代码，并输出有层次感的函数调用链跟踪命令行工具。
自动获取所跟踪函数的函数名
这里，我给出了新版 Trace 函数的实现以及它的使用方法，我们先看一下：
// trace1/trace.go
func Trace() func() {
pc, _, _, ok := runtime.Caller(1)
if !ok {
panic("not found caller")
}
fn := runtime.FuncForPC(pc)
name := fn.Name()
println("enter:", name)
return func() { println("exit:", name) }
}
func foo() {
defer Trace()()
bar()
}
func bar() {
defer Trace()()
}
func main() {
defer Trace()()
foo()
}
在这一版 Trace 函数中，我们通过 runtime.Caller 函数获得当前 Goroutine 的函数调用栈上的信息，runtime.Caller 的参数标识的是要获取的是哪一个栈帧的信息。当参数为 0 时，返回的是 Caller 函数的调用者的函数信息，在这里就是 Trace 函数。但我们需要的是 Trace 函数的调用者的信息，于是我们传入 1。
Caller 函数有四个返回值：第一个返回值代表的是程序计数（pc）；第二个和第三个参数代表对应函数所在的源文件名以及所在行数，这里我们暂时不需要；最后一个参数代表是否能成功获取这些信息，如果获取失败，我们抛出 panic。
接下来，我们通过 runtime.FuncForPC 函数和程序计数器（PC）得到被跟踪函数的函数名称。我们运行一下改造后代码：
enter: main.main
enter: main.foo
enter: main.bar
exit: main.bar
exit: main.foo
exit: main.main
接下来，我们来解决第二个问题，也就是当程序中有多 Goroutine 时，Trace 输出的跟踪信息混杂在一起难以分辨的问题。
增加 Goroutine 标识
那么，接下来我们还继续对 Trace 函数进行改造，让它支持多 Goroutine 函数调用链的跟踪。我们的方案就是在输出的函数出入口信息时，带上一个在程序每次执行时能唯一区分 Goroutine 的 Goroutine ID。
Go 核心团队为了避免Goroutine ID 的滥用，故意没有将 Goroutine ID 暴露给开发者。 但在 Go 标准库的 h2_bundle.go 中，我们却发现了一个获取 Goroutine ID 的标准方法，看下面代码：
// $GOROOT/src/net/http/h2_bundle.go
var http2goroutineSpace = []byte("goroutine ")
func http2curGoroutineID() uint64 {
bp := http2littleBuf.Get().(*[]byte)
defer http2littleBuf.Put(bp)
b := *bp
b = b[:runtime.Stack(b, false)]
// Parse the 4707 out of "goroutine 4707 ["
b = bytes.TrimPrefix(b, http2goroutineSpace)
i := bytes.IndexByte(b, ' ')
if i < 0 {
panic(fmt.Sprintf("No space found in %q", b))
}
b = b[:i]
n, err := http2parseUintBytes(b, 10, 64)
if err != nil {
panic(fmt.Sprintf("Failed to parse goroutine ID out of %q: %v", b, err))
}
return n
}
不过，由于 http2curGoroutineID 不是一个导出函数，我们无法直接使用。我们可以把它复制出来改造一下：
// trace2/trace.go
var goroutineSpace = []byte("goroutine ")
func curGoroutineID() uint64 {
b := make([]byte, 64)
b = b[:runtime.Stack(b, false)]
// Parse the 4707 out of "goroutine 4707 ["
b = bytes.TrimPrefix(b, goroutineSpace)
i := bytes.IndexByte(b, ' ')
if i < 0 {
panic(fmt.Sprintf("No space found in %q", b))
}
b = b[:i]
n, err := strconv.ParseUint(string(b), 10, 64)
if err != nil {
panic(fmt.Sprintf("Failed to parse goroutine ID out of %q: %v", b, err))
}
return n
}
这里，我们改造了两个地方。一个地方是通过直接创建一个 byte 切片赋值给 b，替代原 http2curGoroutineID 函数中从一个 pool 池获取 byte 切片的方式，另外一个是使用 strconv.ParseUint 替代了原先的 http2parseUintBytes。改造后，我们就可以直接使用 curGoroutineID 函数来获取 Goroutine 的 ID 信息了。
好，接下来，我们在 Trace 函数中添加 Goroutine ID 信息的输出：
// trace2/trace.go
func Trace() func() {
pc, _, _, ok := runtime.Caller(1)
if !ok {
panic("not found caller")
}
fn := runtime.FuncForPC(pc)
name := fn.Name()
gid := curGoroutineID()
fmt.Printf("g[%05d]: enter: [%s]\n", gid, name)
return func() { fmt.Printf("g[%05d]: exit: [%s]\n", gid, name) }
}
从上面代码看到，我们在出入口输出的跟踪信息中加入了 Goroutine ID 信息，我们输出的 Goroutine ID 为 5 位数字，如果 ID 值不足 5 位，则左补零，这一切都是 Printf 函数的格式控制字符串“%05d”帮助我们实现的。这样对齐 Goroutine ID 的位数，为的是输出信息格式的一致性更好。如果你的 Go 程序中 Goroutine 的数量超过了 5 位数可以表示的数值范围，也可以自行调整控制字符串。
接下来，我们也要对示例进行一些调整，将这个程序由单 Goroutine 改为多 Goroutine 并发的，这样才能验证支持多 Goroutine 的新版 Trace 函数是否好用：
// trace2/trace.go
func A1() {
defer Trace()()
B1()
}
func B1() {
defer Trace()()
C1()
}
func C1() {
defer Trace()()
D()
}
func D() {
defer Trace()()
}
func A2() {
defer Trace()()
B2()
}
func B2() {
defer Trace()()
C2()
}
func C2() {
defer Trace()()
D()
}
func main() {
var wg sync.WaitGroup
wg.Add(1)
go func() {
A2()
wg.Done()
}()
A1()
wg.Wait()
}
新示例程序共有两个 Goroutine，main groutine 的调用链为A1 -> B1 -> C1 -> D，而另外一个 Goroutine 的函数调用链为A2 -> B2 -> C2 -> D。我们来看一下这个程序的执行结果是否和原代码中两个 Goroutine 的调用链一致：
g[00001]: enter: [main.A1]
g[00001]: enter: [main.B1]
g[00018]: enter: [main.A2]
g[00001]: enter: [main.C1]
g[00001]: enter: [main.D]
g[00001]: exit: [main.D]
g[00001]: exit: [main.C1]
g[00001]: exit: [main.B1]
g[00001]: exit: [main.A1]
g[00018]: enter: [main.B2]
g[00018]: enter: [main.C2]
g[00018]: enter: [main.D]
g[00018]: exit: [main.D]
g[00018]: exit: [main.C2]
g[00018]: exit: [main.B2]
g[00018]: exit: [main.A2]
让输出的跟踪信息更具层次感
对于程序员来说，缩进是最能体现出“层次感”的方法，如果我们将上面示例中 Goroutine 00001 的函数调用跟踪信息以下面的形式展示出来，函数的调用顺序是不是更加一目了然了呢？
g[00001]: ->main.A1
g[00001]: ->main.B1
g[00001]: ->main.C1
g[00001]: ->main.D
g[00001]: <-main.D
g[00001]: <-main.C1
g[00001]: <-main.B1
g[00001]: <-main.A1
那么我们就以这个形式为目标，考虑如何实现输出这种带缩进的函数调用跟踪信息。我们还是直接上代码吧：
// trace3/trace.go
func printTrace(id uint64, name, arrow string, indent int) {
indents := ""
for i := 0; i < indent; i++ {
indents += " "
}
fmt.Printf("g[%05d]:%s%s%s\n", id, indents, arrow, name)
}
var mu sync.Mutex
var m = make(map[uint64]int)
func Trace() func() {
pc, _, _, ok := runtime.Caller(1)
if !ok {
panic("not found caller")
}
fn := runtime.FuncForPC(pc)
name := fn.Name()
gid := curGoroutineID()
mu.Lock()
indents := m[gid] // 获取当前gid对应的缩进层次
m[gid] = indents + 1 // 缩进层次+1后存入map
mu.Unlock()
printTrace(gid, name, "->", indents+1)
return func() {
mu.Lock()
indents := m[gid] // 获取当前gid对应的缩进层次
m[gid] = indents - 1 // 缩进层次-1后存入map
mu.Unlock()
printTrace(gid, name, "<-", indents)
}
}
在上面这段代码中，我们使用了一个 map 类型变量 m 来保存每个 Goroutine 当前的缩进信息：m 的 key 为 Goroutine 的 ID，值为缩进的层次。然后，考虑到 Trace 函数可能在并发环境中运行，根据我们在第 16 讲中提到的“map 不支持并发写”的注意事项，我们增加了一个 sync.Mutex 实例 mu 用于同步对 m 的写操作。
除了增加缩进层次信息外，在这一版的 Trace 函数实现中，我们也把输出出入口跟踪信息的操作提取到了一个独立的函数 printTrace 中，这个函数会根据传入的 Goroutine ID、函数名、箭头类型与缩进层次值，按预定的格式拼接跟踪信息并输出。
g[00001]: ->main.A1
g[00001]: ->main.B1
g[00001]: ->main.C1
g[00001]: ->main.D
g[00001]: <-main.D
g[00001]: <-main.C1
g[00001]: <-main.B1
g[00001]: <-main.A1
g[00018]: ->main.A2
g[00018]: ->main.B2
g[00018]: ->main.C2
g[00018]: ->main.D
g[00018]: <-main.D
g[00018]: <-main.C2
g[00018]: <-main.B2
g[00018]: <-main.A2
利用代码生成自动注入 Trace 函数
将 Trace 函数放入一个独立的 module 中
我们创建一个名为 instrument_trace 的目录，进入这个目录后，通过 go mod init 命令创建一个名为 github.com/bigwhite/instrument_trace 的 module：
$mkdir instrument_trace
$cd instrument_trace
$go mod init github.com/bigwhite/instrument_trace
go: creating new go.mod: module github.com/bigwhite/instrument_trace
接下来，我们将最新版的 trace.go 放入到该目录下，将包名改为 trace，并仅保留 Trace 函数、Trace 使用的函数以及包级变量，其他函数一律删除掉。这样，一个独立的 trace 包就提取完毕了。
接下来，我们将最新版的 trace.go 放入到该目录下，将包名改为 trace，并仅保留 Trace 函数、Trace 使用的函数以及包级变量，其他函数一律删除掉。这样，一个独立的 trace 包就提取完毕了。
// instrument_trace/example_test.go
package trace_test
import (
trace "github.com/bigwhite/instrument_trace"
)
func a() {
defer trace.Trace()()
b()
}
func b() {
defer trace.Trace()()
c()
}
func c() {
defer trace.Trace()()
d()
}
func d() {
defer trace.Trace()()
}
func ExampleTrace() {
a()
// Output:
// g[00001]: ->github.com/bigwhite/instrument_trace_test.a
// g[00001]: ->github.com/bigwhite/instrument_trace_test.b
// g[00001]: ->github.com/bigwhite/instrument_trace_test.c
// g[00001]: ->github.com/bigwhite/instrument_trace_test.d
// g[00001]: <-github.com/bigwhite/instrument_trace_test.d
// g[00001]: <-github.com/bigwhite/instrument_trace_test.c
// g[00001]: <-github.com/bigwhite/instrument_trace_test.b
// g[00001]: <-github.com/bigwhite/instrument_trace_test.a
}
在 example_test.go 文件中，我们用 ExampleXXX 形式的函数表示一个示例，go test 命令会扫描 example_test.go 中的以 Example 为前缀的函数并执行这些函数。
自动注入 Trace 函数
现在，我们在 instrument_trace module 下面增加一个命令行工具，这个工具可以以一个 Go 源文件为单位，自动向这个 Go 源文件中的所有函数注入 Trace 函数。
变化后的 instrument_trace module 的目录结构：
$tree ./instrument_trace -F
./instrument_trace
├── Makefile
├── cmd/
│ └── instrument/
│ └── main.go # instrument命令行工具的main包
├── example_test.go
├── go.mod
├── go.sum
├── instrumenter/ # 自动注入逻辑的相关结构
│ ├── ast/
│ │ └── ast.go
│ └── instrumenter.go
└── trace.go
我们先来看一下 cmd/instrument/main.go 源码，然后自上而下沿着 main 函数的调用逻辑逐一看一下这个功能的实现。下面是 main.go 的源码：
// instrument_trace/cmd/instrument/main.go
... ...
var (
wrote bool
)
func init() {
flag.BoolVar(&wrote, "w", false, "write result to (source) file instead of stdout")
}
func usage() {
fmt.Println("instrument [-w] xxx.go")
flag.PrintDefaults()
}
func main() {
fmt.Println(os.Args)
flag.Usage = usage
flag.Parse() // 解析命令行参数
if len(os.Args) < 2 { // 对命令行参数个数进行校验
usage()
return
}
var file string
if len(os.Args) == 3 {
file = os.Args[2]
}
if len(os.Args) == 2 {
file = os.Args[1]
}
if filepath.Ext(file) != ".go" { // 对源文件扩展名进行校验
usage()
return
}
var ins instrumenter.Instrumenter // 声明instrumenter.Instrumenter接口类型变量
// 创建以ast方式实现Instrumenter接口的ast.instrumenter实例
ins = ast.New("github.com/bigwhite/instrument_trace", "trace", "Trace")
newSrc, err := ins.Instrument(file) // 向Go源文件所有函数注入Trace函数
if err != nil {
panic(err)
}
if newSrc == nil {
// add nothing to the source file. no change
fmt.Printf("no trace added for %s\n", file)
return
}
if !wrote {
fmt.Println(string(newSrc)) // 将生成的新代码内容输出到stdout上
return
}
// 将生成的新代码内容写回原Go源文件
if err = ioutil.WriteFile(file, newSrc, 0666); err != nil {
fmt.Printf("write %s error: %v\n", file, err)
return
}
fmt.Printf("instrument trace for %s ok\n", file)
}
作为命令行工具，instrument 使用标准库的 flag 包实现对命令行参数（这里是 -w）的解析，通过 os.Args 获取待注入的 Go 源文件路径。在完成对命令行参数个数与值的校验后，instrument 程序声明了一个 instrumenter.Instrumenter 接口类型变量 ins，然后创建了一个实现了 Instrumenter 接口类型的 ast.instrumenter 类型的实例，并赋值给变量 ins。
instrumenter.Instrumenter 接口类型的声明放在了 instrumenter/instrumenter.go 中：
type Instrumenter interface {
Instrument(string) ([]byte, error)
}
这里我们看到，这个接口类型的方法列表中只有一个方法 Instrument，这个方法接受一个 Go 源文件路径，返回注入了 Trace 函数的新源文件内容以及一个 error 类型值，作为错误状态标识。我们之所以要抽象出一个接口类型，考虑的就是注入 Trace 函数的实现方法不一，为后续的扩展做好预留。
在这个例子中，我们默认提供了一种自动注入 Trace 函数的实现，那就是 ast.instrumenter，它注入 Trace 的实现原理是这样的：
从原理图中我们可以清楚地看到，在这一实现方案中，我们先将传入的 Go 源码转换为抽象语法树。
在计算机科学中，抽象语法树（abstract syntax tree，AST）是源代码的抽象语法结构的树状表现形式，树上的每个节点都表示源代码中的一种结构。因为 Go 语言是开源编程语言，所以它的抽象语法树的操作包也和语言一起开放给了 Go 开发人员，我们可以基于 Go 标准库以及Go 实验工具库提供的 ast 相关包，快速地构建基于 AST 的应用，这里的 ast.instrumenter 就是一个应用 AST 的典型例子。
一旦我们通过 ast 相关包解析 Go 源码得到相应的抽象语法树后，我们便可以操作这棵语法树，并按我们的逻辑在语法树中注入我们的 Trace 函数，最后我们再将修改后的抽象语法树转换为 Go 源码，就完成了整个自动注入的工作了。
下面是 ast.instrumenter 的 Instructment 方法的代码：
// instrument_trace/instrumenter/ast/ast.go
func (a instrumenter) Instrument(filename string) ([]byte, error) {
fset := token.NewFileSet()
curAST, err := parser.ParseFile(fset, filename, nil, parser.ParseComments) // 解析Go源码，得到AST
if err != nil {
return nil, fmt.Errorf("error parsing %s: %w", filename, err)
}
if !hasFuncDecl(curAST) { // 如果整个源码都不包含函数声明，则无需注入操作，直接返回。
return nil, nil
}
// 在AST上添加包导入语句
astutil.AddImport(fset, curAST, a.traceImport)
// 向AST上的所有函数注入Trace函数
a.addDeferTraceIntoFuncDecls(curAST)
buf := &bytes.Buffer{}
err = format.Node(buf, fset, curAST) // 将修改后的AST转换回Go源码
if err != nil {
return nil, fmt.Errorf("error formatting new code: %w", err)
}
return buf.Bytes(), nil // 返回转换后的Go源码
}
通过代码，我们看到 Instrument 方法的基本步骤与上面原理图大同小异。Instrument 首先通过 go/paser 的 ParserFile 函数对传入的 Go 源文件中的源码进行解析，并得到对应的抽象语法树 AST，然后向 AST 中导入 Trace 函数所在的包，并向这个 AST 的所有函数声明注入 Trace 函数调用。
实际的注入操作发生在 instrumenter 的 addDeferTraceIntoFuncDecls 方法中，我们来看一下这个方法的实现：
// instrument_trace/instrumenter/ast/ast.go
func (a instrumenter) addDeferTraceIntoFuncDecls(f *ast.File) {
for _, decl := range f.Decls { // 遍历所有声明语句
fd, ok := decl.(*ast.FuncDecl) // 类型断言：是否为函数声明
if ok {
// 如果是函数声明，则注入跟踪设施
a.addDeferStmt(fd)
}
}
}
这个方法的逻辑十分清晰，就是遍历语法树上所有声明语句，如果是函数声明，就调用 instrumenter 的 addDeferStmt 方法进行注入，如果不是，就直接返回。addDeferStmt 方法的实现如下：
// instrument_trace/instrumenter/ast/ast.go
func (a instrumenter) addDeferStmt(fd *ast.FuncDecl) (added bool) {
stmts := fd.Body.List
// 判断"defer trace.Trace()()"语句是否已经存在
for _, stmt := range stmts {
ds, ok := stmt.(*ast.DeferStmt)
if !ok {
// 如果不是defer语句，则继续for循环
continue
}
// 如果是defer语句，则要进一步判断是否是defer trace.Trace()()
ce, ok := ds.Call.Fun.(*ast.CallExpr)
if !ok {
continue
}
se, ok := ce.Fun.(*ast.SelectorExpr)
if !ok {
continue
}
x, ok := se.X.(*ast.Ident)
if !ok {
continue
}
if (x.Name == a.tracePkg) && (se.Sel.Name == a.traceFunc) {
// defer trace.Trace()()已存在，返回
return false
}
}
// 没有找到"defer trace.Trace()()"，注入一个新的跟踪语句
// 在AST上构造一个defer trace.Trace()()
ds := &ast.DeferStmt{
Call: &ast.CallExpr{
Fun: &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{
Name: a.tracePkg,
},
Sel: &ast.Ident{
Name: a.traceFunc,
},
},
},
},
}
newList := make([]ast.Stmt, len(stmts)+1)
copy(newList[1:], stmts)
newList[0] = ds // 注入新构造的defer语句
fd.Body.List = newList
return true
}
虽然 addDeferStmt 函数体略长，但逻辑也很清晰，就是先判断函数是否已经注入了 Trace，如果有，则略过；如果没有，就构造一个 Trace 语句节点，并将它插入到 AST 中。
Instrument 的最后一步就是将注入 Trace 后的 AST 重新转换为 Go 代码，这就是我们期望得到的带有 Trace 特性的 Go 代码了。
利用 instrument 工具注入跟踪代码
这里，我在 instrument_trace 项目的 examples 目录下建立了一个名为 demo 的项目，我们就来看看如何使用 instrument 工具为 demo 项目下的 demo.go 文件自动注入跟踪设施。demo.go 文件内容很简单：
// instrument_trace/examples/demo/demo.go
package main
func foo() {
bar()
}
func bar() {
}
func main() {
foo()
}
我们首先构建一下 instrument_trace 下的 instrument 工具：
$cd instrument_trace
$go build github.com/bigwhite/instrument_trace/cmd/instrument
$instrument version
[instrument version]
instrument [-w] xxx.go
-w write result to (source) file instead of stdout
接下来，我们使用 instrument 工具向 examples/demo/demo.go 源文件中的函数自动注入跟踪设施：
$instrument -w examples/demo/demo.go
[instrument -w examples/demo/demo.go]
instrument trace for examples/demo/demo.go ok
注入后的 demo.go 文件变为了下面这个样子：
// instrument_trace/examples/demo/demo.go
package main
import "github.com/bigwhite/instrument_trace"
func foo() {
defer trace.Trace()()
bar()
}
func bar() {
defer trace.Trace()()
}
func main() {
defer trace.Trace()()
foo()
}
此时，如果我们再对已注入 Trace 函数的 demo.go 执行一次 instrument 命令，由于 instrument 会判断 demo.go 各个函数已经注入了 Trace，demo.go 的内容将保持不变。
由于 github.com/bigwhite/instrument_trace 并没有真正上传到 github.com 上，所以如果你要运行 demo.go，我们可以为它配置一个下面这样的 go.mod：
// instrument_trace/examples/demo/go.mod
module demo
go 1.17
require github.com/bigwhite/instrument_trace v1.0.0
replace github.com/bigwhite/instrument_trace v1.0.0 => ../../
这样运行 demo.go 就不会遇到障碍了：
$go run demo.go
g[00001]: ->main.main
g[00001]: ->main.foo
g[00001]: ->main.bar
g[00001]: <-main.bar
g[00001]: <-main.foo
g[00001]: <-main.main
此文章为3月Day18学习笔记，内容来源于极客时间《Tony Bai · Go 语言第一课》。
