xdsresolver: use the refcounted utility for cluster refcounting

internal/xds/resolver/cluster_specifier_plugin_test.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/google/uuid"
@@ -342,6 +343,103 @@ func (s) TestXDSResolverDelayedOnCommittedCSP(t *testing.T) {
 	verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
 }
 
+// TestResolverClusterSpecifierPluginRefCountRace verifies that a cluster
+// specifier plugin is handled correctly when its last in-flight RPC is
+// committed at the same time as an xDS update that names it again. Whichever
+// happens first, the plugin must end up present in the service config: if the
+// commit lands first the entry is torn down and a fresh one replaces it, and if
+// the update lands first the existing entry is simply reused.
+func (s) TestResolverClusterSpecifierPluginRefCountRace(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	nodeID := uuid.New().String()
+	mgmtServer, _, _, bc := setupManagementServerForTest(t, nodeID)
+
+	routeConfigForPlugin := func(name, value string) []*v3routepb.RouteConfiguration {
+		return []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
+			RouteConfigName:              defaultTestRouteConfigName,
+			ListenerName:                 defaultTestServiceName,
+			ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
+			ClusterSpecifierPluginName:   name,
+			ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: value}),
+		})}
+	}
+	wantConfigForPlugin := func(name, value string) string {
+		return fmt.Sprintf(`{
+			"loadBalancingConfig": [{
+				"xds_cluster_manager_experimental": {
+					"children": {
+						"cluster_specifier_plugin:%s": {
+							"childPolicy": [{"csp_experimental": {"arbitrary_field": "%s"}}]
+						}
+					}
+				}
+			}]
+		}`, name, value)
+	}
+
+	listeners := []*v3listenerpb.Listener{e2e.DefaultClientListener(defaultTestServiceName, defaultTestRouteConfigName)}
+	configureResources(ctx, t, mgmtServer, nodeID, listeners, routeConfigForPlugin("cspA", "anythingA"), nil, nil)
+
+	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)
+	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantConfigForPlugin("cspA", "anythingA"))
+
+	// Start an RPC on cspA and leave it uncommitted, so cspA stays referenced.
+	res, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
+	if err != nil {
+		t.Fatalf("cs.SelectConfig(): %v", err)
+	}
+	if got, want := clustermanager.PickedCluster(res.Context), "cluster_specifier_plugin:cspA"; got != want {
+		t.Fatalf("Config selector returned cluster %q, want %q", got, want)
+	}
+
+	// Move the route to cspB. cspA is now held only by the in-flight RPC, so
+	// both plugins stay in the service config. Waiting for that config matters:
+	// it confirms the client processed this update before the race below, and
+	// without it the next update could coalesce with this one and leave the
+	// route config unchanged from the client's point of view.
+	configureResources(ctx, t, mgmtServer, nodeID, listeners, routeConfigForPlugin("cspB", "anythingB"), nil, nil)
+	verifyUpdateFromResolver(ctx, t, stateCh, `{
+		"loadBalancingConfig": [{
+			"xds_cluster_manager_experimental": {
+				"children": {
+					"cluster_specifier_plugin:cspA": {
+						"childPolicy": [{"csp_experimental": {"arbitrary_field": "anythingA"}}]
+					},
+					"cluster_specifier_plugin:cspB": {
+						"childPolicy": [{"csp_experimental": {"arbitrary_field": "anythingB"}}]
+					}
+				}
+			}
+		}]
+	}`)
+
+	// Commit the RPC, dropping cspA's last reference, while an update naming
+	// cspA again is pushed concurrently. The two orderings exercise different
+	// paths through the refcounted entry, and both must converge on cspA being
+	// in the service config.
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		res.OnCommitted()
+	}()
+	configureResources(ctx, t, mgmtServer, nodeID, listeners, routeConfigForPlugin("cspA", "anythingA"), nil, nil)
+	wg.Wait()
+
+	cs = waitForServiceConfig(ctx, t, stateCh, wantConfigForPlugin("cspA", "anythingA"))
+
+	// The surviving entry must still be usable for new RPCs.
+	res, err = cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
+	if err != nil {
+		t.Fatalf("cs.SelectConfig() after the race: %v", err)
+	}
+	if got, want := clustermanager.PickedCluster(res.Context), "cluster_specifier_plugin:cspA"; got != want {
+		t.Fatalf("Config selector returned cluster %q, want %q", got, want)
+	}
+	res.OnCommitted()
+}
+
 // TestResolverClusterSpecifierPlugin_WithFilters tests the case where a route
 // configuration containing cluster specifier plugins is sent by the management
 // server, and HTTP filters are configured. The test verifies that the

