github.com/searKing/golang/go@v1.2.117/x/graceful/graceful.go (about)

     1  package graceful
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"github.com/searKing/golang/go/errors"
     9  	"github.com/searKing/golang/go/runtime"
    10  )
    11  
    12  type Handler struct {
    13  	Name         string       // server name
    14  	StartFunc    StartFunc    // func to start the server
    15  	ShutdownFunc ShutdownFunc // func to shutdown the server
    16  }
    17  
    18  // StartFunc is the type of the function invoked by Graceful to start the server
    19  type StartFunc func(context.Context) error
    20  
    21  // ShutdownFunc is the type of the function invoked by Graceful to shutdown the server
    22  type ShutdownFunc func(context.Context) error
    23  
    24  // Graceful sets up graceful handling of context done, typically for an HTTP server.
    25  // When context is done, the shutdown handler will be invoked with a context.
    26  // Example:
    27  //
    28  //	ctx is wrapped WithShutdownSignal(ctx)
    29  //	When signal is trapped, the shutdown handler will be invoked with a context.
    30  func Graceful(ctx context.Context, handlers ...Handler) (err error) {
    31  	if len(handlers) == 0 {
    32  		return nil
    33  	}
    34  	defer runtime.LogPanic.Recover()
    35  	var wg sync.WaitGroup
    36  
    37  	var mu sync.Mutex
    38  	var errs []error
    39  
    40  	for _, h := range handlers {
    41  		start := h.StartFunc
    42  		shutdown := h.ShutdownFunc
    43  		if start == nil {
    44  			start = func(ctx context.Context) error { return nil }
    45  		}
    46  		if shutdown == nil {
    47  			shutdown = func(ctx context.Context) error { return nil }
    48  		}
    49  
    50  		wg.Add(1)
    51  		go func() {
    52  			defer wg.Done()
    53  			select {
    54  			case <-ctx.Done():
    55  				err = shutdown(ctx)
    56  				if err != nil {
    57  					err = fmt.Errorf("graceful shutdown %s: %w", h.Name, err)
    58  					mu.Lock()
    59  					defer mu.Unlock()
    60  					errs = append(errs, err)
    61  				}
    62  			}
    63  		}()
    64  
    65  		// Start the server
    66  		if err := start(ctx); err != nil {
    67  			return fmt.Errorf("graceful start %s: %w", h.Name, err)
    68  		}
    69  	}
    70  
    71  	wg.Wait()
    72  	return errors.Multi(errs...)
    73  }