Add value to a context without losing type safety

Are you unsure how to pass trace IDs (or other request-scoped data) through your application stack? Or are your fingers sore from typing type assertions for context values?

You’re in the right place!

This article will show you how to store and retrieve values from contexts:

1. In a type safe way.
2. Without littering your code base with keys and type assertions.

But, let’s start at the beginning: how do you actually add a value to a context?

Adding a value to a context

The context package in the standard library contains functionality to store and retrieve values inside of a context.Context.

It’s not possible to set a value in a context directly. You need to create a new context using the WithValue function. This function accepts both a key and a value of type any.

This newly created context will be derived from the existing context and contain the provided key/value pair.

For example:

bg := context.Background()
ctx := context.WithValue(bg, "k", "zabba")

Derives a new context from a background context and sets a value "zabba" for key "k". You can visualize this as ctx wrapping bg:

Getting a value from a context

To retrieve values from a context.Context, there is the Value method.

This method accepts a key of type any and returns a value of type any.

Below you can see how we use this method to retrieve the value for key "k" in both bg and ctx:

package main

import (
	"fmt"
	"context"
)

func main() {
	bg := context.Background()
	ctx := context.WithValue(bg, "k", "zabba")
	
	fmt.Println(bg.Value("k"))
	fmt.Println(ctx.Value("k"))
}

If you ran the example, you should have seen only the derived context ctx has a value. For bg, <nil> was printed.

Type assertions

As you saw, the WithValue and Value methods work with keys and values of type any, the so called empty interface.

The compiler won’t let you do much with this type. Accessing fields, or calling methods is not possible for example.

However, often you will want to do something with the result of Value. Often something type-specific that any does not allow.

Luckily, since any is an interface, we can use type assertions.

Building on our earlier example, we can use a type assertion to convert the result of a Value call to its underlying type.

v := ctx.Value("k") // v is of type any.
under := v.(string) // under is of type string.

This can be simplified to one line:

under := ctx.Value("k").(string)

Be aware though, type assertions can fail if the underlying type of the interface does not match the type that is being asserted for.

For example, this next line fails since the underlying value returned for "k" ("zabba") is not an int.

under := ctx.Value("k").(int)

Attempting to run it will result in a runtime panic:

panic: interface conversion: interface {} is string, not int

Instead of letting the type-assertion panic, we can also choose to check for it ourselves by using the following syntax.

under, ok := ctx.Value("k").(int)

The second variable ok is a bool that indicates if the type assertion succeeded.

The demo below shows a complete example that accesses the underlying value and uppercases it. This would not be possible if v were of type any.

Try to modify it so that it panics.

package main

import (
	"fmt"
	"context"
	"strings"
)

func main() {
	bg := context.Background()
	ctx := context.WithValue(bg, "k", "zabba")
	
	v := ctx.Value("k").(string)

	fmt.Println(strings.ToUpper(v))
}

Isolating the any’s

For values that you might consider storing in a context there is often at most “one per context” or “one per request”.

It’s useful to wrap context.WithValue and Value in functions when dealing with such values. This way you:

• Isolate the code that deals with any away from the rest of your code.
• Remove the need for callers to care about keys.

Below we implement these wrapper functions for a “trace ID”.

Let’s say a trace ID is a value of type string.

We can then create two functions in a separate myctx package.

package main

import (
	"context"
	"fmt"
	"play/myctx"
)

func main() {
	bg := context.Background()
	ctx := myctx.WithTraceID(bg, "zabba")

	// Later...
	traceID, ok := myctx.TraceID(ctx)
	fmt.Printf("traceID: %v, ok: %v\n", traceID, ok)
}

package myctx

import "context"

const key = "traceID"

func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, key, traceID)
}

func TraceID(ctx context.Context) (string, bool) {
	traceID, ok := ctx.Value(key).(string)
	return traceID, ok
}

As you can see in main.go, callers won’t have to bother with type assertions and/or keys.

This implementation has one downside: The key that was used is a string. This means that code outside of the myctx package can also access or inadvertently overwrite our trace ID.

Let’s show this in main.go.

package main

import (
	"context"
	"fmt"
	"play/myctx"
)

func main() {
	bg := context.Background()
	ctx := myctx.WithTraceID(bg, "zabba")

	// Somewhere in a different part of the system...
	subCtx := context.WithValue(ctx, "traceID", "zoo")

	// Even later...
	traceID, ok := myctx.TraceID(subCtx)
	fmt.Printf("traceID: %v, ok: %v\n", traceID, ok)
}

package myctx

import "context"

const key = "traceID"

func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, key, traceID)
}

func TraceID(ctx context.Context) (string, bool) {
	traceID, ok := ctx.Value(key).(string)
	return traceID, ok
}

As you can see, traceID is now "zoo" instead of "zabba" when accessed from the subCtx.

It was probably not intended for subCtx to overwrite our existing trace ID (otherwise WithTraceID would probably have been used). But prevents it from happening.

Luckily we can fix this by using a custom type for our keys in myctx.

package main

import (
	"context"
	"fmt"
	"play/myctx"
)

func main() {
	bg := context.Background()
	ctx := myctx.WithTraceID(bg, "zabba")

	subCtx := context.WithValue(ctx, "traceID", "zoo")

	// Later...
	traceID, ok := myctx.TraceID(subCtx)
	fmt.Printf("traceID: %v, ok: %v\n", traceID, ok)
}

package myctx

import "context"

type ctxKey string

const key ctxKey = "traceID"

func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, key, traceID)
}

func TraceID(ctx context.Context) (string, bool) {
	traceID, ok := ctx.Value(key).(string)
	return traceID, ok
}

Now “zabba” is printed again!

Also, other packages are now unable to access our ctxKey type, because it is not exported from the myctx package.

This ensures that our wrapper functions must be used to access trace IDs from contexts.

Outro

Depending on your requirements you can modify the above wrapping functions to your needs. There is also potential to play around with generics here, but I haven’t gotten around to that yet.

That’s it for now. I hope this was useful :)