/
gproc_signal.go
134 lines (121 loc) · 3.19 KB
/
gproc_signal.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
// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.
package gproc
import (
"context"
"github.com/gogf/gf/v2/internal/intlog"
"github.com/gogf/gf/v2/util/gutil"
"os"
"os/signal"
"sync"
"syscall"
)
// SigHandler defines a function type for signal handling.
type SigHandler func(sig os.Signal)
var (
// Use internal variable to guarantee concurrent safety
// when multiple Listen happen.
listenOnce = sync.Once{}
waitChan = make(chan struct{})
signalChan = make(chan os.Signal, 1)
signalHandlerMu sync.Mutex
signalHandlerMap = make(map[os.Signal][]SigHandler)
shutdownSignalMap = map[os.Signal]struct{}{
syscall.SIGINT: {},
syscall.SIGQUIT: {},
syscall.SIGKILL: {},
syscall.SIGTERM: {},
syscall.SIGABRT: {},
}
)
func init() {
for sig := range shutdownSignalMap {
signalHandlerMap[sig] = make([]SigHandler, 0)
}
}
// AddSigHandler adds custom signal handler for custom one or more signals.
func AddSigHandler(handler SigHandler, signals ...os.Signal) {
signalHandlerMu.Lock()
defer signalHandlerMu.Unlock()
for _, sig := range signals {
signalHandlerMap[sig] = append(signalHandlerMap[sig], handler)
}
notifySignals()
}
// AddSigHandlerShutdown adds custom signal handler for shutdown signals:
// syscall.SIGINT,
// syscall.SIGQUIT,
// syscall.SIGKILL,
// syscall.SIGTERM,
// syscall.SIGABRT.
func AddSigHandlerShutdown(handler ...SigHandler) {
signalHandlerMu.Lock()
defer signalHandlerMu.Unlock()
for _, h := range handler {
for sig := range shutdownSignalMap {
signalHandlerMap[sig] = append(signalHandlerMap[sig], h)
}
}
notifySignals()
}
// Listen blocks and does signal listening and handling.
func Listen() {
listenOnce.Do(func() {
go listen()
})
<-waitChan
}
func listen() {
defer close(waitChan)
var (
ctx = context.Background()
wg = sync.WaitGroup{}
sig os.Signal
)
for {
sig = <-signalChan
intlog.Printf(ctx, `signal received: %s`, sig.String())
if handlers := getHandlersBySignal(sig); len(handlers) > 0 {
for _, handler := range handlers {
wg.Add(1)
var (
currentHandler = handler
currentSig = sig
)
gutil.TryCatch(ctx, func(ctx context.Context) {
defer wg.Done()
currentHandler(currentSig)
}, func(ctx context.Context, exception error) {
intlog.Errorf(ctx, `execute signal handler failed: %+v`, exception)
})
}
}
// If it is shutdown signal, it exits this signal listening.
if _, ok := shutdownSignalMap[sig]; ok {
intlog.Printf(
ctx,
`receive shutdown signal "%s", waiting all signal handler done`,
sig.String(),
)
// Wait until signal handlers done.
wg.Wait()
intlog.Print(ctx, `all signal handler done, exit process`)
return
}
}
}
func notifySignals() {
var signals = make([]os.Signal, 0)
for s := range signalHandlerMap {
signals = append(signals, s)
}
signal.Notify(signalChan, signals...)
}
func getHandlersBySignal(sig os.Signal) []SigHandler {
signalHandlerMu.Lock()
defer signalHandlerMu.Unlock()
return signalHandlerMap[sig]
}