add early signal for use in health checks

README.md

@@ -33,7 +33,18 @@ srv := graceful.DualServerConfig{
 	// Logger optionally adds the ability to log messages, both errors and not.
 	Logger: log,
 }
+intMux.HandleFunc("/health/alive", func(w http.ResponseWriter, _ *http.Request) {
+	w.WriteHeader(http.StatusOK)
+})
+intMux.HandleFunc("/health/ready", func(w http.ResponseWriter, _ *http.Request) {
+	select {
+		case <-srv.ShuttingDown():
+			w.WriteHeader(http.StatusServiceUnavailable)
+		default:
+			w.WriteHeader(http.StatusOK)
+	}
+})
 if err := srv.ListenAndServe(ctx, *intAddr, *extAddr); err != nil {
 	log.Error(err, "Serving failed")
 }
-```
+```

context.go

@@ -16,26 +16,31 @@ import (
 	"github.com/go-logr/logr"
 )
 
-// Contexts returns two contexts which respectively serve as soft and hard
-// shutdown signals. They are cancelled after TERM or INT signals are received.
+// Contexts returns three contexts which respectively serve as warning, soft,
+// and hard shutdown signals. They are cancelled after TERM or INT signals are
+// received.
 //
-// If delay is positive, it will wait that duration after receiving a signal
-// before cancelling the first context. This is useful to allow loadbalancer
-// updates before the server stops accepting new requests.
+// When a shutdown signal is received, the warning context is cancelled. This is
+// useful to start failing health checks while other traffic is still served.
 //
-// If grace is positive, it will wait that duration after cancelling the first
-// context before cancelling the second context. This is useful to set a maximum
-// time to allow pending requests to complete.
+// If delay is positive, the soft context will be cancelled after that duration.
+// This is useful to allow loadbalancer updates before the server stops accepting
+// new requests.
+//
+// If grace is positive, the hard context will be cancelled that duration after
+// the soft context is cancelled. This is useful to set a maximum time to allow
+// pending requests to complete.
 //
 // Repeated TERM or INT signals will bypass any delay or grace time.
-func Contexts(ctx context.Context, log logr.Logger, delay, grace time.Duration) (context.Context, context.Context) {
+func Contexts(ctx context.Context, log logr.Logger, delay, grace time.Duration) (warn, soft, hard context.Context) {
 	if log == nil {
 		log = logr.Discard()
 	}
 	sigCh := make(chan os.Signal, 3)
 	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
 	hardCtx, hardCancel := context.WithCancel(ctx)
 	softCtx, softCancel := context.WithCancel(hardCtx)
+	warnCtx, warnCancel := context.WithCancel(softCtx)
 	go func() {
 		defer signal.Stop(sigCh)
 		defer hardCancel()
@@ -47,6 +52,7 @@ func Contexts(ctx context.Context, log logr.Logger, delay, grace time.Duration)
 		}
 		if delay > 0 {
 			log.Info("Shutdown starting after delay period", "duration", delay)
+			warnCancel()
 			select {
 			case <-time.After(delay):
 				log.Info("Shutdown delay period ended")
@@ -69,5 +75,5 @@ func Contexts(ctx context.Context, log logr.Logger, delay, grace time.Duration)
 			}
 		}
 	}()
-	return softCtx, hardCtx
+	return warnCtx, softCtx, hardCtx
 }

server.go

@@ -71,20 +71,31 @@ type ServerConfig struct {
 	// Logger optionally adds the ability to log messages, both errors and not.
 	Logger logr.Logger
 
-	initOnce sync.Once
+	initOnce   sync.Once
+	warnCtx    context.Context
+	warnCancel context.CancelFunc
 }
 
 func (cfg *ServerConfig) init() {
 	cfg.initOnce.Do(func() {
 		if cfg.Logger == nil {
 			cfg.Logger = logr.Discard()
 		}
+		cfg.warnCtx, cfg.warnCancel = context.WithCancel(context.TODO())
 	})
 }
 
