From 58beff180d7a0627144f39d1bdd4286c6cb9ee0a Mon Sep 17 00:00:00 2001
From: Easwar Swaminathan <easwars@google.com>
Date: Mon, 29 Nov 2021 14:28:18 -0800
Subject: [PATCH] balancergroup: add method to exitIdle a sub-balancer (#4994)

This is required for the RLS LB policy. At pick time, if the RLS picker
finds one of its child policies in IDLE, it needs to be able to ask it
to exit idle.
---
 internal/balancergroup/balancergroup.go      | 10 ++++++++
 internal/balancergroup/balancergroup_test.go | 26 ++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/internal/balancergroup/balancergroup.go b/internal/balancergroup/balancergroup.go
index d516f215323..9776158dd98 100644
--- a/internal/balancergroup/balancergroup.go
+++ b/internal/balancergroup/balancergroup.go
@@ -506,3 +506,13 @@ func (bg *BalancerGroup) ExitIdle() {
 	}
 	bg.outgoingMu.Unlock()
 }
+
+// ExitIdleOne instructs the sub-balancer `id` to exit IDLE state, if
+// appropriate and possible.
+func (bg *BalancerGroup) ExitIdleOne(id string) {
+	bg.outgoingMu.Lock()
+	if config := bg.idToBalancerConfig[id]; config != nil {
+		config.exitIdle()
+	}
+	bg.outgoingMu.Unlock()
+}
diff --git a/internal/balancergroup/balancergroup_test.go b/internal/balancergroup/balancergroup_test.go
index ef11e402ec2..4942f8a7da8 100644
--- a/internal/balancergroup/balancergroup_test.go
+++ b/internal/balancergroup/balancergroup_test.go
@@ -509,3 +509,29 @@ func (s) TestBalancerGroupBuildOptions(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func (s) TestBalancerExitIdleOne(t *testing.T) {
+	const balancerName = "stub-balancer-test-balancergroup-exit-idle-one"
+	exitIdleCh := make(chan struct{}, 1)
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(*stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+	cc := testutils.NewTestClientConn(t)
+	bg := New(cc, balancer.BuildOptions{}, nil, nil)
+	bg.Start()
+	defer bg.Close()
+
+	// Add the stub balancer build above as a child policy.
+	builder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], builder)
+
+	// Call ExitIdle on the child policy.
+	bg.ExitIdleOne(testBalancerIDs[0])
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("Timeout when waiting for ExitIdle to be invoked on child policy")
+	case <-exitIdleCh:
+	}
+}