meain/blog

Jun 09, 2024 . 7 min

Putting Go's Context package into context

Hello Gophers and other beings of the world wide web. The context package is something all Gophers would have used at some point in their life, but do you know what it looks like inside(because apparently the real beauty is on the inside)?

If you like this, you can checkout my previous blog post about how sync.WaitGroup work internally here.

This is intended more of a quick look at what is going on inside and so I'll be simplifying things a bit. If you want to see the full gory details, or if you prefer reading source code instead of a blog, here is the implementation. It is actually a pretty simple implementation.

Basic usage #

All of you know the basic usage of the context package, and if you don't, this is not the blog for you. At least not yet, but just so that this is not jarring, let's see a simple example of using context package.

I really don't know what a good example to showcase would be. After some thinking I decided I would just put in a function that takes in a context and passes it into another function because to most people that is all context is, just something you pass around if the function expects it.

func main() {
	bigFunc(context.Background())
}

func bigFunc(ctx context.Context) {
	smallFunc(ctx)
}

func smallFunc(ctx context.Context) {
	// I don't know what to do with it, let' just print it
	fmt.Println(ctx)
}

If you run this, you will print out context.Background. This is because the thing that context.Background returns, satisfies a Stringer interface which just returns this when calling String.

Now with the foreplay out of the way, lets get into it.

The Context interface #

Let’s start with the basics. The type that you use, context.Context is an interface, and below is the definition for it.

type Context interface {
	Deadline() (deadline time.Time, ok bool) // get the deadline time
	Done() <-chan struct{}                   // get a channel which is closed when cancelled
	Err() error                              // returns non-nil if Done channel is closed
	Value(key any) any                       // get a value from the context store
}

Any struct that satisfies this interface is a valid context object. Let’s have a quick look at what each of these are in case the comments did not give them off.

These are all you need if you want to build a "Context" and you can easily create them. That said, stdlib does provide us with some useful ones.

The emptyCtx struct #

This is a struct which satisfies the bare minimum that is needed to be a Context. Here is what the code looks like:

type emptyCtx struct{}

func (emptyCtx) Deadline() (deadline time.Time, ok bool) {
	return
}

func (emptyCtx) Done() <-chan struct{} {
	return nil
}

func (emptyCtx) Err() error {
	return nil
}

func (emptyCtx) Value(key any) any {
	return nil
}

As you can see, it does nothing, but that is mostly all that is in context.Background and context.TODO.

context.Background and context.TODO #

Both of them are just emptyCtx plus a String method to satisfy the Stringer interface. These provide you a way to create an empty base context. The only difference between them are the name.

You use context.Background when you know that you need an empty context, like in main where you are just starting and you use context.TODO when you don’t know what context to use or haven’t wired things up.

You can think of context.TODO as similar to adding a // TODO comment in your code.

Here is the code for context.Background:

type backgroundCtx struct{ emptyCtx }

func (backgroundCtx) String() string {
	return "context.Background"
}

func Background() Context {
	return backgroundCtx{}
}

and for context.TODO:

type todoCtx struct{ emptyCtx }

func (todoCtx) String() string {
	return "context.TODO"
}

func TODO() Context {
	return todoCtx{}
}

Simple enough, right?

context.WithValue #

Now we get into more useful usecases of the context package. You can use context.WithValue if you want to use context to pass in a value. You might have seen logging or web frameworks make use of this.

Let’s see how this look on the inside:

type valueCtx struct {
	Context
	key, val any
}

func WithValue(parent Context, key, val any) Context {
	return &valueCtx{parent, key, val}
}

It is just returning a struct which contains the parent context, and a key and a value.

If you notice, the instance can only hold one key and one value, but you might have seen in web frameworks where they pull out multiple values from the ctx argument. Since you are embedding the parent into the newer contexts, you can recursively search upwards to get any of the other values.

Let’s say you create something like below:

bgCtx := context.Background()
v1Ctx := context.WithValue(bgCtx, "one", "uno")
v2Ctx := context.WithValue(v1Ctx, "two", "dos")

Now if we have to get the value of “one” from v2Ctx, we can call v2Ctx.Value("one"). This will first check if the key in v2Ctx is “one” and since it is not, it will check if the key of the parent(v1Ctx) is “one”. Now since the key in v1Ctx is "one", we return the value in the context.

