Understanding the context package in golang

 Posted on May 31, 2018  |  12 minutes  |  Parikshit Agnihotry
https://res.cloudinary.com/pagnihotry/image/upload/v1532622064/pagnihotry/Go-Logo_Blue-trimmed.jpg

The context package in go can come in handy while interacting with APIs and slow processes, especially in production-grade systems that serve web requests. Where, you might want to notify all the goroutines to stop work and return. Here is a basic tutorial on how you can use it in your projects with some best practices and gotchas.

Prerequisites

For understanding the context package there are 2 concepts that you should be familiar with.

  1. goroutine
  2. channels

I will try to briefly cover these before moving on to context, if you are already familiar with these, you can move on directly to the Context section.

Goroutine

From the official go documentation: “A goroutine is a lightweight thread of execution”. Goroutines are lighter than a thread so managing them is comparatively less resource intensive.

Playground: https://play.golang.org/p/-TDMgnkJRY6

  package main

  import "fmt"
  
  //function to print hello
  func printHello() {
    fmt.Println("Hello from printHello")
  }

  func main() {
    //inline goroutine. Define a function inline and then call it.
    go func(){fmt.Println("Hello inline")}()
    //call a function as goroutine
    go printHello()
    fmt.Println("Hello from main")
  }

If you run the above program, you may only see it print out Hello from main that is because it fires up couple goroutines and the main function exits before any of them complete. To make sure main waits for the goroutines to complete, you will need some way for the goroutines to tell it that they are done executing, that’s where channels can help us.

Channels

These are the communication channels between goroutines. Channels are used when you want to pass in results or errors or any other kind of information from one goroutine to another. Channels have types, there can be a channel of type int to receive integers or error to receive errors, etc.

Say there is a channel ch of type int If you want to send something to a channel, the syntax is ch <- 1 if you want to receive something from the channel it will be var := <- ch. This recieves from the channel and stores the value in var.

The following program illustrates the use of channels to make sure the goroutines complete and return a value from them to main.

Note: Wait groups (https://golang.org/pkg/sync/#WaitGroup) can also be used for synchronization, but since we discuss channels later on in the context section, I am picking them in my code samples for this blog post

Playground: https://play.golang.org/p/3zfQMox5mHn

  package main

  import "fmt"
  
  //prints to stdout and puts an int on channel
  func printHello(ch chan int) {
    fmt.Println("Hello from printHello")
    //send a value on channel
    ch <- 2
  }

  func main() {
    //make a channel. You need to use the make function to create channels.
    //channels can also be buffered where you can specify size. eg: ch := make(chan int, 2)
    //that is out of the scope of this post.
    ch := make(chan int)
    //inline goroutine. Define a function and then call it.
    //write on a channel when done
    go func(){
      fmt.Println("Hello inline")
      //send a value on channel
      ch <- 1
    }()
    //call a function as goroutine
    go printHello(ch)
    fmt.Println("Hello from main")

    //get first value from channel.
    //and assign to a variable to use this value later
    //here that is to print it
    i := <- ch
    fmt.Println("Recieved ",i)
    //get the second value from channel
    //do not assign it to a variable because we dont want to use that
    <- ch
  }

Context

A way to think about context package in go is that it allows you to pass in a “context” to your program. Context like a timeout or deadline or a channel to indicate stop working and return. For instance, if you are doing a web request or running a system command, it is usually a good idea to have a timeout for production-grade systems. Because, if an API you depend on is running slow, you would not want to back up requests on your system, because, it may end up increasing the load and degrading the performance of all the requests you serve. Resulting in a cascading effect. This is where a timeout or deadline context can come in handy.

Creating context

The context package allows creating and deriving context in following ways:

context.Background() ctx Context

This function returns an empty context. This should be only used at a high level (in main or the top level request handler). This can be used to derive other contexts that we discuss later.

ctx, cancel := context.Background()

context.TODO() ctx Context

This function also creates an empty context. This should also be only used at a high level or when you are not sure what context to use or if the function has not been updated to receive a context. Which means you (or the maintainer) plans to add context to the function in future.

ctx, cancel := context.TODO()

Interestingly, looking at the code (https://golang.org/src/context/context.go), it is exactly same as background. The difference is, this can be used by static analysis tools to validate if the context is being passed around properly, which is an important detail, as the static analysis tools can help surface potential bugs early on, and can be hooked up in a CI/CD pipeline.

From https://golang.org/src/context/context.go:

var (
   background = new(emptyCtx)
   todo       = new(emptyCtx)
)

context.WithValue(parent Context, key, val interface{}) (ctx Context, cancel CancelFunc)

This function takes in a context and returns a derived context where value val is associated with key and flows through the context tree with the context. This means that once you get a context with value, any context that derives from this gets this value. It is not recommended to pass in critical parameters using context value, instead, functions should accept those values in the signature making it explicit.

ctx := context.WithValue(context.Background(), key, "test")

context.WithCancel(parent Context) (ctx Context, cancel CancelFunc)

This is where it starts to get a little interesting. This function creates a new context derived from the parent context that is passed in. The parent can be a background context or a context that was passed into the function.

This returns a derived context and the cancel function. Only the function that creates this should call the cancel function to cancel this context. You can pass around the cancel function if you wanted to, but, that is highly not recommended. This can lead to the invoker of cancel not realizing what the downstream impact of canceling the context may be. There may be other contexts that are derived from this which may cause the program to behave in an unexpected fashion. In short, NEVER pass around the cancel function.

ctx, cancel := context.WithCancel(context.Background())

context.WithDeadline(parent Context, d time.Time) (ctx Context, cancel CancelFunc)

This function returns a derived context from its parent that gets cancelled when the deadline exceeds or cancel function is called. For example, you can create a context that will automatically get canceled at a certain time in future and pass that around in child functions. When that context gets canceled because of deadline running out, all the functions that got the context get notified to stop work and return.

ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2 * time.Second))

