Simplifying context in Go

How can you simplify the Context package in Go, and what is its purpose? In this article, I'll guide you step by step and explain how to use the Context package.

Simplifying context in Go

In Go, we use the Context package for cancellations and
passing values between different application services. Basic knowledge of goroutines would be helpful for a better understanding of this article.

When I created my first goroutine in Go, I stumbled upon a question. What if this goroutine has a bug and never finishes? The rest of my program might keep on running oblivious of a goroutine that never finishes. The simplest example of this situation is a simple “throwaway” goroutine that runs an infinite loop.

package main

import "fmt"

func main() {
  // some "heavy" computation data, 3 entries are only
  // as an example
  dataSet := []string{"apple", "orange", "peach"}

  go func(data []string) {
    // heavy computation that might take a long time and contain a bug
  }(dataSet)

  // the rest of the program goes here and continues
  // to run independent of our goroutine above

  fmt.Println("The rest of our program runs and runs")
}

The above example is not complete but I hope you understand what I'm trying to do. Our “throwaway” goroutine might or might not process data successfully. It might enter an infinite loop or cause an error. The rest of our code would not know what happened.

There are multiple ways to solve this problem. One of them is to use a channel to send a signal to our main thread that this goroutine is taking too long and that it should be cancelled.

package main

import "fmt"
import "time"

func main() {
  stopCh := make(chan bool)

  go func(stopCh chan bool) {
    // simulate long processing
    for {
      select {
        case <-time.After(2 * time.Second):
          fmt.Println("This operation is taking too long. Cancelling...")
          stopCh<- true
      }
    }
  }(stopCh)

  <-stopCh
}

Pretty straightforward. We are using a channel to signal to our main thread that this goroutine is taking too long. But the same thing can be done with context and that is exactly why the context package exists.

package main

import "fmt"
import "context"
import "time"

func main() {
  ctx, cancel := context.WithTimeout(context.Background(), 2 * time.Second)
  // it is always a good practice to cancel the context when
  // we are done with it
  defer cancel()

  go func() {
    // simulate long processing
    for {
      select {
        case <-time.After(2 * time.Second):
          fmt.Println("This operation is taking too long. Cancelling...")
          cancel()
          return
      }
    }
  }()

  select {
    case <-ctx.Done():
        fmt.Println("Context has been cancelled")
  }
}

If you concluded that the context package is a wrapper around a channel to which you react, you would be right. You could easily recreate the context package on your own, but the context package is part of the Go SDK and it is preferable to use it. Also, this package is used by many libraries out there including the httppackage, among others.

Context interface

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key interface{}) interface{}
}

Every time you create a new context you get a type that conforms to this interface. The real implementation of context is hidden in this package and behind this interface. These are the factory types of contexts that you can create:

  1. context.TODO
  2. context.Background
  3. context.WithCancel
  4. context.WithValue
  5. context.WithTimeout
  6. context.WithDeadline

The first context types that we will look at are the context.TODO and context.Background

context.TODO and context.Background

These types of contexts do nothing. It's as if you haven't even created them. Their only use is that you pass some function a context but you don't plan to do anything with that context. A throwaway context. Use this context when you don't plan to do anything with it but some API that you are using requires that you pass it a context.

It is also very important to say that, if a function requires you to pass a context, it is for a good reason. This could be an HTTP request, a connection to a database or a query to the database. HTTP can hang and database queries could take some time. It is good practice that you pass a context in these situations since they protect you from breaking your program.

For example:

func main() {
	ctx := context.TODO()
	req, _ := http.NewRequest(http.MethodGet, "<http://google.com>", nil)

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		fmt.Println("Request failed:", err)
		return
	}
  
	fmt.Println("Status: ", res.StatusCode)
}

The Question from this piece of code is, how long will this request take? Sure, we're calling Google but even Google is not immune from some downtime. A better solution is to create a context that will tell us when this request is taking too long (or at least what we think “too long” is) so we can react to it. A better solution is to implement a timeout context and react when that timeout exceeds. We will go back to this example when we talk about some other types of contexts.

Back to context.TODO(). Under the hood, this context is of type *emptyCtx and it returns empty values for every function in the Context interface. But this context uses a different purpose. It serves as the parent context for some more useful context types. context.Background() is equal to context.TODO. When researching for this blog post, the only difference is semantics. With context.Background(), you are signaling to other developers that you should do something with this context, but from the context package source code, they are the same. If you disagree, leave a comment.

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

// Background returns a non-nil, empty Context. It is never canceled, has no
// values, and has no deadline. It is typically used by the main function,
// initialization, and tests, and as the top-level Context for incoming
// requests.
func Background() Context {
	return background
}

// TODO returns a non-nil, empty Context. Code should use context.TODO when
// it's unclear which Context to use or it is not yet available (because the
// surrounding function has not yet been extended to accept a Context
// parameter).
func TODO() Context {
	return todo
}

Let's dive into some more interesting uses of the context package.

Parent context and canceling

context.TODO and context.Background are types that are used as parent context for other, more useful types of contexts. The first one we will take a look at is the cancel context and its uses.

context.WithCancel()

Let's say you create a goroutine worker that does something very expensive for your CPU. If the work that it is doing is no longer required, you would like it to stop so it does not waste any resources. Normally, you would do this with regular channels.

package main

import (
	"fmt"
	"time"
)

