【“骚”操作】AST实现代码调用跟踪

# 思考 当你接手一个新的Go语言项目时，如果前任开发者没有留下设计文档，或者文档已经过时、代码注释也不够清楚，那你可能会觉得很难快速上手。 这就跟旅行一样，如果你有一张地图，旅途就会顺利很多。因此我想到了一个叫go-callvis的工具，它可以通过分析项目的语法树来展示所有函数之间的调用关系，并以网页的形式呈现出来，方便查看。 不过，在实际使用中我发现这个工具生成网页图的速度有点慢。 为了解决这个问题，我自己开发了两个小工具来辅助。 # goanalysis 此前，我曾开发过这一工具，但遗憾的是，当时的代码并未妥善保存，也未上传至GitHub。因此，本次决定将该工具公开分享。此工具具备以下两大功能： 1. analysis: 通过项目生成相关调用链图，如：  2. track: 将项目内所有函数都内置上defer trace()，在运行过程中，将goroutine中所有调用的函数以及参数输出，比如： ```go time=2024-10-28T23:12:07.398+08:00 level=INFO msg=->example/inner/B.init.0 gid=1 params="[[]], " time=2024-10-28T23:12:07.415+08:00 level=INFO msg=*<-example/inner/B.init.0 gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=->example/inner/A.init.0 gid=1 params="[[]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*->example/inner/B.CalledA gid=1 params="[[]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**<-example/inner/B.CalledA gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*<-example/inner/A.init.0 gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=->main.main gid=1 params="[[]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*->example/inner/A.NewCallA gid=1 params="[[tly]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**<-example/inner/A.NewCallA gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*->example/inner/A.CallA.PrintB gid=1 params="[[]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**->example/inner/B.NewCallB gid=1 params="[[levi]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=***<-example/inner/B.NewCallB gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**->example/inner/B.CallB.PrintB gid=1 params="[[]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=***<-example/inner/B.CallB.PrintB gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**<-example/inner/A.CallA.PrintB gid=1 time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*->example/inner/A.RecursionA gid=1 params="[[%!s(int=1) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=**->example/inner/A.RecursionA gid=1 params="[[%!s(int=2) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=***->example/inner/A.RecursionA gid=1 params="[[%!s(int=3) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=****->example/inner/A.RecursionA gid=1 params="[[%!s(int=4) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*****->example/inner/A.RecursionA gid=1 params="[[%!s(int=5) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=******->example/inner/A.RecursionA gid=1 params="[[%!s(int=6) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*******->example/inner/A.RecursionA gid=1 params="[[%!s(int=7) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=********->example/inner/A.RecursionA gid=1 params="[[%!s(int=8) %!s(int=10)]], " time=2024-10-28T23:12:07.416+08:00 level=INFO msg=*********->example/inner/A.RecursionA gid=1 params="[[%!s(int=9) %!s(int=10)]], " time=2024-10-28T23:12:07.417+08:00 level=INFO msg=**********->example/inner/A.RecursionA gid=1 params="[[%!s(int=10) %!s(int=10)]], " time=2024-10-28T23:12:07.417+08:00 level=INFO msg=***********<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=**********<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=*********<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=********<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=*******<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=******<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=*****<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=****<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=***<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=**<-example/inner/A.RecursionA gid=1 time=2024-10-28T23:12:07.417+08:00 level=INFO msg=*<-main.main gid=1 ``` # 原理说明 ## analysis ### 通过标准库获取编译后callGraph ```go // 1. 加载项目代码所有的package名称 initial, _ := packages.Load(&packages.Config{}, p.Dir+"/...") if packages.PrintErrors(initial) > 0 { return fmt.Errorf("packages contain errors") } // 2. 基于指定的package名称，创建SSA项目（包含所有引用的包） prog, _ := ssautil.AllPackages(initial, 0) prog.Build() // 3. 通过不同的算法，获取相关的调用图 switch p.algo { case CallGraphTypeStatic: p.callGraph = static.CallGraph(prog) case CallGraphTypeCha: p.callGraph = cha.CallGraph(prog) case CallGraphTypeRta: mains, err := p.GetMainPackage(prog.AllPackages()) if err != nil { return err } var roots []*ssa.Function for _, main := range mains { roots = append(roots, main.Func("main")) } p.callGraph = rta.Analyze(roots, true).CallGraph case CallGraphTypeVta: p.callGraph = vta.CallGraph(ssautil.AllFunctions(prog), cha.CallGraph(prog)) } ``` ### 通过callgraph生成项目整体的语法树 ```go type FuncNode struct { Key string `json:"key"` // 唯一表示 Pkg string `json:"pkg"` Name string `json:"name"` Parent []string `json:"parent"` // 通过key来索引 Children []string `json:"children"` // 通过key来索引 } err := callgraph.GraphVisitEdges(p.callGraph, func(edge *callgraph.Edge) error { // 获取调用方 caller := edge.Caller // 获取被调用方 callee := edge.Callee if isSynthetic(edge) { return nil } // 排除标准库 if inStd(caller) || inStd(callee) { return nil } if isInter(edge) { return nil } // 排除需要忽略的库, 在启动前可以指定 if inIgnores(caller) || inIgnores(callee) { return nil } // caller是否存在 var pNode, qNode *FuncNode var ok bool // 如果不存在, 则创建 if pNode, ok = p.tree.Nodes[caller.String()]; !ok { pNode = &FuncNode{ Key: caller.String(), Pkg: caller.Func.Pkg.Pkg.Path(), Name: caller.Func.RelString(caller.Func.Pkg.Pkg), } p.tree.Nodes[pNode.Key] = pNode } if qNode, ok = p.tree.Nodes[callee.String()]; !ok { qNode = &FuncNode{ Key: callee.String(), Pkg: callee.Func.Pkg.Pkg.Path(), Name: callee.Func.RelString(callee.Func.Pkg.Pkg), } p.tree.Nodes[qNode.Key] = qNode } if strings.HasSuffix(caller.String(), "main") && p.tree.MainKey == "" { p.tree.MainKey = caller.String() } pNode.Children = append(pNode.Children, qNode.Key) qNode.Parent = append(qNode.Parent, pNode.Key) fmt.Printf("%s to %s \n", caller, callee) return nil }) ``` 通过以上操作将所有函数都保存到内存以及缓存文件中，这样就可以进行各种相关链路图片的生成了。 ### 举例 在github：[https://github.com/toheart/goanalysis](https://github.com/toheart/goanalysis) 中，存在测试的example，获取其中`B.PrintB`的相关链路关系，可以执行一下命令： ```bash go run . analysis "D:\code\goanalysis\example" -p "n44:(example/inner/B.CallB).PrintB" ``` 其中`B.PrintB`的key可以通过缓存文件来查询：  ## track 想快速的分析代码运行时如何执行，在函数中添加入口函数以及defer函数是最快的方式。 按照一般思路，我们需要手动在代码内部一行行的加入相同代码。 对于程序员来说，”懒惰“是优良品质，所有的重复性劳动都需要思考如何自动化完成。 那么自带干电池的”go“也提供这种方式。 在上面我们了解了ast相关库的使用，其中`parser.ParseFile`可以用来静态解析文件，通过分析出当前文件语法树包含的所有函数。 代码：[https://github.com/toheart/goanalysis/blob/main/pkg/track/rewrite.go](https://github.com/toheart/goanalysis/blob/main/pkg/track/rewrite.go) 核心流程： ```go func (r *Rewrite) RewriteFile() { flag := false // 插入defer函数 for _, item := range r.f.Decls { funcDel, ok := item.(*ast.FuncDecl) if !ok { continue } // 判断是否需要插入defer函数 if r.HasSameDefer(funcDel) { continue } elts := r.genTraceParams(funcDel.Type) deferStmt := r.genDefer(elts) // 将defer语句添加到函数体的开头 funcDel.Body.List = append([]ast.Stmt{deferStmt}, funcDel.Body.List...) flag = true } if flag { // 如果上面插入了defer函数, 那么就说明需要插入import r.ImportFunctrace() } buf := &bytes.Buffer{} err := format.Node(buf, r.fset, r.f) if err != nil { return } if debug { fmt.Println(buf.String()) return } // 重写文件 if err = os.WriteFile(r.fullPath, buf.Bytes(), 0666); err != nil { fmt.Printf("write %s error: %v\n", r.fullPath, err) return } } ``` # 总结 对于这个小工具，还有一个“装逼”功能没有完成：通过gitlab的merge，输出当前合并所有改动影响。（后续补坑吧）。 如果对于你后续阅读源码有帮助，一键三连，支持一下。 # 参考说明 1. [https://tonybai.com/2020/12/10/a-kind-of-thinking-about-how-to-trace-function-call-chain/](https://tonybai.com/2020/12/10/a-kind-of-thinking-about-how-to-trace-function-call-chain/) 2. https://mattermost.com/blog/instrumenting-go-code-via-ast/ 3. [https://yuroyoro.github.io/goast-viewer/](https://yuroyoro.github.io/goast-viewer/)