Contexts - Cancellation, Timeout and Propagation

What is a context?

Context is a package in the standard library which is mainly used to propagate cancellation signals from one function to another or even from micro service to another. Let’s consider the example of a user sending a get request to a web server to download 50 images, zip it and send the zipped response back. The request has been triggered from the browser and let’s say it takes 70 seconds to complete. The user is not patient enough and decides to cancel it when it is still being processed by the server. Wouldn’t it be nice if the server gets to know that the user(client) has cancelled the request so that the server can also terminate the request and save valuable CPU and memory? This is the perfect use case for a context. The context allows the server to know when a request has been cancelled by the client so that it can terminate its resources and move on. We will write the program for this as we progress through the tutorial.

A simple example

Let’s start by writing a simple example where we have a function which has a long running code that is called from main. We must terminate the long running function when a cancellation signal is sent from the caller. Let’s discuss how we can achieve this with the help of a context.

 1package main
 2
 3import (
 4	"fmt"
 5	"time"
 6)
 7
 8func longRunning() int {
 9	count := 0
10	for i := 0; i < 5; i++ {
11		count = count + i
12                fmt.Println("Current value of count:",count)
13		time.Sleep(2 * time.Second)
14
15	}
16	return count
17}
18
19func main() {
20	count := longRunning()
21	fmt.Println("count is ", count)
22}

In the above program, longRunning function runs a for loop for 5 iterations which sleeps for 2 seconds between each iteration. Although this code is contrived, it is perfect for understanding context, so please stay with me on this :). The above code will print

Current value of count: 0
Current value of count: 1
Current value of count: 3
Current value of count: 6
Current value of count: 10
count is  10

after 10 seconds.

What if we want to terminate the long running function after 5 seconds? This is possible with the help of context. The main function will send a cancellation signal to longRunning function after 5 seconds using the context. The longRunning function will receive it and return immediately. Let’s see how we can achieve this.

Context with Cancellation

The first step towards the above goal is to create a new context in the main function. We should also be able to cancel this context after 5 seconds so that the longRunning will return. Please don’t worry if it doesn’t make sense now. As we progress with the tutorial it will be clear. A context with a cancellation signal can be created using the WithCancel function. This function needs a context to be passed as a parameter. You might be confused by this since we are trying to create a new context and this function expects a context :O. This is where the background context has its place. The background is a non-nil empty context. It is usually used as the starting point to create any new context. The WithCancel function returns a context and a CancelFunc. Calling the CancelFunc will send the cancellation signal.

1 
2ctx, cancelFunc := context.WithCancel(context.Background())
3go func() {
4	time.Sleep(2 * time.Second)
5	cancelFunc()
6}()
7count, err := longRunning(ctx)

In the above code we do what we just described. We create context with a cancellation using the WithCancel function by passing it the background context as a parameter. We create a new goroutine which will call the cancelFunc after 2 seconds. This sends the cancellation signal. The complete main function with these changes is provided below.

 1func main() {
 2	ctx, cancelFunc := context.WithCancel(context.Background())
 3	go func() {
 4		time.Sleep(2 * time.Second)
 5		cancelFunc()
 6	}()
 7	count, err := longRunning(ctx)
 8	if err != nil {
 9		fmt.Println("long running task exited with error", err)
10		return
11	}
12
13	fmt.Println("count is ", count)
14}

Now that we are done with the main function changes, let’s go ahead and finish the changes for the longRunning function. The first change needed is the addition of the context parameter which is passed from main. The next change needed in the longRunning function is it needs to know when the context is cancelled from the main function. The context has a done channel which is used to notify context cancellation. The context’s done channel which will be closed when the cancelFunc is called from main. This can be used in the longRunning function and it will know when the context is cancelled. The changes described above are incorporated in the longRunning below.

 1func longRunning(ctx context.Context) (int, error) {
 2	count := 0
 3	for i := 0; i < 5; i++ {
 4		select {
 5		case <-ctx.Done():
 6			return 0, ctx.Err()
 7		default:
 8			count = count + i
 9            fmt.Println("Current value of count:",count)
10			time.Sleep(2 * time.Second)
11		}
12	}
13	return count, nil
14}

 1package main
 2
 3import (
 4	"context"
 5	"fmt"
 6	"time"
 7)
 8
 9func longRunning(ctx context.Context) (int, error) {
10	count := 0
11	for i := 0; i < 5; i++ {
12		select {
13		case <-ctx.Done():
14			return 0, ctx.Err()
15		default:
16			count = count + i
17            fmt.Println("Current value of count:",count)
18			time.Sleep(2 * time.Second)
19		}
20	}
21	return count, nil
22}
23
24func main() {
25	ctx, cancelFunc := context.WithCancel(context.Background())
26	go func() {
27		time.Sleep(2 * time.Second)
28		cancelFunc()
29	}()
30	count, err := longRunning(ctx)
31	if err != nil {
32		fmt.Println("long running task exited with error", err)
33		return
34	}
35
36	fmt.Println("count is ", count)
37}

