ronantakizawa/codereview-bench huggingface dataset:
Repository: grpc/grpc-go
File: examples/features/gracefulstop/client/main.go (Go)

Code before:
var addr = flag.String("addr", "localhost:50052", "the address to connect to")

func main() {
	flag.Parse()

	conn, err := grpc.NewClient(*addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("Failed to create new client: %v", err)
	}
	defer conn.Close()
	c := pb.NewEchoClient(conn)

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	stream, err := c.ServerStreamingEcho(ctx, &pb.EchoRequest{})
	if err != nil {
		log.Fatalf("Error starting stream: %v", err)
	}

	for {
		r, err := stream.Recv()
		if err != nil {
			// Handle the error and close the stream gracefully
			log.Printf("Error sending request: %v\n", err)
			err := stream.CloseSend()
			if err != nil {
				log.Fatalf("Error closing stream: %v", err)
			}
			log.Println("Stream closed gracefully")
			break
		}
		log.Printf(r.Message)
	}

	log.Printf("Client finished interaction with server.")
}


Review comment (suggestion):
This doesn't make sense.  The stream ended (by virtue of receiving an error), so there is no reason to `CloseSend`.  Probably `CloseSend` should also return an error in this scenario, but it's also probably not important.

Code after:

var addr = flag.String("addr", "localhost:50052", "the address to connect to")

func main() {
	flag.Parse()

	conn, err := grpc.NewClient(*addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("Failed to create new client: %v", err)
	}
	defer conn.Close()
	c := pb.NewEchoClient(conn)

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	// Start a client stream and keep calling the `c.UnaryEcho` until receiving
	// an error. Error will indicate that server graceful stop is initiated and
	// it won't accept any new requests.
	stream, err := c.ClientStreamingEcho(ctx)
	if err != nil {
		log.Fatalf("Error starting stream: %v", err)
	}

	// Keep track of successful unary requests which can be compared later to
	// the successful unary requests reported by the server.
	unaryRequests := 0
	for {
		r, err := c.UnaryEcho(ctx, &pb.EchoRequest{Message: "Hello"})
		if err != nil {
			log.Printf("Error calling `UnaryEcho`. Server graceful stop initiated: %v", err)
			break
		}
		unaryRequests++
		time.Sleep(200 * time.Millisecond)
		log.Printf(r.Message)
	}
	log.Printf("Successful unary requests made by client: %d", unaryRequests)

	r, err := stream.CloseAndRecv()
	if err != nil {
		log.Fatalf("Error closing stream: %v", err)
	}
	if fmt.Sprintf("%d", unaryRequests) != r.Message {
		log.Fatalf("Got %s successful unary requests processed from server, want: %d", r.Message, unaryRequests)
	}
	log.Printf("Successful unary requests processed by server and made by client are same.")
}
Repository: grpc/grpc-go
File: examples/features/gracefulstop/server/main.go (Go)

Code before:
	return &pb.EchoResponse{Message: fmt.Sprintf("Request Processed: %d", s.unaryRequests)}, nil
}

func main() {
	flag.Parse()

	address := fmt.Sprintf(":%v", *port)
	lis, err := net.Listen("tcp", address)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	s := grpc.NewServer()
	ss := &server{streamStart: make(chan struct{})}
	pb.RegisterEchoServer(s, ss)

	go func() {
		<-ss.streamStart // wait until server streaming starts
		time.Sleep(1 * time.Second)
		log.Println("Initiating graceful shutdown...")
		s.GracefulStop() // gracefully stop server after in-flight server streaming rpc finishes
		log.Println("Server stopped gracefully.")
	}()

	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}


Review comment (suggestion):
This still should do a hard Stop after N seconds (10?), probably via a `time.AfterFunc`.  This is what users will need to do, so we need to show it.

Code after:
	return &pb.EchoResponse{Message: req.Message}, nil
}

func main() {
	flag.Parse()

	address := fmt.Sprintf(":%v", *port)
	lis, err := net.Listen("tcp", address)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	s := grpc.NewServer()
	ss := &server{streamStart: make(chan struct{})}
	pb.RegisterEchoServer(s, ss)

	go func() {
		<-ss.streamStart // wait until server streaming starts
		time.Sleep(1 * time.Second)
		log.Println("Initiating graceful shutdown...")
		timer := time.AfterFunc(10*time.Second, func() {
			log.Println("Server couldn't stop gracefully in time. Doing force stop.")
			s.Stop()
		})
		defer timer.Stop()
		s.GracefulStop() // gracefully stop server after in-flight server streaming rpc finishes
		log.Println("Server stopped gracefully.")
	}()

	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}
