Repository: grpc/grpc-go

Issue:
Flaky Test: Balancer_TwoAddresses_WeightExpiration
https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506

```
--- FAIL: Test (6.19s)
    --- FAIL: Test/Balancer_TwoAddresses_WeightExpiration (1.19s)
        tlogger.go:116: INFO server.go:685 [core] [Server #328]Server created  (t=+1.018725ms)
        tlogger.go:116: INFO server.go:685 [core] [Server #330]Server created  (t=+1.759362ms)
        balancer_test.go:625: Marshaled service config: {"loadBalancingConfig": [ {"weighted_round_robin":{"enableOobLoadReport":true,"oobReportingPeriod":"60s","blackoutPeriod":"0s","weightExpirationPeriod":"60s","weightUpdatePeriod":".050s","errorUtilizationPenalty":0}} ] }
        tlogger.go:116: INFO clientconn.go:1678 [core] original dial target is: "whatever:///127.0.0.1:42173"  (t=+2.267587ms)
        tlogger.go:116: INFO server.go:881 [core] [Server #328 ListenSocket #329]ListenSocket created  (t=+2.984119ms)
        tlogger.go:116: INFO server.go:881 [core] [Server #330 ListenSocket #331]ListenSocket created  (t=+3.153685ms)
        tlogger.go:116: INFO clientconn.go:314 [core] [Channel #332]Channel created  (t=+3.606116ms)
        tlogger.go:116: INFO clientconn.go:193 [core] [Channel #332]parsed dial target is: resolver.Target{URL:url.URL{Scheme:"whatever", Opaque:"", User:(*url.Userinfo)(nil), Host:"", Path:"/127.0.0.1:42173", RawPath:"", OmitHost:false, ForceQuery:false, RawQuery:"", Fragment:"", RawFragment:""}}  (t=+3.86649ms)
        tlogger.go:116: INFO clientconn.go:194 [core] [Channel #332]Channel authority set to "127.0.0.1:42173"  (t=+4.005739ms)
        tlogger.go:116: INFO clientconn.go:314 [core] [Channel #332]Channel exiting idle mode  (t=+4.413988ms)
        tlogger.go:116: INFO resolver_wrapper.go:200 [core] [Channel #332]Resolver state updated: {
              "Addresses": [
                {
                  "Addr": "127.0.0.1:42173",
                  "ServerName": "",
                  "Attributes": null,
                  "BalancerAttributes": null,
                  "Metadata": null
                }
              ],
              "Endpoints": [
                {
                  "Addresses": [
                    {
                      "Addr": "127.0.0.1:42173",
                      "ServerName": "",
                      "Attributes": null,
                      "BalancerAttributes": null,
                      "Metadata": null
                    }
                  ],
                  "Attributes": null
                }
              ],
              "ServiceConfig": null,
              "Attributes": null
            } (resolver returned new addresses)  (t=+4.704539ms)
        tlogger.go:116: INFO balancer_wrapper.go:107 [core] [Channel #332]Channel switches to new LB policy "weighted_round_robin"  (t=+4.983848ms)
        tlogger.go:116: INFO balancer.go:105 [weighted-round-robin] [0x40000b8460] Created  (t=+5.168582ms)
        tlogger.go:116: INFO balancer.go:166 [weighted-round-robin] [0x40000b8460] UpdateCCS: {{[{Addr: "127.0.0.1:42173", ServerName: "", }] [{[{Addr: "127.0.0.1:42173", ServerName: "", }] <nil>}] <nil> <nil>} 0x400043e400}  (t=+5.486323ms)
        tlogger.go:116: INFO clientconn.go:842 [core] [Channel #332 SubChannel #333]Subchannel created  (t=+5.655888ms)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #333]Subchannel Connectivity change to CONNECTING  (t=+6.037838ms)
        tlogger.go:116: INFO clientconn.go:1320 [core] [Channel #332 SubChannel #333]Subchannel picks a new address "127.0.0.1:42173" to connect  (t=+6.201913ms)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #333]Subchannel Connectivity change to READY  (t=+7.214086ms)
        tlogger.go:116: INFO clientconn.go:544 [core] [Channel #332]Channel Connectivity change to CONNECTING  (t=+7.529052ms)
        tlogger.go:116: INFO clientconn.go:544 [core] [Channel #332]Channel Connectivity change to READY  (t=+8.971524ms)
        tlogger.go:116: INFO resolver_wrapper.go:200 [core] [Channel #332]Resolver state updated: {
              "Addresses": [
                {
                  "Addr": "127.0.0.1:42173",
                  "ServerName": "",
                  "Attributes": null,
                  "BalancerAttributes": null,
                  "Metadata": null
                },
                {
                  "Addr": "127.0.0.1:41821",
                  "ServerName": "",
                  "Attributes": null,
                  "BalancerAttributes": null,
                  "Metadata": null
                }
              ],
              "Endpoints": [
                {
                  "Addresses": [
                    {
                      "Addr": "127.0.0.1:42173",
                      "ServerName": "",
                      "Attributes": null,
                      "BalancerAttributes": null,
                      "Metadata": null
                    }
                  ],
                  "Attributes": null
                },
                {
                  "Addresses": [
                    {
                      "Addr": "127.0.0.1:41821",
                      "ServerName": "",
                      "Attributes": null,
                      "BalancerAttributes": null,
                      "Metadata": null
                    }
                  ],
                  "Attributes": null
                }
              ],
              "ServiceConfig": null,
              "Attributes": null
            } ()  (t=+9.309473ms)
        tlogger.go:116: INFO balancer.go:166 [weighted-round-robin] [0x40000b8460] UpdateCCS: {{[{Addr: "127.0.0.1:42173", ServerName: "", } {Addr: "127.0.0.1:41821", ServerName: "", }] [{[{Addr: "127.0.0.1:42173", ServerName: "", }] <nil>} {[{Addr: "127.0.0.1:41821", ServerName: "", }] <nil>}] <nil> <nil>} 0x400043e400}  (t=+9.711981ms)
        tlogger.go:116: INFO clientconn.go:842 [core] [Channel #332 SubChannel #336]Subchannel created  (t=+9.874784ms)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #336]Subchannel Connectivity change to CONNECTING  (t=+34.939867ms)
        tlogger.go:116: INFO clientconn.go:1320 [core] [Channel #332 SubChannel #336]Subchannel picks a new address "127.0.0.1:41821" to connect  (t=+35.241417ms)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #336]Subchannel Connectivity change to READY  (t=+38.[113](https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506#step:8:114)778ms)
        balancer_test.go:637: Failed iteration 0; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 1; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 2; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 3; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 4; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 5; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 6; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 7; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 8; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed iteration 9; counts: map[127.0.0.1:41821:50 127.0.0.1:42173:50]; want [{srv:0x4000598780 w:9} {srv:0x4000598d50 w:90}]
        balancer_test.go:637: Failed to route RPCs with proper ratio
        tlogger.go:[116](https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506#step:8:117): INFO server.go:817 [core] [Server #330 ListenSocket #331]ListenSocket deleted  (t=+1.177759206s)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #336]Subchannel Connectivity change to IDLE  (t=+1.178468345s)
        tlogger.go:116: INFO clientconn.go:544 [core] [Channel #332]Channel Connectivity change to SHUTDOWN  (t=+1.178716236s)
        tlogger.go:116: INFO resolver_wrapper.go:100 [core] [Channel #332]Closing the name resolver  (t=+1.178940063s)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #336]Subchannel Connectivity change to CONNECTING  (t=+1.179136017s)
        tlogger.go:116: INFO balancer_wrapper.go:145 [core] [Channel #332]ccBalancerWrapper: closing  (t=+1.179310612s)
        tlogger.go:116: INFO clientconn.go:1320 [core] [Channel #332 SubChannel #336]Subchannel picks a new address "127.0.0.1:41821" to connect  (t=+1.179727447s)
        tlogger.go:116: INFO clientconn.go:[120](https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506#step:8:121)2 [core] [Channel #332 SubChannel #333]Subchannel Connectivity change to SHUTDOWN  (t=+1.180289782s)
        tlogger.go:116: INFO clientconn.go:1551 [core] [Channel #332 SubChannel #333]Subchannel deleted  (t=+1.180489033s)
        tlogger.go:116: INFO clientconn.go:1202 [core] [Channel #332 SubChannel #336]Subchannel Connectivity change to SHUTDOWN  (t=+1.180963806s)
        tlogger.go:116: INFO clientconn.go:1551 [core] [Channel #332 SubChannel #336]Subchannel deleted  (t=+1.181167394s)
        tlogger.go:116: INFO clientconn.go:314 [core] [Channel #332]Channel deleted  (t=+1.181333333s)
        tlogger.go:116: INFO server.go:817 [core] [Server #328 ListenSocket #329]ListenSocket deleted  (t=+1.181606[121](https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506#step:8:122)s)
        tlogger.go:116: WARNING clientconn.go:1382 [core] [Channel #332 SubChannel #336]grpc: addrConn.createTransport failed to connect to {Addr: "127.0.0.1:41821", ServerName: "127.0.0.1:42173", }. Err: connection error: desc = "transport: Error while dialing: dial tcp [127](https://github.com/grpc/grpc-go/actions/runs/10377601227/job/28732125169?pr=7506#step:8:128).0.0.1:41821: operation was canceled"  (t=+1.184049915s)
FAIL
FAIL	google.golang.org/grpc/balancer/weightedroundrobin	11.550s
```