The above program will print

Current value of count: 0
long running task exited with error context canceled

and will terminate after 2 seconds. This is exactly what we wanted. We have sent a termination signal to a long running program after 2 seconds and have successfully stopped the long running function. Perfect :)

Context with timeout

In the above program, we are creating a new goroutine which will call the cancelFunc after two seconds. There is one more way to achieve the same in context. The WithTimeout function of the context package creates a context which will be automatically cancelled after the timeout. We can use this in our code and cancel the context after 2 seconds.

The modified main function which uses WithTimeout is provided below.

 1func main() {
 2	ctx, cancelFunc := context.WithTimeout(context.Background(), 2*time.Second)
 3	defer cancelFunc()
 4	count, err := longRunning(ctx)
 5	if err != nil {
 6		fmt.Println("long running task exited with error", err)
 7		return
 8	}
 9
10	fmt.Println("count is ", count)
11}

In line no. 2 of the main function above, we create a context which will timeout after 2 seconds. The rest of the program remains the same. The WithTimeout returns a cancelfunc which needs to be called to prevent resource leak in case longRunning completes before the 2 second timeout. As a best practice we can call it every time using defer. Running this program will print,

Current value of count: 0
long running task exited with error context deadline exceeded

Context propagation

One of the other uses of context is to terminate all goroutines spawned by a function once the context is cancelled. This is done by passing/propagating the context to all child goroutines.

Let’s say we run a web server and it has to contact two DB servers to process an incoming request. It spawns two goroutines, waits for both of them to return and then it computes the output based on the result of those two goroutines and returns it. If the user decided to cancel the request when it’s in the middle of processing, we need to terminate the goroutines spawned by the web server to prevent goroutine leak and also to save valuable CPU compute time. This can be done with the help of context. An example will make things clear.

 1package main
 2
 3import (
 4	"context"
 5	"fmt"
 6	"sync"
 7	"time"
 8)
 9
10type output struct {
11	count int
12	err   error
13}
14
15func dbTask1(ctx context.Context, wg *sync.WaitGroup) (int, error) {
16	defer wg.Done()
17	select {
18	case <-ctx.Done():
19		fmt.Println("access DB task1 error:", ctx.Err())
20		return 0, fmt.Errorf("access DB task1 error: %s", ctx.Err())
21	case <-time.After(7 * time.Second):
22		return 20, nil
23	}
24}
25
26func dbTask2(ctx context.Context, wg *sync.WaitGroup) (int, error) {
27	defer wg.Done()
28	select {
29	case <-ctx.Done():
30		fmt.Println("access DB task2 error:", ctx.Err())
31		return 0, fmt.Errorf("access DB task2 error: %s", ctx.Err())
32	case <-time.After(5 * time.Second):
33		return 30, nil
34	}
35}
36
37func webApi(ctx context.Context) (int, error) {
38	wg := sync.WaitGroup{}
39	outputCh1 := make(chan output, 1)
40	outputCh2 := make(chan output, 1)
41	wg.Add(1)
42	go func() {
43		count1, err := dbTask1(ctx, &wg)
44		o := output{
45			count: count1,
46			err:   err,
47		}
48		outputCh1 <- o
49
50	}()
51	wg.Add(1)
52	go func() {
53		count2, err := dbTask2(ctx, &wg)
54		o := output{
55			count: count2,
56			err:   err,
57		}
58		outputCh2 <- o
59	}()
60	wg.Wait()
61	output1 := <-outputCh1
62	if output1.err != nil {
63		return 0, output1.err
64	}
65	output2 := <-outputCh2
66	if output2.err != nil {
67		return 0, output2.err
68	}
69	output := output1.count + output2.count
70	return output, nil
71}
72
73func main() {
74	ctx, cancelFunc := context.WithTimeout(context.Background(), 2*time.Second)
75	defer cancelFunc()
76	count, err := webApi(ctx)
77	if err != nil {
78		fmt.Println("long running task exited with error", err)
79		return
80	}
81
82	fmt.Println("count is", count)
83}

In the above program functions dbTask1 and dbTask2 are the two expensive DB operations. dbTask1 returns a successful response after 7 seconds and dbTask2 returns a successful response after 5 seconds. These two functions accept a context as parameter. If the context is cancelled, the done channel is closed and these functions return error in line 20 and line 31 respectively.

