github.com/m3db/m3@v1.5.0/src/x/os/interrupt.go (about)

     1  // Copyright (c) 2019 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package xos
    22  
    23  import (
    24  	"fmt"
    25  	"os"
    26  	"os/signal"
    27  	"sync"
    28  	"syscall"
    29  
    30  	"go.uber.org/zap"
    31  )
    32  
    33  // InterruptOptions are options to use when waiting for an interrupt.
    34  type InterruptOptions struct {
    35  	// InterruptChannel is an existing interrupt channel, if none
    36  	// specified one will be created.
    37  	InterruptCh <-chan error
    38  
    39  	// InterruptedChannel is a channel that will be closed once an
    40  	// interrupt has been seen. Use this to pass to goroutines who
    41  	// want to be notified about interruptions that you don't want
    42  	// consuming from the main interrupt channel.
    43  	InterruptedCh chan struct{}
    44  }
    45  
    46  // InterruptError is an error representing an interrupt.
    47  type InterruptError struct {
    48  	interrupt string
    49  }
    50  
    51  // ErrInterrupted is an error indicating that the interrupted channel was closed,
    52  // meaning an interrupt was received on the main interrupt channel.
    53  var ErrInterrupted = NewInterruptError("interrupted")
    54  
    55  // NewInterruptOptions creates InterruptOptions with sane defaults.
    56  func NewInterruptOptions() InterruptOptions {
    57  	return InterruptOptions{
    58  		InterruptCh:   NewInterruptChannel(1),
    59  		InterruptedCh: make(chan struct{}),
    60  	}
    61  }
    62  
    63  // NewInterruptError creates a new InterruptError.
    64  func NewInterruptError(interrupt string) error {
    65  	return &InterruptError{interrupt: interrupt}
    66  }
    67  
    68  func (i *InterruptError) Error() string {
    69  	return i.interrupt
    70  }
    71  
    72  // WatchForInterrupt watches for interrupts in a non-blocking fashion and closes
    73  // the interrupted channel when an interrupt is found. Use this method to
    74  // watch for interrupts while the caller continues to execute (e.g. during
    75  // server startup). To ensure child goroutines get properly closed, pass them
    76  // the interrupted channel. If the interrupted channel is closed, then the
    77  // goroutine knows to stop its work. This method returns a function that
    78  // can be used to stop the watch.
    79  func WatchForInterrupt(logger *zap.Logger, opts InterruptOptions) func() {
    80  	interruptCh := opts.InterruptCh
    81  	closed := make(chan struct{})
    82  	go func() {
    83  		select {
    84  		case err := <-interruptCh:
    85  			logger.Warn("interrupt", zap.Error(err))
    86  			close(opts.InterruptedCh)
    87  		case <-closed:
    88  			logger.Info("interrupt watch stopped")
    89  			return
    90  		}
    91  	}()
    92  
    93  	var doOnce sync.Once
    94  	return func() {
    95  		doOnce.Do(func() {
    96  			close(closed)
    97  		})
    98  	}
    99  }
   100  
   101  // WaitForInterrupt will wait for an interrupt to occur and return when done.
   102  func WaitForInterrupt(logger *zap.Logger, opts InterruptOptions) {
   103  	// Handle interrupts.
   104  	interruptCh := opts.InterruptCh
   105  	if interruptCh == nil {
   106  		// Need to catch our own interrupts.
   107  		interruptCh = NewInterruptChannel(1)
   108  		logger.Info("registered new interrupt handler")
   109  	} else {
   110  		logger.Info("using registered interrupt handler")
   111  	}
   112  
   113  	logger.Warn("interrupt", zap.Error(<-interruptCh))
   114  
   115  	if opts.InterruptedCh != nil {
   116  		close(opts.InterruptedCh)
   117  	}
   118  }
   119  
   120  // NewInterruptChannel will return an interrupt channel useful with multiple
   121  // listeners.
   122  func NewInterruptChannel(numListeners int) <-chan error {
   123  	interruptCh := make(chan error, numListeners)
   124  	go func() {
   125  		err := NewInterruptError(fmt.Sprintf("%v", <-interrupt()))
   126  		for i := 0; i < numListeners; i++ {
   127  			interruptCh <- err
   128  		}
   129  	}()
   130  	return interruptCh
   131  }
   132  
   133  func interrupt() <-chan os.Signal {
   134  	c := make(chan os.Signal)
   135  	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
   136  	return c
   137  }