Discussion:
@easwars: https://github.com/grpc/grpc-go/actions/runs/10394308167/job/28783810539?pr=7498@purnesh42H: https://github.com/grpc/grpc-go/actions/runs/10465483190/job/28980791860?pr=7462@easwars: https://github.com/grpc/grpc-go/actions/runs/10685201154/job/29617436939?pr=7579@dfawley: I believe this was supposed to be fixed by #7523.  Maybe that wasn't effective enough?@zasweq: https://github.com/grpc/grpc-go/actions/runs/10855329479/job/30127706645?pr=7614 argh this shows up on a branch with  the fix merged. So it still persists.@purnesh42H: @eshitachandwani to forge test and confirm the flakiness percentage@purnesh42H: https://github.com/grpc/grpc-go/actions/runs/10858951486/job/30137631363?pr=7609@eshitachandwani: Failed 0/100k (in all 7 iterations) test runs on forge@zasweq: https://github.com/grpc/grpc-go/actions/runs/11524690877/job/32085409887

Solution patch:
diff --git a/Documentation/encoding.md b/Documentation/encoding.md
index 64358f10..cce9384f 100644
--- a/Documentation/encoding.md
+++ b/Documentation/encoding.md
@@ -22,7 +22,7 @@ package proto
 import "google.golang.org/grpc/encoding"
 
 func init() {
-	encoding.RegisterCodec(protoCodec{})
+    encoding.RegisterCodec(protoCodec{})
 }
 
 // ... implementation of protoCodec ...
