DEV Community

Arash Outadi
Arash Outadi

Posted on • Edited on

2 1

Adding Contexts via Go AST (Code Instrumentation)

Problem

You changed one function to require a ctx context.Context and now you have to change the function signatures of the
all the upstream functions in your codebase 😠!

Do you have to make all these changes manually or can you automate the process somehow?

We can automate it using Go's AST.

Demo

Check out this playground link to see how to source code can be programmatically changed to include the proper arguments/parameters: PLAYGROUND
We'll describe how it all works throughout the tutorial

Illustration

We want to change this function to what is commented out in the TODO:

func changedFn() {
fmt.Println("Nothing to do here")
}
// TODO: Change this function to make a downstream call which needs a context.Context.
// func changedFn(ctx context.Context) {
// fmt.Println("Do some important work...")
// // Now also make a DB call
// makeDBCall(ctx, "Some important data!")
// }
// NOTE:
// This Context must flow down from the very start of the program!
// Which means the functions that call this function must also have a context.Context parameter now

And imagine the function sits inside the fictitious file below, with many other functions calling it.

Again the TODO:s indicate what needs to change to make this file compile.

NOTE: The other functions are not important, only there to illustrate we have to change all the functions that call our changed function

Fictitious File Notice all the TODOs

package test
import (
"fmt"
"context"
)
// Added new context.Context parameter for downstream call
func changedFn(ctx context.Context) {
fmt.Println("Do some important work...")
// Now also make a downstream call
makeDownstreamRequest(ctx, "Some important data!")
}
// TODO: func needsctx1(ctx context.Context, n int)
func needsctx1(n int) {
if true {
// TODO: changedFn(ctx)
changedFn()
}
}
// TODO: func needsctx2(ctx context.Context) bool
func needsctx2() bool {
for index := 0; index < 3; index++ {
needsctx1(ctx, 1)
}
return true
}
// TODO: func needsctx3(ctx context.Context)
func needsctx3() {
if needsctx2(ctx) {
changedFn(ctx)
}
}
type SS struct{}
// TODO: func (rec *SS) save(ctx context.Context, s string, n int)
func (rec *SS) save(s string, n int) {
// TODO: needsctx1(ctx, 2)
needsctx1(2)
}
view raw example.go hosted with ❤ by GitHub

Manual Solution

Fixing this manually, usually involves a trial and error solution process:

  1. Adjust the function body to correctly call your new function with ctx arguments
  2. Adjust the function definition/signature to include a ctx context.Context parameter
  3. Iterate until the compiler no longer complains

Programmatic Solution (Pseudo-Code)

The algorithm for the manual solution is very simple, so on a high level to automate this process we can:

  1. Parse the Go code
  2. Generate an AST(Abstract Syntax Tree) [3]
  3. Programmatically determine where we need "inject" code, specifically:
    • ctx argument to a function call
    • ctx context.Context function parameter to function declaration
  4. Edit the AST
  5. Iterate until there are no more places where an "injection" is required
  6. Convert the AST back into the text representation -> Our new Golang source code

Pre-Requisites

In this tutorial I assume that:

  • You are proficient with Go
  • Somewhat familiar with what an AST is

Tutorial Conventions

Playground Links

After a code example, I will provide a full working example via the Go Playground see you can run the code for yourself.
Look out for PLAYGROUND

Brevity

I will often only include the minimal amount of code to demonstrate a new concept in the code examples and will generally cut out any boiler-plate code.
Look for the Playground links for full working examples.

Setup

Libraries

Note in this tutorial we will be using these libraries:

  • github.com/dave/dst (Alternative to go/ast)

You can download them with the go get command:

go get github.com/dave/dst

github.com/dave/dst is a fork of the official go/ast package that is meant specifically for instrumenting go code.

In contrast go/ast was primarily meant for code generation.

This is an important difference because in code instrumentation, we only want to change a very specific region of code and leave the rest of the AST exactly as it was. go/ast has a difficulties achieving this, especially with comments [4]

Code

Step 0 - Visualize the AST

Using this go-ast visualizer, we can get an idea of what the Go AST looks like and what we need to look for when injecting new code [5].

So go to http://goast.yuroyoro.net/ and play around with these above gists

Our primary focus should be on the on the FuncDecl and CallExpr Nodes since our injection points will be either when we are:

  • Defining a function <-- Might need to add a ctx context.Context parameter
  • Calling a function <-- Might need to add a ctx argument

