@@ -3,8 +3,10 @@ package protocol
33import (
44 "context"
55 "fmt"
6+ healthpb "google.golang.org/grpc/health/grpc_health_v1"
67 "io"
78 "sync"
9+ "time"
810
911 "google.golang.org/grpc"
1012 "google.golang.org/grpc/codes"
@@ -30,6 +32,10 @@ type Protocol struct {
3032 openerMutex sync.Mutex
3133
3234 closeChan chan struct {}
35+
36+ // errorChan is to send an error messsage to restart the connection
37+ reconnectErrorChan chan error
38+ checkInterval time.Duration
3339}
3440
3541var (
@@ -137,6 +143,10 @@ func (p *Protocol) OpenInbound(ctx context.Context) error {
137143 }
138144 }()
139145
146+ if p .reconnectErrorChan != nil {
147+ go p .healthCheck (ctx )
148+ }
149+
140150 // Wait until external or internal context done
141151 select {
142152 case <- ctx .Done ():
@@ -164,3 +174,38 @@ func (p *Protocol) Close(ctx context.Context) error {
164174 close (p .closeChan )
165175 return nil
166176}
177+
178+ // healthCheck check status of the server
179+ func (p * Protocol ) healthCheck (ctx context.Context ) {
180+ logger := cecontext .LoggerFrom (ctx )
181+ healthClient := healthpb .NewHealthClient (p .clientConn )
182+ stream , err := healthClient .Watch (ctx , & healthpb.HealthCheckRequest {Service : "" })
183+ if err != nil {
184+ p .reconnectErrorChan <- err
185+ return
186+ }
187+
188+ for {
189+ timeout := time .NewTimer (p .checkInterval )
190+ done := make (chan struct {})
191+
192+ go func () {
193+ resp , err := stream .Recv ()
194+ if err != nil {
195+ p .reconnectErrorChan <- err
196+ close (done )
197+ return
198+ }
199+ logger .Infof ("Received server health status %s" , resp .Status )
200+ close (done )
201+ }()
202+
203+ select {
204+ case <- timeout .C :
205+ p .reconnectErrorChan <- fmt .Errorf ("timeout waiting for health check" )
206+ return
207+ case <- done :
208+ timeout .Stop ()
209+ }
210+ }
211+ }
0 commit comments