internal/xds/resolver/helpers_test.go

@@ -173,6 +173,36 @@ func verifyUpdateFromResolver(ctx context.Context, t *testing.T, stateCh chan re
 	return cs
 }
 
+// waitForServiceConfig drains updates from the resolver until one carries a
+// service config matching wantSC, and fails if none does before ctx expires.
+// Use this instead of verifyUpdateFromResolver when the resolver is expected to
+// publish intermediate configs on the way to the wanted one.
+//
+// Returns the config selector from the matching update.
+func waitForServiceConfig(ctx context.Context, t *testing.T, stateCh chan resolver.State, wantSC string) iresolver.ConfigSelector {
+	t.Helper()
+
+	want := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(wantSC)
+	for {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("Timeout waiting for the resolver to publish service config:\n%s", wantSC)
+		case state := <-stateCh:
+			if err := state.ServiceConfig.Err; err != nil {
+				t.Fatalf("Received error in service config: %v", err)
+			}
+			if !internal.EqualServiceConfigForTesting(state.ServiceConfig.Config, want.Config) {
+				continue
+			}
+			cs := iresolver.GetConfigSelector(state)
+			if cs == nil {
+				t.Fatal("Received nil config selector in update from resolver")
+			}
+			return cs
+		}
+	}
+}
+
 // verifyNoUpdateFromResolver verifies that no update is pushed on stateCh.
 // Calls t.Fatal() if an update is received before defaultTestShortTimeout
 // expires.

internal/xds/resolver/serviceconfig.go

