Closing Go Worker Pools

Introduction

Following on from an earlier post on the [Worker Pool design pattern in Go]({% post_url 2024-04-20-go-worker-pools %}), here is an approach to closing worker pools when we are done with them. Its essence is to use the standard library’s sync.WaitGroup to keep track of how many of the workers are running.

Code Example Without Closing

The following code example shows a worker pool that does not close the output channel:

// Spin up as many workers as needed
for i := 0; i < noOfWorkers; i++ {
    // This goroutine is the worker
    go func(in <-chan int, out chan<- int) {
        // Actually do the work
        foo()
    }()
}

Code Example Where the Output Channel is Closed

// Outer goroutine manages the worker pool
go func() {
    // Prep waitgroup
    wg := sync.WaitGroup{}
    wg.Add(noOfWorkers)

    // Spin up as many workers as needed
    for i := 0; i < noOfWorkers; i++ {
        // Inner goroutine is the worker
        go func(in <-chan int, out chan<- int) {
            defer wg.Done()
            // Actually do the work
            foo()
        }()
    }

    // Wait at the end of the outer goroutine
    wg.Wait()
    // All workers finished. Safe to close channel.
    close(out)
}

Explanation

There is a fairly typical for loop, creating noOfWorkers workers. Each worker is a goroutine that runs foo(). Outside this, there is another layer of goroutine. There is just one of these, and all the worker goroutines are opened within it. It waits until they have all finished, then it closes their output channel.

The waitgroup is set up at the top of the orchestration goroutine, which calls wg.Add(noOfWorkers) at the start. Deferring the wg.Done() call in each worker ensures they will all decrement the waitgroup counter when they have finished their work. This may happen, for example, when their input channel closes.

Because all of this is happening within an outer goroutine, subsequent execution is not blocked while the workers are finishing their tasks.

Summary

Goroutines are cheap, and adding one more to get the channel closed shouldn’t add noticeable overhead to your code. Although closing the output channel can be challenging with multiple goroutines writing to it, the outer orchestrating goroutine can handle it without issue. This technique is most valuable when there are other workers doing something else downstream, which need to know when to stop. Downstream workers can use the familiar pattern for task := range(channel) may rely on the closure of the channel to indicate that their are no more tasks to come.