From 30df270f98c87ddf836331adc1b20726fd4172d8 Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Tue, 19 Mar 2024 08:38:08 -0700
Subject: [PATCH 1/2] channelz: re-add state for subchannels

---
 clientconn.go         |  2 ++
 test/channelz_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/clientconn.go b/clientconn.go
index 91508fb4e48e..d0116fb3fed7 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -1204,12 +1204,14 @@ func (ac *addrConn) updateConnectivityState(s connectivity.State, lastErr error)
 	close(ac.stateChan)
 	ac.stateChan = make(chan struct{})
 	ac.state = s
+	ac.channelz.ChannelMetrics.State.Store(&s)
 	if lastErr == nil {
 		channelz.Infof(logger, ac.channelz, "Subchannel Connectivity change to %v", s)
 	} else {
 		channelz.Infof(logger, ac.channelz, "Subchannel Connectivity change to %v, last error: %s", s, lastErr)
 	}
 	ac.acbw.updateState(s, lastErr)
+
 }
 
 // adjustParams updates parameters used to create transports upon
diff --git a/test/channelz_test.go b/test/channelz_test.go
index e223c5cce1eb..5e883d5342d3 100644
--- a/test/channelz_test.go
+++ b/test/channelz_test.go
@@ -131,6 +131,43 @@ func (s) TestCZGetChannel(t *testing.T) {
 	}
 }
 
+func (s) TestCZGetSubChannel(t *testing.T) {
+	e := tcpClearRREnv
+	e.balancer = ""
+	te := newTest(t, e)
+	te.startServer(&testServer{security: e.security})
+	r := manual.NewBuilderWithScheme("whatever")
+	addrs := []resolver.Address{{Addr: te.srvAddr}}
+	r.InitialState(resolver.State{Addresses: addrs})
+	te.resolverScheme = r.Scheme()
+	te.clientConn(grpc.WithResolvers(r))
+	defer te.tearDown()
+	if err := verifyResultWithDelay(func() (bool, error) {
+		tcs, _ := channelz.GetTopChannels(0, 0)
+		if len(tcs) != 1 {
+			return false, fmt.Errorf("there should only be one top channel, not %d", len(tcs))
+		}
+		scs := tcs[0].SubChans()
+		if len(scs) != 1 {
+			return false, fmt.Errorf("there should be one subchannel, not %d", len(scs))
+		}
+		var scid int64
+		for scid = range scs {
+		}
+		sc := channelz.GetSubChannel(scid)
+		if sc == nil {
+			return false, fmt.Errorf("subchannel with id %v is nil", scid)
+		}
+		state := sc.ChannelMetrics.State.Load()
+		if state == nil || *state != connectivity.Ready {
+			return false, fmt.Errorf("Got subchannel state=%v; want %q", state, connectivity.Ready)
+		}
+		return true, nil
+	}); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func (s) TestCZGetServer(t *testing.T) {
 	e := tcpClearRREnv
 	te := newTest(t, e)

From 4366e7bba207f7483d2dac67285734ad800674cf Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Tue, 19 Mar 2024 10:47:33 -0700
Subject: [PATCH 2/2] Apply suggestions from code review

Co-authored-by: Arvind Bright <arvind.bright100@gmail.com>
---
 clientconn.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/clientconn.go b/clientconn.go
index d0116fb3fed7..524159a985a8 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -1211,7 +1211,6 @@ func (ac *addrConn) updateConnectivityState(s connectivity.State, lastErr error)
 		channelz.Infof(logger, ac.channelz, "Subchannel Connectivity change to %v, last error: %s", s, lastErr)
 	}
 	ac.acbw.updateState(s, lastErr)
-
 }
 
 // adjustParams updates parameters used to create transports upon