Context Cancellation in Go

Originally published at ayada.dev

In Go, we can use context to send cancellation signals to goroutines that is doing some work. The Done method in context returns a channel that acts as a cancellation signal to goroutines using the given context. When the channel is closed, it indicates to the goroutines that they should stop the work they are doing. We can use WithCancel context function to obtain a new context along with the cancel function that can be used to close the channel.

Here we have two examples of WithCancel function. In the first example, we will run two functions, runTicker and anotherTicker, in separate goroutines. The ticker functions periodically prints a string to the console and it will run indefinitely. Only way to stop the goroutines is to cancel the context. In the ticker functions, we wait for the channel returned by ctx.Done() to be closed by running the cancelFunc function that WithCancel provides.

package main

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

func exampleOne() {
    ctx, cancelFunc := context.WithCancel(context.Background())

    wg := new(sync.WaitGroup)
    wg.Add(2) // we will start two ticker goroutines.

    go func() {
        // start the two ticker functions in separate goroutines.
        go runTicker(ctx, wg)
        go anotherTicker(ctx, wg)

        // sleep for 10s for ticker functions to print something to the console.
        time.Sleep(10 * time.Second)

        // after 10s, cancel the context.
        cancelFunc()
    }()

    // wait for the ticker functions to complete.
    wg.Wait()
}

func anotherTicker(ctx context.Context, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-time.After(2 * time.Second):
            sleepContext(ctx, time.Minute) // sleep here indicates some work.
            fmt.Println(time.Now(), "another ticker executed")
        case <-ctx.Done(): // will execute if cancel func is called.
            return
        }
    }

}

func runTicker(ctx context.Context, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-time.After(2 * time.Second):
            fmt.Println(time.Now(), "ticker executed")
        case <-ctx.Done(): // will execute if cancel func is called.
            return
        }
    }
}

// a sleep function that honors context cancellation.
func sleepContext(ctx context.Context, delay time.Duration) {
    select {
    case <-ctx.Done(): // will execute if cancel func is called.
    case <-time.After(delay):
    }
}

In the second example, we will see that once the context is cancelled, even though execute function is run multiple times, it will return immediately without any sleep. This is because the channel returned by ctx.Done() in sleepContext function is already closed.

package main

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

func exampleTwo() {
    ctx, cancelFunc := context.WithCancel(context.Background())
    cancelFunc() // cancel the context immediately.

    for i := 0; i < 5; i++ {
        execute(ctx) // execute function 5 times, each time there is no sleep.
    }
}

func execute(ctx context.Context) {
    fmt.Println(time.Now(), "executing func")

    // since ctx was cancelled immediately, sleepContext will return immediately without any sleep.
    sleepContext(ctx, 5*time.Second)
}

package main

import (
    "flag"
    "fmt"
)

func main() {
    example := flag.Int("example", 1, "the example to run")
    flag.Parse()

    switch *example {
    case 1:
        fmt.Printf("===== Running example 1 =====\n\n")
        exampleOne()

    case 2:
        fmt.Printf("===== Running example 2 =====\n\n")
        exampleTwo()
    }
}

The output for both exampleOne and exampleTwo will look like this:

$ go build -o bin .

$ ./bin -example 1
===== Running example 1 =====

2021-12-29 16:25:26.166643 +0530 IST m=+2.001564014 ticker executed
2021-12-29 16:25:28.168533 +0530 IST m=+4.003442497 ticker executed
2021-12-29 16:25:30.169012 +0530 IST m=+6.003912087 ticker executed
2021-12-29 16:25:32.169898 +0530 IST m=+8.004789488 ticker executed
2021-12-29 16:25:34.167073 +0530 IST m=+10.001956888 another ticker executed

$ ./bin -example 2
===== Running example 2 =====

2021-12-29 16:25:42.014868 +0530 IST m=+0.000194393 executing func
2021-12-29 16:25:42.015062 +0530 IST m=+0.000387651 executing func
2021-12-29 16:25:42.015068 +0530 IST m=+0.000393697 executing func
2021-12-29 16:25:42.015071 +0530 IST m=+0.000397068 executing func
2021-12-29 16:25:42.015074 +0530 IST m=+0.000400107 executing func