github.com/blend/go-sdk@v1.20220411.3/graceful/background.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     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 OptBackgroundSkipStopOnSignal(skipStopOnSignal bool) BackgroundOption {
    28  	return func(bo *BackgroundOptions) { bo.SkipStopOnSignal = skipStopOnSignal }
    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  	Signals          []os.Signal
    37  	Log              Logger
    38  	SkipStopOnSignal bool
    39  }
    40  
    41  // Background yields a context that will signal `<-ctx.Done()` when
    42  // a signal is sent to the process (as specified in `DefaultShutdownSignals`).
    43  func Background(opts ...BackgroundOption) context.Context {
    44  	options := BackgroundOptions{
    45  		Signals:          DefaultShutdownSignals,
    46  		SkipStopOnSignal: false,
    47  	}
    48  	for _, opt := range opts {
    49  		opt(&options)
    50  	}
    51  
    52  	ctx, cancel := context.WithCancel(context.Background())
    53  	shutdown := Notify(options.Signals...)
    54  
    55  	go func() {
    56  		MaybeDebugf(options.Log, "graceful background; waiting for shutdown signal")
    57  		<-shutdown
    58  
    59  		MaybeDebugf(options.Log, "graceful background; shutdown signal received, canceling context")
    60  		cancel()
    61  
    62  		if !options.SkipStopOnSignal {
    63  			MaybeDebugf(options.Log, "graceful background; stopping shutdown signal channel")
    64  			signal.Stop(shutdown)
    65  		}
    66  	}()
    67  	return ctx
    68  }