Using contexts to avoid leaking goroutines
Table of Contents
The context package makes it possible to manage a chain of calls within the same call path by signaling the context’s Done channel.
In this article, we will examine how to use the context package to avoid leaking goroutines.
The Problem: Goroutine Leaks #
Assume you have a function that starts a goroutine internally. Once this function is called, the caller may not be able to terminate the goroutine started by the function.
// gen is a broken generator that will leak a goroutine.
func gen() <-chan int {
ch := make(chan int)
go func() {
var n int
for {
ch <- n
n++
}
}()
return ch
}
The generator above starts a goroutine with an infinite loop, but the caller consumes the values until n equals 5.
// The call site of gen doesn't have a way to stop the goroutine
for n := range gen() {
fmt.Println(n)
if n == 5 {
break
}
}
Once the caller is done with the generator (when it breaks the loop), the goroutine will run forever executing the infinite loop. Our code will leak a goroutine.
The Solution: Cancellable Contexts #
We can avoid the problem by signaling the internal goroutine with a stop channel, but there is a better solution: cancellable contexts. The generator can select on a context’s Done channel, and once the context is done, the internal goroutine can be cancelled.
// gen is a generator that can be cancelled by cancelling the ctx.
func gen(ctx context.Context) <-chan int {
ch := make(chan int)
go func() {
var n int
for {
select {
case <-ctx.Done():
return // avoid leaking this goroutine when ctx is done
case ch <- n:
n++
}
}
}()
return ch
}
Usage Example #
Now, the caller can signal the generator when it is done consuming. Once the cancel function is called, the internal goroutine will return.
ctx, cancel := context.WithCancel(context.Background())
defer cancel() // make sure all paths cancel the context to avoid context leak
for n := range gen(ctx) {
fmt.Println(n)
if n == 5 {
cancel()
break
}
}
Key Benefits #
- Prevents goroutine leaks: The context cancellation ensures goroutines are properly terminated
- Clean resource management: Using
defer cancel()
ensures cleanup happens even if errors occur - Flexible cancellation: Contexts can be cancelled from anywhere in the call chain
The full program is available as a gist.