priority: implement changes to child policy caching specified in A115 (grpc#8997)

Summary of changes:
- Add a new API to the `balancergroup` that allows removal of child
policies with immediate effect, even when a cache is configured.
- Make the `priority` LB use this new API when removing child policies
that have been removed from its configuration.
- This change is protected with an environment variable that turns on
the new behavior by default.
- Update the name generation algorithm to give higher preference to
reuse names of existing priority (if there is a locality match) over
lower priorities.
- Fix a bunch of tests in the `priority` LB policy to ensure that they
work with the *old* and *new* behavior.
- Fix a couple of aggregate cluster tests such that they don't
continuously handle the same NACKed resource from the management server.
- Other minor cleanups in tests that I had to look at as part of making
this change.

RELEASE NOTES:
* xds/priority: Stop caching child LB policies removed from the
configuration. This will help reduce memory and cpu usage when
localities are constantly switching between priorities.

Author: goingforstudying-ctrl

balancer/weightedtarget/weightedtarget_test.go

@@ -1348,7 +1348,7 @@ func addressesToAddrs(as []resolver.Address) []string {
 	return ret
 }
 
-const initIdleBalancerName = "test-init-Idle-balancer"
+const initIdleBalancerName = "test-init-idle-balancer"
 
 var errTestInitIdle = fmt.Errorf("init Idle balancer error 0")
 
@@ -1388,7 +1388,7 @@ func (s) TestInitialIdle(t *testing.T) {
   "targets": {
     "cluster_1": {
       "weight":1,
-      "childPolicy": [{"test-init-Idle-balancer": ""}]
+      "childPolicy": [{"test-init-idle-balancer": ""}]
     }
   }
 }`))

internal/balancergroup/balancergroup.go

@@ -338,6 +338,16 @@ func (bg *BalancerGroup) Add(id string, builder balancer.Builder) {
 // closed after timeout. Cleanup work (closing sub-balancer and removing
 // subconns) will be done after timeout.
 func (bg *BalancerGroup) Remove(id string) {
+	bg.removeInternal(id, true)
+}
+
+// RemoveImmediately removes and closes the balancer with id from the group
+// immediately.
+func (bg *BalancerGroup) RemoveImmediately(id string) {
+	bg.removeInternal(id, false)
+}
+
+func (bg *BalancerGroup) removeInternal(id string, withCaching bool) {
 	bg.logger.Infof("Removing child policy for child %q", id)
 
 	bg.outgoingMu.Lock()
@@ -356,32 +366,40 @@ func (bg *BalancerGroup) Remove(id string) {
 	// Unconditionally remove the sub-balancer config from the map.
 	delete(bg.idToBalancerConfig, id)
 
-	if bg.deletedBalancerCache != nil {
-		if bg.logger.V(2) {
-			bg.logger.Infof("Adding child policy for child %q to the balancer cache", id)
-			bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
-		}
-
-		bg.deletedBalancerCache.Add(id, sbToRemove, func() {
+	if withCaching {
+		if bg.deletedBalancerCache != nil {
+			if bg.logger.V(2) {
+				bg.logger.Infof("Adding child policy for child %q to the balancer cache", id)
+			}
+			bg.deletedBalancerCache.Add(id, sbToRemove, func() {
+				if bg.logger.V(2) {
+					bg.logger.Infof("Removing child policy for child %q from the balancer cache after timeout", id)
+					bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+				}
+
+				// A sub-balancer evicted from the timeout cache needs to closed
+				// and its subConns need to removed, unconditionally. There is a
+				// possibility that a sub-balancer might be removed (thereby
+				// moving it to the cache) around the same time that the
+				// balancergroup is closed, and by the time we get here the
+				// balancergroup might be closed.  Check for `outgoingStarted ==
+				// true` at that point can lead to a leaked sub-balancer.
+				bg.outgoingMu.Lock()
+				sbToRemove.stopBalancer()
+				bg.outgoingMu.Unlock()
+				bg.cleanupSubConns(sbToRemove)
+			})
 			if bg.logger.V(2) {
-				bg.logger.Infof("Removing child policy for child %q from the balancer cache after timeout", id)
 				bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
 			}
