Join me next Friday, at 3:00pm EDT (19:00 UTC), for a live stream where I’ll be building a new feature in my open-source project. Subscribe on YouTube and hit the Notify me button so you don’t miss it!

Yesterday I introduced the sync.WaitGroup type. Today I want to look at some possible ways to improve the same code.

First, the code as we left off yesterday:

wg := &sync.WaitGroup{}

for i := range 10 {
	wg.Add(1)
	go func(i int) {
		fmt.Println(i)
		wg.Done() // Signal that this goroutine is done
	}(i)
}

// Wait for all goroutines to finish
wg.Wait()

The first improvement we can make here is to call wg.Add only once. This doesn’t work in all cases. You may not know how many goroutines you want to launch ahead of time, for example. But in this case, we know we’re going to launch exactly 10 goroutines, so we can call wg.Add(10) before the loop:

wg := &sync.WaitGroup{}

wg.Add(10)
for i := range 10 {
	go func(i int) {
		fmt.Println(i)
		wg.Done() // Signal that this goroutine is done
	}(i)
}

// Wait for all goroutines to finish
wg.Wait()

The second, and more substantial improvement we can make is to defer the call to wg.Done(). In this particular case, it doesn’t matter, so let’s first make a small change. Let’s imagine we only want to print the even numbers. Here’s one possible (buggy) implementation:

wg := &sync.WaitGroup{}

wg.Add(10)
for i := range 10 {
	go func(i int) {
		if i%2 != 0 {
			return // Skip odd numbers
		}
		fmt.Println(i)
		wg.Done() // Signal that this goroutine is done
	}(i)
}

// Wait for all goroutines to finish
wg.Wait()

This code will succeed in only printing even numbers, but it will then hang forever, because the wg.Done() call is never reached for odd numbers. The WaitGroup counter will never reach zero, and the program will block on wg.Wait().

To solve this, we can call wg.Done() using a defer statement as the first line in the anonymous function:

wg := &sync.WaitGroup{}

wg.Add(10)
for i := range 10 {
	go func(i int) {
		defer wg.Done() // Signal that this goroutine is done
		if i%2 != 0 {
			return // Skip odd numbers
		}
		fmt.Println(i)
	}(i)
}

// Wait for all goroutines to finish
wg.Wait()

Now, no matter whether the function returns early or not, wg.Done() will always be called, and the program will not hang.