github.com/searKing/golang/go@v1.2.74/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  // 	ctx is wrapped WithShutdownSignal(ctx)
    28  // 	When signal is trapped, the shutdown handler will be invoked with a context.
    29  func Graceful(ctx context.Context, handlers ...Handler) (err error) {
    30  	if len(handlers) == 0 {
    31  		return nil
    32  	}
    33  	defer runtime.LogPanic.Recover()
    34  	var wg sync.WaitGroup
    35  
    36  	var mu sync.Mutex
    37  	var errs []error
    38  
    39  	for _, h := range handlers {
    40  		start := h.StartFunc
    41  		shutdown := h.ShutdownFunc
    42  		if start == nil {
    43  			start = func(ctx context.Context) error { return nil }
    44  		}
    45  		if shutdown == nil {
    46  			shutdown = func(ctx context.Context) error { return nil }
    47  		}
    48  
    49  		wg.Add(1)
    50  		go func() {
    51  			defer wg.Done()
    52  			select {
    53  			case <-ctx.Done():
    54  				err = shutdown(ctx)
    55  				if err != nil {
    56  					err = fmt.Errorf("graceful shutdown %s: %w", h.Name, err)
    57  					mu.Lock()
    58  					defer mu.Unlock()
    59  					errs = append(errs, err)
    60  				}
    61  			}
    62  		}()
    63  
    64  		// Start the server
    65  		if err := start(ctx); err != nil {
    66  			return fmt.Errorf("graceful start %s: %w", h.Name, err)
    67  		}
    68  	}
    69  
    70  	wg.Wait()
    71  	return errors.Multi(errs...)
    72  }