Step 1 - Use "dst" to parse Go code

Before we start jumping into the logic of the program, let's just see a quick demo of how we parse Go code using the "dave/dst" package and print out the AST representation of the FuncDecl nodes.

package main
import (
"go/parser"
"go/token"
"github.com/dave/dst"
"github.com/dave/dst/decorator"
"github.com/dave/dst/dstutil"
)
// Utility error checking function for when you don't need to gracefully handle errors
func must(err error) {
if err != nil {
panic(err)
}
}
func main() {
file, err := decorator.Parse(srcCodeString)
must(err)
// Notice that we have to define our own function examining/editting a node during AST traversal
applyFunc := func(c *dstutil.Cursor) bool {
node := c.Node()
// Use a switch-case construct based on the node "type"
// This is a very useful of navigating the AST
switch n := node.(type) {
case (*dst.FuncDecl):
// Pretty print the Node AST
dst.Print(n)
}
return true
}
// We traverse the Go AST via the Apply function
// If the node is "nil" or the return value is "false" traversal stops
// Lastly, it's possible to edit the AST while doing the traversal and return the result
_ = dstutil.Apply(file, applyFunc, nil)
}
view raw parse_simple.go hosted with ❤ by GitHub

PLAYGROUND

I've left comments in the above code snippet, so be sure to read those before moving on.

Step 2 - Helper Functions

Remember that we either need to

  • Add a ctx context.Context parameter to a function declaration (FuncDecl Node)
  • Add a ctx argument to a function call (CallExpr) The Go AST structs for these two actions can be defined as follows:
    func newCtxParam() dst.Field {
    return dst.Field{
    Names: []*dst.Ident{&dst.Ident{Name: "ctx"}},
    Type: &dst.SelectorExpr{
    X: &dst.Ident{Name: "context"},
    Sel: &dst.Ident{Name: "Context"},
    },
    }
    }
    func newCtxArg() dst.Ident {
    return dst.Ident{Name: "ctx"}
    }
    view raw helper1.go hosted with ❤ by GitHub

Step 3 - Editing the AST

To demonstrate how to add arguments and parameters using our helper functions, observe the following "naive" applyFunc that adds an additional ctx argument and ctx context.Context parameter to every function call and function definition respectively.

applyFunc := func(c *dstutil.Cursor) bool {
node := c.Node()
switch n := node.(type) {
// Add an (additional) "ctx context.Context" parameter to EVERY function definition
case (*dst.FuncDecl):
ctxP := newCtxParam()
// We "prepend" the "ctx" parameter so it is the first parameter in the function call
n.Type.Params.List = append([]*dst.Field{&ctxP}, n.Type.Params.List...)
// Add an (additional) "ctx" argument to EVERY function call
case (*dst.CallExpr):
ctxA := newCtxArg()
// We "prepend" the "ctx" argument so it is the first argument in the function call
n.Args = append([]dst.Expr{&ctxA}, n.Args...)
}
return true
}

Let's use our applyFunc and print out the source code! Notice I added another utility function for converting the AST representation into actual Go code.
package main
import (
"bytes"
"fmt"
"github.com/dave/dst"
"github.com/dave/dst/decorator"
"github.com/dave/dst/dstutil"
)
...
// FormatNode converts the ast representation into it's textual format (Basically actual Go code)
func FormatNode(file dst.File) string {
var buf bytes.Buffer
decorator.Fprint(&buf, &file)
return buf.String()
}
func main() {
file, err := decorator.Parse(srcCodeString)
must(err)
applyFunc := func(c *dstutil.Cursor) bool {
node := c.Node()
switch n := node.(type) {
// Add an (additional) "ctx context.Context" parameter to EVERY function definition
case (*dst.FuncDecl):
ctxP := newCtxParam()
// We "prepend" the "ctx" parameter so it is the first parameter in the function call
n.Type.Params.List = append([]*dst.Field{&ctxP}, n.Type.Params.List...)
// Add an (additional) "ctx" argument to EVERY function call
case (*dst.CallExpr):
ctxA := newCtxArg()
// We "prepend" the "ctx" argument so it is the first argument in the function call
n.Args = append([]dst.Expr{&ctxA}, n.Args...)
}
return true
}
// We traverse the Go AST via the Apply function
// If the node is "nil" or the return value is "false" traversal stops
// Lastly, it's possible to edit the AST while doing the traversal and return the result
_ = dstutil.Apply(file, applyFunc, nil)
fmt.Println(FormatNode(*file))
}

