-
Notifications
You must be signed in to change notification settings - Fork 0
/
server.go
100 lines (82 loc) · 2.3 KB
/
server.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
package grpchealth
import (
"context"
"time"
"google.golang.org/grpc/health"
healthpb "google.golang.org/grpc/health/grpc_health_v1"
)
// ServerOption is server option
type ServerOption func(*serverOptions)
// WithChecksFunc set checks function
func WithChecksFunc(checksFunc func() healthpb.HealthCheckResponse_ServingStatus) ServerOption {
return func(o *serverOptions) {
o.checksFunc = checksFunc
}
}
// WithRegularInterval set checking health interval
func WithRegularInterval(d time.Duration) ServerOption {
return func(o *serverOptions) {
o.interval = d
}
}
type serverOptions struct {
checksFunc func() healthpb.HealthCheckResponse_ServingStatus
interval time.Duration
}
func applyServerOptions(opts ...ServerOption) *serverOptions {
o := &serverOptions{}
for _, opt := range opts {
opt(o)
}
return o
}
// Server implements service Health.
type Server struct {
service string
hs *health.Server
checkFn func() healthpb.HealthCheckResponse_ServingStatus
interval time.Duration
}
// NewServer returns a health Server.
func NewServer(service string, opts ...ServerOption) *Server {
o := applyServerOptions(opts...)
//checkFn: fn,
s := &Server{
service: service,
hs: health.NewServer(),
checkFn: o.checksFunc,
interval: o.interval,
}
if s.interval > 0 {
go s.checksAtRegularInterval()
}
return s
}
// Check implements `service Health`.
func (s *Server) Check(ctx context.Context, in *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
// regular checks
if s.interval > 0 {
return s.hs.Check(ctx, in)
}
if s.checkFn != nil {
servingStatus := s.checkFn()
s.hs.SetServingStatus(s.service, servingStatus)
}
return s.hs.Check(ctx, in)
}
// Watch implements `service Health`.
func (s *Server) Watch(in *healthpb.HealthCheckRequest, stream healthpb.Health_WatchServer) error {
return s.hs.Watch(in, stream)
}
// SetServingStatus is called when need to reset the serving status of a service
// or insert a new service entry into the statusMap.
func (s *Server) SetServingStatus(servingStatus healthpb.HealthCheckResponse_ServingStatus) {
s.hs.SetServingStatus(s.service, servingStatus)
}
func (s *Server) checksAtRegularInterval() {
for {
servingStatus := s.checkFn()
s.hs.SetServingStatus(s.service, servingStatus)
time.Sleep(s.interval)
}
}