Repository: grpc/grpc-go
File: examples/features/gracefulstop/client/main.go (Go)

Code before:
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	// Start a client stream and keep calling the `c.UnaryEcho` until receiving
	// an error. Error will indicate that server graceful stop is initiated and
	// it won't accept any new requests.
	stream, err := c.ClientStreamingEcho(ctx)
	if err != nil {
		log.Fatalf("Error starting stream: %v", err)
	}

	for {
		r, err := c.UnaryEcho(ctx, &pb.EchoRequest{Message: "Hello"})
		if err != nil {
			log.Printf("Error calling `UnaryEcho`. Server graceful stop initiated: %v", err)
			break
		}
		time.Sleep(10 * time.Millisecond)
		log.Printf(r.Message)
	}

	r, err := stream.CloseAndRecv()
	if err != nil {
		log.Fatalf("Error closing stream: %v", err)
	}
	log.Printf(r.Message)
}


Review comment (suggestion):
The client should keep its own count of requests and confirm that the value returned in this message matches its count.  Then it should output a "happy" message, which is how we can test it in `examples_test.sh`

Code after:

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	// Start a client stream and keep calling the `c.UnaryEcho` until receiving
	// an error. Error will indicate that server graceful stop is initiated and
	// it won't accept any new requests.
	stream, err := c.ClientStreamingEcho(ctx)
	if err != nil {
		log.Fatalf("Error starting stream: %v", err)
	}

	// Keep track of successful unary requests which can be compared later to
	// the successful unary requests reported by the server.
	unaryRequests := 0
	for {
		r, err := c.UnaryEcho(ctx, &pb.EchoRequest{Message: "Hello"})
		if err != nil {
			log.Printf("Error calling `UnaryEcho`. Server graceful stop initiated: %v", err)
			break
		}
		unaryRequests++
		time.Sleep(200 * time.Millisecond)
		log.Printf(r.Message)
	}
	log.Printf("Successful unary requests made by client: %d", unaryRequests)

	r, err := stream.CloseAndRecv()
	if err != nil {
		log.Fatalf("Error closing stream: %v", err)
	}
	if fmt.Sprintf("%d", unaryRequests) != r.Message {
		log.Fatalf("Got %s successful unary requests processed from server, want: %d", r.Message, unaryRequests)
	}
	log.Printf("Successful unary requests processed by server and made by client are same.")
}
Repository: grpc/grpc-go
File: stats/stats_test.go (Go)

Code before:
type (
	connCtxKey struct{}
	rpcCtxKey  struct{}
)