@@ -50,14 +50,14 @@ On the client-side, to specify a `Codec` to use for message transmission, the
 `CallOption` `CallContentSubtype` should be used as follows:
 
 ```go
-	response, err := myclient.MyCall(ctx, request, grpc.CallContentSubtype("mycodec"))
+    response, err := myclient.MyCall(ctx, request, grpc.CallContentSubtype("mycodec"))
 ```
 
 As a reminder, all `CallOption`s may be converted into `DialOption`s that become
 the default for all RPCs sent through a client using `grpc.WithDefaultCallOptions`:
 
 ```go
-	myclient := grpc.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mycodec")))
+    myclient := grpc.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mycodec")))
 ```
 
 When specified in either of these ways, messages will be encoded using this
@@ -98,7 +98,7 @@ package gzip
 import "google.golang.org/grpc/encoding"
 
 func init() {
-	encoding.RegisterCompressor(compressor{})
+    encoding.RegisterCompressor(compressor{})
 }
 
 // ... implementation of compressor ...
@@ -125,14 +125,14 @@ On the client-side, to specify a `Compressor` to use for message transmission,
 the `CallOption` `UseCompressor` should be used as follows:
 
 ```go
-	response, err := myclient.MyCall(ctx, request, grpc.UseCompressor("gzip"))
+    response, err := myclient.MyCall(ctx, request, grpc.UseCompressor("gzip"))
 ```
 
 As a reminder, all `CallOption`s may be converted into `DialOption`s that become
 the default for all RPCs sent through a client using `grpc.WithDefaultCallOptions`:
 
 ```go
-	myclient := grpc.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
+    myclient := grpc.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
 ```
 
 When specified in either of these ways, messages will be compressed using this
