Wait groups

A wait group waits for a collection of things, usually goroutines, to finish. It is essentially a thread-safe counter that allows you to wait until the counter reaches zero. A common pattern for their usage is this:

// Create a waitgroup
wg := sync.WaitGroup{}
for i := 0; i < 10; i++ {
    // Add to the wait group **before** creating the 
    //goroutine
    wg.Add(1)
    go func() {
        // Make sure the waitgroup knows about
        // goroutine completion
         defer wg.Done()
         // Do work
    }()
}
// Wait until all goroutines are done
wg.Wait()

When you create a WaitGroup, it is initialized to zero, so a call to Wait will not wait for anything. So, you have to add the number of things it has to wait for before calling Wait. To do this, we call Add(n), where n is the number of things to add for waiting. It makes it easier for the reader to call Add(1) just before creating the thing to wait, which is, in this case, a goroutine. The main goroutine then calls Wait, which will wait until the wait group counter reaches zero. For that to happen, we have to make sure that the Done method is called for each goroutine that returns. Using a defer statement is the easiest way to ensure that.

A common use for a WaitGroup is in an orchestrator service that calls multiple services and collects the results. The orchestrator service has to wait for all the services to return to continue computation.

Look at the following example:

func orchestratorService() (Result1, Result2) {
    wg := sync.WaitGroup{}  // Create a WaitGroup
    wg.Add(1)     // Add the first goroutine
    var result1 Result1
    go func() {
         defer wg.Done() // Make sure waitgroup
                        // knows completion
         result1 = callService1() // Call service1
    }()
    wg.Add(1)     // Add the second goroutine
    var result2 Result2
    go func() {
        defer wg.Done()  // Make sure waitgroup
                         // knows completion
         result2 = callService2()  // Call service2
    }()
    wg.Wait()     // Wait for both services to return
    return result1, result2    // Return results
}

A common mistake when working with a WaitGroup is calling Add or Done at the wrong place. There are two points to keep in mind:

• Add must be called before the program has a chance to run Wait. That implies that you cannot call Add inside the goroutine you are waiting for using the WaitGroup. There is no guarantee that the goroutine will run before Wait is called.
• Done must be called eventually. The safest way to do it is to use a defer statement inside the goroutine, so if the goroutine logic changes in time or it returns in an unexpected way (such as a panic), Done is called.

Sometimes using a wait group and channels together can cause some chicken or the egg problems: you have to close a channel after Wait, but Wait will not terminate unless you close the channel. Look at the following program:

 1: func main() {
 2:    ch := make(chan int)
 3:    var wg sync.WaitGroup
 4:    for i := 0; i < 10; i++ {
 5:        wg.Add(1)
 6:        go func(i int) {
 7:            defer wg.Done()
 8:            ch <- i
 9:        }(i)
10:    }
11:    // There is no goroutine reading from ch
12:    // None of the goroutines will return
13:    // so this will deadlock at Wait below
14:    wg.Wait()
15:    close(ch)
16:    for i := range ch {
17:        fmt.Println(i)
18:    }
19: }

One possible solution is to put the for loop at lines 16-18 into a separate goroutine before Wait, so there will be a goroutine reading from the channels. Since the channels will be read, all goroutines will terminate, which will release the wg.Wait, and close the channel, terminating the reader for loop:

go func() {
     for i := range ch {
           fmt.Println(i)
     }
}()
wg.Wait()
close(ch)

Another solution is as follows:

go func() {
     wg.Wait()
     close(ch)
}()
for i := range ch {
     fmt.Println(i)
}

The wait group is now waiting inside another goroutine, and after all the waited-for goroutines return, it closes the channel.