Signals via Context

In addition to the Notify function, the os/signal package also includes the NotifyContext function. Rather than accepting a channel, this function accepts a context and returns a new context that includes everything in the old context, but this new context will receive a value from the ctx.Done() channel when a matching signal is received.

NotifyContext also returns a second argument - a function that we can call to stop listening for our specified signals.

Let’s look at an example. We will use the same example we had at the end of the Using Signals with Go writeup, except we will use NotifyContext instead of Notify and Stop.

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop() // Ensure that this gets in case we return early.
	// ctx.Done() returns a channel that will have a message
	// when the context is cancelled. We wait for that signal, which means
	// we received the signal, or our context was cancelled for some other reason.
	<-ctx.Done()
	// Stop receiving signals
	stop()

	fmt.Println("Shutting down...")
	// Simulate a slow shutdown
	time.Sleep(3 * time.Second)
	// Shutdown has completed
	fmt.Println("Shutdown complete")
}

The major upside to using NotifyContext is that we often want to work with a context anyway, as a large portion of Go code is designed to work with the context object. NotifyContext is a nice way to simplify that process rather than needing to code a custom context object on our own.

A quick example of where this might be useful is if our program needed to end in a certain amount of time. We can use the context.WithTimeout function to create a context that will be cancelled if it doesn’t complete within that time frame, and this can easily be combined with our signal context.

func main() {
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
	defer stop()

	// Wait for the context to be cancelled via timeout or signal.
	<-ctx.Done()
	stop()
	fmt.Println("Shutting down...")
	time.Sleep(3 * time.Second)
	fmt.Println("Shutdown complete")
}

You can find this in the cmd/ctx directory of the accompanying repository for this article.

Now our program will terminate either when it times out at after 3 seconds, or when it receives an interrupt signal, whichever comes first. After that it will start shutting down (which coincidentally introduces another 3 second timeout in this example).

There are two notable downsides to using NotifyContext:

1. We won’t know which specific signal was received, so we likely only want to use this for a small set of signals that we are going to respond to in the same way.
2. We will only be notified of the first signal received.

In many situations these two restrictions won’t matter, but they are worth noting.

Be Sure to Call the stop Function!

One last thing to cover before we wrap this up - it is important to ensure that the stop function returned by NotifyContext is being called after we receive a signal. In the last example we saw this in two places:

1. We called defer stop() immediately after creating our context.
2. We called stop() immediately after receiving a message via the ctx.Done() channel.

Technically we only needed to call stop() once, and that was immediately after we received a message via the ctx.Done() channel. We want to call stop() here as otherwise we will have the same issue we saw when using Notify and Stop - our program would continue to capture signals and would prevent an interrupt signal from terminating the program before the 3 second sleep. We can verify this by removing the call to the stop function.

func main() {
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
	defer stop()

	// Wait for the context to be cancelled via timeout or signal.
	<-ctx.Done()
	// stop() // comment this out to show it is necessary.
	fmt.Println("Shutting down...")
	time.Sleep(3 * time.Second)
	fmt.Println("Shutdown complete")
}

With the stop() call commented out, try running this program and pressing Ctrl+C multiple times. No matter how many times you send the SIGINT signal, it will still take 3 seconds to terminate.

This happens because our defer stop() function call will never happen until our main() function exits, and despite the fact that our context can only receive one signal, additional signals will be captured until we call stop. We need to explicitly tell our program to stop listening for signals.

Then why do we even bother with the deferred call to stop? Technically, we don’t need to, but I have found it is better to add it just to be sure it gets called. Code can change over time, and we might find ourselves with another branching path that allows a function to exit early. In those cases, we want to be sure stop is eventually called. Plus, stop is an idempotent function, which means calling it multiple times won’t have any adverse side effects.