-
-			// A sub-balancer evicted from the timeout cache needs to closed
-			// and its subConns need to removed, unconditionally. There is a
-			// possibility that a sub-balancer might be removed (thereby
-			// moving it to the cache) around the same time that the
-			// balancergroup is closed, and by the time we get here the
-			// balancergroup might be closed.  Check for `outgoingStarted ==
-			// true` at that point can lead to a leaked sub-balancer.
-			bg.outgoingMu.Lock()
-			sbToRemove.stopBalancer()
 			bg.outgoingMu.Unlock()
-			bg.cleanupSubConns(sbToRemove)
-		})
-		bg.outgoingMu.Unlock()
-		return
+			return
+		}
+
+		// Fall through to remove the sub-balancer with immediate effect if we are not caching.
+		if bg.logger.V(2) {
+			bg.logger.Infof("Child policy for child %q was requested to be cached before eventual removal. No such cache exists. Removing right away.", id)
+		}
 	}
 
 	// Remove the sub-balancer with immediate effect if we are not caching.
@@ -481,7 +499,7 @@ func (bg *BalancerGroup) ResolverError(err error) {
 // from map. Delete sc from the map only when state changes to Shutdown. Since
 // it's just forwarding the action, there's no need for a removeSubConn()
 // wrapper function.
-func (bg *BalancerGroup) newSubConn(config *subBalancerWrapper, addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+func (bg *BalancerGroup) newSubConn(sbw *subBalancerWrapper, addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
 	// NOTE: if balancer with id was already removed, this should also return
 	// error. But since we call balancer.stopBalancer when removing the balancer, this
 	// shouldn't happen.
@@ -493,12 +511,12 @@ func (bg *BalancerGroup) newSubConn(config *subBalancerWrapper, addrs []resolver
 	var sc balancer.SubConn
 	oldListener := opts.StateListener
 	opts.StateListener = func(state balancer.SubConnState) { bg.updateSubConnState(sc, state, oldListener) }
-	sc, err := bg.cc.NewSubConn(addrs, opts)
+	sc, err := sbw.ClientConn.NewSubConn(addrs, opts)
 	if err != nil {
 		bg.incomingMu.Unlock()
 		return nil, err
 	}
-	bg.scToSubBalancer[sc] = config
+	bg.scToSubBalancer[sc] = sbw
 	bg.incomingMu.Unlock()
 	return sc, nil
 }

internal/balancergroup/balancergroup_ext_test.go

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package balancergroup_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+// Tests verifies that the RemoveImmediately method of balancergroup closes the
+// child balancer immediately, while the Remove method does not close the child
+// balancer immediately, when the cache is enabled.
+func (s) TestBalancerGroup_RemoveImmediately(t *testing.T) {
+	// Channels to track child balancer creation and closure.
+	childLBCreated := make(chan string, 1)
+	childLBClosed := make(chan string, 1)
+
+	childLBName1 := strings.ToLower(t.Name()) + "-child-1"
+	t.Logf("Registering a child balancer with name %q", childLBName1)
+	stub.Register(childLBName1, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName1
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName1
+		},
+	})
+
+	childLBName2 := strings.ToLower(t.Name()) + "-child-2"
+	t.Logf("Registering a child balancer with name %q", childLBName2)
+	stub.Register(childLBName2, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName2
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName2
+		},
+	})
+
+	t.Logf("Creating a balancergroup with cache enabled")
+	tcc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      tcc,
+		BuildOpts:               balancer.BuildOptions{},
+		SubBalancerCloseTimeout: defaultTestTimeout,
+	})
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-1")
+	bg.AddWithClientConn("child-1", childLBName1, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-2")
+	bg.AddWithClientConn("child-2", childLBName2, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with immediate effect", "child-1")
+	bg.RemoveImmediately("child-1")
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with caching", "child-2")
+	bg.Remove("child-2")
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-childLBClosed:
+		t.Fatalf("Child LB closed when expected to be cached")
+	case <-sCtx.Done():
+	}
+
+	t.Logf("Closing the balancergroup, which should close the second child balancer")
+	bg.Close()
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+}

internal/envconfig/envconfig.go

