github.com/anycable/anycable-go@v1.5.1/utils/graceful_signals.go (about)

     1  package utils
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"os/signal"
     7  	"sync"
     8  	"syscall"
     9  	"time"
    10  )
    11  
    12  type signalHandler func(ctx context.Context) error
    13  
    14  // The struct implementing logic for graceful shutdown in response to OS signals.
    15  type GracefulSignals struct {
    16  	handlers              []signalHandler
    17  	forceTerminateHandler func()
    18  	timeout               time.Duration
    19  	executed              bool
    20  
    21  	ch chan os.Signal
    22  	mu sync.Mutex
    23  }
    24  
    25  // Create new GracefulSignals struct.
    26  func NewGracefulSignals(timeout time.Duration) *GracefulSignals {
    27  	return &GracefulSignals{
    28  		timeout:               timeout,
    29  		forceTerminateHandler: func() { os.Exit(0) },
    30  		handlers:              make([]signalHandler, 0),
    31  		ch:                    make(chan os.Signal, 1),
    32  	}
    33  }
    34  
    35  func (s *GracefulSignals) Handle(handler signalHandler) {
    36  	s.mu.Lock()
    37  	defer s.mu.Unlock()
    38  
    39  	s.handlers = append(s.handlers, handler)
    40  }
    41  
    42  func (s *GracefulSignals) HandleForceTerminate(handler func()) {
    43  	s.mu.Lock()
    44  	defer s.mu.Unlock()
    45  
    46  	s.forceTerminateHandler = handler
    47  }
    48  
    49  func (s *GracefulSignals) Listen() {
    50  	signal.Notify(s.ch, syscall.SIGINT, syscall.SIGTERM)
    51  	go s.listen()
    52  }
    53  
    54  func (s *GracefulSignals) listen() {
    55  	for { // nolint:gosimple
    56  		select {
    57  		case <-s.ch:
    58  			s.exec()
    59  		}
    60  	}
    61  }
    62  
    63  func (s *GracefulSignals) exec() {
    64  	s.mu.Lock()
    65  
    66  	if s.executed {
    67  		s.mu.Unlock()
    68  		return
    69  	}
    70  
    71  	shutdown := make(chan struct{})
    72  	s.executed = true
    73  
    74  	terminateCtx, terminateImmediately := context.WithCancel(context.Background())
    75  
    76  	timeoutCtx, cancelTimeout := context.WithTimeout(terminateCtx, s.timeout)
    77  	defer cancelTimeout()
    78  
    79  	go func() {
    80  		termSig := make(chan os.Signal, 1)
    81  		signal.Notify(termSig, syscall.SIGINT, syscall.SIGTERM)
    82  		<-termSig
    83  
    84  		terminateImmediately()
    85  
    86  		// Wait for handlers to interrupt
    87  		// NOTE: it's the responsibility of handlers to react on the context cancellation
    88  		<-shutdown
    89  
    90  		s.mu.Lock()
    91  		defer s.mu.Unlock()
    92  
    93  		if s.forceTerminateHandler != nil {
    94  			s.forceTerminateHandler()
    95  		}
    96  	}()
    97  
    98  	handlers := s.handlers[:] // nolint:gocritic
    99  	s.mu.Unlock()
   100  
   101  	for _, handler := range handlers {
   102  		handler(timeoutCtx) // nolint:errcheck
   103  	}
   104  
   105  	close(shutdown)
   106  }