diff --git a/light/client.go b/light/client.go index dfb4fccf7..f38e9d59d 100644 --- a/light/client.go +++ b/light/client.go @@ -1034,7 +1034,12 @@ func (c *Client) findNewPrimary(ctx context.Context, height int64, remove bool) // process all the responses as they come in for i := 0; i < cap(witnessResponsesC); i++ { - response := <-witnessResponsesC + var response witnessResponse + select { + case response = <-witnessResponsesC: + case <-ctx.Done(): + return nil, ctx.Err() + } switch response.err { // success! We have found a new primary case nil: @@ -1063,10 +1068,6 @@ func (c *Client) findNewPrimary(ctx context.Context, height int64, remove bool) // return the light block that new primary responded with return response.lb, nil - // catch canceled contexts or deadlines - case context.Canceled, context.DeadlineExceeded: - return nil, response.err - // process benign errors by logging them only case provider.ErrNoResponse, provider.ErrLightBlockNotFound, provider.ErrHeightTooHigh: lastError = response.err