github.com/blend/go-sdk@v1.20220411.3/graceful/shutdown_by_signal.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  	"os/signal"
    12  	"sync"
    13  
    14  	"github.com/blend/go-sdk/ex"
    15  )
    16  
    17  // ShutdownBySignal gracefully stops a set hosted processes based on a set of variadic options.
    18  // A "Graceful" processes *must* block on start.
    19  // Fatal errors will be returned, that is, errors that are returned by either .Start() or .Stop().
    20  // Panics are not caught by graceful, and it is assumed that your .Start() or .Stop methods will catch relevant panics.
    21  func ShutdownBySignal(hosted []Graceful, opts ...ShutdownOption) error {
    22  	var options ShutdownOptions
    23  	for _, opt := range opts {
    24  		opt(&options)
    25  	}
    26  
    27  	shouldShutdown := make(chan struct{})
    28  	serverExited := make(chan struct{})
    29  
    30  	waitShutdownComplete := sync.WaitGroup{}
    31  	waitShutdownComplete.Add(len(hosted))
    32  
    33  	waitServerExited := sync.WaitGroup{}
    34  	waitServerExited.Add(len(hosted))
    35  
    36  	errors := make(chan error, 2*len(hosted))
    37  
    38  	for _, hostedInstance := range hosted {
    39  		// start the instance
    40  		go func(instance Graceful) {
    41  			defer func() {
    42  				_ = safely(func() { close(serverExited) }) // close the server exited channel, but do so safely
    43  				waitServerExited.Done()                    // signal the normal exit process is done
    44  			}()
    45  			if err := instance.Start(); err != nil {
    46  				errors <- err
    47  			}
    48  		}(hostedInstance)
    49  
    50  		// wait to stop the instance
    51  		go func(instance Graceful) {
    52  			defer waitShutdownComplete.Done()
    53  			<-shouldShutdown // tell the hosted process to stop "gracefully"
    54  			if err := instance.Stop(); err != nil {
    55  				errors <- err
    56  			}
    57  		}(hostedInstance)
    58  	}
    59  
    60  	select {
    61  	case <-options.ShutdownSignal: // if we've issued a shutdown, wait for the server to exit
    62  		signal.Stop(options.ShutdownSignal) // unhook the process signal redirects, the next ^c will crash the process etc.
    63  		close(shouldShutdown)
    64  		waitShutdownComplete.Wait()
    65  		waitServerExited.Wait()
    66  
    67  	case <-serverExited: // if any of the servers exited on their own
    68  		close(shouldShutdown) // quit the signal listener
    69  		waitShutdownComplete.Wait()
    70  	}
    71  	if len(errors) > 0 {
    72  		return <-errors
    73  	}
    74  	return nil
    75  }
    76  
    77  func safely(action func()) (err error) {
    78  	defer func() {
    79  		if r := recover(); r != nil {
    80  			err = ex.New(r)
    81  		}
    82  	}()
    83  	action()
    84  	return
    85  }