PLAYGROUND

Notice that this applyFunc doesn't check if it is actually necessary to inject additional code. It just does it.

To see why this applyFunc is not adequate try running changing the srcCodeString to:

package test

import (
    "fmt"
    "context"
)

func alreadyHasContext(ctx context.Context) {
    fmt.Println("Do some important work...")
    makeDownstreamRequest(ctx, "Some important data!") 
}

Or if you prefer go to playground link below.

PLAYGROUND

Step 4 - Being Selective

Instead of indiscriminately adding ctx everywhere like in the naive implementation, this time we will examine the nodes in the AST to determine where we need to inject a ctx argument or ctx content.Context function parameter.

FuncDecl Node

Let's look at the FuncDecl, node for the following code.

func changedFn(ctx context.Context) {
fmt.Println("Do some important work...")
// Now also make a downstream call
makeDBCall(ctx, "Some important data!")
}
view raw new_func.go hosted with ❤ by GitHub

Function Signature

Specifically let's start the with FuncType, which essentially describes the function signature
changedFn signature

The feature that we care about most is SelectorExpr that contains the Ident for the Context parameter. With this in mind, we can construct a function to check if the FuncDecl contains a context.Context as a parameter in the function signature

func hasContextParam(fd *dst.FuncDecl) bool {
// 1. Check if a context is already passed as parameter, if so return early
for _, p := range fd.Type.Params.List {
// If it's not a *SelectorExpr, then skip it (e.g. "a" and "b" in func(a, b, c string))
se, ok := p.Type.(*dst.SelectorExpr)
if !ok {
continue
}
if se.Sel.Name == "Context" {
return true
}
}
return false
}

If it already has a context.Context then we don't need to do anything and can move on.

Function Body

To determine if we need to add a context.Context to the function signature, we need to examine the function body to check if there are any calls to functions which require a context.Context parameter.

At this point, you might be asking how we can know if a function requires a Context in the first place?

Below are two possible ways of determining that:

  1. If there is a function call to a function we have already determined needs a Context (Infinite Recursion 😆)
    • Obviously, we would need a mechanism for recording which functions have Context parameters
    • Also we need a "seed" function which already has Context parameter, so we have at-least one function which we KNOW requires Context
  2. We can naively look for function calls that have a argument named ctx

In this tutorial, we'll go over method 2 as it slightly simpler and doesn't require an initial scan.

Functions To Examine The FuncDecl & CallExpr

var needsContextFuncs = make(map[string]bool)
// looks into function body to see if call like fn(ctx, ...) if so add "ctx context.Context" parameter
func doesFuncDeclRequireCtx(fd *dst.FuncDecl) bool {
// 1. Check if a context is already passed as parameter, if so return early
for _, p := range fd.Type.Params.List {
se, ok := p.Type.(*dst.SelectorExpr)
if !ok {
continue
}
if se.Sel.Name == "Context" {
return false
}
}
// 2. If it doesn't check if the function body has a CallExpr with "ctx" argument
needsCtx := false
dst.Inspect(fd.Body, func(node dst.Node) bool {
switch n := node.(type) {
case (*dst.CallExpr):
if hasCtxArg(n) {
needsCtx = true
return false
}
}
return true
})
return needsCtx
}
func hasCtxArg(ce *dst.CallExpr) bool {
for _, arg := range ce.Args {
switch v := arg.(type) {
case (*dst.Ident):
if v.Name == "ctx" {
return true
}
}
}
return false
}

Note: the global map needsContextsFuncs (However, we use the map as a Set)

Step 5 - Selective ApplyFunc

To make the our ApplyFunc function more selective we will have to examine the FuncDecl and CallExpr nodes

Selective Apply Function

applyFunc := func(c *dstutil.Cursor) bool {
node := c.Node()
switch n := node.(type) {
// Look into function body to see if call like fn(ctx, ...) if so add "ctx context.Context" parameter
case (*dst.FuncDecl):
if doesFuncDeclRequireCtx(n) {
needsContextFuncs[n.Name.Name] = true
ctxP := newCtxParam()
n.Type.Params.List = append([]*dst.Field{&ctxP}, n.Type.Params.List...)
}
// From populated map based on previous case, check if a function NOW needs a ctx argument
case (*dst.CallExpr):
ident, ok := n.Fun.(*dst.Ident)
if ok && !hasCtxArg(n) && needsContextFuncs[ident.Name] {
ctxA := newCtxArg()
n.Args = append([]dst.Expr{&ctxA}, n.Args...)
}
}
return true
}