func main() {
	// a channel to tell the goroutine to stop
	// what it is doing
	stopCh := make(chan bool)
	collectIntegers := func(stop chan bool) <-chan int {
		dst := make(chan int)
		n := 1
		go func() {
			for {
				select {
				case <-stop:
					close(dst)
					return
				case dst <- n:
					n++
				}
			}
		}()

		return dst
	}

	ciCh := collectIntegers(stopCh)
	time.AfterFunc(2*time.Second, func() {
		i := <-ciCh
		stopCh <- true
		close(stopCh)
		fmt.Println(fmt.Sprintf("%d integers collected", i))
	})

	// block here until AfterFunc runs
	for _ = range ciCh {}
}

The code is pretty straightforward. We are using a channel to signal to the goroutine to stop working after 2 seconds. Now let's try this with a canceller context.

package main

import (
  "context"
  "fmt"
  "time"
)

func main() {
	collectIntegers := func(ctx context.Context) <-chan int {
		dst := make(chan int)
		n := 1
		go func() {
			for {
				select {
				// listen here to react when cancel() function is called
				case <-ctx.Done():
					close(dst)
					return
				case dst <- n:
					n++
				}
			}
		}()

		return dst
	}

	ctx, cancel := context.WithCancel(context.Background())
	ciCh := collectIntegers(ctx)
	time.AfterFunc(2*time.Second, func() {
		i := <-ciCh
		// when cancel() is called, done channel inside the context is closed
		cancel()
		fmt.Println(fmt.Sprintf("%d integers collected", i))
	})

	// block here until AfterFunc runs
	for _ = range ciCh {
	}
}

From the documentation:

Calling the CancelFunc cancels the child and its children, removes the parent's reference to the child, and stops any associated timers. Failing to call the CancelFunc leaks the child and its children until the parent is canceled or the timer fires.

So, when we called cancel(), select case for ctx.Done()was fulfilled and we could return from the goroutine.

A very important thing to say here is that you always defer cancel(). From the official Go blog:

Always defer a call to the cancel function that’s returned when you create a new Context with a timeout or deadline. This releases resources held by the new Context when the containing function exits

Parent <-> Child relationship

Contexts work based on a parent-child relationship. When you create a context from another context, that created context is said to be derived from the parent context. If you cancel the parent context, all of its children are canceled as well. You can create as many derived contexts as you like. Here, we are creating 2 derived contexts and canceling the parent context. After the parent is canceled, the child is canceled as well.

package main

import (
  "context"
  "fmt"
  "sync"
  "time"
)

func main() {
	parent := context.Background()

	ctx2, cancelCtx2 := context.WithCancel(parent)
	ctx3, _ := context.WithCancel(ctx2)

	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		for {
			select {
			case <-ctx2.Done():
				fmt.Println("ctx2 is done")
				wg.Done()
				return
			}
		}
	}()

	go func() {
		for {
			select {
			case <-ctx3.Done():
				fmt.Println("ctx3 is done")
				wg.Done()
				return
			}
		}
	}()

	time.AfterFunc(1*time.Second, func() {
		cancelCtx2()
	})

	wg.Wait()
}

As you can see, the first ones to be cancelled are the children. The last one to be cancelled is the parent.

Timers -> context.WithTimeout and context.WithDeadline

WithTimeout and WithDeadline and contexts that are set to automatically cancel when some time expires or reaches, depending on the type. They are essentially the same, but WithDeadline receives time. Time while WithTimeout accepts time. Duration but returns a WithDeadline context.

From the source code:

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

WithDeadline can be set to expire at some future date and time. For example, the code below takes whatever date and time you are reading this and expires in 2 seconds in the future.

package main

import (
  "context"
  "fmt"
  "time"
)

func main() {
	now := time.Now()
	fmt.Println(fmt.Sprintf("Current date and time is: %s", now.Format("02.01.2006 15:04:05")))
	ctx, cancel := context.WithDeadline(context.Background(), now.Add(2*time.Second))
	defer cancel()

	wait := make(chan bool)
	go func(ctx context.Context, wait chan bool) {
		for {
			select {
			case <-ctx.Done():
				wait <- false
				fmt.Println(fmt.Sprintf("Deadline is reached on: %s", time.Now().Format("02.01.2006 15:04:05")))
				return
			}
		}
	}(ctx, wait)

	<-wait
	fmt.Println("Main thread has waited long enough!")
}

As I said, WithTimeout is the same as WithDeadline, but it only accepts a time. Duration time. The code above could be written in the same way WithTimeout:

package main

import (
  "context"
  "fmt"
  "time"
)

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

	wait := make(chan bool)
	go func(ctx context.Context, wait chan bool) {
		for {
			select {
			case <-ctx.Done():
				wait <- false
				fmt.Println("Deadline reached!")
				return
			}
		}
	}(ctx, wait)

	<-wait
	fmt.Println("Main thread has waited long enough!")
}

Conclusion

Context package is great for concurrency patterns and essential to learn well since most of the libraries that are dealing with some kind of timers or future resolutions use context. This includes database connections, HTTP requests, etc…

Check out this article on the Rebel Source website where you can run and edit the code examples.

If you would like to know more about this package here are a few useful links:

Thank you for your the time to read this blog! Feel free to share your thoughts about this topic and drop us an email at hello@prototyp.digital.