335 lines
8.3 KiB
Go
335 lines
8.3 KiB
Go
|
// Copyright 2012 The Go Authors. All rights reserved.
|
||
|
// Use of this source code is governed by a BSD-style
|
||
|
// license that can be found in the LICENSE file.
|
||
|
|
||
|
package signal
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"os"
|
||
|
"sync"
|
||
|
)
|
||
|
|
||
|
var handlers struct {
|
||
|
sync.Mutex
|
||
|
// Map a channel to the signals that should be sent to it.
|
||
|
m map[chan<- os.Signal]*handler
|
||
|
// Map a signal to the number of channels receiving it.
|
||
|
ref [numSig]int64
|
||
|
// Map channels to signals while the channel is being stopped.
|
||
|
// Not a map because entries live here only very briefly.
|
||
|
// We need a separate container because we need m to correspond to ref
|
||
|
// at all times, and we also need to keep track of the *handler
|
||
|
// value for a channel being stopped. See the Stop function.
|
||
|
stopping []stopping
|
||
|
}
|
||
|
|
||
|
type stopping struct {
|
||
|
c chan<- os.Signal
|
||
|
h *handler
|
||
|
}
|
||
|
|
||
|
type handler struct {
|
||
|
mask [(numSig + 31) / 32]uint32
|
||
|
}
|
||
|
|
||
|
func (h *handler) want(sig int) bool {
|
||
|
return (h.mask[sig/32]>>uint(sig&31))&1 != 0
|
||
|
}
|
||
|
|
||
|
func (h *handler) set(sig int) {
|
||
|
h.mask[sig/32] |= 1 << uint(sig&31)
|
||
|
}
|
||
|
|
||
|
func (h *handler) clear(sig int) {
|
||
|
h.mask[sig/32] &^= 1 << uint(sig&31)
|
||
|
}
|
||
|
|
||
|
// Stop relaying the signals, sigs, to any channels previously registered to
|
||
|
// receive them and either reset the signal handlers to their original values
|
||
|
// (action=disableSignal) or ignore the signals (action=ignoreSignal).
|
||
|
func cancel(sigs []os.Signal, action func(int)) {
|
||
|
handlers.Lock()
|
||
|
defer handlers.Unlock()
|
||
|
|
||
|
remove := func(n int) {
|
||
|
var zerohandler handler
|
||
|
|
||
|
for c, h := range handlers.m {
|
||
|
if h.want(n) {
|
||
|
handlers.ref[n]--
|
||
|
h.clear(n)
|
||
|
if h.mask == zerohandler.mask {
|
||
|
delete(handlers.m, c)
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
action(n)
|
||
|
}
|
||
|
|
||
|
if len(sigs) == 0 {
|
||
|
for n := 0; n < numSig; n++ {
|
||
|
remove(n)
|
||
|
}
|
||
|
} else {
|
||
|
for _, s := range sigs {
|
||
|
remove(signum(s))
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Ignore causes the provided signals to be ignored. If they are received by
|
||
|
// the program, nothing will happen. Ignore undoes the effect of any prior
|
||
|
// calls to Notify for the provided signals.
|
||
|
// If no signals are provided, all incoming signals will be ignored.
|
||
|
func Ignore(sig ...os.Signal) {
|
||
|
cancel(sig, ignoreSignal)
|
||
|
}
|
||
|
|
||
|
// Ignored reports whether sig is currently ignored.
|
||
|
func Ignored(sig os.Signal) bool {
|
||
|
sn := signum(sig)
|
||
|
return sn >= 0 && signalIgnored(sn)
|
||
|
}
|
||
|
|
||
|
var (
|
||
|
// watchSignalLoopOnce guards calling the conditionally
|
||
|
// initialized watchSignalLoop. If watchSignalLoop is non-nil,
|
||
|
// it will be run in a goroutine lazily once Notify is invoked.
|
||
|
// See Issue 21576.
|
||
|
watchSignalLoopOnce sync.Once
|
||
|
watchSignalLoop func()
|
||
|
)
|
||
|
|
||
|
// Notify causes package signal to relay incoming signals to c.
|
||
|
// If no signals are provided, all incoming signals will be relayed to c.
|
||
|
// Otherwise, just the provided signals will.
|
||
|
//
|
||
|
// Package signal will not block sending to c: the caller must ensure
|
||
|
// that c has sufficient buffer space to keep up with the expected
|
||
|
// signal rate. For a channel used for notification of just one signal value,
|
||
|
// a buffer of size 1 is sufficient.
|
||
|
//
|
||
|
// It is allowed to call Notify multiple times with the same channel:
|
||
|
// each call expands the set of signals sent to that channel.
|
||
|
// The only way to remove signals from the set is to call Stop.
|
||
|
//
|
||
|
// It is allowed to call Notify multiple times with different channels
|
||
|
// and the same signals: each channel receives copies of incoming
|
||
|
// signals independently.
|
||
|
func Notify(c chan<- os.Signal, sig ...os.Signal) {
|
||
|
if c == nil {
|
||
|
panic("os/signal: Notify using nil channel")
|
||
|
}
|
||
|
|
||
|
handlers.Lock()
|
||
|
defer handlers.Unlock()
|
||
|
|
||
|
h := handlers.m[c]
|
||
|
if h == nil {
|
||
|
if handlers.m == nil {
|
||
|
handlers.m = make(map[chan<- os.Signal]*handler)
|
||
|
}
|
||
|
h = new(handler)
|
||
|
handlers.m[c] = h
|
||
|
}
|
||
|
|
||
|
add := func(n int) {
|
||
|
if n < 0 {
|
||
|
return
|
||
|
}
|
||
|
if !h.want(n) {
|
||
|
h.set(n)
|
||
|
if handlers.ref[n] == 0 {
|
||
|
enableSignal(n)
|
||
|
|
||
|
// The runtime requires that we enable a
|
||
|
// signal before starting the watcher.
|
||
|
watchSignalLoopOnce.Do(func() {
|
||
|
if watchSignalLoop != nil {
|
||
|
go watchSignalLoop()
|
||
|
}
|
||
|
})
|
||
|
}
|
||
|
handlers.ref[n]++
|
||
|
}
|
||
|
}
|
||
|
|
||
|
if len(sig) == 0 {
|
||
|
for n := 0; n < numSig; n++ {
|
||
|
add(n)
|
||
|
}
|
||
|
} else {
|
||
|
for _, s := range sig {
|
||
|
add(signum(s))
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Reset undoes the effect of any prior calls to Notify for the provided
|
||
|
// signals.
|
||
|
// If no signals are provided, all signal handlers will be reset.
|
||
|
func Reset(sig ...os.Signal) {
|
||
|
cancel(sig, disableSignal)
|
||
|
}
|
||
|
|
||
|
// Stop causes package signal to stop relaying incoming signals to c.
|
||
|
// It undoes the effect of all prior calls to Notify using c.
|
||
|
// When Stop returns, it is guaranteed that c will receive no more signals.
|
||
|
func Stop(c chan<- os.Signal) {
|
||
|
handlers.Lock()
|
||
|
|
||
|
h := handlers.m[c]
|
||
|
if h == nil {
|
||
|
handlers.Unlock()
|
||
|
return
|
||
|
}
|
||
|
delete(handlers.m, c)
|
||
|
|
||
|
for n := 0; n < numSig; n++ {
|
||
|
if h.want(n) {
|
||
|
handlers.ref[n]--
|
||
|
if handlers.ref[n] == 0 {
|
||
|
disableSignal(n)
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Signals will no longer be delivered to the channel.
|
||
|
// We want to avoid a race for a signal such as SIGINT:
|
||
|
// it should be either delivered to the channel,
|
||
|
// or the program should take the default action (that is, exit).
|
||
|
// To avoid the possibility that the signal is delivered,
|
||
|
// and the signal handler invoked, and then Stop deregisters
|
||
|
// the channel before the process function below has a chance
|
||
|
// to send it on the channel, put the channel on a list of
|
||
|
// channels being stopped and wait for signal delivery to
|
||
|
// quiesce before fully removing it.
|
||
|
|
||
|
handlers.stopping = append(handlers.stopping, stopping{c, h})
|
||
|
|
||
|
handlers.Unlock()
|
||
|
|
||
|
signalWaitUntilIdle()
|
||
|
|
||
|
handlers.Lock()
|
||
|
|
||
|
for i, s := range handlers.stopping {
|
||
|
if s.c == c {
|
||
|
handlers.stopping = append(handlers.stopping[:i], handlers.stopping[i+1:]...)
|
||
|
break
|
||
|
}
|
||
|
}
|
||
|
|
||
|
handlers.Unlock()
|
||
|
}
|
||
|
|
||
|
// Wait until there are no more signals waiting to be delivered.
|
||
|
// Defined by the runtime package.
|
||
|
func signalWaitUntilIdle()
|
||
|
|
||
|
func process(sig os.Signal) {
|
||
|
n := signum(sig)
|
||
|
if n < 0 {
|
||
|
return
|
||
|
}
|
||
|
|
||
|
handlers.Lock()
|
||
|
defer handlers.Unlock()
|
||
|
|
||
|
for c, h := range handlers.m {
|
||
|
if h.want(n) {
|
||
|
// send but do not block for it
|
||
|
select {
|
||
|
case c <- sig:
|
||
|
default:
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Avoid the race mentioned in Stop.
|
||
|
for _, d := range handlers.stopping {
|
||
|
if d.h.want(n) {
|
||
|
select {
|
||
|
case d.c <- sig:
|
||
|
default:
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// NotifyContext returns a copy of the parent context that is marked done
|
||
|
// (its Done channel is closed) when one of the listed signals arrives,
|
||
|
// when the returned stop function is called, or when the parent context's
|
||
|
// Done channel is closed, whichever happens first.
|
||
|
//
|
||
|
// The stop function unregisters the signal behavior, which, like signal.Reset,
|
||
|
// may restore the default behavior for a given signal. For example, the default
|
||
|
// behavior of a Go program receiving os.Interrupt is to exit. Calling
|
||
|
// NotifyContext(parent, os.Interrupt) will change the behavior to cancel
|
||
|
// the returned context. Future interrupts received will not trigger the default
|
||
|
// (exit) behavior until the returned stop function is called.
|
||
|
//
|
||
|
// The stop function releases resources associated with it, so code should
|
||
|
// call stop as soon as the operations running in this Context complete and
|
||
|
// signals no longer need to be diverted to the context.
|
||
|
func NotifyContext(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc) {
|
||
|
ctx, cancel := context.WithCancel(parent)
|
||
|
c := &signalCtx{
|
||
|
Context: ctx,
|
||
|
cancel: cancel,
|
||
|
signals: signals,
|
||
|
}
|
||
|
c.ch = make(chan os.Signal, 1)
|
||
|
Notify(c.ch, c.signals...)
|
||
|
if ctx.Err() == nil {
|
||
|
go func() {
|
||
|
select {
|
||
|
case <-c.ch:
|
||
|
c.cancel()
|
||
|
case <-c.Done():
|
||
|
}
|
||
|
}()
|
||
|
}
|
||
|
return c, c.stop
|
||
|
}
|
||
|
|
||
|
type signalCtx struct {
|
||
|
context.Context
|
||
|
|
||
|
cancel context.CancelFunc
|
||
|
signals []os.Signal
|
||
|
ch chan os.Signal
|
||
|
}
|
||
|
|
||
|
func (c *signalCtx) stop() {
|
||
|
c.cancel()
|
||
|
Stop(c.ch)
|
||
|
}
|
||
|
|
||
|
type stringer interface {
|
||
|
String() string
|
||
|
}
|
||
|
|
||
|
func (c *signalCtx) String() string {
|
||
|
var buf []byte
|
||
|
// We know that the type of c.Context is context.cancelCtx, and we know that the
|
||
|
// String method of cancelCtx returns a string that ends with ".WithCancel".
|
||
|
name := c.Context.(stringer).String()
|
||
|
name = name[:len(name)-len(".WithCancel")]
|
||
|
buf = append(buf, "signal.NotifyContext("+name...)
|
||
|
if len(c.signals) != 0 {
|
||
|
buf = append(buf, ", ["...)
|
||
|
for i, s := range c.signals {
|
||
|
buf = append(buf, s.String()...)
|
||
|
if i != len(c.signals)-1 {
|
||
|
buf = append(buf, ' ')
|
||
|
}
|
||
|
}
|
||
|
buf = append(buf, ']')
|
||
|
}
|
||
|
buf = append(buf, ')')
|
||
|
return string(buf)
|
||
|
}
|