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