+// ShuttingDown returns a channel that is closed when the server encounters
+// an error or it receives its first signal to begin shutting down.
+func (cfg *ServerConfig) ShuttingDown() <-chan struct{} {
+	cfg.init()
+	return cfg.warnCtx.Done()
+}
+
 // ListenAndServe listens on the given address and calls Serve.
 func (cfg *ServerConfig) ListenAndServe(ctx context.Context, addr string) error {
 	cfg.init()
+	defer cfg.warnCancel()
 
 	lis, err := net.Listen("tcp", addr)
 	if err != nil {
@@ -101,9 +112,10 @@ func (cfg *ServerConfig) ListenAndServe(ctx context.Context, addr string) error
 // cancelled a hard shutdown is initiated.
 func (cfg *ServerConfig) Serve(ctx context.Context, lis net.Listener) error {
 	cfg.init()
+	defer cfg.warnCancel()
 
 	g, ctx := errgroup.WithContext(ctx)
-	softCtx, hardCtx := Contexts(ctx, cfg.Logger, cfg.ShutdownDelay, cfg.ShutdownGrace)
+	warnCtx, softCtx, hardCtx := Contexts(ctx, cfg.Logger, cfg.ShutdownDelay, cfg.ShutdownGrace)
 	// Serve.
 	g.Go(func() error {
 		defer lis.Close()
@@ -113,6 +125,12 @@ func (cfg *ServerConfig) Serve(ctx context.Context, lis net.Listener) error {
 		}
 		return nil
 	})
+	// Signal shutting down.
+	g.Go(func() error {
+		<-warnCtx.Done()
+		cfg.warnCancel()
+		return nil
+	})
 	// Watch for shutdown signal.
 	g.Go(func() error {
 		<-softCtx.Done()
@@ -150,20 +168,31 @@ type DualServerConfig struct {
 	// Logger optionally adds the ability to log messages, both errors and not.
 	Logger logr.Logger
 
-	initOnce sync.Once
+	initOnce   sync.Once
+	warnCtx    context.Context
+	warnCancel context.CancelFunc
 }
 
 func (cfg *DualServerConfig) init() {
 	cfg.initOnce.Do(func() {
 		if cfg.Logger == nil {
 			cfg.Logger = logr.Discard()
 		}
+		cfg.warnCtx, cfg.warnCancel = context.WithCancel(context.TODO())
 	})
 }
 
+// ShuttingDown returns a channel that is closed when the server encounters
+// an error or it receives its first signal to begin shutting down.
+func (cfg *DualServerConfig) ShuttingDown() <-chan struct{} {
+	cfg.init()
+	return cfg.warnCtx.Done()
+}
+
 // ListenAndServe listens on the given addresses and calls Serve.
 func (cfg *DualServerConfig) ListenAndServe(ctx context.Context, intAddr, extAddr string) error {
 	cfg.init()
+	defer cfg.warnCancel()
 
 	intLis, err := net.Listen("tcp", intAddr)
 	if err != nil {
@@ -188,12 +217,13 @@ func (cfg *DualServerConfig) ListenAndServe(ctx context.Context, intAddr, extAdd
 // a hard shutdown is initiated.
 func (cfg *DualServerConfig) Serve(ctx context.Context, intLis, extLis net.Listener) error {
 	cfg.init()
+	defer cfg.warnCancel()
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	g, ctx := errgroup.WithContext(ctx)
-	softCtx, hardCtx := Contexts(ctx, cfg.Logger, cfg.ShutdownDelay, cfg.ShutdownGrace)
+	warnCtx, softCtx, hardCtx := Contexts(ctx, cfg.Logger, cfg.ShutdownDelay, cfg.ShutdownGrace)
 	extShutdownDone := make(chan struct{})
 
 	// Serve internal.
@@ -214,6 +244,12 @@ func (cfg *DualServerConfig) Serve(ctx context.Context, intLis, extLis net.Liste
 		}
 		return nil
 	})
+	// Signal shutting down.
+	g.Go(func() error {
+		<-warnCtx.Done()
+		cfg.warnCancel()
+		return nil
+	})
 	// Watch for signal and shutdown external first.
 	g.Go(func() error {
 		defer close(extShutdownDone)