var (
	// For headers sent to server:
	testMetadata = metadata.MD{
		"key1":       []string{"value1"},
		"key2":       []string{"value2"},
		"user-agent": []string{fmt.Sprintf("test/0.0.1 grpc-go/%s", grpc.Version)},
	}
	// For headers sent from server:
	testHeaderMetadata = metadata.MD{
		"hkey1": []string{"headerValue1"},
		"hkey2": []string{"headerValue2"},
	}
	// For trailers sent from server:
	testTrailerMetadata = metadata.MD{
		"tkey1": []string{"trailerValue1"},
		"tkey2": []string{"trailerValue2"},
	}
	// The id for which the service handler should return error.
	errorID int32 = 32202
    // Server Stats
	// Ensure that Unary RPC server stats events are logged in the correct order.
	expectedUnarySequence = []string{
		"ConnStats",
		"InHeader",
		"Begin",
		"InPayload",
		"OutHeader",
		"OutPayload",
		"OutTrailer",
		"End",
	}
    // Server Stats
	// Ensure that the sequence of server-side stats events for a 
	// client stream RPC matches the expected flow.
	expectedClientStreamSequence = []string{
		"ConnStats",
		"InHeader",
		"Begin",
		"OutHeader",
		"InPayload",
		"InPayload",
		"InPayload",
		"InPayload",
		"InPayload",
		"OutPayload",
		"OutTrailer",

Review comment (suggestion):
variable name is self explanatory. If at all we want to put comment here, something like this is more suitable, `To verify if the Unary RPC server stats events are logged in the correct order.`

Code after:
type connCtxKey struct{}
type rpcCtxKey struct{}

var (
	// For headers sent to server:
	testMetadata = metadata.MD{
		"key1":       []string{"value1"},
		"key2":       []string{"value2"},
		"user-agent": []string{fmt.Sprintf("test/0.0.1 grpc-go/%s", grpc.Version)},
	}
	// For headers sent from server:
	testHeaderMetadata = metadata.MD{
		"hkey1": []string{"headerValue1"},
		"hkey2": []string{"headerValue2"},
	}
	// For trailers sent from server:
	testTrailerMetadata = metadata.MD{
		"tkey1": []string{"trailerValue1"},
		"tkey2": []string{"trailerValue2"},
	}
	// The id for which the service handler should return error.
	errorID int32 = 32202
)

func idToPayload(id int32) *testpb.Payload {
	return &testpb.Payload{Body: []byte{byte(id), byte(id >> 8), byte(id >> 16), byte(id >> 24)}}
}

func payloadToID(p *testpb.Payload) int32 {
	if p == nil || len(p.Body) != 4 {
		panic("invalid payload")
	}
	return int32(p.Body[0]) + int32(p.Body[1])<<8 + int32(p.Body[2])<<16 + int32(p.Body[3])<<24
}

func setIncomingStats(ctx context.Context, mdKey string, b []byte) context.Context {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}
	md.Set(mdKey, string(b))
	return metadata.NewIncomingContext(ctx, md)
}

func getOutgoingStats(ctx context.Context, mdKey string) []byte {
	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		return nil
	}
	tagValues := md.Get(mdKey)
	if len(tagValues) == 0 {
Repository: grpc/grpc-go
File: stats/stats_test.go (Go)

Code before:

type connCtxKey struct{}
type rpcCtxKey struct{}

var (
	// For headers sent to server:
	testMetadata = metadata.MD{
		"key1":       []string{"value1"},
		"key2":       []string{"value2"},
		"user-agent": []string{fmt.Sprintf("test/0.0.1 grpc-go/%s", grpc.Version)},
	}
	// For headers sent from server:
	testHeaderMetadata = metadata.MD{
		"hkey1": []string{"headerValue1"},
		"hkey2": []string{"headerValue2"},
	}
	// For trailers sent from server:
	testTrailerMetadata = metadata.MD{
		"tkey1": []string{"trailerValue1"},
		"tkey2": []string{"trailerValue2"},
	}
	// The id for which the service handler should return error.
	errorID int32 = 32202
	// To verify if the Unary RPC server stats events are logged in the
	// correct order.
	expectedUnarySequence = []string{
		"ConnStats",
		"InHeader",
		"Begin",
		"InPayload",
		"OutHeader",
		"OutPayload",
		"OutTrailer",
		"End",
	}
	// To verify if the Client Stream RPC server stats events are logged in the
	// correct order.
	expectedClientStreamSequence = []string{
		"ConnStats",
		"InHeader",
		"Begin",
		"OutHeader",
		"InPayload",
		"InPayload",
		"InPayload",
		"InPayload",
		"InPayload",
		"OutPayload",
		"OutTrailer",
		"End",
	}

Review comment (suggestion):
this can be initialized within the test itself because we are not using it anywhere else. Also, convention is to use `want/wanted` instead of `expected`

Code after:

type connCtxKey struct{}
type rpcCtxKey struct{}

var (
	// For headers sent to server:
	testMetadata = metadata.MD{
		"key1":       []string{"value1"},
		"key2":       []string{"value2"},
		"user-agent": []string{fmt.Sprintf("test/0.0.1 grpc-go/%s", grpc.Version)},
	}
	// For headers sent from server:
	testHeaderMetadata = metadata.MD{
		"hkey1": []string{"headerValue1"},
		"hkey2": []string{"headerValue2"},
	}
	// For trailers sent from server:
	testTrailerMetadata = metadata.MD{
		"tkey1": []string{"trailerValue1"},
		"tkey2": []string{"trailerValue2"},
	}
	// The id for which the service handler should return error.
	errorID int32 = 32202
)

func idToPayload(id int32) *testpb.Payload {
	return &testpb.Payload{Body: []byte{byte(id), byte(id >> 8), byte(id >> 16), byte(id >> 24)}}
}

func payloadToID(p *testpb.Payload) int32 {
	if p == nil || len(p.Body) != 4 {
		panic("invalid payload")
	}
	return int32(p.Body[0]) + int32(p.Body[1])<<8 + int32(p.Body[2])<<16 + int32(p.Body[3])<<24
}

func setIncomingStats(ctx context.Context, mdKey string, b []byte) context.Context {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}
	md.Set(mdKey, string(b))
	return metadata.NewIncomingContext(ctx, md)
}

func getOutgoingStats(ctx context.Context, mdKey string) []byte {
	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		return nil
	}
	tagValues := md.Get(mdKey)
Repository: grpc/grpc-go
File: balancer_wrapper.go (Go)

Code before:
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		curHD := acbw.healthData
		if curHD != hd {
			return
		}
		if !healthCheckEnabled {
			if listener != nil {
				listener(balancer.SubConnState{ConnectivityState: connectivity.Ready})
			}
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				curHD := acbw.healthData
				acbw.healthMu.Unlock()
				if curHD != hd {
					return
				}
				listener(scs)
			})
		}

		if listener == nil {
			listenerWrapper = nil
		}

		healthOpts := HealthCheckOptions{
			HealthServiceName: cfg.ServiceName,
			Listener:          listenerWrapper,
		}
		regHealthLisFn.(func(context.Context, balancer.SubConn, HealthCheckOptions))(ctx, acbw, healthOpts)
	})
}


