From 08b938982bd4006c6edc83d6845e68f767ef4dee Mon Sep 17 00:00:00 2001 From: Nishant Das Date: Thu, 25 Feb 2021 01:10:25 +0800 Subject: [PATCH] Propagate Cancellation To Beacon Node Server (#8512) Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com> --- validator/rpc/health.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/validator/rpc/health.go b/validator/rpc/health.go index 6d0e0725b..b1ef903de 100644 --- a/validator/rpc/health.go +++ b/validator/rpc/health.go @@ -57,7 +57,12 @@ func (s *Server) GetVersion(ctx context.Context, _ *ptypes.Empty) (*pb.VersionRe // StreamBeaconLogs from the beacon node via a gRPC server-side stream. func (s *Server) StreamBeaconLogs(req *ptypes.Empty, stream pb.Health_StreamBeaconLogsServer) error { - client, err := s.beaconNodeHealthClient.StreamBeaconLogs(s.ctx, req) + // Wrap service context with a cancel in order to propagate the exiting of + // this method properly to the beacon node server. + ctx, cancel := context.WithCancel(s.ctx) + defer cancel() + + client, err := s.beaconNodeHealthClient.StreamBeaconLogs(ctx, req) if err != nil { return err } @@ -67,6 +72,8 @@ func (s *Server) StreamBeaconLogs(req *ptypes.Empty, stream pb.Health_StreamBeac return status.Error(codes.Canceled, "Context canceled") case <-stream.Context().Done(): return status.Error(codes.Canceled, "Context canceled") + case <-client.Context().Done(): + return status.Error(codes.Canceled, "Context canceled") default: resp, err := client.Recv() if err != nil {