Using contexts to avoid leaking goroutines

Tue, Oct 25, 2016

The context package makes it possible to manage a chain of calls within the same call path by signaling context’s Done channel.

In this article, we will examine how to use the context package to avoid leaking goroutines.

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 is equal to 5.

// The call site of gen doesn't have a 
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.

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 cancellable 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 of this goroutine when ctx is done.
			case ch <- n:
				n++
			}
		}
	}()
	return ch
}

Now, the caller can signal the generator when it is done consuming. Once cancel function is called, the internal goroutine will be returned.

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
    }
}

// ...

The full program is available as a gist.

rakyll.org

Using contexts to avoid leaking goroutines