Review comment (refactor):
Can we make a helper method to check if a new health listener has been registered, and therefore we don't have to process updates from this one. This same logic is currently being done in two places.

Maybe something like:

```
func (acbw *acBalancerWrapper) isHealthDataCurrent(hd *healthData) bool { ... }
```

Code after:
func (acbw *acBalancerWrapper) RegisterHealthListener(listener func(balancer.SubConnState)) {
	acbw.healthMu.Lock()
	defer acbw.healthMu.Unlock()
	acbw.healthData.closeHealthProducer()
	// listeners should not be registered when the connectivity state
	// isn't Ready. This may happen when the balancer registers a listener
	// after the connectivityState is updated, but before it is notified
	// of the update.
	if acbw.healthData.connectivityState != connectivity.Ready {
		return
	}
	// Replace the health data to stop sending updates to any previously
	// registered health listeners.
	hd := newHealthData(connectivity.Ready)
	acbw.healthData = hd
	if listener == nil {
		return
	}

	registerFn := acbw.healthListenerRegFn()
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		if acbw.healthData != hd {
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				defer acbw.healthMu.Unlock()
				if acbw.healthData != hd {
					return
				}
				listener(scs)
			})
		}

		hd.closeHealthProducer = registerFn(ctx, listenerWrapper)
	})
}
Repository: grpc/grpc-go
File: balancer_wrapper.go (Go)

Code before:
		if !healthCheckEnabled {
			if listener != nil {
				listener(balancer.SubConnState{ConnectivityState: connectivity.Ready})
			}
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				curHD := acbw.healthData
				acbw.healthMu.Unlock()
				if curHD != hd {
					return
				}
				listener(scs)
			})
		}

		if listener == nil {
			listenerWrapper = nil
		}

		healthOpts := HealthCheckOptions{
			HealthServiceName: cfg.ServiceName,
			Listener:          listenerWrapper,
		}
		regHealthLisFn.(func(context.Context, balancer.SubConn, HealthCheckOptions))(ctx, acbw, healthOpts)
	})
}


Review comment (suggestion):
If I'm reading this right, the only reason to pass a nil `Listener` in `HealthCheckOptions` to the health checking function is to stop the previous running health checking stream.

What are your thoughts on exposing functionality to cancel the running stream from the producer, via the `internal` package (just like how we do to register the listener)? I feel that would make this code more straightforward to read and understand.

Code after:
	// of the update.
	if acbw.healthData.connectivityState != connectivity.Ready {
		return
	}
	// Replace the health data to stop sending updates to any previously
	// registered health listeners.
	hd := newHealthData(connectivity.Ready)
	acbw.healthData = hd
	if listener == nil {
		return
	}

	registerFn := acbw.healthListenerRegFn()
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		if acbw.healthData != hd {
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				defer acbw.healthMu.Unlock()
				if acbw.healthData != hd {
					return
				}
				listener(scs)
			})
		}

		hd.closeHealthProducer = registerFn(ctx, listenerWrapper)
	})
}
Repository: grpc/grpc-go
File: balancer_wrapper.go (Go)

Code before:
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				curHD := acbw.healthData
				acbw.healthMu.Unlock()
				if curHD != hd {
					return
				}
				listener(scs)
			})
		}

		if listener == nil {
			listenerWrapper = nil
		}

		healthOpts := HealthCheckOptions{
			HealthServiceName: cfg.ServiceName,
			Listener:          listenerWrapper,
		}
		regHealthLisFn.(func(context.Context, balancer.SubConn, HealthCheckOptions))(ctx, acbw, healthOpts)
	})
}


