Copying the Context

Go nicely formalized the idea of a “request context” with the context package. context.Context is a lightweight object that you can pass around to control processes started by a request, and to pass request related data downstream. Since context.Context is tied to a request, it is usually canceled after the request handling is done. But sometimes, a request simply starts a long running process that continues long after the completion of the request. For these cases, you have to create a new context, but then how can you copy the values in the original context into this new one? This is a common question for HTTP API handlers that start long-running goroutines. This article is about a simple method to “copy” one context into another for such cases.

But first, a few words about the Context.

The primary purpose of a context is to control the timeout and cancellation of long-running goroutines. At the same time, it can serve as a key-value store for relatively few request-scoped objects. I say “relatively few”, because every time you add a new key-value to a context, you create a new one wrapping the old. With each added value, context adds a new layer.

Take a look at the HTTP middleware below. It authenticates the user, creates a request id, and adds both the user id and request id to the context:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

func authMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // Authenticate the user and get user id
        userId, err := authenticateUser(r)
        if err!=nil {
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }
        // Create a new request id
        requestId := uuid.New()
        // Add user id to the context
        ctx := context.WithValue(r.Context(), userIdKey, userId)
        // Add request id to the context
        ctx = context.WithValue(ctx, requestIdKey, requestId)
        next.ServeHTTP(w,r.WithContext(ctx))
	})
}

After adding the userId, the context looks like this:

After adding the requestId, it becomes:

As you can see, a Context has an onion-like layered structure. When you query a context instance for a value, the Context instance at the outermost layer checks if it recognizes the key. If not, it passes the control to the next layer enclosed in it, until the requested key is found, or a nil is returned. If you wrap the Context n times, finding a key will take O(n).

There is a reason for this design. Every new context wraps another one without touching the values there, so you can create multiple contexts for different uses that are based on the same context. For instance, let’s say our HTTP API starts several new goroutines, and for one of these goroutines, we need to impersonate another user (an admin user, for instance). We can simply override the userId value for that context, and pass it into the goroutines that need the admin user while the base context can be sent to the other goroutines that use the logged in user. The context that overrides the userId property looks like this:

To see how this is implemented, let’s take a look at the simplified version of context.WithValue:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

type valueCtx struct {
	Context
	key, val any
}

func WithValue(parent Context, key, val any) Context {
    // Error checking removed
	return &valueCtx{parent, key, val}
}

func (c *valueCtx) Value(key any) any {
	if c.key == key {
		return c.val
	}
    // The standard library implementetion checks many other cases here,
    // but conceptually, this is what it is doing:
	return c.Context.Value(key)
}

As you can see, the WithValue implementation simply wraps the previous context with a valueCtx containing the key-value pair. And valueCtx.Value is implemented to check the key in the currecnt context, and if not found, pass it down the layers.

We can use the same principle to copy a context:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

type copyContext struct {
   // This context contains the cancel/timeout channel
   context.Context
   // This context contains the values
   values context.Context
}

func NewCopyContext(baseContext, requestContext context.Context) context.Context {
   return &copyContext{
      context: baseContext,
      values: requestContext,
  }
}

This is illustrated below:

The copyContext contains two contexts, one being the base context that will be used in the new goroutines, and the request context that will be canceled once the HTTP handler returns. The requestContext will only be lept because of its key-value pairs. The following implementation of Value will use the request context for values:

1
2
3

func (c *copyContext) Value(key any) any {
   return c.values.Value(key)
}

As you can see, this implementation disregards all the values of the base context, and uses the request context to answer all value queries. The context cancellation or timeout will happen using the embedded context.

Now we can implement the scenario we described earlier. A goroutine starts with a new context that is a copy of the request context, and another goroutine starts with a new context that impoersonates an admin user:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

func handler(w http.ResponseWriter, r *http.Request) {
   // Do some processing
   
   // Create a new context by copying the request context
   newCtx:=NewCopyContext(context.Background(), r.Context())
   // Create an admin context
   adminCtx:=context.WithValue(newCtx,userIdKey, adminUserId)
   go longRunningTask(newCtx)
   go adminTask(adminCtx)
}

Above, the handler creates two goroutines that continue to run after the handler returns. Both goroutines are started with a context that is a copy of the request context. One of the contexts overrides the userId property with an admin user id.

So, you can “copy” the contents of one context into another by simple composition.