Using the "context" Package in Go

By Silviu, on 2017-04-28

Samples of code, showing how to use the context package to handle timeouts, early cancellations and context-related data.

Manual Cancelation

The following example starts three concurrent "racing pets" go routines. Each of them arrives to the finish line after a designated amount of time. The first to arrive explicitly invokes the cancelation function of the shared context instance.

To allow any of the "pet" go routines to easily send a cancelation signal, we are embedding the cancel function as a value inside ctx, using the context.WithValue function.

Alternatively, you can orchestrate the cancelation from the "parent" go routine (in our case, the main function, and call cancel() from there, upon receipt of a custom channel value.

package main

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

// CancelFuncKey is a custom type that prevents potential string key conflicts
// inside the context key value map.
type CancelFuncKey struct{}

var wg sync.WaitGroup

func main() {

	// Construct a cancelable context
	ctx, cancel := context.WithCancel(context.Background())

	// Wrap the cancellable context into a context with value to store the cancel function.
	// This way, it can be invoked by the the pet routine that finishes first.
	ctxVal := context.WithValue(ctx, CancelFuncKey{}, cancel)

	// The Dog pet routine finishes first, in about 3 seconds, so it will call cancel.
	// The other routines will receive the cancelation notfication, and will perform
	// the necessary cleanups (stop the ticker, decrease the main wait group delta, etc)
	startPetRoutine(ctxVal, "Dog", 3)
	startPetRoutine(ctxVal, "Cat", 6)
	startPetRoutine(ctxVal, "Mouse", 9)

	// Block until all pet routines are safely finished
	wg.Wait()
	fmt.Println("==== Done. Exiting. ====")
}

func startPetRoutine(ctx context.Context, pet string, finishLine int) {
	wg.Add(1)
	ticker := time.NewTicker(time.Millisecond * 1000)
	i := 0
	go func() {
		for {
			select {
			case <-ticker.C:
				i++
				if i == finishLine {
					fmt.Printf("%s pet routine: finish line reached.Clean up and cancel context.\n", pet)
					ticker.Stop()
					cancel := ctx.Value(CancelFuncKey{})
					(cancel).(context.CancelFunc)()
					wg.Done()
					return
				}
			case <-ctx.Done():
				fmt.Printf("%s pet routine: cancelation notice received via context.Clean up and exit.\n", pet)
				ticker.Stop()
				wg.Done()
				return
			}
		}
	}()
}

Timeout Cancelation

The following example has two scenario, again involving our three "racing pets" go routines. Both scenarios use context.WithTimeout to construct a context the times out after a specified duration. In the first scenario, one of the "racing pets" finishes ahead of the timeout, and sends an explicit cancel invocation. In the second scenario the "pets" are not able to finish in time, so the timeout gets triggered and the go routines are notified through the ctx.Done() channel.

package main

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

// CancelFuncKey is a custom type that prevents potential string key conflicts
// inside the context key value map.
type CancelFuncKey struct{}

var wg sync.WaitGroup

func main() {
	TimeoutNotTriggered()
	TimeoutTriggered()
}

// The Dog pet routine finishes first, in about 3 seconds, so it will call cancel.
// The timeout will not have a chance to get triggered.
// The other go routines will receive the cancelation notfication, via the Done channel,
// so they are able to perform the necessary cleanups (stop the ticker,
// decrease the main wait group delta, etc).
func TimeoutNotTriggered() {

	fmt.Println("==== Begin: Scenario where timeout does not get triggered ====")
	// Construct a timeout context to expire after about 5 seconds
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

	// Wrap the cancellable context into a context with value to store the cancel function.
	// This way, it can be invoked by the the pet routine that finishes first.
	ctxVal := context.WithValue(ctx, CancelFuncKey{}, cancel)

	startPetRoutine(ctxVal, "Dog", 3)
	startPetRoutine(ctxVal, "Cat", 4)
	startPetRoutine(ctxVal, "Mouse", 9)

	// Block until all pet routines are safely finished
	wg.Wait()
	fmt.Println("==== End: Scenario where timeout does not get triggered ====\n")
}

// Since the timeout in this example is only 2 seconds, and the earliest "pet"
// routine finish line is after 3 seconds, the timeout gets triggered first.
// The other routines will receive the timeout notfication, via the Done channel,
// so they are able to perform the necessary cleanups (stop the ticker,
// decrease the main wait group delta, etc).
func TimeoutTriggered() {

	fmt.Println("==== Begin: Scenario where timeout does get triggered ====")
	// Construct a timeout context to expire after about 2 seconds
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)

	// Wrap the cancellable context into a context with value to store the cancel function.
	// This way, it can be invoked by the the pet routine that finishes first.
	ctxVal := context.WithValue(ctx, CancelFuncKey{}, cancel)

	startPetRoutine(ctxVal, "Dog", 3)
	startPetRoutine(ctxVal, "Cat", 4)
	startPetRoutine(ctxVal, "Mouse", 9)

	// Block until all pet routines are safely finished
	wg.Wait()
	fmt.Println("==== End: Scenario where timeout does get triggered ====")
}

func startPetRoutine(ctx context.Context, pet string, finishLine int) {
	wg.Add(1)
	ticker := time.NewTicker(time.Millisecond * 1000)
	i := 0
	fmt.Printf("Starting \"%s\" for %v seconds.\n", pet, finishLine)
	go func() {
		for {
			select {
			case <-ticker.C:
				i++
				if i == finishLine {
					fmt.Printf("%s pet routine: finish line reached.Clean up and cancel context.\n", pet)
					ticker.Stop()
					cancel := ctx.Value(CancelFuncKey{})
					(cancel).(context.CancelFunc)()
					wg.Done()
					return
				}
			case <-ctx.Done():
				fmt.Printf("%s pet routine: cancelation or timeout notice received.Clean up and exit.\n", pet)
				ticker.Stop()
				wg.Done()
				return
			}
		}
	}()
}