blob: 419f06040d1d3d37f8d966abfad8293ee9f92420 [file] [log] [blame]
khenaidooab1f7bd2019-11-14 14:00:27 -05001/*
2 *
3 * Copyright 2018 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19package health
20
21import (
22 "context"
23 "fmt"
24 "io"
25 "time"
26
27 "google.golang.org/grpc"
28 "google.golang.org/grpc/codes"
29 "google.golang.org/grpc/connectivity"
30 healthpb "google.golang.org/grpc/health/grpc_health_v1"
31 "google.golang.org/grpc/internal"
32 "google.golang.org/grpc/internal/backoff"
33 "google.golang.org/grpc/status"
34)
35
Andrea Campanella3614a922021-02-25 12:40:42 +010036var (
37 backoffStrategy = backoff.DefaultExponential
38 backoffFunc = func(ctx context.Context, retries int) bool {
39 d := backoffStrategy.Backoff(retries)
40 timer := time.NewTimer(d)
41 select {
42 case <-timer.C:
43 return true
44 case <-ctx.Done():
45 timer.Stop()
46 return false
47 }
khenaidooab1f7bd2019-11-14 14:00:27 -050048 }
Andrea Campanella3614a922021-02-25 12:40:42 +010049)
khenaidooab1f7bd2019-11-14 14:00:27 -050050
51func init() {
52 internal.HealthCheckFunc = clientHealthCheck
53}
54
55const healthCheckMethod = "/grpc.health.v1.Health/Watch"
56
57// This function implements the protocol defined at:
58// https://github.com/grpc/grpc/blob/master/doc/health-checking.md
59func clientHealthCheck(ctx context.Context, newStream func(string) (interface{}, error), setConnectivityState func(connectivity.State), service string) error {
60 tryCnt := 0
61
62retryConnection:
63 for {
64 // Backs off if the connection has failed in some way without receiving a message in the previous retry.
65 if tryCnt > 0 && !backoffFunc(ctx, tryCnt-1) {
66 return nil
67 }
68 tryCnt++
69
70 if ctx.Err() != nil {
71 return nil
72 }
73 setConnectivityState(connectivity.Connecting)
74 rawS, err := newStream(healthCheckMethod)
75 if err != nil {
76 continue retryConnection
77 }
78
79 s, ok := rawS.(grpc.ClientStream)
80 // Ideally, this should never happen. But if it happens, the server is marked as healthy for LBing purposes.
81 if !ok {
82 setConnectivityState(connectivity.Ready)
83 return fmt.Errorf("newStream returned %v (type %T); want grpc.ClientStream", rawS, rawS)
84 }
85
86 if err = s.SendMsg(&healthpb.HealthCheckRequest{Service: service}); err != nil && err != io.EOF {
87 // Stream should have been closed, so we can safely continue to create a new stream.
88 continue retryConnection
89 }
90 s.CloseSend()
91
92 resp := new(healthpb.HealthCheckResponse)
93 for {
94 err = s.RecvMsg(resp)
95
96 // Reports healthy for the LBing purposes if health check is not implemented in the server.
97 if status.Code(err) == codes.Unimplemented {
98 setConnectivityState(connectivity.Ready)
99 return err
100 }
101
102 // Reports unhealthy if server's Watch method gives an error other than UNIMPLEMENTED.
103 if err != nil {
104 setConnectivityState(connectivity.TransientFailure)
105 continue retryConnection
106 }
107
Andrea Campanella3614a922021-02-25 12:40:42 +0100108 // As a message has been received, removes the need for backoff for the next retry by resetting the try count.
khenaidooab1f7bd2019-11-14 14:00:27 -0500109 tryCnt = 0
110 if resp.Status == healthpb.HealthCheckResponse_SERVING {
111 setConnectivityState(connectivity.Ready)
112 } else {
113 setConnectivityState(connectivity.TransientFailure)
114 }
115 }
116 }
117}