/
signals.go
164 lines (148 loc) · 5.28 KB
/
signals.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
// Copyright 2015 The Vanadium 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 signals implements utilities for managing process shutdown with
// support for signal-handling.
package signals
// TODO(caprita): Rename the function to Shutdown() and the package to shutdown
// since it's not just signals anymore.
import (
"os"
"os/signal"
"sync"
"syscall"
"time"
"v.io/v23/context"
)
// SameSignalTimeWindow specifies the time window during which multiple
// deliveries of the same signal are counted as one signal. If set to zero, no
// such de-duping occurs. This is useful in situations where a process receives
// a signal explicitly sent by its parent when the parent receives the signal,
// but also receives it independently by virtue of being part of the same
// process group.
//
// This is a variable, so that it can be set appropriately. Note, there is no
// locking around it, the assumption being that it's set during initialization
// and never reset afterwards.
var SameSignalTimeWindow time.Duration
const (
DoubleStopExitCode = 1
)
// Default returns a set of platform-specific signals that applications are
// encouraged to listen on.
func Default() []os.Signal {
return []os.Signal{syscall.SIGTERM, syscall.SIGINT}
}
// ShutdownOnSignals registers signal handlers for the specified signals, or, if
// none are specified, the default signals. The first signal received will be
// made available on the returned channel; upon receiving a second signal, the
// process will exit if that signal differs from the first, or if the same it
// arrives more than a second after the first.
func ShutdownOnSignals(ctx *context.T, signals ...os.Signal) <-chan os.Signal {
if len(signals) == 0 {
signals = Default()
}
// At least a buffer of length two so that we don't drop the first two
// signals we get on account of the channel being full.
ch := make(chan os.Signal, 2)
signal.Notify(ch, signals...)
// At least a buffer of length one so that we don't block on ret <- sig.
ret := make(chan os.Signal, 1)
go func() {
// First signal received.
var sig os.Signal
var sigTime time.Time
select {
case sig = <-ch:
sigTime = time.Now()
ret <- sig
case <-ctx.Done():
ret <- ContextDoneSignal(ctx.Err().Error())
return
}
// Wait for a second signal, and force an exit if the process is
// still executing cleanup code.
for {
secondSig := <-ch
// If signal de-duping is enabled, ignore the signal if
// it's the same signal and has occurred within the
// specified time window.
if SameSignalTimeWindow <= 0 || secondSig.String() != sig.String() || sigTime.Add(SameSignalTimeWindow).Before(time.Now()) {
os.Exit(DoubleStopExitCode)
}
}
}()
return ret
}
// Handler represents a signal handler that can be used to wait for signal
// reception or context cancelation as per NotifyWithCancel. In addition
// it can be used to register additional cancel functions to be invoked
// on signal reception or context cancelation.
type Handler struct {
ctx *context.T
cancel context.CancelFunc
ch <-chan os.Signal
mu sync.Mutex
cancelList []func() // GUARDED_BY(mu)
}
// RegisterCancel registers one or more cancel functions to be invoked
// when a signal is received or the original context is canceled.
func (h *Handler) RegisterCancel(fns ...func()) {
h.mu.Lock()
defer h.mu.Unlock()
h.cancelList = append(h.cancelList, fns...)
}
func (h *Handler) cancelAll() {
h.mu.Lock()
defer h.mu.Unlock()
h.cancel()
for _, cancel := range h.cancelList {
cancel()
}
}
// WaitForSignal will wait for a signal to be received. Context cancelation
// is translated into a ContextDoneSignal signal.
func (h *Handler) WaitForSignal() os.Signal {
select {
case sig := <-h.ch:
h.cancelAll()
return sig
case <-h.ctx.Done():
h.cancelAll()
return ContextDoneSignal(h.ctx.Err().Error())
}
}
// ShutdownOnSignalsWithCancel is like ShutdownOnSignals except it forks the
// supplied context to obtain a cancel function which is called by the returned
// function when a signal is received. The returned function can be called to
// wait for the signal to be received or for the context to be canceled.
// Typical usage would be:
//
// func main() {
// ctx, shutdown := v23.Init()
// defer shutdown()
// ctx, handler := ShutdownOnSignalsWithCancel(ctx)
// defer handler.WaitForSignal()
//
// _, srv, err := v23.WithNewServer(ctx, ...)
//
// }
//
// waitForInterrupt will wait for a signal to be received at which point it
// will cancel the context and thus the server created by WithNewServer to
// initiate its internal shutdown. The deferred shutdown returned by v23.Init()
// will then wait for that the server to complete its shutdown.
// Canceling the context is treated as receipt of a custom signal,
// ContextDoneSignal, in terms of its returns value.
func ShutdownOnSignalsWithCancel(ctx *context.T, signals ...os.Signal) (*context.T, *Handler) {
ctx, cancel := context.WithCancel(ctx)
handler := &Handler{
ctx: ctx,
cancel: cancel,
ch: ShutdownOnSignals(ctx, signals...),
}
return ctx, handler
}
type ContextDoneSignal string
func (ContextDoneSignal) Signal() {}
func (s ContextDoneSignal) String() string { return string(s) }