【“骚”操作】AST实现代码调用跟踪思考当你接手一个新的Go语言项目时，如果前任开发者没有留下设计文档，或者文档已经过

思考

当你接手一个新的Go语言项目时，如果前任开发者没有留下设计文档，或者文档已经过时、代码注释也不够清楚，那你可能会觉得很难快速上手。

这就跟旅行一样，如果你有一张地图，旅途就会顺利很多。因此我想到了一个叫go-callvis的工具，它可以通过分析项目的语法树来展示所有函数之间的调用关系，并以网页的形式呈现出来，方便查看。
不过，在实际使用中我发现这个工具生成网页图的速度有点慢。

为了解决这个问题，我自己开发了两个小工具来辅助。

goanalysis

此前，我曾开发过这一工具，但遗憾的是，当时的代码并未妥善保存，也未上传至GitHub。因此，本次决定将该工具公开分享。此工具具备以下两大功能：

analysis: 通过项目生成相关调用链图，如：

track: 将项目内所有函数都内置上defer trace()，在运行过程中，将goroutine中所有调用的函数以及参数输出，比如：

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

	// 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生成项目整体的语法树

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：github.com/toheart/goa… 中，存在测试的example，获取其中B.PrintB的相关链路关系，可以执行一下命令：

go run . analysis "D:\code\goanalysis\example" -p "n44:(example/inner/B.CallB).PrintB"

其中B.PrintB的key可以通过缓存文件来查询：

B.PrintB

track

想快速的分析代码运行时如何执行，在函数中添加入口函数以及defer函数是最快的方式。按照一般思路，我们需要手动在代码内部一行行的加入相同代码。

对于程序员来说，”懒惰“是优良品质，所有的重复性劳动都需要思考如何自动化完成。那么自带干电池的”go“也提供这种方式。

在上面我们了解了ast相关库的使用，其中parser.ParseFile可以用来静态解析文件，通过分析出当前文件语法树包含的所有函数。

代码：github.com/toheart/goa…

核心流程：

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，输出当前合并所有改动影响。（后续补坑吧）。

如果对于你后续阅读源码有帮助，一键三连，支持一下。

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

思考