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