context.WithTimeout(parent Context, timeout time.Duration) (ctx Context, cancel CancelFunc)

This function is similar to context.WithDeadline. The difference is that it takes in time duration as an input instead of the time object. This function returns a derived context that gets canceled if the cancel function is called or the timeout duration is exceeded.

ctx, cancel := context.WithTimeout(context.Background(), time.Duration(150)*time.Millisecond)

Accepting and using contexts in your functions

Now that we know how to create the contexts (Background and TODO) and how to derive contexts (WithValue, WithCancel, Deadline, and Timeout), let’s discuss how to use them.

In the following example, you can see a function accepting context starts a goroutine and waits for that goroutine to return or that context to cancel. The select statement helps us to pick whatever case happens first and return.

<-ctx.Done() once the Done channel is closed, the case <-ctx.Done(): is selected. Once this happens, the function should abandon work and prepare to return. That means you should close any open pipes, free resources and return form the function. There are cases when freeing up resources can hold up the return, like doing some clean up that hangs, etc. You should look out for any such possibilities while handling the context return.

An example that follows this section has a full go program that illustrates timeout and the cancel functions.

//Function that does slow processing with a context
//Note that context is the first argument
func sleepRandomContext(ctx context.Context, ch chan bool) {

  //Cleanup tasks
  //There are no contexts being created here
  //Hence, no canceling needed
  defer func() {
    fmt.Println("sleepRandomContext complete")
    ch <- true
  }()

  //Make a channel
  sleeptimeChan := make(chan int)

  //Start slow processing in a goroutine
  //Send a channel for communication
  go sleepRandom("sleepRandomContext", sleeptimeChan)

  //Use a select statement to exit out if context expires
  select {
  case <-ctx.Done():
    //If context expires, this case is selected
    //Free up resources that may no longer be needed because of aborting the work
    //Signal all the goroutines that should stop work (use channels)
    //Usually, you would send something on channel,
    //wait for goroutines to exit and then return
    //Or, use wait groups instead of channels for synchronization
    fmt.Println("Time to return")
  case sleeptime := <-sleeptimeChan:
    //This case is selected when processing finishes before the context is cancelled
    fmt.Println("Slept for ", sleeptime, "ms")
  }
}

Example

So far we have seen that using contexts you can set a deadline, timeout, or call the cancel function to notify all the functions that use any derived context to stop work and return. Here is an example of how it may work:

  • main function:
    • Creates a context with cancel
    • Calls cancel function after a random timeout
  • doWorkContext function
    • Derives a timeout context
    • This context will be canceled when
      • main calls cancelFunction or
      • Timeout elapses or
      • doWorkContext calls its cancelFunction
    • Starts a goroutine to do some slow processing passing the derived context
    • Waits for the goroutine to complete or context to be canceled by main whichever happens first
  • sleepRandomContext function
    • Starts a goroutine to do the slow processing
    • Waits for the goroutine to complete or,
    • Waits for the context to be canceled by main, timeout or a cancel call to its own cancelFunction
  • sleepRandom function
    • Sleeps for random amount of time
    • This example uses sleep to simulate random processing times, in a real-world example you may use channels to signal this function to start cleanup and wait on a channel for it to confirm that cleanup is complete.

Playground: https://play.golang.org/p/grQAUN3MBlg (Looks like random seed I use, time, in playground is not really changing. You may have to executing this in your local machine to see randomness)

Github: https://github.com/pagnihotry/golang_samples/blob/master/go_context_sample.go

package main

import (
  "context"
  "fmt"
  "math/rand"
  "time"
)