@@ -119,6 +119,13 @@ var (
 	// A future release will remove this environment variable, enabling strict
 	// path checking behavior unconditionally.
 	DisableStrictPathChecking = boolFromEnv("GRPC_GO_EXPERIMENTAL_DISABLE_STRICT_PATH_CHECKING", false)
+
+	// EnablePriorityLBChildPolicyCache controls whether the priority balancer
+	// should cache child balancers that are removed from the LB policy config,
+	// for a period of 15 minutes. This is disabled by default, but can be
+	// enabled by setting the env variable
+	// GRPC_EXPERIMENTAL_ENABLE_PRIORITY_LB_CHILD_POLICY_CACHE to true.
+	EnablePriorityLBChildPolicyCache = boolFromEnv("GRPC_EXPERIMENTAL_ENABLE_PRIORITY_LB_CHILD_POLICY_CACHE", false)
 )
 
 func boolFromEnv(envVar string, def bool) bool {

internal/xds/balancer/cdsbalancer/configbuilder_childname.go

@@ -31,9 +31,10 @@ import (
 // struct keeps state between generate() calls, and a later generate() might
 // return names returned by the previous call.
 type nameGenerator struct {
-	existingNames map[clients.Locality]string
-	prefix        uint64
-	nextID        uint64
+	prevLocalitiesToChildNames map[clients.Locality]string // locality to child name mapping generated for the previous update
+	prevChildNames             []string                    // prioritized list of child names generated for the previous update
+	prefix                     uint64
+	nextID                     uint64
 }
 
 func newNameGenerator(prefix uint64) *nameGenerator {
@@ -53,36 +54,58 @@ func newNameGenerator(prefix uint64) *nameGenerator {
 // - update 3: [[L1, L2], [L3]] --> ["0", "2"]   (Two priorities were merged)
 // - update 4: [[L1], [L4]] --> ["0", "3",]      (A priority was split, and a new priority was added)
 func (ng *nameGenerator) generate(priorities [][]xdsresource.Locality) []string {
-	var ret []string
+	ret := make([]string, len(priorities))
 	usedNames := make(map[string]bool)
 	newNames := make(map[clients.Locality]string)
-	for _, priority := range priorities {
-		var nameFound string
+
+	// Pass 1: Same priority index match.
+	for i, priority := range priorities {
+		if i >= len(ng.prevChildNames) {
+			continue
+		}
+		targetName := ng.prevChildNames[i]
+		for _, locality := range priority {
+			if name, ok := ng.prevLocalitiesToChildNames[locality.ID]; ok && name == targetName {
+				ret[i] = targetName
+				usedNames[targetName] = true
+				break
+			}
+		}
+	}
+
+	// Pass 2: Greedy reuse.
+	for i, priority := range priorities {
+		if ret[i] != "" {
+			continue
+		}
 		for _, locality := range priority {
-			if name, ok := ng.existingNames[locality.ID]; ok {
+			if name, ok := ng.prevLocalitiesToChildNames[locality.ID]; ok {
 				if !usedNames[name] {
-					nameFound = name
-					// Found a name to use. No need to process the remaining
-					// localities.
+					ret[i] = name
+					usedNames[name] = true
 					break
 				}
 			}
 		}
+	}
 
-		if nameFound == "" {
-			// No appropriate used name is found. Make a new name.
-			nameFound = fmt.Sprintf("priority-%d-%d", ng.prefix, ng.nextID)
+	// Pass 3: New name.
+	for i, name := range ret {
+		if name == "" {
+			newID := fmt.Sprintf("priority-%d-%d", ng.prefix, ng.nextID)
 			ng.nextID++
+			ret[i] = newID
+			usedNames[newID] = true
 		}
+	}
 
-		ret = append(ret, nameFound)
-		// All localities in this priority share the same name. Add them all to
-		// the new map.
+	// Update state.
+	for i, priority := range priorities {
 		for _, l := range priority {
-			newNames[l.ID] = nameFound
+			newNames[l.ID] = ret[i]
 		}
-		usedNames[nameFound] = true
 	}
-	ng.existingNames = newNames
+	ng.prevLocalitiesToChildNames = newNames
+	ng.prevChildNames = ret
 	return ret
 }