Go: Writing Our Very Own WaitGroup

September 29, 2022

When we want to wait for the completion of multiple goroutines, we typically use WaitGroup, provided by the sync standard library package. Here is a quick sample that demonstrates how it is generally used:

package main

import (
        "sync"
)

type task struct{}

func worker(wg *sync.WaitGroup, t task) {
        defer wg.Done()
        // perform task
}

func main() {
        tasks := []task{{}, {}, {}}
        var wg sync.WaitGroup
        wg.Add(len(tasks)) // one goroutine for each task
        for _, task := range tasks {
                go worker(&wg, task)
        }
        wg.Wait() // wait for all goroutines to finish execution

	// continue
}

Using WaitGroup is more convenient than using channels for this particular use-case—were we to take the channel route, we would have had to create a buffered channel and decrement a counter each time we received from it.

Let us build our own simplified WaitGroup that supports Add(), Done(), and Wait(), as an exercise. We have our requirements already: the three foregoing functions are what we want to support. An implicit requirement is that our waitGroup must be thread-safe, for obvious reasons.

type waitGroup struct {
	counter int
	mu sync.Mutex
}

sync.Mutex could have been an embedded struct, but that would mean that waitGroup would expose Lock() and Unlock() to its users—we do not want this. Since mu is unexported, users importing waitGroup will not have access to the underlying mutex. counter is unexported for similar reasons.

Next, let us define Add and Done on a pointer receiver type. Were we to define these methods on a non-pointer receiver, the methods would have no effect—Add and Done would increment and decrement different counters respectively!

func (wg *waitGroup) Add(n int) {
	wg.mu.Lock()
	wg.counter += n
	wg.mu.Unlock()
}

func (wg *waitGroup) Done() {
	wg.mu.Lock()
	wg.counter--
	wg.mu.Unlock()
}

Easy enough, wasn't it? Now let us define Wait().

func (wg *waitGroup) Wait() {
	for {
		wg.mu.Lock()
		if wg.counter == 0 {
			wg.mu.Unlock()
			return
		}
		wg.mu.Unlock()
	}
}

Here, we are "busy waiting" for counter to become 0. Until it does, Wait will not return.

And there we have it! Our very own custom waitGroup.