Problem

You changed one function to require actx 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 theTODO:

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

Again theTODO: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

Manual Solution

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

1. Adjust the function body to correctly call your new function withctx arguments
2. Adjust the function definition/signature to include actx 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 anAST(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
  • If you want to learn more about Go's AST, I recommend thispost from Eno Compton[7]

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 forPLAYGROUND

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 togo/ast)

You can download them with thego get command:

go get github.com/dave/dst

github.com/dave/dst is a fork of the officialgo/ast package that is meant specifically forinstrumenting go code.

In contrastgo/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 difficultiesachieving this, especially with comments [4]

Code

Step 0 - Visualize the AST

Using thisgo-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 tohttp://goast.yuroyoro.net/ and play around with theseabove gists

Our primary focus should be on the on theFuncDecl andCallExpr Nodes since our injection points will be either when we are:

• Defining a function <-- Might need to add actx context.Context parameter
• Calling a function <-- Might need to add actx 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 theFuncDecl nodes.

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 actx context.Context parameter to a function declaration (FuncDecl Node)
• Add actx argument to a function call (CallExpr)The Go AST structs for these two actions can be defined as follows:

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 additionalctx argument andctx context.Context parameter to every function call and function definition respectively.

PLAYGROUND

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

To see why thisapplyFunc is not adequate try running changing thesrcCodeString to:

packagetestimport("fmt""context")funcalreadyHasContext(ctxcontext.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 addingctx everywhere like in the naive implementation, this time we will examine the nodes in the AST to determine where we need to inject actx argument orctx content.Context function parameter.

FuncDecl Node

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

Function Signature

Specifically let's start the withFuncType, which essentially describes the function signature

The feature that we care about most isSelectorExpr that contains theIdent for theContext parameter. With this in mind, we can construct a function to check if theFuncDecl contains acontext.Context as a parameter in the function signature

Function Body

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

At this point, you might be asking how we can know if a function requires aContext 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 aContext (Infinite Recursion 😆)
   • Obviously, we would need a mechanism for recording which functions haveContext parameters
   • Also we need a "seed" function which already hasContext parameter, so we have at-least one function which we KNOW requiresContext
2. We can naively look for function calls that have a argument namedctx

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

Functions To Examine TheFuncDecl &CallExpr

Step 5 - Selective ApplyFunc

To make the ourApplyFunc function more selective we will have to examine theFuncDecl andCallExpr nodes

Selective Apply Function

Step 6 - Is this it?

So now we have a SelectiveApplyFunc will it be able to add all theContexts 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.
RememberStep 5 of our pseudo code algorithm

Infinite For Loop?

To make sure that all the ancestors are updated we could simply run theApply with ourApplyFunc over and over again inside afor {} 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 theApplyFunc deems that there are no new areas to addctx arguments orctx context.Context parameters.

Code

Optional - Add comments describing iteration

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

PLAYGROUND

Output:

packagetestimport("context""fmt")// Added on 'ctx context.Context' parameter on iteration 0funcchangedFn(ctxcontext.Context){fmt.Println("Do some important work...")// Now also make a downstream callmakeDownstreamRequest(ctx,"Some important data!")}// Added on 'ctx context.Context' parameter on iteration 1funcneedsctx1(ctxcontext.Context,nint){iftrue{changedFn(ctx)// Added on ctx arg on iteration 0}}// Added on 'ctx context.Context' parameter on iteration 2funcneedsctx2(ctxcontext.Context)bool{forindex:=0;index<3;index++{needsctx1(ctx,1)// Added on ctx arg on iteration 1}returntrue}// Added on 'ctx context.Context' parameter on iteration 1funcneedsctx3(ctxcontext.Context){ifneedsctx2(ctx){// Added on ctx arg on iteration 2changedFn(ctx)// Added on ctx arg on iteration 0}}typeSSstruct{}// Added on 'ctx context.Context' parameter on iteration 2func(rec*SS)save(ctxcontext.Context,sstring,nint){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 viaParseFile and entire directories withParseDir

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 toDave for making the awesomedst package!

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

Related Reading

• I recommend this article aboutInstrumenting 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/