The code looks something like below. I’ve removed some bits about how timeout/cancel values are handled. You can see the full code here.

func (c *valueCtx) Value(key any) any {
	// If it this one, just return it
	if c.key == key {
		return c.val
	}
	return value(c.Context, key)
}

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			// If the parent is a `valueCtx`, check its key
			if key == ctx.key {
				return ctx.val
			}
			c = ctx.Context
		case backgroundCtx, todoCtx:
			// If we have reached the top, ie base context
			// Return as we did not find anything
			return nil
		default:
			// If it is some other context,
			// just calls its `.Value` method
			return c.Value(key)
		}
	}
}

The code, as you can see is just recursively searching up it parent contexts to see of any of them match the key and of so returning its value.

context.WithCancel #

Let's look into something more useful. You can use the context package to create a ctx that can be used to signal a cancellation to downstream functions.

Let’s see an example of how this could be used:

func doWork(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		time.Sleep(1 * time.Second)
		fmt.Println("doing work...")
	}
}

func main() {
	bgCtx := context.Background()
	innerCtx, cancel := context.WithCancel(bgCtx)

	go doWork(innerCtx) // call goroutine
	time.Sleep(3 * time.Second) // do work in main

	// well, if `doWork` is still not done, just cancel it
	cancel()
}

In this case you can signal the doWork function to stop doing work by calling cancel in the main function.

Now to how this works. Let’s start with the function definition(we’ll get to the struct def soon):

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := &cancelCtx{}
	c.propagateCancel(parent, c)
	return c, func() { c.cancel(true, Canceled, nil) }
}

When you call context.WithCancel, it returns you an instance of cancelCtx and a function that you can call to cancel the context. Just from this, we can infer that a cancelCtx is a context that has a cancel function which can be used to "cancel" the context.

In case you forgot, cancelling a context just means you close the channel returned by Done().

BTW, the propagateCancel function's primary role in this context here is to create a cancelCtx, and also to do things like making sure that the parent is not already cancelled before creating it.

OK, now let’s look at the struct and after that we will go into how this works (source).

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     atomic.Value          // chan struct{} for Done
	children map[canceler]struct{}
	err      error
	cause    error
}

OK, what do we have here:

Oh BTW, this is the cancel function(source):

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	if cause == nil {
		cause = err
	}

	c.err = err
	c.cause = cause

	// load the chan struct{} from atomic.Value and close it
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		// if it does not exist, store a closed channel
		c.done.Store(closedchan)
	} else {
		// if it exists, close it
		close(d)
	}

	// call cancel on all children to propagate the cancellation
	for child := range c.children {
		child.cancel(false, err, cause)
	}
	c.children = nil

	// remove itself from the parent
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

First we set the cause and err on the instance and then we close the channel returned in Done. After that it cancelles all of its children and finally removes itself from the parent.

And finally, context.WithDeadline and context.WithTimeout #

These are useful when you want to create a context that automatically cancels itself when a deadline is reached. This can be very useful for enforcing things like server timeouts.

Just to get it out of the way, context.WithTimeout is just computing the deadline and calling context.WithDeadline. In fact, this is the entire code for it:

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

Now to get into the details. As some of you might have guessed, WithDeadline is basically just a normal WithCancel context, but wit the context package handling the cancellation.

Let’s look at what the code does. Here I have the code for the function WithDeadlineCause which is a variant of WithDeadline, but with the added ability to pass in a "cause" for the cancellation. FYI, Cause variants are available for other context package functions as well and the non Cause one like WithDeadline is just calling Cause variants with a nil for cause.

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// if parent deadline earlier than child,
		// just return a cancelCtx
		return WithCancel(parent)
	}

	// create a new timerCtx
	c := &timerCtx{
		deadline: d,
	}
	c.cancelCtx.propagateCancel(parent, c)

	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
		return c, func() { c.cancel(false, Canceled, nil) }
	}

	// if all good, setup a new timer and return a cancel func
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

Let’s walk through what it does:

Oh BTW, timerCtx is just a cancelCtx with a timer:

type timerCtx struct {
	cancelCtx
	timer *time.Timer
	deadline time.Time
}

Aaaaand, that is it. As I mentioned in the beginning, I have cut out parts from the original code to make things a little easier to explain. The actual source is very readable I would encourage you to go through it.

← Home