-
Notifications
You must be signed in to change notification settings - Fork 2
/
gracegrpc.go
159 lines (138 loc) · 3.32 KB
/
gracegrpc.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
package gracegrpc
/*
*
* Created by 0x5010 on 2018/02/01.
* gracegrpc
* https://github.com/0x5010/gracegrpc
*
* Copyright 2018 0x5010.
* Licensed under the MIT license.
*
*/
import (
"fmt"
"log"
"net"
"os"
"os/signal"
"syscall"
"github.com/facebookgo/grace/gracenet"
"google.golang.org/grpc"
)
var (
// Used to indicate a graceful restart in the new process.
envKey = "LISTEN_FDS"
ppid = os.Getppid()
)
type logger interface {
Printf(string, ...interface{})
}
// GraceGrpc is used to wrap a grpc server that can be gracefully terminated & restarted
type GraceGrpc struct {
server *grpc.Server
grace *gracenet.Net
listener net.Listener
errors chan error
pidPath string
logger logger
}
// New is used to construct a new GraceGrpc
func New(s *grpc.Server, net, addr, pidPath string, l logger) (*GraceGrpc, error) {
if l == nil {
l = log.New(os.Stderr, "", log.LstdFlags)
}
gr := &GraceGrpc{
server: s,
grace: &gracenet.Net{},
//for StartProcess error.
errors: make(chan error),
pidPath: pidPath,
logger: l,
}
listener, err := gr.grace.Listen(net, addr)
if err != nil {
return nil, err
}
gr.listener = listener
return gr, nil
}
func (gr *GraceGrpc) startServe() {
if err := gr.server.Serve(gr.listener); err != nil {
gr.errors <- err
}
}
func (gr *GraceGrpc) handleSignal() <-chan struct{} {
terminate := make(chan struct{})
go func() {
ch := make(chan os.Signal, 10)
signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR2)
for {
sig := <-ch
switch sig {
case syscall.SIGINT, syscall.SIGTERM:
signal.Stop(ch)
gr.server.GracefulStop()
close(terminate)
return
case syscall.SIGUSR2:
if _, err := gr.grace.StartProcess(); err != nil {
gr.errors <- err
}
}
}
}()
return terminate
}
// storePid is used to write out PID to pidPath
func (gr *GraceGrpc) storePid(pid int) error {
pidPath := gr.pidPath
if pidPath == "" {
return fmt.Errorf("No pid file path: %s", pidPath)
}
pidFile, err := os.OpenFile(pidPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
if err != nil {
return fmt.Errorf("Could not open pid file: %v", err)
}
defer pidFile.Close()
_, err = pidFile.WriteString(fmt.Sprintf("%d", pid))
if err != nil {
return fmt.Errorf("Could not write to pid file: %s", err)
}
return nil
}
// Serve is used to start grpc server.
// Serve will gracefully terminated or restarted when handling signals.
func (gr *GraceGrpc) Serve() error {
if gr.listener == nil || gr.logger == nil {
return fmt.Errorf("gracegrpc must construct by new")
}
inherit := os.Getenv(envKey) != ""
pid := os.Getpid()
addrString := gr.listener.Addr().String()
if inherit {
if ppid == 1 {
gr.logger.Printf("Listening on init activated %s\n", addrString)
} else {
gr.logger.Printf("Graceful handoff of %s with new pid %d replace old pid %d\n", addrString, pid, ppid)
}
} else {
gr.logger.Printf("Serving %s with pid %d\n", addrString, pid)
}
if err := gr.storePid(pid); err != nil {
return err
}
go gr.startServe()
if inherit && ppid != 1 {
if err := syscall.Kill(ppid, syscall.SIGTERM); err != nil {
return fmt.Errorf("failed to close parent: %s", err)
}
}
terminate := gr.handleSignal()
select {
case err := <-gr.errors:
return err
case <-terminate:
gr.logger.Printf("Exiting pid %d.", os.Getpid())
return nil
}
}