Review comment (bug):
How about doing the type assertion when `regHealthLisFn` is first set to `internal.RegisterClientHealthCheckListener` (probably after the `nil` check that is done as part of the other checks to see if client side health checking needs to be enabled) and maybe throwing an error log and/or returning if the type assertion fails.

Code after:
	hd := newHealthData(connectivity.Ready)
	acbw.healthData = hd
	if listener == nil {
		return
	}

	registerFn := acbw.healthListenerRegFn()
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		if acbw.healthData != hd {
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				defer acbw.healthMu.Unlock()
				if acbw.healthData != hd {
					return
				}
				listener(scs)
			})
		}

		hd.closeHealthProducer = registerFn(ctx, listenerWrapper)
	})
}
Repository: grpc/grpc-go
File: test/healthcheck_test.go (Go)

Code before:
	_, r := setupClient(t, nil)
	svcCfg := fmt.Sprintf(`{
	    "healthCheckConfig": {
		    "serviceName": "foo"
	    },
	    "loadBalancingConfig": [{"%s":{}}]
    }`, t.Name())
	r.UpdateState(resolver.State{
		Addresses:     []resolver.Address{{Addr: lis.Addr().String()}},
		ServiceConfig: parseServiceConfig(t, r, svcCfg)})

	var sc balancer.SubConn
	select {
	case sc = <-scChan:
	case <-ctx.Done():
		t.Fatal("Context timed out waiting for SubConn creation")
	}

	// Wait for the SubConn to enter READY.
	select {
	case <-readyUpdateReceivedCh:
	case <-ctx.Done():
		t.Fatalf("Context timed out waiting for SubConn to enter READY")
	}

	// Register a health listener and verify it receives updates.
	select {
	case <-hcEnterChan:
		t.Fatalf("Health service client created prematurely.")
	case <-time.After(defaultTestShortTimeout):
	}

	healthChan := make(chan balancer.SubConnState, 1)
	sc.RegisterHealthListener(func(scs balancer.SubConnState) {
		healthChan <- scs
	})

	select {
	case <-hcEnterChan:
	case <-ctx.Done():
		t.Fatalf("Context timed out waiting for health check to begin.")
	}

	for readyReceived := false; !readyReceived; {
		select {
		case scs := <-healthChan:
			readyReceived = scs.ConnectivityState == connectivity.Ready
		case <-ctx.Done():
			t.Fatalf("Context timed out waiting for healthy backend.")
		}
	}

Review comment (suggestion):
This comment should probably move to the block where we register the health listener, right?

And here we should probably comment that we need to verify that health check is not started before registering a listener or something like that.

Code after:
	r.UpdateState(resolver.State{
		Addresses:     []resolver.Address{{Addr: lis.Addr().String()}},
		ServiceConfig: parseServiceConfig(t, r, svcCfg)})

	var sc balancer.SubConn
	select {
	case sc = <-scChan:
	case <-ctx.Done():
		t.Fatal("Context timed out waiting for SubConn creation")
	}

	// Wait for the SubConn to enter READY.
	select {
	case <-readyUpdateReceivedCh:
	case <-ctx.Done():
		t.Fatalf("Context timed out waiting for SubConn to enter READY")
	}

	// Health check should start only after a health listener is registered.
	select {
	case <-hcEnterChan:
		t.Fatalf("Health service client created prematurely.")
	case <-time.After(defaultTestShortTimeout):
	}

	// Register a health listener and verify it receives updates.
	healthChan := make(chan balancer.SubConnState, 1)
	sc.RegisterHealthListener(func(scs balancer.SubConnState) {
		healthChan <- scs
	})

	select {
	case <-hcEnterChan:
	case <-ctx.Done():
		t.Fatalf("Context timed out waiting for health check to begin.")
	}

	for readyReceived := false; !readyReceived; {
		select {
		case scs := <-healthChan:
			t.Logf("Received health update: %v", scs)
			readyReceived = scs.ConnectivityState == connectivity.Ready
		case <-ctx.Done():
			t.Fatalf("Context timed out waiting for healthy backend.")
		}
	}

	// Registering a nil listener should invalidate the previously registered
	// listener and close the health service client.
	sc.RegisterHealthListener(nil)
	select {