GoLang exiting from multiple go routines with context and wait group

context_cancel.go

package main

// Here's a simple example to show how to properly terminate multiple go routines by using a context.
// Thanks to the WaitGroup we'll be able to end all go routines gracefully before the main function ends.

import (
	"context"
	"fmt"
	"math/rand"
	"os"
	"sync"
	"time"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

// it prints all values pushed into "ch" ("ch" here is read only)
func reader(ctx context.Context, wg *sync.WaitGroup, ch <-chan int) {
	defer wg.Done() // decrements the WaitGroup counter by one when the function returns

	for {
		select {
		case <-ctx.Done(): // Done returns a channel that's closed when work done on behalf of this context is canceled
			fmt.Println("Exiting from reading go routine")
			return
		case v, ok := <-ch:
			if !ok {
				fmt.Println("Channel has been closed")
				return
			}

			fmt.Println(v)
		}
	}
}

// it writes a random integer into "ch" every second ("ch" here is write only)
func writer(ctx context.Context, wg *sync.WaitGroup, ch chan<- int) {
	defer wg.Done() // decrements the WaitGroup counter by one when the function returns

	for {
		select {
		case <-ctx.Done(): // Done returns a channel that's closed when work done on behalf of this context is canceled
			fmt.Println("Exiting from writing go routine")
			return
		case ch <- rand.Intn(100): // pushes a random integer from 0 to 100 into the channel
			time.Sleep(1 * time.Second) // sleeps one second
		}
	}
}

func main() {
	channel := make(chan int)
	defer close(channel)

	// a WaitGroup waits for a collection of goroutines to finish, pass this by address
	waitGroup := sync.WaitGroup{}

	// context.WithCancel returns a copy of parent with a new Done channel.
	// The returned context's Done channel is closed when the returned cancel function is called or when the parent
	// context's Done channel is closed, whichever happens first.
	ctx, cancel := context.WithCancel(context.Background())

	waitGroup.Add(2)                    // adds delta, if the counter becomes zero, all goroutines blocked on Wait are released
	go reader(ctx, &waitGroup, channel) // go routine that prints all values pushed into "channel"
	go writer(ctx, &waitGroup, channel) // go routine that writes a random integer into "channel" every second

	// go routine that listens for an Enter keystroke to terminate the program
	go func() {
		os.Stdin.Read(make([]byte, 1)) // wait for Enter keystroke
		cancel()                       // cancel the associated context
	}()

	waitGroup.Wait() // it blocks until the WaitGroup counter is zero

	fmt.Println("Exiting from main")
}

It's a big gotcha to pass the WaitGroup pointer reference. Thanks!

It is unsafe to put wg.Add(1) inside the goroutine. There are chances that this runs after waitGroup.Wait(). So it should be safer to move wg.Add(1) before the goroutine like this:

	waitGroup.Add(1)
	go reader(&waitGroup, ctx, channel) // go routine that prints all values pushed into "channel"
	waitGroup.Add(1)
	go writer(&waitGroup, ctx, channel) // go routine that writes a random integer into "channel" every second

Good spot @xianghuzhao I've updated the gist 👍

Thanks, this gist is really helpful.

In the main function waitGroup variable should have name "wg" or "waitGroup" otherwise this will not compile.

Amazing. Thanks for sharing this.

Do you want me to open merge request for this? It is really nice example, pitty to not have it workable...

Do you want me to open merge request for this? It is really nice example, pitty to not have it workable...

@landrisek Busy times 😅 I updated the gist now it works 👍
@johnmackenzie91 glad it's been useful 👍