-
Notifications
You must be signed in to change notification settings - Fork 0
/
handler.go
executable file
·208 lines (168 loc) · 5.4 KB
/
handler.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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
package healthcheck
import (
"encoding/json"
"fmt"
"net/http"
"sync"
)
const (
// LivenessHandlerPath path to process liveness probe.
LivenessHandlerPath = "/live"
// ReadinessHandlerPath path to process readiness probe.
ReadinessHandlerPath = "/ready"
successCheckerResultString = "OK"
)
// Handler is a wrapper over http.Handler,
// allowing you to add liveness and readiness checks
type Handler interface {
// Handler is http.Handler, so it can be exposed directly and processed
// /live and /ready endpoints.
http.Handler
// AddLivenessCheck adds a check indicating that this instance
// of the application should be destroyed or restarted. A failed liveness check
// indicates that this instance is not running.
// Each liveness check is also included as a readiness check.
AddLivenessCheck(name string, check Check)
// AddReadinessCheck adds a check indicating that this
// application instance is currently unable to serve requests due to an external
// dependency or some kind of temporary failure. If the readiness check fails, this instance
// should no longer receive requests, but it should not be restarted or destroyed.
AddReadinessCheck(name string, check Check)
// LiveEndpoint is an HTTP handler for the /live endpoint only, which
// is useful if you need to add it to your own HTTP handler tree.
LiveEndpoint(http.ResponseWriter, *http.Request)
//ReadyEndpoint is an HTTP handler for the /ready endpoint only, which
// is useful if you need to add it to your own HTTP handler tree.
ReadyEndpoint(http.ResponseWriter, *http.Request)
// AddCheckErrorHandler adds a callback to process a failed check (in order to log errors, etc.).
AddCheckErrorHandler(handler ErrorHandler)
}
// Check signature of check proccess function
type Check func() error
// ErrorHandler error handler's signature for failed checks.
type ErrorHandler func(name string, err error)
// NewHandler creates a new basic Handler
func NewHandler() Handler {
h := &basicHandler{
livenessChecks: make(map[string]Check),
readinessChecks: make(map[string]Check),
}
h.Handle("/live", http.HandlerFunc(h.LiveEndpoint))
h.Handle("/ready", http.HandlerFunc(h.ReadyEndpoint))
return h
}
// basicHandler implementation of Handler.
type basicHandler struct {
http.ServeMux
checksMutex sync.RWMutex
livenessChecks map[string]Check
readinessChecks map[string]Check
errorHandler ErrorHandler
}
func (s *basicHandler) LiveEndpoint(w http.ResponseWriter, r *http.Request) {
s.handle(w, r, s.livenessChecks)
}
func (s *basicHandler) ReadyEndpoint(w http.ResponseWriter, r *http.Request) {
s.handle(w, r, s.readinessChecks, s.livenessChecks)
}
func (s *basicHandler) AddLivenessCheck(name string, check Check) {
s.checksMutex.Lock()
defer s.checksMutex.Unlock()
s.livenessChecks[name] = check
}
func (s *basicHandler) AddReadinessCheck(name string, check Check) {
s.checksMutex.Lock()
defer s.checksMutex.Unlock()
s.readinessChecks[name] = check
}
func (s *basicHandler) AddCheckErrorHandler(handler ErrorHandler) {
s.errorHandler = handler
}
type result struct {
name string
result string
}
func (s *basicHandler) collectChecks(checks map[string]Check, resultsOut map[string]string) (status int) {
s.checksMutex.RLock()
defer s.checksMutex.RUnlock()
status = http.StatusOK
if len(checks) == 0 {
return
}
var (
wg = sync.WaitGroup{}
results = make(chan result)
)
for name, check := range checks {
wg.Add(1)
go func(name string, check Check) {
defer func() {
wg.Done()
// check panic error
if r := recover(); r != nil {
results <- result{
name: name,
result: fmt.Sprintf("checker panic recovered: %v", r),
}
if s.errorHandler != nil {
s.errorHandler(name, fmt.Errorf("checker panic recovered: %v", r))
}
}
}()
var val = successCheckerResultString
if err := check(); err != nil {
val = err.Error()
if s.errorHandler != nil {
s.errorHandler(name, err)
}
}
results <- result{
name: name,
result: val,
}
}(name, check)
}
// wait for all checks to be made
// then close the results channel
go func() {
wg.Wait()
close(results)
}()
for res := range results {
resultsOut[res.name] = res.result
if res.result != successCheckerResultString {
status = http.StatusServiceUnavailable
}
}
return status
}
func (s *basicHandler) handle(w http.ResponseWriter, r *http.Request, checks ...map[string]Check) {
if r.Method != http.MethodGet {
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
return
}
checkResults := make(map[string]string)
status := http.StatusOK
for _, m := range checks {
if s := s.collectChecks(m, checkResults); s != http.StatusOK {
status = s
}
}
// Set response code and content header
w.Header().Set("Content-Type", "application/json; charset=utf-8")
w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
w.Header().Set("Pragma", "no-cache")
w.Header().Set("Expires", "0")
w.WriteHeader(status)
// If not ?full=1, we return an empty body. Kubernetes only cares about
// HTTP status codes, so we won't waste bytes on the full request body.
if r.URL.Query().Get("full") != "1" {
_, _ = w.Write([]byte("{}\n"))
return
}
// Write the JSON body, ignoring any encoding errors (which
// are actually not possible because we encode map[string]string).
encoder := json.NewEncoder(w)
encoder.SetIndent("", " ")
_ = encoder.Encode(checkResults)
}