Save request-scoped values to parent context in Golang

Context in Golang can be use to carry deadline, cancellation signals and other request-scoped values across API boundaries and between process. Context can be use to save request-scoped variables but one should not use it for passing optional parameters to functions.

Context with values

However context could be very useful to store any data that is uniform only within the scope of the request. With the context provided by Golang standard library we can save request-scoped values using context.WithValue.

The parent function save the request-scoped value by creating a new context from the parent context.

type ctxKey string

var key = ctxKey("key")

func parentFn(parentCtx context.Context) {
	childCtx := context.WithValue(parentCtx, key, "request value")
	childFn(childCtx)
}

Custom context

By design any values that the child function save in the context will not be accessible by the parent function. As context are treated as immutable, new context needs to be created from the parent context.

There may be scenarios where it make sense to save request-scoped values in a context. To save request-scoped values and make it accessible to parents function, we can extend the functionality of context by implementing your own custom context. Take note that this method is not recommended and should only be use when it is absolutely necessary.

import (
	"context"
)

type CustomContext struct {
	context.Context

	username string
}

func NewCustomContext(ctx context.Context) *CustomContext {
	return &CustomContext{
		Context:  ctx,
		username: "",
	}
}

func (customCtx *CustomContext) SetUsername(username string) {
	customCtx.username = username
}

func (customCtx CustomContext) GetUsername() string {
	return customCtx.username
}

We can now use the saved values from the child function.

func parentFn(parentCtx context.Context) {
	childCtx := NewCustomContext(parentCtx)
	childFn(childCtx)

	fmt.Println(childCtx.GetUsername())
}

func childFn(ctx *CustomContext) {
	ctx.SetUsername("user1")
}

A different child function can also access the saved value

func parentFn(parentCtx context.Context) {
	childCtx := NewCustomContext(parentCtx)
	retrieveUsername(childCtx)
	printUsername(childCtx)
}

func retrieveUsername(ctx *CustomContext) {
	ctx.SetUsername("user1")
}

func printUsername(ctx *CustomContext) {
	fmt.Println(ctx.GetUsername())
}

Do note that the above implementation is not concurrency safe but can be easily adapted to ensure it is safe for simultaneous use by multiple goroutines.

Compatibility with Go context

We can further extend the functionality of your custom context to also implement the same interface as context.Context, providing the flexibility to use context with deadline or cancellation. You can also use your custom context as normal context.Context on third party packages since the custom context implements context.Context.

func (customCtx *CustomContext) Deadline() (deadline time.Time, ok bool) {
	return customCtx.Context.Deadline()
}

func (customCtx *CustomContext) Done() <-chan struct{} {
	return customCtx.Context.Done()
}

func (customCtx *CustomContext) Err() error {
	return customCtx.Context.Err()
}

func (customCtx *CustomContext) Value(key any) any {
	return customCtx.Context.Value(key)
}

With the interface fulfilled, we can essentially treat our custom context as the usual context provided by context.Context

func runSomething(parentCtx context.Context) {
	childCtx := NewCustomContext(parentCtx)

	retrieveUsername(childCtx)

	ctx, cancel := context.WithCancel(childCtx)

	go doSomething(ctx)

	time.Sleep(time.Second)

	cancel()
}

func doSomething(ctx context.Context) {
	// Do something here...
}

Here is an sample code at Go Playground

Conclusion

The functionality of context can be easily extend. Such use of extensibility should be prudent and use only when absolutely necessary. Clear guidelines of the limitations of a custom context must be clearly written as it potentially could defy the usual expectation of what a context provides.