WaitGroup in Go

By Gavin     @2020-06-22     1114 views

Go package sync provides some helpful synchronization primitives like Mutex, Once and WaitGroup. Let’s try WaitGroup today.

About WaitGroup

https://golang.org/pkg/sync/#WaitGroup

WaitGroup waits for a collection of goroutines to finish. Which means we have a main goroutine, and bunch of other goroutines, we could use WaitGroup to make sure our main goroutine will wait for other goroutines to finish their jobs.

Add(delta int) in which delta can be negative. Goroutines blocked by Wait() will be released if the counter becomes 0. It causes panic if the counter goes negative.

Done() decreases the counter by 1, it equals to Add(-1). And actually Done() invokes Add(-1) if you check out the source code.

// Done decrements the WaitGroup counter by one.
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

Wait() will block until the counter we mentioned before goes 0.

WaitGroup Usage

As the code below, we call Add(1) in the main goroutine, then print something and call Done() in another goroutine. Wait() will block the main goroutine until Done().

wg := sync.WaitGroup{}
wg.Add(1)  // we guarantee wg.Add() executed first
go func() {
	defer wg.Done()
	println("1. goroutine execute first")
}()
wg.Wait() // this line will wait for wg.Done()
println("2. then wait return")

We can get an output like this:

1. goroutine execute first
2. then wait return

The output will not change if you add a time.Sleep(time.Second) or some other works between line 4 and line 5. And WaitGroup can be defined like this, and we don’t have to initialize it explicitly.

var wg sync.WaitGroup

Why my goroutines executed after Wait()

Sometimes we found our goroutines executed unexpected, like this:

wg := sync.WaitGroup{}
go func() {
	wg.Add(1) // this goroutine may be not scheduled before wg.Wait()
	defer wg.Done()
	println("2. then goroutine execute")
}()
wg.Wait()
println("1. wait return first")

And the output is

1. wait return first
2. then goroutine execute

Obviously we put the wg.Add(1) at a wrong place (in the new goroutine). As we know the new goroutine may not execute immediately, so the counter still is 0 when we reach wg.Wait()

How to verify the theory? runtime.Gosched() can make the main goroutine yields the processor, so another goroutine could have a chance to run. (Thanks to a friend on Reddit, runtime.Gosched() doesn’t guarantee other goroutine will run, and the case below is just for verify our guess. We should always use wg.Add(1) in main goroutine to make sure it run before wg.Wait() just like our first case in WaitGroup Usage)

wg := sync.WaitGroup{}
go func() {
	wg.Add(1)
	defer wg.Done()
	println("1. goroutine execute first")
}()
runtime.Gosched() // yields the processor, our goroutine will get a chance to run
wg.Wait()
println("2. then wait return")

And the output looks good. Cool, That’s it.

1. goroutine execute first
2. then wait return

So we have to make sure the wg.Add() execute before wg.Wait(), or we’ll get some unexpected bugs.

And you can check out the test code on GitHub.

Gavin's Daily

© 2019-2020 Gavin's GoBlog