diff --git a/balancer_wrapper.go b/balancer_wrapper.go
index 8ad6ce2f..efdbe7cf 100644
--- a/balancer_wrapper.go
+++ b/balancer_wrapper.go
@@ -262,7 +262,7 @@ type acBalancerWrapper struct {
 
 // updateState is invoked by grpc to push a subConn state update to the
 // underlying balancer.
-func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolver.Address, err error) {
+func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolver.Address, err error, readyChan chan struct{}) {
 	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
 		if ctx.Err() != nil || acbw.ccb.balancer == nil {
 			return
@@ -278,12 +278,11 @@ func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolve
 		acbw.ac.mu.Lock()
 		defer acbw.ac.mu.Unlock()
 		if s == connectivity.Ready {
-			// When changing states to READY, reset stateReadyChan.  Wait until
+			// When changing states to READY, close stateReadyChan.  Wait until
 			// after we notify the LB policy's listener(s) in order to prevent
 			// ac.getTransport() from unblocking before the LB policy starts
 			// tracking the subchannel as READY.
-			close(acbw.ac.stateReadyChan)
-			acbw.ac.stateReadyChan = make(chan struct{})
+			close(readyChan)
 		}
 	})
 }
diff --git a/clientconn.go b/clientconn.go
index e5615e71..a680fefc 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -1193,6 +1193,14 @@ func (ac *addrConn) updateConnectivityState(s connectivity.State, lastErr error)
 	if ac.state == s {
 		return
 	}
+	if ac.state == connectivity.Ready {
+		// When leaving ready, re-create the ready channel.
+		ac.stateReadyChan = make(chan struct{})
+	}
+	if s == connectivity.Shutdown {
+		// Wake any producer waiting to create a stream on the transport.
+		close(ac.stateReadyChan)
+	}
 	ac.state = s
 	ac.channelz.ChannelMetrics.State.Store(&s)
 	if lastErr == nil {
@@ -1200,7 +1208,7 @@ func (ac *addrConn) updateConnectivityState(s connectivity.State, lastErr error)
 	} else {
 		channelz.Infof(logger, ac.channelz, "Subchannel Connectivity change to %v, last error: %s", s, lastErr)
 	}
-	ac.acbw.updateState(s, ac.curAddr, lastErr)
+	ac.acbw.updateState(s, ac.curAddr, lastErr, ac.stateReadyChan)
 }
 
 // adjustParams updates parameters used to create transports upon