Step 6 - Is this it?

So now we have a Selective ApplyFunc will it be able to add all the Contexts now?

PLAYGROUND

Spoilers... It doesn't

Step 7 - Iterative Solution

Our new ApplyFunc only manages to change the immediate ancestors, but if want the changes to propagate up further we need an iterative solution.
Remember Step 5 of our pseudo code algorithm

Infinite For Loop?

To make sure that all the ancestors are updated we could simply run the Apply with our ApplyFunc over and over again inside a for {} loop, but when should we stop?
A simple idea that will work is to stop when the previous code is the same as the current code generated.

That is when the ApplyFunc deems that there are no new areas to add ctx arguments or ctx context.Context parameters.

Code

prev := startCode
i := 0
file, err := decorator.Parse(prev)
curN := dstutil.Apply(file, applyFunc, nil)
must(err)
for {
i++
curN = dstutil.Apply(curN, applyFunc, nil)
cur = FormatNode(*file)
if cur == prev {
break
}
prev = cur
}
view raw iteration.go hosted with ❤ by GitHub

PLAYGROUND

Optional - Add comments describing iteration

Lastly, as a illustrative exercise, below is a playground link for adding comments that describe in which iteration the ctx/ctx context.Context was added

PLAYGROUND

Output:

package test

import (
    "context"
    "fmt"
)

// Added on 'ctx context.Context' parameter on iteration 0
func changedFn(ctx context.Context) {
    fmt.Println("Do some important work...")
    // Now also make a downstream call
    makeDownstreamRequest(ctx, "Some important data!")
}

// Added on 'ctx context.Context' parameter on iteration 1
func needsctx1(ctx context.Context, n int) {
    if true {
        changedFn(ctx) // Added on ctx arg on iteration 0
    }
}

// Added on 'ctx context.Context' parameter on iteration 2
func needsctx2(ctx context.Context) bool {
    for index := 0; index < 3; index++ {
        needsctx1(ctx, 1) // Added on ctx arg on iteration 1
    }
    return true
}

// Added on 'ctx context.Context' parameter on iteration 1
func needsctx3(ctx context.Context) {
    if needsctx2(ctx) { // Added on ctx arg on iteration 2

        changedFn(ctx) // Added on ctx arg on iteration 0
    }
}

type SS struct{}

// Added on 'ctx context.Context' parameter on iteration 2
func (rec *SS) save(ctx context.Context, s string, n int) {
    needsctx1(ctx, 2) // Added on ctx arg on iteration 1
}

Future Work

It's not to hard to extend this to changing the AST of a file via ParseFile and entire directories with ParseDir

If there is any interest, I'm happy to make a second tutorial covering these topics and additionally how to handle imported packages.

Thanks

Thanks for reading this tutorial about using Go's AST I hope you found it useful.

Let me know if you liked it or hated it!

Also huge thanks to Dave for making the awesome dst package!

As well as all the other sources I used to make this tutorial

Related Reading

  • I recommend this article about Instrumenting Go code via AST which served as the basis of this post. In some ways, they are solving a simpler problem because the function signature that they need to change is always the same, where in our case it can vary.

Sources

  1. https://faiface.github.io/post/context-should-go-away-go2/
  2. https://medium.com/@cep21/how-to-correctly-use-context-context-in-go-1-7-8f2c0fafdf39
  3. https://en.wikipedia.org/wiki/Abstract_syntax_tree
  4. https://github.com/golang/go/issues/20744 "Free-floating comments are single-biggest issue when manipulating the AST"
  5. http://goast.yuroyoro.net/ Golang AST visualizer
  6. https://godoc.org/github.com/dave/dst GoDoc for github.com/dave/dst package
  7. https://commandercoriander.net/blog/2016/12/30/reading-go-ast/

Sentry image

See why 4M developers consider Sentry, “not bad.”

Fixing code doesn’t have to be the worst part of your day. Learn how Sentry can help.

Learn more

Top comments (0)

Billboard image

The Next Generation Developer Platform

Coherence is the first Platform-as-a-Service you can control. Unlike "black-box" platforms that are opinionated about the infra you can deploy, Coherence is powered by CNC, the open-source IaC framework, which offers limitless customization.

Learn more