github.com/go-playground/pkg/v5@v5.29.1/app/context.go (about)

     1  package appext
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"os"
     7  	"os/signal"
     8  	"syscall"
     9  	"time"
    10  )
    11  
    12  type contextBuilder struct {
    13  	signals   []os.Signal
    14  	timeout   time.Duration
    15  	exitFn    func(int)
    16  	forceExit bool
    17  }
    18  
    19  // Context returns a new context builder, with sane defaults, that can be overridden. Calling `Build()` finalizes
    20  // the new desired context and returns the configured `context.Context`.
    21  func Context() *contextBuilder {
    22  	return &contextBuilder{
    23  		signals: []os.Signal{
    24  			os.Interrupt,
    25  			syscall.SIGTERM,
    26  			syscall.SIGQUIT,
    27  		},
    28  		timeout:   30 * time.Second,
    29  		forceExit: true,
    30  		exitFn:    os.Exit,
    31  	}
    32  }
    33  
    34  // Signals sets the signals to listen for. Defaults to `os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT`.
    35  func (c *contextBuilder) Signals(signals ...os.Signal) *contextBuilder {
    36  	c.signals = signals
    37  	return c
    38  }
    39  
    40  // Timeout sets the timeout for graceful shutdown before forcing the issue exiting with exit code 1.
    41  // Defaults to 30 seconds.
    42  //
    43  // A timeout of <= 0, not recommended, disables the timeout and will wait forever for a seconds signal or application
    44  // shuts down.
    45  func (c *contextBuilder) Timeout(timeout time.Duration) *contextBuilder {
    46  	c.timeout = timeout
    47  	return c
    48  }
    49  
    50  // ForceExit sets whether to force terminate ungracefully upon receipt of a second signal. Defaults to true.
    51  func (c *contextBuilder) ForceExit(forceExit bool) *contextBuilder {
    52  	c.forceExit = forceExit
    53  	return c
    54  }
    55  
    56  // ExitFn sets the exit function to use. Defaults to `os.Exit`.
    57  //
    58  // This is used in the unit tests but can be used to intercept the exit call and do something else as needed also.
    59  func (c *contextBuilder) ExitFn(exitFn func(int)) *contextBuilder {
    60  	c.exitFn = exitFn
    61  	return c
    62  }
    63  
    64  // Build finalizes the context builder and returns the configured `context.Context`.
    65  //
    66  // This will spawn another goroutine listening for the configured signals and will cancel the context when received with
    67  // the configured settings.
    68  func (c *contextBuilder) Build() context.Context {
    69  	var sig = make(chan os.Signal, 1)
    70  	signal.Notify(sig, c.signals...)
    71  
    72  	ctx, cancel := context.WithCancel(context.Background())
    73  
    74  	go listen(sig, cancel, c.exitFn, c.timeout, c.forceExit)
    75  
    76  	return ctx
    77  }
    78  
    79  func listen(sig <-chan os.Signal, cancel context.CancelFunc, exitFn func(int), timeout time.Duration, forceExit bool) {
    80  	s := <-sig
    81  	cancel()
    82  	log.Printf("received shutdown signal %q\n", s)
    83  
    84  	if timeout > 0 {
    85  		select {
    86  		case s := <-sig:
    87  			if forceExit {
    88  				log.Printf("received second shutdown signal %q, forcing exit\n", s)
    89  				exitFn(1)
    90  			}
    91  		case <-time.After(timeout):
    92  			log.Printf("timeout of %s reached, forcing exit\n", timeout)
    93  			exitFn(1)
    94  		}
    95  	} else {
    96  		s = <-sig
    97  		if forceExit {
    98  			log.Printf("received second shutdown signal %q, forcing exit\n", s)
    99  			exitFn(1)
   100  		}
   101  	}
   102  }