Decryption Go: WaitGroup

It may be the most straightforward sync library in Golang

Note: Non-members can read the full story in this link.

sync.WaitGroup, I believe all gophers have used it. It is a concurrency primitive in the package sync used for task coordination. It solves the problem of concurrency-waiting: when a goroutine A is waiting at a checkpoint for a group of goroutines to complete. Without this synchronization primitive, how can we achieve this functionality? There are many ways, let’s first try using channels to implement it.

package main 
 ​ 
 import ( 
   "fmt" 
   "time" 
 ) 
 ​ 
 func main() { 
   channels := make(chan struct{}, 10) 
   for i := 0; i < 10; i++ { 
     go func() { 
       println("hello,world") 
       time.Sleep(time.Second * 10) 
       channels <- struct{}{} 
     }() 
   } 
   for i := 0; i < 10; i++ { 
     <-channels 
     fmt.Println(fmt.Sprintf("done:%d", i)) 
   } 
   fmt.Println("all done") 
 }

Although this code can fulfill the functionality, it is cumbersome. We not only need to create a channel but also create a goroutine to wait for this channel. Moreover, if the number of goroutines we need to wait for is uncertain, this approach is unsuitable. Therefore, sync.WaitGroup comes into play.

Now, let’s implement the above functionality differently using sync.WaitGroup go.dev link.

package main 
 ​ 
 import ( 
   "fmt" 
   "sync" 
   "time" 
 ) 
 ​ 
 func main() { 
   wg := sync.WaitGroup{} 
   wg.Add(10) 
   for i := 0; i < 10; i++ { 
     go func() { 
       defer wg.Done() 
       println("hello,world") 
       time.Sleep(time.Second * 10) 
     }() 
   } 
   for i := 0; i < 10; i++ { 
     go func() { 
       wg.Wait() 
       fmt.Println(fmt.Sprintf("done,%d", i)) 
     }() 
   } 
   wg.Wait() 
   fmt.Println("all done") 
   time.Sleep(time.Second * 1) 
 }

Many operating systems and programming languages provide similar concurrency primitives. For example, the barrier in Linux, the barrier in Pthread (POSIX threads), std::barrier in C++, CyclicBarrier and CountDownLatch in Java, and so on. This concurrency primitive is, therefore, a very fundamental type of concurrency.

Basic Usage of WaitGroup

The primary usage of WaitGroup is straightforward, with only three methods:

func (wg *WaitGroup) Add(delta int) 
func (wg *WaitGroup) Done() 
func (wg *WaitGroup) Wait()

The Add method increases the number of goroutines to wait for, the Done method decreases the number of goroutines to wait for, and the Wait method waits for all goroutines to complete.

Implementation of WaitGroup

This article is based on go1.21.4

First, let’s take a look at the data structure of WaitGroup:

type WaitGroup struct { 
   noCopy noCopy 
   state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count. 
   sema  uint32 
 }

The WaitGroup data structure is very simple, with only two fields: state and sema. State is an atomic Uint64 type used to store the counter and the number of waiting goroutines, and sema is a semaphore used for waiting.

Add Method

Let’s look at the implementation of the Add method, which removes the race detection.

func (wg *WaitGroup) Add(delta int) { 
   state := wg.state.Add(uint64(delta) << 32) 
   v := int32(state >> 32) // counter 
   w := uint32(state)      // waiter count 
   if v < 0 { 
     panic("sync: negative WaitGroup counter") 
   } 
   if w != 0 && delta > 0 && v == int32(delta) { 
     panic("sync: WaitGroup misuse: Add called concurrently with Wait") 
   } 
   if v > 0 || w == 0 { 
     return 
   } 
   // This goroutine has set counter to 0 when waiters > 0. 
   // Now there can't be concurrent mutations of state: 
   // - Adds must not happen concurrently with Wait, 
   // - Wait does not increment waiters if it sees counter == 0. 
   // Still do a cheap sanity check to detect WaitGroup misuse. 
   if wg.state.Load() != state { 
     panic("sync: WaitGroup misuse: Add called concurrently with Wait") 
   } 
   // Reset waiters count to 0. Wake up all waiters. 
   wg.state.Store(0) 
   for ; w != 0; w-- { 
     runtime_Semrelease(&wg.sema, false, 0) 
   } 
 }

The implementation of the Add method is straightforward. It left-shifts the delta by 32 bits and adds it to the state. Then, it checks if the counter is less than 0, and if so, it panics. It also checks if the waiter count is not 0, and if the delta is greater than 0 and the counter is equal to the delta, it panics. It returns if the counter is greater than 0 or the waiter count is 0. If other goroutines have modified the state, it panics. Finally, it sets the waiter count to 0 and wakes up all waiters.

Done Method

The implementation of the Done method is straightforward. It calls the Add method with a parameter of -1.

Wait Method

The logic of the Wait method is as follows: it continuously checks the state's value. If the counter becomes 0, all tasks have been completed, and the caller does not need to wait any longer, so it returns directly. If the counter is greater than 0, there are still tasks that have not been completed, so the caller becomes a waiter, needs to join the waiter queue, and blocks itself.

func (wg *WaitGroup) Wait() { 
   for { 
     state := wg.state.Load() 
     v := int32(state >> 32) 
     w := uint32(state) 
     if v == 0 { 
       return 
     } 
     // Increment waiters count. 
     if wg.state.CompareAndSwap(state, state+1) { 
       runtime_Semacquire(&wg.sema) // block until signalled 
       if wg.state.Load() != 0 { 
         panic("sync: WaitGroup is reused before previous Wait has returned") 
       } 
       return 
     } 
   } 
 }

Common Errors when Using WaitGroup

1. The high 32 bits of the state are the counter, and the low 32 bits are the waiter count. If the value of the counter exceeds 2³¹-1, it will panic. This problem rarely occurs in practical use because this value is too large.
2. Add and Done do not appear in pairs, which means that if you add n times, you should also call Done n times. If you add n times but call Done m times, where m < n, it will panic or cause deadlock.
3. The timing of Add is unexpected. If Add is called after Wait, it will panic. This will cause the counter to become 0, but there are still goroutines waiting.
4. Reusing WaitGroup.

Bugs in WaitGroup usage in real projects

1. In Golang issue 28123, the biggest problem in this code is that line 9 copies the WaitGroup instance w. Although this code can be executed successfully, it does violate the rule of not copying the WaitGroup instance after use. In projects, we can use the vet tool to detect such errors.
2. Docker issue 28161 and issue 27011 are errors caused by reusing WaitGroup without waiting for the previous Wait to finish before calling Add.
3. Etcd issue 6534 is also a bug in reusing WaitGroup, where Add is called without waiting for the previous Wait to finish.
4. Kubernetes issue 59574 is a bug in which the Wait is forgotten before increasing the count. This bug is considered almost impossible to occur in our usual understanding.

Conclusion

sync.WaitGroup is a very basic concurrency primitive. Its implementation is straightforward, but it has many pitfalls. When using it, you must pay attention to not copying the WaitGroup instance, not reusing the WaitGroup instance, not calling Add after Wait, and ensuring that Add and Done appear in pairs. We can use the ‘go vet’ tool to check for these issues in real projects.