Fix minor issues not covered by PR #9137 (#9147)

This PR is a follow up of already merged PR #9137

Delivers minor fixes for nit comments unaddressed in the previous PR

RELEASE NOTES: N/A

Author: mswierq

balancer/rls/balancer_test.go

@@ -999,6 +999,9 @@ func waitForConnectivityState(ctx context.Context, t *testing.T, ch *buffer.Unbo
 		select {
 		case gotState := <-ch.Get():
 			ch.Load()
+			if gotState == nil {
+				t.Fatalf("Channel closed while waiting for RLS control channel to become one of %v", wants)
+			}
 			got := gotState.(connectivity.State)
 			for _, want := range wants {
 				if got == want {
@@ -1048,8 +1051,10 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	wrappedSubscriber := &wrappingConnectivityStateSubscriber{connStateCh: buffer.NewUnbounded()}
 	origConnectivityStateSubscriber := newConnectivityStateSubscriber
 	newConnectivityStateSubscriber = func(delegate grpcsync.Subscriber) grpcsync.Subscriber {
-		wrappedSubscriber.delegate = delegate
-		return wrappedSubscriber
+		return &wrappingConnectivityStateSubscriber{
+			delegate:    delegate,
+			connStateCh: wrappedSubscriber.connStateCh,
+		}
 	}
 	defer func() { newConnectivityStateSubscriber = origConnectivityStateSubscriber }()
 
@@ -1116,10 +1121,10 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 		// and move it to TRANSIENT_FAILURE.
 		ctxFailed := metadata.AppendToOutgoingContext(ctx, "n1", "v1")
 		makeTestRPCAndVerifyError(ctxFailed, t, cc, codes.Unavailable, nil)
-	}
 
-	// Wait for the control channel to move to TRANSIENT_FAILURE.
-	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.TransientFailure)
+		// Wait for the control channel to move to TRANSIENT_FAILURE.
+		waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.TransientFailure)
+	}
 
 	// Restart the RLS server.
 	lis.Restart()
@@ -1191,8 +1196,10 @@ func (s) TestControlChannelIdleTransitionNoBackoffReset(t *testing.T) {
 	wrappedSubscriber := &wrappingConnectivityStateSubscriber{connStateCh: buffer.NewUnbounded()}
 	origConnectivityStateSubscriber := newConnectivityStateSubscriber
 	newConnectivityStateSubscriber = func(delegate grpcsync.Subscriber) grpcsync.Subscriber {
-		wrappedSubscriber.delegate = delegate
-		return wrappedSubscriber
+		return &wrappingConnectivityStateSubscriber{
+			delegate:    delegate,
+			connStateCh: wrappedSubscriber.connStateCh,
+		}
 	}
 	defer func() { newConnectivityStateSubscriber = origConnectivityStateSubscriber }()
 

balancer/rls/control_channel.go

@@ -62,11 +62,11 @@ type controlChannel struct {
 	// hammering the RLS service while it is overloaded or down.
 	throttler adaptiveThrottler
 
-	cc                   *grpc.ClientConn
-	client               rlsgrpc.RouteLookupServiceClient
-	logger               *internalgrpclog.PrefixLogger
-	unsubscribe          func()
-	seenTransientFailure bool
+	cc                      *grpc.ClientConn
+	client                  rlsgrpc.RouteLookupServiceClient
+	logger                  *internalgrpclog.PrefixLogger
+	dropConnStateSubscriber func()
+	seenTransientFailure    bool
 }
 
 // newControlChannel creates a controlChannel to rlsServerName and uses
@@ -90,7 +90,8 @@ func newControlChannel(rlsServerName, serviceConfig string, rpcTimeout time.Dura
 	}
 	// Subscribe to connectivity state before connecting to avoid missing initial
 	// updates, which are only delivered to active subscribers.
-	ctrlCh.unsubscribe = internal.SubscribeToConnectivityStateChanges.(func(cc *grpc.ClientConn, s grpcsync.Subscriber) func())(ctrlCh.cc, newConnectivityStateSubscriber(ctrlCh))
+	subscribe := internal.SubscribeToConnectivityStateChanges.(func(cc *grpc.ClientConn, s grpcsync.Subscriber) func())
+	ctrlCh.dropConnStateSubscriber = subscribe(ctrlCh.cc, newConnectivityStateSubscriber(ctrlCh))
 	ctrlCh.cc.Connect()
 	ctrlCh.client = rlsgrpc.NewRouteLookupServiceClient(ctrlCh.cc)
 	ctrlCh.logger.Infof("Control channel created to RLS server at: %v", rlsServerName)
@@ -180,7 +181,7 @@ func (cc *controlChannel) dialOpts(bOpts balancer.BuildOptions, serviceConfig st
 }
 
 func (cc *controlChannel) close() {
-	cc.unsubscribe()
+	cc.dropConnStateSubscriber()
 	cc.cc.Close()
 	cc.logger.Infof("Shutdown")
 }