New function signal.NotifyContext in GoLang 1.16

os/signal package in GoLang may not be frequently used but it provides some good features like Shutdown() which can be used to gracefully shutdown a running HTTP server.

func (srv *Server) Shutdown(ctx context.Context) error

With this function, there is no need to use third party library to gracefully shutdown HTTP server. How is it being used?

package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"time"
)

func main() {
	server := http.Server{
		Addr: ":8080",
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(time.Second * 10)
		fmt.Fprint(w, "Hello world!")
	})

	go server.ListenAndServe()

	// Listen to interrupt signal（CTRL + C）
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c

	// Reset os.Interrupt default behavior
	signal.Reset(os.Interrupt)

	fmt.Println("shutting down gracefully, press Ctrl+C again to force")

	// Give 5s more for the system to process existing requests
	timeoutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := server.Shutdown(timeoutCtx); err != nil {
		fmt.Println(err)
	}
}

In above code, it uses os/signal to listen to the OS interrupt signal, after receiving the signal, the channel will be unblocked and the code will proceed to create the context with cancel and shutdown the server. In order to allow force exit when Shutdown is being ran, the Reset function will reset os.Interrupt default behavior(this is not necessary though).

The key to graceful shutdown are:

• No new request can get in
• Finish processing existing queued requests

In Go 1.16, a new function called NotifyContext() is introduced.

func NotifyContext(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc)

Its function is similar to Notify but with a different way to write the code. To change above code with NotifyContext function.

package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"time"
)

func main() {
	server := http.Server{
		Addr: ":8080",
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(time.Second * 10)
		fmt.Fprint(w, "Hello world!")
	})

	go server.ListenAndServe()

	// Listen to interrupt signal（CTRL + C）
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	<-ctx.Done()

	// Reset os.Interrupt default behavior, similar to signal.Reset
	stop()
	fmt.Println("shutting down gracefully, press Ctrl+C again to force")

	// Gievn 5s more to process existing requests
	timeoutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := server.Shutdown(timeoutCtx); err != nil {
		fmt.Println(err)
	}
}

Indeed, internal implementation of NotifyContext uses Notify.

func NotifyContext(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc) {
	ctx, cancel := context.WithCancel(parent)
	c := &signalCtx{
		Context: ctx,
		cancel:  cancel,
		signals: signals,
	}
	c.ch = make(chan os.Signal, 1)
	Notify(c.ch, c.signals...)
	if ctx.Err() == nil {
		go func() {
			select {
			case <-c.ch:
				c.cancel()
			case <-c.Done():
			}
		}()
	}
	return c, c.stop
}

When stop() is called, it will call Stop() function in os/signal package. This function has similar function to Reset(). 

From its implementation, NotifyContext does a better encapsulation of the function and can also take Context which would provide more useful value passing at different places.

Reference: Go1.16 中的新函数 signal.NotifyContext 怎么用？ (qq.com)