Proper Usage of sync.WaitGroup in Go Programming
The sync.WaitGroup type provides a mechanism to wait for a set number of operations to complete. Typically, it's used to await the completion of multiple goroutines. Let's first examine its public interface and then analyze a common mistake that leads to non-deterministic behavior.
A WaitGroup can be instantiated using its zero value:
wg := sync.WaitGroup{}
Internally, sync.WaitGroup maintains an internal counter initialized to zero. The Add(int) method increments this counter, while Done() or Add with a negative value decrements it. To block until the counter reaches zero, the Wait() method must be called.
Note: The counter cannot go below zero, which would cause a panic in the goroutine.
In the following example, we initialize a WaitGroup, launch three goroutines that automatically update the counter, and then wait for their completion. We expect all three goroutines to increment a shared value (which should be 3). Can you identify the issue in this code?
wg := sync.WaitGroup{}
var v uint64
for i := 0; i < 3; i++ {
go func() {
wg.Add(1)
atomic.AddUint64(&v, 1)
wg.Done()
}()
}
wg.Wait()
fmt.Println(v)
Running this snippet produces a non-deterministic result—any value from 0 to 3 may be printed. Furthermore, with the -race flag enabled, Go detects data races. Given that we're using the sync/atomic package to modify v, how could this happen?
The problem lies in calling wg.Add(1) inside the newly created goroutine rather than in the parent goroutine. As a result, there's no guarantee that the parent goroutine will wait for all three goroutines before proceeding.
The diagram below illustrates a scenario where the output is 2. In this case, the main goroutine initiates three goroutines. However, the last one executes after the first two have already called wg.Done(), thus allowing the parent goroutine to proceed prematurely. Consequently, when the main goroutine reads v, it equals 2. Race detection can also identify unsafe access to v.
When working with goroutines, remember that execution order is not guaranteed without synchronization. For instance, the following code might print either "ab" or "ba":
go func() {
fmt.Print("a")
}()
go func() {
fmt.Print("b")
}()
Both goroutines may run on different threads, and there's no guarantee about execution order.
CPU memory barriers (also known as memory fences) are required to enforce ordering. Go offers various synchronization techniques to achieve this, such as sync.WaitGroup ensuring a "happens-before" relationship between wg.Add and wg.Wait.
Returning to our example, there are two viable solutions. First, call wg.Add before entering the loop:
wg := sync.WaitGroup{}
var v uint64
wg.Add(3)
for i := 0; i < 3; i++ {
go func() {
// ...
}()
}
// ...
Alternatively, invoke wg.Add within each loop iteration before starting the child goroutine:
wg := sync.WaitGroup{}
var v uint64
for i := 0; i < 3; i++ {
wg.Add(1)
go func() {
// ...
}()
}
// ...
Both approaches are valid. If the total count for the WaitGroup is known beforehand, the first solution avoids repeated calls to wg.Add. However, it requires careful handling to ensure consistency across all usages to prevent subtle errors.
Avoid replicating this frequent mistake made by Go developers. When using sync.WaitGroup, the Add operation must occur before launching goroutines in the parent goroutine, and Done must be called within each goroutine.