go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/graceful/background.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package graceful
     9  
    10  import (
    11  	"context"
    12  	"os"
    13  	"os/signal"
    14  )
    15  
    16  // OptBackgroundSignals sets the signals.
    17  func OptBackgroundSignals(signals ...os.Signal) BackgroundOption {
    18  	return func(bo *BackgroundOptions) { bo.Signals = signals }
    19  }
    20  
    21  // OptBackgroundLog sets the logger.
    22  func OptBackgroundLog(log Logger) BackgroundOption {
    23  	return func(bo *BackgroundOptions) { bo.Log = log }
    24  }
    25  
    26  // OptBackgroundSkipStopOnSignal sets if we should stop the signal channel on stop.
    27  func OptBackgroundContext(ctx context.Context) BackgroundOption {
    28  	return func(bo *BackgroundOptions) { bo.Context = ctx }
    29  }
    30  
    31  // BackgroundOption mutates background options
    32  type BackgroundOption func(*BackgroundOptions)
    33  
    34  // BackgroundOptions are options for the background context.
    35  type BackgroundOptions struct {
    36  	// Context is a root context, if unset `context.Background()` is used.
    37  	Context context.Context
    38  	// Signals are the specific os signals to intercept.
    39  	Signals []os.Signal
    40  	// Log holds an reference to a graceful logger.
    41  	Log Logger
    42  }
    43  
    44  // Background yields a context that will signal `<-ctx.Done()` when
    45  // a signal is sent to the process (as specified in `DefaultShutdownSignals`).
    46  //
    47  // This context will cancel only (1) time.
    48  func Background(opts ...BackgroundOption) context.Context {
    49  	options := BackgroundOptions{
    50  		Context: context.Background(),
    51  		Signals: DefaultShutdownSignals,
    52  	}
    53  	for _, opt := range opts {
    54  		opt(&options)
    55  	}
    56  
    57  	ctx, cancel := context.WithCancel(options.Context)
    58  	shutdown := SignalNotify(options.Signals...)
    59  	go func() {
    60  		defer func() {
    61  			signal.Stop(shutdown)
    62  		}()
    63  		MaybeDebugf(options.Log, "graceful background; waiting for shutdown signal")
    64  		select {
    65  		case <-options.Context.Done():
    66  			return
    67  		case <-shutdown:
    68  			MaybeDebugf(options.Log, "graceful background; shutdown signal received, canceling context")
    69  			cancel()
    70  			return
    71  		}
    72  
    73  	}()
    74  	return ctx
    75  }