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

     1  /*
     2  
     3  Copyright (c) 2024 - 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  	"errors"
    13  	"os"
    14  	"sync"
    15  
    16  	"go.charczuk.com/sdk/errutil"
    17  )
    18  
    19  // Graceful is the main entrypoint for hosting graceful processes.
    20  type Graceful struct {
    21  	Hosted          []Service
    22  	ShutdownSignals []os.Signal
    23  	RestartSignals  []os.Signal
    24  	Log             Logger
    25  }
    26  
    27  // StartForShutdown starts, and prepares to gracefully stop, a set hosted
    28  // processes based on a provided context's cancellation.
    29  //
    30  // The context is used to stop the goroutines this function spawns,
    31  // as well as call `Stop(...)` on the hosted processes when it cancels.
    32  func (g Graceful) StartForShutdown(ctx context.Context) error {
    33  	shouldShutdown := make(chan struct{})
    34  	shouldRestart := make(chan struct{})
    35  	serverExited := make(chan struct{})
    36  
    37  	if len(g.ShutdownSignals) > 0 {
    38  		notifyShutdown := SignalNotify(g.ShutdownSignals...)
    39  		go func() {
    40  			MaybeDebugf(g.Log, "graceful background; waiting for shutdown signal")
    41  			select {
    42  			case <-ctx.Done():
    43  				return
    44  			case <-notifyShutdown:
    45  				MaybeDebugf(g.Log, "graceful background; shutdown signal received, canceling context")
    46  				_ = safelyClose(shouldShutdown)
    47  			}
    48  		}()
    49  	}
    50  
    51  	if len(g.RestartSignals) > 0 {
    52  		restart := SignalNotify(g.RestartSignals...)
    53  		go func() {
    54  			for {
    55  				MaybeDebugf(g.Log, "graceful background; waiting for restart signal")
    56  				select {
    57  				case <-ctx.Done():
    58  					return
    59  				case <-restart:
    60  					// NOTE(wc): we don't close here because we may need to do this more than once!
    61  					shouldRestart <- struct{}{}
    62  					MaybeDebugf(g.Log, "graceful background; shutdown signal received, canceling context")
    63  				}
    64  			}
    65  		}()
    66  	}
    67  
    68  	var waitShutdownComplete sync.WaitGroup
    69  	waitShutdownComplete.Add(len(g.Hosted))
    70  
    71  	var waitServerExited sync.WaitGroup
    72  	waitServerExited.Add(len(g.Hosted))
    73  
    74  	hostedErrors := make(chan error, 2*len(g.Hosted))
    75  
    76  	for _, hostedInstance := range g.Hosted {
    77  		// start the instance
    78  		go func(instance Service) {
    79  			defer func() {
    80  				_ = safelyClose(serverExited)
    81  				waitServerExited.Done() // signal the normal exit process is done
    82  			}()
    83  			if err := instance.Start(ctx); err != nil {
    84  				hostedErrors <- err
    85  			}
    86  		}(hostedInstance)
    87  
    88  		// wait to restart the instance
    89  		go func(instance Service) {
    90  			<-shouldRestart
    91  			if err := instance.Restart(ctx); err != nil {
    92  				hostedErrors <- err
    93  			}
    94  		}(hostedInstance)
    95  
    96  		// wait to stop the instance
    97  		go func(instance Service) {
    98  			defer waitShutdownComplete.Done()
    99  			<-shouldShutdown // the hosted process has been told to stop "gracefully"
   100  			if err := instance.Stop(ctx); err != nil && !errors.Is(err, context.Canceled) {
   101  				hostedErrors <- err
   102  			}
   103  		}(hostedInstance)
   104  	}
   105  
   106  	select {
   107  	case <-ctx.Done(): // if we've issued a shutdown, wait for the server to exit
   108  		_ = safelyClose(shouldShutdown)
   109  		waitShutdownComplete.Wait()
   110  		waitServerExited.Wait()
   111  	case <-serverExited:
   112  		// if any of the servers exited on their
   113  		// own, we should crash the whole party
   114  		_ = safelyClose(shouldShutdown)
   115  		waitShutdownComplete.Wait()
   116  	}
   117  	if errorCount := len(hostedErrors); errorCount > 0 {
   118  		var err error
   119  		for x := 0; x < errorCount; x++ {
   120  			err = errutil.AppendFlat(err, <-hostedErrors)
   121  		}
   122  		return err
   123  	}
   124  	return nil
   125  }
   126  
   127  func safelyClose(c chan struct{}) (err error) {
   128  	defer func() {
   129  		if r := recover(); r != nil {
   130  			err = errutil.New(r)
   131  		}
   132  	}()
   133  	close(c)
   134  	return
   135  }
   136  
   137  func safely(action func()) (err error) {
   138  	defer func() {
   139  		if r := recover(); r != nil {
   140  			err = errutil.New(r)
   141  		}
   142  	}()
   143  	action()
   144  	return
   145  }