Add common Go App shutdown logic via Context

CHANGELOG.md

@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [5.24.0] - 2024-01-21
+### Added
+- `appext` package for application level helpers. Specifically added setting up os signal trapping and cancellation of context.Context.
+
 ## [5.23.0] - 2024-01-14
 ### Added
 - `And` and `AndThen` functions to `Option` & `Result` types.
@@ -87,6 +91,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `timext.NanoTime` for fast low level monotonic time with nanosecond precision.
 
 [Unreleased]: https://github.com/go-playground/pkg/compare/v5.23.0...HEAD
+[5.24.0]: https://github.com/go-playground/pkg/compare/v5.23.0..v5.24.0
 [5.23.0]: https://github.com/go-playground/pkg/compare/v5.22.0..v5.23.0
 [5.22.0]: https://github.com/go-playground/pkg/compare/v5.21.3..v5.22.0
 [5.21.3]: https://github.com/go-playground/pkg/compare/v5.21.2..v5.21.3

README.md

@@ -1,6 +1,6 @@
 # pkg
 
-![Project status](https://img.shields.io/badge/version-5.22.0-green.svg)
+![Project status](https://img.shields.io/badge/version-5.24.0-green.svg)
 [![Lint & Test](https://github.com/go-playground/pkg/actions/workflows/go.yml/badge.svg)](https://github.com/go-playground/pkg/actions/workflows/go.yml)
 [![Coverage Status](https://coveralls.io/repos/github/go-playground/pkg/badge.svg?branch=master)](https://coveralls.io/github/go-playground/pkg?branch=master)
 [![GoDoc](https://godoc.org/github.com/go-playground/pkg?status.svg)](https://pkg.go.dev/mod/github.com/go-playground/pkg/v5)

app/context.go

@@ -0,0 +1,102 @@
+package appext
+
+import (
+	"context"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+type contextBuilder struct {
+	signals   []os.Signal
+	timeout   time.Duration
+	exitFn    func(int)
+	forceExit bool
+}
+
+// Context returns a new context builder, with sane defaults, that can be overridden. Calling `Build()` finalizes
+// the new desired context and returns the configured `context.Context`.
+func Context() *contextBuilder {
+	return &contextBuilder{
+		signals: []os.Signal{
+			os.Interrupt,
+			syscall.SIGTERM,
+			syscall.SIGQUIT,
+		},
+		timeout:   30 * time.Second,
+		forceExit: true,
+		exitFn:    os.Exit,
+	}
+}
+
+// Signals sets the signals to listen for. Defaults to `os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT`.
+func (c *contextBuilder) Signals(signals ...os.Signal) *contextBuilder {
+	c.signals = signals
+	return c
+}
+
+// Timeout sets the timeout for graceful shutdown before forcing the issue exiting with exit code 1.
+// Defaults to 30 seconds.
+//
+// A timeout of <= 0, not recommended, disables the timeout and will wait forever for a seconds signal or application
+// shuts down.
+func (c *contextBuilder) Timeout(timeout time.Duration) *contextBuilder {
+	c.timeout = timeout
+	return c
+}
+
+// ForceExit sets whether to force terminate ungracefully upon receipt of a second signal. Defaults to true.
+func (c *contextBuilder) ForceExit(forceExit bool) *contextBuilder {
+	c.forceExit = forceExit
+	return c
+}
+
+// ExitFn sets the exit function to use. Defaults to `os.Exit`.
+//
+// This is used in the unit tests but can be used to intercept the exit call and do something else as needed also.
+func (c *contextBuilder) ExitFn(exitFn func(int)) *contextBuilder {
+	c.exitFn = exitFn
+	return c
+}
+
+// Build finalizes the context builder and returns the configured `context.Context`.
+//
+// This will spawn another goroutine listening for the configured signals and will cancel the context when received with
+// the configured settings.
+func (c *contextBuilder) Build() context.Context {
+	var sig = make(chan os.Signal, 1)
+	signal.Notify(sig, c.signals...)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go listen(sig, cancel, c.exitFn, c.timeout, c.forceExit)
+
+	return ctx
+}
+
+func listen(sig <-chan os.Signal, cancel context.CancelFunc, exitFn func(int), timeout time.Duration, forceExit bool) {
+	s := <-sig
+	cancel()
+	log.Printf("received shutdown signal %q\n", s)
+
+	if timeout > 0 {
+		select {
+		case s := <-sig:
+			if forceExit {
+				log.Printf("received second shutdown signal %q, forcing exit\n", s)
+				exitFn(1)
+			}
+		case <-time.After(timeout):
+			log.Printf("timeout of %s reached, forcing exit\n", timeout)
+			exitFn(1)
+		}
+	} else {
+		s = <-sig
+		if forceExit {
+			log.Printf("received second shutdown signal %q, forcing exit\n", s)
+			exitFn(1)
+		}
+	}
+}

app/context_test.go

@@ -0,0 +1,83 @@
+package appext
+
+import (
+	"context"
+	. "github.com/go-playground/assert/v2"
+	"os"
+	"os/signal"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestForceExitWithNoTimeout(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	exitFn := func(code int) {
+		defer wg.Done()
+		Equal(t, 1, code)
+	}
+
+	c := Context().Timeout(0).ExitFn(exitFn)
+
+	// copy of Build for testing
+	var sig = make(chan os.Signal, 1)
+	signal.Notify(sig, c.signals...)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go listen(sig, cancel, c.exitFn, c.timeout, c.forceExit)
+
+	sig <- os.Interrupt
+	sig <- os.Interrupt
+	wg.Wait()
+	Equal(t, context.Canceled, ctx.Err())
+}
+
+func TestForceExitWithTimeout(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	exitFn := func(code int) {
+		defer wg.Done()
+		Equal(t, 1, code)
+	}
+
+	c := Context().Timeout(time.Hour).ExitFn(exitFn)
+
+	// copy of Build for testing
+	var sig = make(chan os.Signal, 1)
+	signal.Notify(sig, c.signals...)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go listen(sig, cancel, c.exitFn, c.timeout, c.forceExit)
+
+	sig <- os.Interrupt
+	sig <- os.Interrupt
+	wg.Wait()
+	Equal(t, context.Canceled, ctx.Err())
+}
+
+func TestTimeoutWithNoForceExit(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	exitFn := func(code int) {
+		defer wg.Done()
+		Equal(t, 1, code)
+	}
+
+	c := Context().Timeout(time.Millisecond * 200).ForceExit(false).ExitFn(exitFn)
+
+	// copy of Build for testing
+	var sig = make(chan os.Signal, 1)
+	signal.Notify(sig, c.signals...)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go listen(sig, cancel, c.exitFn, c.timeout, c.forceExit)
+
+	// only sending one, timeout must be reached for test to finish
+	sig <- os.Interrupt
+	wg.Wait()
+	Equal(t, context.Canceled, ctx.Err())
+}