@@ -1510,18 +1518,20 @@ func (ac *addrConn) getReadyTransport() transport.ClientTransport {
 func (ac *addrConn) getTransport(ctx context.Context) (transport.ClientTransport, error) {
 	for ctx.Err() == nil {
 		ac.mu.Lock()
-		t, state, sc := ac.transport, ac.state, ac.stateReadyChan
+		t, state, readyChan := ac.transport, ac.state, ac.stateReadyChan
 		ac.mu.Unlock()
-		if state == connectivity.Ready {
-			return t, nil
-		}
 		if state == connectivity.Shutdown {
+			// Return an error immediately in only this case since a connection
+			// will never occur.
 			return nil, status.Errorf(codes.Unavailable, "SubConn shutting down")
 		}
 
 		select {
 		case <-ctx.Done():
-		case <-sc:
+		case <-readyChan:
+			if state == connectivity.Ready {
+				return t, nil
+			}
 		}
 	}
 	return nil, status.FromContextError(ctx.Err()).Err()
diff --git a/examples/gotutorial.md b/examples/gotutorial.md
index b32649e6..b1b75468 100644
--- a/examples/gotutorial.md
+++ b/examples/gotutorial.md
@@ -151,13 +151,13 @@ func (s *routeGuideServer) RouteChat(stream pb.RouteGuide_RouteChatServer) error
 
 ```go
 func (s *routeGuideServer) GetFeature(ctx context.Context, point *pb.Point) (*pb.Feature, error) {
-	for _, feature := range s.savedFeatures {
-		if proto.Equal(feature.Location, point) {
-			return feature, nil
-		}
-	}
-	// No feature was found, return an unnamed feature
-	return &pb.Feature{"", point}, nil
+    for _, feature := range s.savedFeatures {
+        if proto.Equal(feature.Location, point) {
+            return feature, nil
+        }
+    }
+    // No feature was found, return an unnamed feature
+    return &pb.Feature{"", point}, nil
 }
 ```
 
@@ -168,14 +168,14 @@ Now let's look at one of our streaming RPCs. `ListFeatures` is a server-side str
 
 ```go
 func (s *routeGuideServer) ListFeatures(rect *pb.Rectangle, stream pb.RouteGuide_ListFeaturesServer) error {
-	for _, feature := range s.savedFeatures {
-		if inRange(feature.Location, rect) {
-			if err := stream.Send(feature); err != nil {
-				return err
-			}
-		}
-	}
-	return nil
+    for _, feature := range s.savedFeatures {
+        if inRange(feature.Location, rect) {
+            if err := stream.Send(feature); err != nil {
+                return err
+            }
+        }
+    }
+    return nil
 }
 ```
 
@@ -188,34 +188,34 @@ Now let's look at something a little more complicated: the client-side streaming
 
 ```go
 func (s *routeGuideServer) RecordRoute(stream pb.RouteGuide_RecordRouteServer) error {
-	var pointCount, featureCount, distance int32
-	var lastPoint *pb.Point
-	startTime := time.Now()
-	for {
-		point, err := stream.Recv()
-		if err == io.EOF {
-			endTime := time.Now()
-			return stream.SendAndClose(&pb.RouteSummary{
-				PointCount:   pointCount,
-				FeatureCount: featureCount,
-				Distance:     distance,
-				ElapsedTime:  int32(endTime.Sub(startTime).Seconds()),
-			})
-		}
-		if err != nil {
-			return err
-		}
-		pointCount++
-		for _, feature := range s.savedFeatures {
-			if proto.Equal(feature.Location, point) {
-				featureCount++
-			}
-		}
-		if lastPoint != nil {
-			distance += calcDistance(lastPoint, point)
-		}
-		lastPoint = point
-	}
+    var pointCount, featureCount, distance int32
+    var lastPoint *pb.Point
+    startTime := time.Now()
+    for {
+        point, err := stream.Recv()
+        if err == io.EOF {
+            endTime := time.Now()
+            return stream.SendAndClose(&pb.RouteSummary{
+                PointCount:   pointCount,
+                FeatureCount: featureCount,
+                Distance:     distance,
+                ElapsedTime:  int32(endTime.Sub(startTime).Seconds()),
+            })
+        }
+        if err != nil {
+            return err
+        }
+        pointCount++
+        for _, feature := range s.savedFeatures {
+            if proto.Equal(feature.Location, point) {
+                featureCount++
+            }
+        }
+        if lastPoint != nil {
+            distance += calcDistance(lastPoint, point)
+        }
+        lastPoint = point
+    }
 }
 ```
 
@@ -226,22 +226,22 @@ Finally, let's look at our bidirectional streaming RPC `RouteChat()`.
 
 ```go
 func (s *routeGuideServer) RouteChat(stream pb.RouteGuide_RouteChatServer) error {
-	for {
-		in, err := stream.Recv()
-		if err == io.EOF {
-			return nil
-		}
-		if err != nil {
-			return err
-		}
-		key := serialize(in.Location)
+    for {
+        in, err := stream.Recv()
+        if err == io.EOF {
+            return nil
+        }
+        if err != nil {
+            return err
+        }
+        key := serialize(in.Location)
                 ... // look for notes to be sent to client
-		for _, note := range s.routeNotes[key] {
-			if err := stream.Send(note); err != nil {
-				return err
-			}
-		}
-	}
+        for _, note := range s.routeNotes[key] {
+            if err := stream.Send(note); err != nil {
+                return err
+            }
+        }
+    }
 }
 ```
 
@@ -353,21 +353,21 @@ r := rand.New(rand.NewSource(time.Now().UnixNano()))
 pointCount := int(r.Int31n(100)) + 2 // Traverse at least two points
 var points []*pb.Point
 for i := 0; i < pointCount; i++ {
-	points = append(points, randomPoint(r))
+    points = append(points, randomPoint(r))
 }
 log.Printf("Traversing %d points.", len(points))
 stream, err := client.RecordRoute(ctx)
 if err != nil {
-	log.Fatalf("%v.RecordRoute(_) = _, %v", client, err)
+    log.Fatalf("%v.RecordRoute(_) = _, %v", client, err)
 }
 for _, point := range points {
-	if err := stream.Send(point); err != nil {
-		log.Fatalf("%v.Send(%v) = %v", stream, point, err)
-	}
+    if err := stream.Send(point); err != nil {
+        log.Fatalf("%v.Send(%v) = %v", stream, point, err)
+    }
 }
 reply, err := stream.CloseAndRecv()
 if err != nil {
-	log.Fatalf("%v.CloseAndRecv() got error %v, want %v", stream, err, nil)
+    log.Fatalf("%v.CloseAndRecv() got error %v, want %v", stream, err, nil)
 }
 log.Printf("Route summary: %v", reply)
 ```
@@ -382,23 +382,23 @@ Finally, let's look at our bidirectional streaming RPC `RouteChat()`. As in the
 stream, err := client.RouteChat(ctx)
 waitc := make(chan struct{})
 go func() {
-	for {
-		in, err := stream.Recv()
-		if err == io.EOF {
-			// read done.
-			close(waitc)
-			return
-		}
-		if err != nil {
-			log.Fatalf("Failed to receive a note : %v", err)
-		}
-		log.Printf("Got message %s at point(%d, %d)", in.Message, in.Location.Latitude, in.Location.Longitude)
-	}
+    for {
+        in, err := stream.Recv()
+        if err == io.EOF {
+            // read done.
+            close(waitc)
+            return
+        }
+        if err != nil {
+            log.Fatalf("Failed to receive a note : %v", err)
+        }
+        log.Printf("Got message %s at point(%d, %d)", in.Message, in.Location.Latitude, in.Location.Longitude)
+    }
 }()
 for _, note := range notes {
-	if err := stream.Send(note); err != nil {
-		log.Fatalf("Failed to send a note: %v", err)
-	}
+    if err := stream.Send(note); err != nil {
+        log.Fatalf("Failed to send a note: %v", err)
+    }
 }
 stream.CloseSend()
 <-waitc



Test patch:
diff --git a/producer_ext_test.go b/producer_ext_test.go
index 2da36d8d..628a1185 100644
--- a/producer_ext_test.go
+++ b/producer_ext_test.go
@@ -51,7 +51,7 @@ func (*producerBuilder) Build(cci any) (balancer.Producer, func()) {
 	}
 }
 
-func (p *producer) TestStreamStart(t *testing.T, streamStarted chan<- struct{}) {
+func (p *producer) testStreamStart(t *testing.T, streamStarted chan<- struct{}) {
 	go func() {
 		defer close(p.stopped)
 		ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
@@ -68,8 +68,11 @@ var producerBuilderSingleton = &producerBuilder{}
 // TestProducerStreamStartsAfterReady ensures producer streams only start after
 // the subchannel reports as READY to the LB policy.
 func (s) TestProducerStreamStartsAfterReady(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
 	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
 	producerCh := make(chan balancer.Producer)
+	var producerClose func()
 	streamStarted := make(chan struct{})
 	done := make(chan struct{})
 	bf := stub.BalancerFuncs{
@@ -90,7 +93,8 @@ func (s) TestProducerStreamStartsAfterReady(t *testing.T) {
 			if err != nil {
 				return err
 			}
-			producer, _ := sc.GetOrBuildProducer(producerBuilderSingleton)
+			var producer balancer.Producer
+			producer, producerClose = sc.GetOrBuildProducer(producerBuilderSingleton)
 			producerCh <- producer
 			sc.Connect()
 			return nil
@@ -119,7 +123,15 @@ func (s) TestProducerStreamStartsAfterReady(t *testing.T) {
 
 	go cc.Connect()
 	p := <-producerCh
-	p.(*producer).TestStreamStart(t, streamStarted)
+	p.(*producer).testStreamStart(t, streamStarted)
 
-	<-done
+	select {
+	case <-done:
+		// Wait for the stream to start before exiting; otherwise the ClientConn
+		// will close and cause stream creation to fail.
+		<-streamStarted
+		producerClose()
+	case <-ctx.Done():
+		t.Error("Timed out waiting for test to complete")
+	}
 }