Cancellation Done Right

cancel.go

// TODO(as): find a production foss program where this issue exists
//
// This trivial program generates integers and prints them to
// standard output until it reaches 1000 or the context is done.
//
// It contains a fix for a difficult-to-find bug caused in many
// Go programs written by authors at all experience levels.
package main

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

func main() {
	in := make(chan int)
	go func() {
		for i := 0; i < 1000; i++ {
			in <- i
		}
		close(in)
	}()

	ctx, fn := context.WithTimeout(context.Background(), time.Second)
	defer fn()

	do(ctx, in)
}

func do(ctx context.Context, in chan int) {
	for {
		select {
		case <-ctx.Done():
      // first cancellation check
			return
		case i, ok := <-in:
			if !ok {
				// consumer closed the channel
				return
			}

			select {
			case <-ctx.Done():
				// edge-triggered cancellation check
				// both channels were ready, but the "in" channel won
				return
			default:
			}
			// The select statement is a random selection, so if both cases in the outer
			// select are ready at the time of selection, a random path is taken. If that
			// random path is taken to this case, there is a chance the context was "done" too
			//
			// Many users assume that when a context (or done channel for that matter) is
			// cancelled/closed that operations in that function cease. If you only have
			// the level-triggered check select you are rolling the dice on that, and may
			// execute in the function after a context is "done".
			//
			// This can lead to nasty issues and when it does, is very difficult
			// to trace. What happens specifically depends on what you do in the
			// function itself. The edge-triggered check is not necessary for all
			// applications, but it's vital to have when your constraints dictate
			// that cancellation must not result in further data processing.
			//
			// (We assume the context is independent of the "in" channel for the purpose
			// of demonstrating this effectively, and that someone can cancel it out-of-band
			// with the producer to the in channel)

			fmt.Println(i)
		}
	}
}

// Note: This do function is an example to demonstrate two issues, in
// reality this function shouldn't even have a context. Closing the
// in channel is a good-enough signal to tell the function it's done
// processing data (in fact you could even use a for range over the
// channel and avoid the select altogether). Real functions like this
// exist though, and the edge/level triggered checks become less apparent
// when they're missing and necessary.