WaitGroups in Golang

31 May 2022 ⏱️ 5 min

Get latest articles directly in your inbox

Concurrency in Golang is the ability for functions to run independently of each other. It refers to the composition of a set of independently executing processes. To support concurrency, Golang provides us with goroutines. These are functions that run concurrently with other functions or methods. Goroutines are quite efficient and light-weight (use only 2kB of stack space)

What are WaitGroups?

WaitGroups in Go allows a program to wait for specified goroutines. This is provided as a part of the sync package mainly because WaitGroups are type of sync mechanism that blocks the execution of a program until the execution of goroutines in the group is completed.

WaitGroups are important because they allow a goroutine to block the thread and execute it. One would have to handle adding delay in the main thread manually to let the goroutines execute. Using wait groups ensures that your main thread continues execution till wait() is reached.

Here’s an example to show the issue without wait group. If you run the code below you’ll a sum that is not 10. The issue here is that the main thread exits while your goroutines couldn’t complete the execution.

package main

import (
	"fmt"
)

func main() {
	sum := 0

	// function to increment a sum variable by 1
	increment := func() {
		sum++
	}

	for i := 0; i < 10; i++ {
		go func() {
			increment()
		}()
	}

	fmt.Printf("Sum is %d\n", sum)
}

As mentioned above one solution to this can be an addition of sleep in the main thread. You will get the sum as 10 but imagine blocking your main thread at multiple places across your big codebase.

package main

import (
	"fmt"
)

func main() {
	sum := 0

	// function to increment a sum variable by 1
	increment := func() {
		sum++
	}

	for i := 0; i < 10; i++ {
		go func() {
			increment()
		}()
	}

	time.Sleep(5 * time.Second)
	fmt.Printf("Sum is %d\n", sum)
}

How do the WaitGroups work?

Waitgroups comes with 3 core methods that help in blocking the execution of a program.

wg.Add(int)
wg.Wait()
wg.Done()

Add() Method

To define the number of goroutines to wait for, WaitGroup maintains an internal counter that is added via wg.Add() method. In simple terms, this function adds a value (positive or negative) to the WaitGroup counter. If this counter becomes 0, the waitgroup releases the goroutines that are blocked on the wait(). (We’ll discuss more about wait below)

Note: If the counter goes negative, the code will panic. When dealing with negative values in Add(), handle the panics.

Wait() Method

The Wait() method blocks the execution of code until the internal counter maintained by the waitGroup reduces to a 0 value.

Done() Method

Once a goroutine has completed its execution, we need to decrease the count parameter defined in Add(int) by 1. This can be achieved using the Done() method. This is usually used with the defer statement inside the executing goroutine method. If you check the go codebase, you’ll see that done method calls Add() with -1 value.

// Ref: https://go.dev/src/sync/waitgroup.go
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

Learning with Example

Let’s use the same example as above to understand the three WaitGroup methods better. Here we create multiple goroutines that increment the sum by 1 and finally print out the sum.

Before diving into the wait group solution we’ll see the problem

package main

import (
	"fmt"
	"sync"
)

func main() {
	sum := 0

 	// function to increment a sum variable by 1
	increment := func() {
		sum++
	}

  // declare a waitgroup
	var wg sync.WaitGroup

  // increment internal counter by 10
	wg.Add(10)

	for i := 0; i < 10; i++ {
		go func() {
			defer wg.Done() // each iteration reduces the internal counter by 1
			increment()
		}()
	}

  // wait for all goroutines to be released.
	wg.Wait()

	fmt.Printf("Sum is %d\n", sum)
}

Here var wg sync.WaitGroup creates a new WaitGroup
Then wg.Add(10) tells that wait group must wait for ten goroutines.
Post that we do defer wg.Done() updating the WaitGroup and decreasing the count by 1 on each execution of goroutine completes.
Finally wg.Wait() blocks the execution until the counter reaches 0 and all the goroutines are released.

Things to Remember

If a WaitGroup is explicitly passed into functions, it should be done by a pointer. The waitgroup struct in waitgroup.go defines a noCopy attribute. Thus, it must not be copied after first use. Thereby we use pointer.
```
type WaitGroup struct {
	noCopy noCopy

	state1 uint64
	state2 uint32
}
```
WaitGoups are just enough if you don’t need any data returned from a goroutine. When data needs to be returned, one can use channels in Go.
Can’t I use a channel to do everything that a WaitGroup does?
The answer is Yes. But channels are generally used when there is some communication required among goroutines whereas WaitGroups have a single task to wait until all independent operations are executed.