signal.go

     1  // Copyright 2012 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package signal
     6  
     7  import (
     8  	"context"
     9  	"os"
    10  	"sync"
    11  )
    12  
    13  var handlers struct {
    14  	sync.Mutex
    15  	// Map a channel to the signals that should be sent to it.
    16  	m map[chan<- os.Signal]*handler
    17  	// Map a signal to the number of channels receiving it.
    18  	ref [numSig]int64
    19  	// Map channels to signals while the channel is being stopped.
    20  	// Not a map because entries live here only very briefly.
    21  	// We need a separate container because we need m to correspond to ref
    22  	// at all times, and we also need to keep track of the *handler
    23  	// value for a channel being stopped. See the Stop function.
    24  	stopping []stopping
    25  }
    26  
    27  type stopping struct {
    28  	c chan<- os.Signal
    29  	h *handler
    30  }
    31  
    32  type handler struct {
    33  	mask [(numSig + 31) / 32]uint32
    34  }
    35  
    36  func (h *handler) want(sig int) bool {
    37  	return (h.mask[sig/32]>>uint(sig&31))&1 != 0
    38  }
    39  
    40  func (h *handler) set(sig int) {
    41  	h.mask[sig/32] |= 1 << uint(sig&31)
    42  }
    43  
    44  func (h *handler) clear(sig int) {
    45  	h.mask[sig/32] &^= 1 << uint(sig&31)
    46  }
    47  
    48  // Stop relaying the signals, sigs, to any channels previously registered to
    49  // receive them and either reset the signal handlers to their original values
    50  // (action=disableSignal) or ignore the signals (action=ignoreSignal).
    51  func cancel(sigs []os.Signal, action func(int)) {
    52  	handlers.Lock()
    53  	defer handlers.Unlock()
    54  
    55  	remove := func(n int) {
    56  		var zerohandler handler
    57  
    58  		for c, h := range handlers.m {
    59  			if h.want(n) {
    60  				handlers.ref[n]--
    61  				h.clear(n)
    62  				if h.mask == zerohandler.mask {
    63  					delete(handlers.m, c)
    64  				}
    65  			}
    66  		}
    67  
    68  		action(n)
    69  	}
    70  
    71  	if len(sigs) == 0 {
    72  		for n := 0; n < numSig; n++ {
    73  			remove(n)
    74  		}
    75  	} else {
    76  		for _, s := range sigs {
    77  			remove(signum(s))
    78  		}
    79  	}
    80  }
    81  
    82  // Ignore causes the provided signals to be ignored. If they are received by
    83  // the program, nothing will happen. Ignore undoes the effect of any prior
    84  // calls to [Notify] for the provided signals.
    85  // If no signals are provided, all incoming signals will be ignored.
    86  func Ignore(sig ...os.Signal) {
    87  	cancel(sig, ignoreSignal)
    88  }
    89  
    90  // Ignored reports whether sig is currently ignored.
    91  func Ignored(sig os.Signal) bool {
    92  	sn := signum(sig)
    93  	return sn >= 0 && signalIgnored(sn)
    94  }
    95  
    96  var (
    97  	// watchSignalLoopOnce guards calling the conditionally
    98  	// initialized watchSignalLoop. If watchSignalLoop is non-nil,
    99  	// it will be run in a goroutine lazily once Notify is invoked.
   100  	// See Issue 21576.
   101  	watchSignalLoopOnce sync.Once
   102  	watchSignalLoop     func()
   103  )
   104  
   105  // Notify causes package signal to relay incoming signals to c.
   106  // If no signals are provided, all incoming signals will be relayed to c.
   107  // Otherwise, just the provided signals will.
   108  //
   109  // Package signal will not block sending to c: the caller must ensure
   110  // that c has sufficient buffer space to keep up with the expected
   111  // signal rate. For a channel used for notification of just one signal value,
   112  // a buffer of size 1 is sufficient.
   113  //
   114  // It is allowed to call Notify multiple times with the same channel:
   115  // each call expands the set of signals sent to that channel.
   116  // The only way to remove signals from the set is to call [Stop].
   117  //
   118  // It is allowed to call Notify multiple times with different channels
   119  // and the same signals: each channel receives copies of incoming
   120  // signals independently.
   121  func Notify(c chan<- os.Signal, sig ...os.Signal) {
   122  	if c == nil {
   123  		panic("os/signal: Notify using nil channel")
   124  	}
   125  
   126  	handlers.Lock()
   127  	defer handlers.Unlock()
   128  
   129  	h := handlers.m[c]
   130  	if h == nil {
   131  		if handlers.m == nil {
   132  			handlers.m = make(map[chan<- os.Signal]*handler)
   133  		}
   134  		h = new(handler)
   135  		handlers.m[c] = h
   136  	}
   137  
   138  	add := func(n int) {
   139  		if n < 0 {
   140  			return
   141  		}
   142  		if !h.want(n) {
   143  			h.set(n)
   144  			if handlers.ref[n] == 0 {
   145  				enableSignal(n)
   146  
   147  				// The runtime requires that we enable a
   148  				// signal before starting the watcher.
   149  				watchSignalLoopOnce.Do(func() {
   150  					if watchSignalLoop != nil {
   151  						go watchSignalLoop()
   152  					}
   153  				})
   154  			}
   155  			handlers.ref[n]++
   156  		}
   157  	}
   158  
   159  	if len(sig) == 0 {
   160  		for n := 0; n < numSig; n++ {
   161  			add(n)
   162  		}
   163  	} else {
   164  		for _, s := range sig {
   165  			add(signum(s))
   166  		}
   167  	}
   168  }
   169  
   170  // Reset undoes the effect of any prior calls to [Notify] for the provided
   171  // signals.
   172  // If no signals are provided, all signal handlers will be reset.
   173  func Reset(sig ...os.Signal) {
   174  	cancel(sig, disableSignal)
   175  }
   176  
   177  // Stop causes package signal to stop relaying incoming signals to c.
   178  // It undoes the effect of all prior calls to [Notify] using c.
   179  // When Stop returns, it is guaranteed that c will receive no more signals.
   180  func Stop(c chan<- os.Signal) {
   181  	handlers.Lock()
   182  
   183  	h := handlers.m[c]
   184  	if h == nil {
   185  		handlers.Unlock()
   186  		return
   187  	}
   188  	delete(handlers.m, c)
   189  
   190  	for n := 0; n < numSig; n++ {
   191  		if h.want(n) {
   192  			handlers.ref[n]--
   193  			if handlers.ref[n] == 0 {
   194  				disableSignal(n)
   195  			}
   196  		}
   197  	}
   198  
   199  	// Signals will no longer be delivered to the channel.
   200  	// We want to avoid a race for a signal such as SIGINT:
   201  	// it should be either delivered to the channel,
   202  	// or the program should take the default action (that is, exit).
   203  	// To avoid the possibility that the signal is delivered,
   204  	// and the signal handler invoked, and then Stop deregisters
   205  	// the channel before the process function below has a chance
   206  	// to send it on the channel, put the channel on a list of
   207  	// channels being stopped and wait for signal delivery to
   208  	// quiesce before fully removing it.
   209  
   210  	handlers.stopping = append(handlers.stopping, stopping{c, h})
   211  
   212  	handlers.Unlock()
   213  
   214  	signalWaitUntilIdle()
   215  
   216  	handlers.Lock()
   217  
   218  	for i, s := range handlers.stopping {
   219  		if s.c == c {
   220  			handlers.stopping = append(handlers.stopping[:i], handlers.stopping[i+1:]...)
   221  			break
   222  		}
   223  	}
   224  
   225  	handlers.Unlock()
   226  }
   227  
   228  // Wait until there are no more signals waiting to be delivered.
   229  // Defined by the runtime package.
   230  func signalWaitUntilIdle()
   231  
   232  func process(sig os.Signal) {
   233  	n := signum(sig)
   234  	if n < 0 {
   235  		return
   236  	}
   237  
   238  	handlers.Lock()
   239  	defer handlers.Unlock()
   240  
   241  	for c, h := range handlers.m {
   242  		if h.want(n) {
   243  			// send but do not block for it
   244  			select {
   245  			case c <- sig:
   246  			default:
   247  			}
   248  		}
   249  	}
   250  
   251  	// Avoid the race mentioned in Stop.
   252  	for _, d := range handlers.stopping {
   253  		if d.h.want(n) {
   254  			select {
   255  			case d.c <- sig:
   256  			default:
   257  			}
   258  		}
   259  	}
   260  }
   261  
   262  // NotifyContext returns a copy of the parent context that is marked done
   263  // (its Done channel is closed) when one of the listed signals arrives,
   264  // when the returned stop function is called, or when the parent context's
   265  // Done channel is closed, whichever happens first.
   266  //
   267  // The stop function unregisters the signal behavior, which, like [signal.Reset],
   268  // may restore the default behavior for a given signal. For example, the default
   269  // behavior of a Go program receiving [os.Interrupt] is to exit. Calling
   270  // NotifyContext(parent, os.Interrupt) will change the behavior to cancel
   271  // the returned context. Future interrupts received will not trigger the default
   272  // (exit) behavior until the returned stop function is called.
   273  //
   274  // The stop function releases resources associated with it, so code should
   275  // call stop as soon as the operations running in this Context complete and
   276  // signals no longer need to be diverted to the context.
   277  func NotifyContext(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc) {
   278  	ctx, cancel := context.WithCancel(parent)
   279  	c := &signalCtx{
   280  		Context: ctx,
   281  		cancel:  cancel,
   282  		signals: signals,
   283  	}
   284  	c.ch = make(chan os.Signal, 1)
   285  	Notify(c.ch, c.signals...)
   286  	if ctx.Err() == nil {
   287  		go func() {
   288  			select {
   289  			case <-c.ch:
   290  				c.cancel()
   291  			case <-c.Done():
   292  			}
   293  		}()
   294  	}
   295  	return c, c.stop
   296  }
   297  
   298  type signalCtx struct {
   299  	context.Context
   300  
   301  	cancel  context.CancelFunc
   302  	signals []os.Signal
   303  	ch      chan os.Signal
   304  }
   305  
   306  func (c *signalCtx) stop() {
   307  	c.cancel()
   308  	Stop(c.ch)
   309  }
   310  
   311  type stringer interface {
   312  	String() string
   313  }
   314  
   315  func (c *signalCtx) String() string {
   316  	var buf []byte
   317  	// We know that the type of c.Context is context.cancelCtx, and we know that the
   318  	// String method of cancelCtx returns a string that ends with ".WithCancel".
   319  	name := c.Context.(stringer).String()
   320  	name = name[:len(name)-len(".WithCancel")]
   321  	buf = append(buf, "signal.NotifyContext("+name...)
   322  	if len(c.signals) != 0 {
   323  		buf = append(buf, ", ["...)
   324  		for i, s := range c.signals {
   325  			buf = append(buf, s.String()...)
   326  			if i != len(c.signals)-1 {
   327  				buf = append(buf, ' ')
   328  			}
   329  		}
   330  		buf = append(buf, ']')
   331  	}
   332  	buf = append(buf, ')')
   333  	return string(buf)
   334  }
   335
View as plain text