@@ -75,14 +75,14 @@ type xdsClusterManagerConfig struct {
 // serviceConfigJSON produces a service config in JSON format that contains LB
 // policy config for the "xds_cluster_manager" LB policy, with entries in the
 // children map for all active clusters.
-func serviceConfigJSON(activeClusters map[string]*clusterInfo, activePlugins map[string]*clusterInfo) []byte {
+func serviceConfigJSON(activeClusters, activePlugins map[string]*grpcsync.RefCounted[*clusterInfo]) []byte {
 	// Generate children (all entries in activeClusters).
 	children := make(map[string]xdsChildConfig)
 	for cluster, ci := range activeClusters {
-		children[cluster] = ci.cfg
+		children[cluster] = ci.Value().cfg
 	}
 	for plugin, ci := range activePlugins {
-		children[plugin] = ci.cfg
+		children[plugin] = ci.Value().cfg
 	}
 
 	sc := serviceConfig{
@@ -109,6 +109,10 @@ type virtualHost struct {
 type routeCluster struct {
 	name        string                       // Name of the cluster.
 	interceptor httpfilter.ClientInterceptor // HTTP filters to run for RPCs matching this route.
+	// info is the resolver-wide entry for this cluster, shared by every route
+	// that references it. An RPC routed here holds a reference on it until the
+	// RPC is committed.
+	info *grpcsync.RefCounted[*clusterInfo]
 }
 
 type route struct {
@@ -159,8 +163,8 @@ type configSelector struct {
 	// Configuration received from the xDS management server.
 	virtualHost      virtualHost
 	routes           []route
-	clusters         map[string]*clusterInfo
-	plugins          map[string]*clusterInfo
+	clusters         map[string]*grpcsync.RefCounted[*clusterInfo]
+	plugins          map[string]*grpcsync.RefCounted[*clusterInfo]
 	httpFilterConfig []xdsresource.HTTPFilter
 	xdsConfig        *xdsresource.XDSConfig
 }
@@ -223,43 +227,16 @@ func (cs *configSelector) SelectConfig(rpcInfo iresolver.RPCInfo) (*iresolver.RP
 	// Add a ref to the selected cluster to keep the interceptors alive until RPC
 	// is committed.
 	rc.Increment()
-	if info, ok := cs.clusters[cluster.name]; ok {
-		// Add a ref to the selected cluster, as this RPC needs this
-		// cluster until it is committed.
-		info.refCount.Add(1)
-		config.OnCommitted = sync.OnceFunc(func() {
-			if v := info.refCount.Add(-1); v == 0 {
-				// We call unsubscribe rather than sendNewServiceConfig to
-				// prevent redundant updates. If the reference count in the
-				// dependency manager drops to zero, it will automatically
-				// trigger a service config update with this cluster
-				// removed. Calling unsubscribe allows the dependency
-				// manager to handle the update flow once and for all.
-				info.unsubscribe()
-			}
-			// Decrement the refcount of the route cluster and close the interceptor
-			// if refcount goes to zero.
-			rc.Decrement()
-		})
-	} else if info, ok := cs.plugins[cluster.name]; ok {
-		// Add a ref to the selected plugin, as this RPC needs this
-		// plugin until it is committed.
-		info.refCount.Add(1)
-		config.OnCommitted = sync.OnceFunc(func() {
-			if v := info.refCount.Add(-1); v == 0 {
-				// This entry will be removed from activePlugins when
-				// producing a new service config update.
-				cs.sendNewServiceConfig()
-			}
-			// Decrement the refcount of the route cluster and close the interceptor
-			// if refcount goes to zero.
-			rc.Decrement()
-		})
-	} else {
-		// This should be unreachable because all route clusters are normalized
-		// into cs.clusters or cs.plugins during config selector creation.
-		panic(fmt.Sprintf("matched cluster %q not found in ConfigSelector", cluster.name))
-	}
+	// Add a ref to the selected cluster or plugin, as this RPC needs it until it
+	// is committed. Releasing the last reference unsubscribes from the cluster
+	// or pushes a new service config for a plugin.
+	cluster.info.Increment()
+	config.OnCommitted = sync.OnceFunc(func() {
+		cluster.info.Decrement()
+		// Decrement the refcount of the route cluster and close the interceptor
+		// if refcount goes to zero.
+		rc.Decrement()
+	})
 
 	if rt.maxStreamDuration != 0 {
 		config.MethodConfig.Timeout = &rt.maxStreamDuration
@@ -364,18 +341,14 @@ func (cs *configSelector) stop() {
 		}
 	}
 
-	// If any reference counts drop to zero, a service config update is required
-	// to remove the clusters. Since the old config selector is stopped
-	// after a new one is active, we must trigger a subsequent update to delete
-	// the now-unused clusters.
+	// Release this config selector's reference on each cluster and plugin. If
+	// any reference count drops to zero, the cleanup registered when the entry
+	// was created removes it from the resolver's active maps and triggers the
+	// service config update needed to drop it from the channel's config.
 	for _, ci := range cs.clusters {
-		if v := ci.refCount.Add(-1); v == 0 {
-			ci.unsubscribe()
-		}
+		ci.Decrement()
 	}
 	for _, ci := range cs.plugins {
-		if v := ci.refCount.Add(-1); v == 0 {
-			cs.sendNewServiceConfig()
-		}
+		ci.Decrement()
 	}
 }

internal/xds/resolver/serviceconfig_test.go

@@ -25,7 +25,6 @@ import (
 	"time"
 
 	xxhash "github.com/cespare/xxhash/v2"
-	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/grpcutil"
 	iresolver "google.golang.org/grpc/internal/resolver"
@@ -44,52 +43,6 @@ func Test(t *testing.T) {
 	grpctest.RunSubTests(t, s{})
 }
 
-func (s) TestPruneActiveClusters(t *testing.T) {
-	newClusterInfo := func(ref int32, unsubscribe func()) *clusterInfo {
-		ci := &clusterInfo{unsubscribe: unsubscribe}
-		ci.refCount.Store(ref)
-		return ci
-	}
-	r := &xdsResolver{
-		activeClusters: map[string]*clusterInfo{
-			"zero":        newClusterInfo(0, func() {}),
-			"one":         newClusterInfo(1, func() {}),
-			"two":         newClusterInfo(2, func() {}),
-			"anotherzero": newClusterInfo(0, func() {}),
-		},
-		activePlugins: map[string]*clusterInfo{
-			"zero":        newClusterInfo(0, nil),
-			"one":         newClusterInfo(1, nil),
-			"two":         newClusterInfo(2, nil),
-			"anotherzero": newClusterInfo(0, nil),
-		},
-	}
-	wantActiveClusters := map[string]int32{
-		"one": 1,
-		"two": 2,
-	}
-	wantActivePlugins := map[string]int32{
-		"one": 1,
-		"two": 2,
-	}
-	r.pruneActiveClustersAndPlugins()
-
-	getRefCounts := func(m map[string]*clusterInfo) map[string]int32 {
-		res := make(map[string]int32)
-		for k, v := range m {
-			res[k] = v.refCount.Load()
-		}
-		return res
-	}
-
-	if d := cmp.Diff(getRefCounts(r.activeClusters), wantActiveClusters); d != "" {
-		t.Fatalf("r.activeClusters refCounts mismatch (-got +want):\n%s", d)
-	}
-	if d := cmp.Diff(getRefCounts(r.activePlugins), wantActivePlugins); d != "" {
-		t.Fatalf("r.activePlugins refCounts mismatch (-got +want):\n%s", d)
-	}
-}
-
 func (s) TestGenerateRequestHash(t *testing.T) {
 	const channelID = 12378921
 	cs := &configSelector{channelID: channelID}