Understanding how cancellation works

cancel.go

package main

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

// simply runs a program to showcase how to properly leverage
// context cancellations to abort goroutines timely.
func main() {
	ctx, cancel := context.WithCancel(context.Background())
	
	// ALWAYS defer the cancel. Even if you plan on cancelling manually at some
	// point, it's safer to defer a cancel because when you invoke a WithX function,
	// WithCancel in this case, a new goroutine is spawned in the background to
	// propagate the cancellation to children contexts as well as keeping track
	// of them. Were you to forget cancelling, that would end up being a leak.
	defer cancel()

	var wg sync.WaitGroup

	// this goroutine thinks it's listening to canceled context, but it isn't.
	wg.Add(1)
	go func() {
		defer wg.Done()
		fmt.Println("starting work 1!")

		// at the time this line is reached, the context hasn't been cancelled yet,
		// so the program continues and doesn't abort.
		if ctx.Err() != nil {
			fmt.Printf("aborting work 1: %s!\n", ctx.Err().Error())
			return
		}
		time.Sleep(time.Millisecond * 5)
		fmt.Println("done work 1!")
	}()

	// this goroutine actually is listening to cancelled contexts and will abort timely.
	wg.Add(1)
	go func() {
		defer wg.Done()
		fmt.Println("starting work 2!")

		select {
		case <-ctx.Done():
			fmt.Printf("aborting work 2: %s!\n", ctx.Err().Error())
			return
		case <-time.After(time.Millisecond * 5):
			fmt.Println("done work 2!")
			return
		}
	}()

	// cancel emits an empty struct through the "done" channel child contexts have,
	// so if any code is listening to ctx.Done(), they'll get the message.
	time.Sleep(time.Millisecond)
	cancel()
	wg.Wait()

	fmt.Println("finished!")
}