dbTask1 and dbTask2 are called from the webApi function as two goroutines in line 42 and line 52. The webApi receives a context from main and the same context is passed to dbTask1 and dbTask2 in line no. 43 and 53. webApi waits for these two functions to return and finally computes the result which is the sum of the return values of both of these goroutines and returns it in line no. 70. This program uses waitgroup. I recommend reading /buffered-channels-worker-pools/ if you are not sure how waitgroups work

If the context is cancelled before dbTask1 and dbTask2 complete processing, both dbTask1 and dbTask2 return error and they terminate. This prevents goroutine leak and it also saves valuable CPU time.

If context is not used, these two groutines will be dangling and the program will leak resources.

main creates a context with a timeout of 2 seconds and hence after 2 seconds both dbTask1 and dbTask2 return an error and are terminated. The rest of the code is self explanatory. This program will print the following output.

access DB task2 error: context deadline exceeded
access DB task1 error: context deadline exceeded
long running task exited with error access DB task1 error: context deadline exceeded

Practical Example - Web server with context cancellation

Now that we have a solid understanding of context, let’s try to use this practically when creating a web server. Let’s write the program for this right away.

 1package main
 2
 3import (
 4	"io"
 5	"log"
 6	"net/http"
 7	"time"
 8)
 9
10func webServer(w http.ResponseWriter, r *http.Request) {
11	timer := time.NewTimer(10 * time.Second)
12	<-timer.C
13	log.Println("writing response....")
14	_, err := io.WriteString(w, "Hello context")
15	if err != nil {
16		log.Println("Error when writing response", err)
17	}
18}
19
20func main() {
21	http.HandleFunc("/", webServer)
22	log.Println("starting web server...")
23	log.Fatal(http.ListenAndServe(":8080", nil))
24}

The above program implements a simple web server listening on port 8080. This web server performs an expensive operation and it’s simulated by a timer which fires after 10 seconds. After the web server receives a request it will respond with the output of 10 seconds. Save this program and run it from your favourite IDE. If you need help running this program, please visit /hello-world-gomod/

After this program is run, we sill see the output starting web server.... Now the web server is ready to serve request. Let’s go ahead and send a request using curl.

curl localhost:8080

10 seconds after the above request is sent, the server logs will print

writing response....

and we get the following response for the curl command.

Hello context

Now let’s take the case where the user is not patient enough and after sending the curl request, he cancels it by pressing ctrl + c in the keyboard. Immediately after pressing these keys, the curl command will return. What do you think will happen to the request which is in process in the web server? Well let’s find out.

If a curl request is sent and if it’s cancelled before 10 seconds using ctrl + c, the server will never know that the request is cancelled and it will still continue processing it. In this case too, the server will print

writing response....

and will write the response. This is not good of course. The request is cancelled and the server is wasting valuable CPU time on a request which is already cancelled. As you have guessed by now, this is the typical use case for a context. What we ought to be doing is watching on the context of the http request and return immediately if the context is cancelled. The program is modified to do this and provided below.

 1package main
 2
 3import (
 4	"io"
 5	"log"
 6	"net/http"
 7	"time"
 8)
 9
10func webServer(w http.ResponseWriter, r *http.Request) {
11	timer := time.NewTimer(10 * time.Second)
12	select {
13	case <-r.Context().Done():
14		log.Println("Error when processing request:", r.Context().Err())
15		return
16	case <-timer.C:
17		log.Println("writing response...")
18		_, err := io.WriteString(w, "Hello context")
19		if err != nil {
20			log.Println("Error when writing response", err)
21		}
22		return
23	}
24}
25
26func main() {
27	http.HandleFunc("/", webServer)
28	log.Println("starting web server...")
29	log.Fatal(http.ListenAndServe(":8080", nil))
30}

In the above program, in line no. 13, we watch on the context of the request and immediately return if it’s cancelled.

If the context is not cancelled, the server will respond with Hello context after 10 seconds.

Let’s try the same scenario which we tried earlier. Let’s run this new server and send a request using curl localhost:8080 and cancel it immediately using ctrl + c. The request will be cancelled and the done channel of the context will be closed. The server will print

Error when processing request: context canceled

and return immediately. Thus the server doesn’t continue to process a request which has been cancelled saving valuable CPU time.

Context can also be used to passed request scoped values such as a request ID. We will discuss this in a separate tutorial of its own.

I hope you liked this tutorial. Please leave your feedback and comments. Please consider sharing this tutorial on twitter or LinkedIn. Have a good day.