//Slow function
func sleepRandom(fromFunction string, ch chan int) {
  //defer cleanup
  defer func() { fmt.Println(fromFunction, "sleepRandom complete") }()

  //Perform a slow task
  //For illustration purpose,
  //Sleep here for random ms
  seed := time.Now().UnixNano()
  r := rand.New(rand.NewSource(seed))
  randomNumber := r.Intn(100)
  sleeptime := randomNumber + 100
  fmt.Println(fromFunction, "Starting sleep for", sleeptime, "ms")
  time.Sleep(time.Duration(sleeptime) * time.Millisecond)
  fmt.Println(fromFunction, "Waking up, slept for ", sleeptime, "ms")

  //write on the channel if it was passed in
  if ch != nil {
    ch <- sleeptime
  }
}

//Function that does slow processing with a context
//Note that context is the first argument
func sleepRandomContext(ctx context.Context, ch chan bool) {

  //Cleanup tasks
  //There are no contexts being created here
  //Hence, no canceling needed
  defer func() {
    fmt.Println("sleepRandomContext complete")
    ch <- true
  }()

  //Make a channel
  sleeptimeChan := make(chan int)

  //Start slow processing in a goroutine
  //Send a channel for communication
  go sleepRandom("sleepRandomContext", sleeptimeChan)

  //Use a select statement to exit out if context expires
  select {
  case <-ctx.Done():
    //If context is cancelled, this case is selected
    //This can happen if the timeout doWorkContext expires or
    //doWorkContext calls cancelFunction or main calls cancelFunction
    //Free up resources that may no longer be needed because of aborting the work
    //Signal all the goroutines that should stop work (use channels)
    //Usually, you would send something on channel, 
    //wait for goroutines to exit and then return
    //Or, use wait groups instead of channels for synchronization
    fmt.Println("sleepRandomContext: Time to return")
  case sleeptime := <-sleeptimeChan:
    //This case is selected when processing finishes before the context is cancelled
    fmt.Println("Slept for ", sleeptime, "ms")
  }
}

//A helper function, this can, in the real world do various things.
//In this example, it is just calling one function.
//Here, this could have just lived in main
func doWorkContext(ctx context.Context) {

  //Derive a timeout context from context with cancel
  //Timeout in 150 ms
  //All the contexts derived from this will returns in 150 ms
  ctxWithTimeout, cancelFunction := context.WithTimeout(ctx, time.Duration(150)*time.Millisecond)

  //Cancel to release resources once the function is complete
  defer func() {
    fmt.Println("doWorkContext complete")
    cancelFunction()
  }()

  //Make channel and call context function
  //Can use wait groups as well for this particular case
  //As we do not use the return value sent on channel
  ch := make(chan bool)
  go sleepRandomContext(ctxWithTimeout, ch)

  //Use a select statement to exit out if context expires
  select {
  case <-ctx.Done():
    //This case is selected when the passed in context notifies to stop work
    //In this example, it will be notified when main calls cancelFunction
    fmt.Println("doWorkContext: Time to return")
  case <-ch:
    //This case is selected when processing finishes before the context is cancelled
    fmt.Println("sleepRandomContext returned")
  }
}

func main() {
  //Make a background context
  ctx := context.Background()
  //Derive a context with cancel
  ctxWithCancel, cancelFunction := context.WithCancel(ctx)

  //defer canceling so that all the resources are freed up 
  //For this and the derived contexts
  defer func() {
    fmt.Println("Main Defer: canceling context")
    cancelFunction()
  }()

  //Cancel context after a random time
  //This cancels the request after a random timeout
  //If this happens, all the contexts derived from this should return
  go func() {
    sleepRandom("Main", nil)
    cancelFunction()
    fmt.Println("Main Sleep complete. canceling context")
  }()
  //Do work
  doWorkContext(ctxWithCancel)
}

Gotchas

If a function takes in a context, make sure you check on how it respects the cancel notification. For instance, the exec.CommandContext does not close the reader pipe till the command has executed all the forks that a process created (Github issue: https://github.com/golang/go/issues/23019), which means that the context cancellation will not make this function return right away if you are waiting on cmd.Wait() until all forks of the external command have completed processing. If you used a timeout or deadline with a max execution time, you may see this not working as expected. If you run into any such issues, you can implement timeouts using time.After.

Best practices

  1. context.Background should be used only at the highest level, as the root of all derived contexts
  2. context.TODO should be used where not sure what to use or if the current function will be updated to use context in future
  3. context cancelations are advisory, the functions may take time to clean up and exit
  4. context.Value should be used very rarely, it should never be used to pass in optional parameters. This makes the API implicit and can introduce bugs. Instead, such values should be passed in as arguments.
  5. Don’t store contexts in a struct, pass them explicitly in functions, preferably, as the first argument.
  6. Never pass nil context, instead, use a TODO if you are not sure what to use.
  7. The Context struct does not have a cancel method because only the function that derives the context should cancel it.
programming  golang