How and when to use sync.WaitGroup in Golang

Go is known for its first-class support for concurrency, or the ability for a program to deal with multiple things at once. Running code concurrently is becoming a more critical part of programming as computers move from running a single code stream faster to running more streams simultaneously.

Goroutines solve the problem of running concurrent code in a program, and channels solve the problem of communicating safely between concurrently running code.

Read details concept of Goroutines and Channels Concurrency in GO 101 - goroutine and channels

In that article, we looked at using channels in Go to get data from one concurrent operation to another. But we also ran across an another concern: How do we know when all the concurrent operations are complete? One answer is that we can use sync.WaitGroup.

Waitgroup

Allows you to block a specific code block to allow a set of goroutines to complete their executions.

WaitGroup ships with 3 methods:

• wg.Add(int): Increase the counter based on the parameter (generally "1").

• wg.Wait(): Blocks the execution of code until the internal counter reduces to 0 value.

• wg.Done(): Decreases the count parameter in Add(int) by 1.

Example:

package main
import (
    "fmt"
    "sync"
    "time"
)

var taskIDs = []int32{1, 2, 3}

func main() {
    var wg sync.WaitGroup
    for _, taskID := range taskIDs {
        wg.Add(1)
        go performTask(&wg, taskID)
    }
    // waiting untill counter value reach to 0
    wg.Wait()
}

func performTask(wg *sync.WaitGroup, taskID int32) {
    defer wg.Done()
    time.Sleep(3 * time.Second)
    fmt.Println(taskID, " - Done")
}

Run code here

In the above example, we have simply initiated 3 goroutines (we can also initiate more goroutines as we want). Before initiating every goroutine wg.Add(1) increases counter by 1. From inside goroutine defer wg.Done() executes after all functionality done and decreases counter value by 1. And finally wg.Wait() releases main block when counter value reached to 0. So, WaitGroup ensures completion of all initiated gorutine's execution.

When to use Waitgroup?

• Major use case is to simply wait for a set of goroutines untill complete their execution.
• Another is to use with channel(s) to achieve better outcome.

For better explanation of 2nd use case lets solve a real world problem. Assume we have 7 order ids. using concurrency we have to collect successful order details of those orders.

Firstly we will go for channel approach to collect individual order details.

package main
import (
    "fmt"
    "time"
)

var orderIDs = []int32{1, 2, 3, 4, 5, 6, 7}
type order struct {
    id   int32
    time string
}

func main() {
    var orders []order
    ch := make(chan order)

    for _, orderID := range orderIDs {
        go getOrderDetails(ch, orderID)
    }

    for i := 0; i < len(orderIDs); i++ {
        orders = append(orders, <-ch)
    }

    fmt.Printf("collected orders: %v", orders)
}

func getOrderDetails(ch chan order, orderID int32) {
    time.Sleep(3 * time.Second)
    // details collection logic
    orderDetails := order{id: orderID, time: time.Now().UTC().Format("15:04:05")}

    ch <- orderDetails
}

Run code here

Here, we have created an unbuffered channel ch. In every goroutine collects order details and sends data to channel. Same time for collecting order details we have iterated over length of expected order ids to receive order data one by one from channel for every order ids. As far the program works fine as expected and output is:

Output:

collected orders: [{3 23:00:03} {1 23:00:03} {7 23:00:03} {4 23:00:03} {2 23:00:03} {6 23:00:03} {5 23:00:03}]

But, suppose one of our goroutine returns an error instead of sending order data to channel on the other hand our channel trying to receive data but data not available in channel. In this situation program will enter deadlock situation. Try in here

Now we will combine channel and waitgroup to collect successful order data.

package main

import (
    "fmt"
    "sync"
    "time"
)

var orderIDs = []int32{1, 2, 3, 4, 5, 6, 7}
type order struct {
    id   int32
    time string
}

func main() {
    var wg sync.WaitGroup
    var orders []order

    // create a buffered channel wit length of orderIDs
    ch := make(chan order, len(orderIDs))

    for _, orderID := range orderIDs {
        wg.Add(1)
        go getOrderDetails(ch, &wg, orderID)
    }
    // here we wait until all jobs done
    wg.Wait()
    // closing channel after all job done
    close(ch)
    for v := range ch {
        orders = append(orders, v)
    }

    fmt.Printf("collected orders: %v", orders)
}

func getOrderDetails(ch chan order, wg *sync.WaitGroup, orderID int32) {
    defer wg.Done()

    time.Sleep(3 * time.Second)
    // details collection logic
    orderDetails := order{id: orderID, time: time.Now().UTC().Format("15:04:05")}

    if orderID == 4 {
        fmt.Println("Something went wrong")
        return
    }

    ch <- orderDetails
}

Run code here

Here we have created an buffered channel having max buffer size of orders ids length. From goroutine order data will be sent to channel(not necessary all order data will be sent, in case of error). Here for order id 4 no data sent to channel but wg.Done() executes as expected. After wg.Wait() reaches to its ends, we have closed channel as all goroutine's job done. Then from channel we can simply receive available orders data by iterate over available channel results.

In our case there was 7 incoming order ids and we finally got 6 successful orders detail from channel without any error.

Output:

collected orders: [{5 23:00:00} {6 23:00:00} {3 23:00:00} {1 23:00:00} {7 23:00:00} {2 23:00:00}]

Final thoughts

Channels being a signature feature of the Go language shouldn't mean that it is idiomatic to use them alone whenever possible. What is idiomatic in Go is to use the simplest and easiest to understand solution. The WaitGroup convey both the meaning (main function is Waiting for workers to be done) and the mechanic (the workers notify when they are Done).