github.com/blend/go-sdk@v1.20220411.3/graceful/shutdown_by_context.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  	"sync"
    13  
    14  	"github.com/blend/go-sdk/ex"
    15  )
    16  
    17  // ShutdownByContext gracefully stops a set hosted processes based on context cancellation.
    18  func ShutdownByContext(ctx context.Context, hosted ...Graceful) error {
    19  	shouldShutdown := make(chan struct{})
    20  	serverExited := make(chan struct{})
    21  
    22  	waitShutdownComplete := sync.WaitGroup{}
    23  	waitShutdownComplete.Add(len(hosted))
    24  
    25  	waitServerExited := sync.WaitGroup{}
    26  	waitServerExited.Add(len(hosted))
    27  
    28  	errors := make(chan error, 2*len(hosted))
    29  
    30  	for _, hostedInstance := range hosted {
    31  		// start the instance
    32  		go func(instance Graceful) {
    33  			defer func() {
    34  				_ = safely(func() { close(serverExited) }) // close the server exited channel, but do so safely
    35  				waitServerExited.Done()                    // signal the normal exit process is done
    36  			}()
    37  			if err := instance.Start(); err != nil {
    38  				errors <- err
    39  			}
    40  		}(hostedInstance)
    41  
    42  		// wait to stop the instance
    43  		go func(instance Graceful) {
    44  			defer waitShutdownComplete.Done()
    45  			<-shouldShutdown // tell the hosted process to stop "gracefully"
    46  			if err := instance.Stop(); err != nil {
    47  				errors <- err
    48  			}
    49  		}(hostedInstance)
    50  	}
    51  
    52  	select {
    53  	case <-ctx.Done(): // if we've issued a shutdown, wait for the server to exit
    54  		close(shouldShutdown)
    55  		waitShutdownComplete.Wait()
    56  		waitServerExited.Wait()
    57  	case <-serverExited: // if any of the servers exited on their own
    58  		close(shouldShutdown) // quit the signal listener
    59  		waitShutdownComplete.Wait()
    60  	}
    61  	if errorCount := len(errors); errorCount > 0 {
    62  		var err error
    63  		for x := 0; x < errorCount; x++ {
    64  			err = ex.Append(err, <-errors)
    65  		}
    66  		return err
    67  	}
    68  	return nil
    69  }