xdsdepmgr: add functionality for cluster subscription

internal/xds/resolver/xds_resolver.go

@@ -151,6 +151,7 @@ func (b *xdsResolverBuilder) Build(target resolver.Target, cc resolver.ClientCon
 	r.logger = prefixLogger(r)
 	r.logger.Infof("Creating resolver for target: %+v", target)
 	r.dm = xdsdepmgr.New(r.ldsResourceName, opts.Authority, r.xdsClient, r)
+	r.dm.Start()
 	return r, nil
 }
 
@@ -329,6 +330,8 @@ func (r *xdsResolver) sendNewServiceConfig(cs stoppableConfigSelector) bool {
 	state := iresolver.SetConfigSelector(resolver.State{
 		ServiceConfig: r.cc.ParseServiceConfig(string(sc)),
 	}, cs)
+	state = xdsresource.SetXDSConfig(state, r.xdsConfig)
+	state = xdsdepmgr.SetDependencyManager(state, r.dm)
 	if err := r.cc.UpdateState(xdsclient.SetClient(state, r.xdsClient)); err != nil {
 		if r.logger.V(2) {
 			r.logger.Infof("Channel rejected new state: %+v with error: %v", state, err)

internal/xds/resolver/xds_resolver_test.go

@@ -36,6 +36,7 @@ import (
 	estats "google.golang.org/grpc/experimental/stats"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/grpcsync"
 	iresolver "google.golang.org/grpc/internal/resolver"
 	iringhash "google.golang.org/grpc/internal/ringhash"
 	"google.golang.org/grpc/internal/testutils"
@@ -1455,3 +1456,128 @@ func (s) TestResolver_AutoHostRewrite(t *testing.T) {
 		})
 	}
 }
+
+// resourcesMatch returns true if the got slice matches resource names in want.
+func resourcesMatch(got, want []string) bool {
+	if len(got) != len(want) {
+		return false
+	}
+	wantMap := make(map[string]bool, len(want))
+	for _, w := range want {
+		wantMap[w] = true
+	}
+	for _, n := range got {
+		if !wantMap[n] {
+			return false
+		}
+	}
+	return true
+}
+
+// TestResolverKeepWatchOpen_ActiveRPCs tests that the dependency manager keeps
+// a cluster watch open when there are active RPCs using that cluster, even if
+// the cluster is no longer referenced by the current route configuration.
+func (s) TestResolverKeepWatchOpen_ActiveRPCs(t *testing.T) {
+	t.Skip("Will be enabled when all the watchers have shifted to dependency manager")
+	gotBothClusterRequest := grpcsync.NewEvent()
+	gotOnlySecondCluster := grpcsync.NewEvent()
+
+	// These are only accessed in the callback, which is executed serially.
+	seenBothClusters := false
+	seenSecondClusterOnly := false
+
+	clusterA := "cluster-A"
+	clusterB := "cluster-B"
+	mgmtServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{
+		OnStreamRequest: func(_ int64, req *v3discoverypb.DiscoveryRequest) error {
+			if req.GetTypeUrl() == version.V3ClusterURL {
+				resourceNames := req.GetResourceNames()
+				if !seenBothClusters {
+					if resourcesMatch(resourceNames, []string{clusterA, clusterB}) {
+						seenBothClusters = true
+						gotBothClusterRequest.Fire()
+					}
+				}
+				if seenBothClusters && !seenSecondClusterOnly {
+					if resourcesMatch(resourceNames, []string{clusterB}) {
+						seenSecondClusterOnly = true
+						gotOnlySecondCluster.Fire()
+					}
+				}
+			}
+			return nil
+		},
+		AllowResourceSubset: true,
+	})
+
+	nodeID := uuid.New().String()
+	bc := e2e.DefaultBootstrapContents(t, nodeID, mgmtServer.Address)
+
+	// Configure initial resources: Route -> ClusterA
+	routeA := e2e.DefaultRouteConfig(defaultTestRouteConfigName, defaultTestServiceName, clusterA)
+	listeners := []*v3listenerpb.Listener{e2e.DefaultClientListener(defaultTestServiceName, defaultTestRouteConfigName)}
+	clusters := []*v3clusterpb.Cluster{
+		e2e.DefaultCluster(clusterA, "endpoint-A", e2e.SecurityLevelNone),
+		e2e.DefaultCluster(clusterB, "endpoint-B", e2e.SecurityLevelNone),
+	}
+	endpoints := []*v3endpointpb.ClusterLoadAssignment{
+		e2e.DefaultEndpoint("endpoint-A", "localhost", []uint32{8080}),
+		e2e.DefaultEndpoint("endpoint-B", "localhost", []uint32{8081}),
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	configureAllResourcesOnManagementServer(ctx, t, mgmtServer, nodeID, listeners, []*v3routepb.RouteConfiguration{routeA}, clusters, endpoints)
+
+	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)
+
+	// Start RPC (Ref Counts ClusterA)
+	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantServiceConfig(clusterA))
+	res, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
+	if err != nil {
+		t.Fatalf("cs.SelectConfig(): %v", err)
+	}
+
+	// Switch Configuration to ClusterB
+	routeB := e2e.DefaultRouteConfig(defaultTestRouteConfigName, defaultTestServiceName, clusterB)
+	configureAllResourcesOnManagementServer(ctx, t, mgmtServer, nodeID, listeners, []*v3routepb.RouteConfiguration{routeB}, clusters, endpoints)
+
+	// Resolver should request BOTH A (due to active RPC) and B (due to new config)
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timeout waiting for updated CDS request including clusters A and B")
+	case <-gotBothClusterRequest.Done():
+	}
+
+	// Verify Service Config has both clusters
+	const wantServiceRaw = `{
+      "loadBalancingConfig": [{
+        "xds_cluster_manager_experimental": {
+          "children": {
+            "cluster:cluster-A": {
+              "childPolicy": [{"cds_experimental": {"cluster": "cluster-A"}}]
+            },
+            "cluster:cluster-B": {
+              "childPolicy": [{"cds_experimental": {"cluster": "cluster-B"}}]
+            }
+          }
+        }
+      }]
+    }`
+	verifyUpdateFromResolver(ctx, t, stateCh, wantServiceRaw)
+
+	// Finish RPC (Drops Ref to ClusterA)
+	res.OnCommitted()
+
+	// ONLY cluster B should be requested now that there are no references to
+	// cluster A.
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timeout waiting for updated CDS request including only cluster B")
+	case <-gotOnlySecondCluster.Done():
+	}
+
+	// ServiceConfig update should also contain only cluster B.
+	verifyUpdateFromResolver(ctx, t, stateCh, wantServiceConfig(clusterB))
+}

internal/xds/xdsdepmgr/xds_dependency_manager.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"net/url"
 	"sync"
+	"sync/atomic"
 
 	"google.golang.org/grpc/grpclog"
 	internalgrpclog "google.golang.org/grpc/internal/grpclog"
@@ -129,6 +130,7 @@ type DependencyManager struct {
 	clusterWatchers         map[string]*xdsResourceState[xdsresource.ClusterUpdate, struct{}]
 	endpointWatchers        map[string]*xdsResourceState[xdsresource.EndpointsUpdate, struct{}]
 	dnsResolvers            map[string]*xdsResourceState[xdsresource.DNSUpdate, dnsExtras]
+	clusterSubscriptions    map[string]*ClusterRef
 }
 
 // New creates a new DependencyManager.
@@ -154,25 +156,27 @@ func New(listenerName, dataplaneAuthority string, xdsClient xdsclient.XDSClient,
 		endpointWatchers:        make(map[string]*xdsResourceState[xdsresource.EndpointsUpdate, struct{}]),
 		dnsResolvers:            make(map[string]*xdsResourceState[xdsresource.DNSUpdate, dnsExtras]),
 		clusterWatchers:         make(map[string]*xdsResourceState[xdsresource.ClusterUpdate, struct{}]),
+		clusterSubscriptions:    make(map[string]*ClusterRef),
 	}
 	dm.logger = prefixLogger(dm)
-
-	// Start the listener watch. Listener watch will start the other resource
-	// watches as needed.
 	dm.listenerWatcher = &xdsResourceState[xdsresource.ListenerUpdate, struct{}]{}
-	lw := &xdsResourceWatcher[xdsresource.ListenerUpdate]{
+	return dm
+}
+
+// Start registers the listener watch on the xDS client which in turn will start
+// watches for other resources.
+func (m *DependencyManager) Start() {
+	m.listenerWatcher.stop = xdsresource.WatchListener(m.xdsClient, m.ldsResourceName, &xdsResourceWatcher[xdsresource.ListenerUpdate]{
 		onUpdate: func(update *xdsresource.ListenerUpdate, onDone func()) {
-			dm.onListenerResourceUpdate(update, onDone)
+			m.onListenerResourceUpdate(update, onDone)
 		},
 		onError: func(err error, onDone func()) {
-			dm.onListenerResourceError(err, onDone)
+			m.onListenerResourceError(err, onDone)
 		},
 		onAmbientError: func(err error, onDone func()) {
-			dm.onListenerResourceAmbientError(err, onDone)
+			m.onListenerResourceAmbientError(err, onDone)
 		},
-	}
-	dm.listenerWatcher.stop = xdsresource.WatchListener(dm.xdsClient, listenerName, lw)
-	return dm
+	})
 }
 
 // Close cancels all registered resource watches.
@@ -236,7 +240,17 @@ func (m *DependencyManager) maybeSendUpdateLocked() {
 	dnsResourcesSeen := make(map[string]bool)
 	clusterResourcesSeen := make(map[string]bool)
 	haveAllResources := true
+	// Get all clusters to be watched: from route config and from cluster
+	// subscriptions, the subscriptions can be from RPC referencing the cluster
+	// or from balancer for cluster specifier plugins.
+	clustersToWatch := make(map[string]bool)
 	for cluster := range m.clustersFromRouteConfig {
+		clustersToWatch[cluster] = true
+	}
+	for cluster := range m.clusterSubscriptions {
+		clustersToWatch[cluster] = true
+	}
+	for cluster := range clustersToWatch {
 		ok, leafClusters, err := m.populateClusterConfigLocked(cluster, 0, config.Clusters, edsResourcesSeen, dnsResourcesSeen, clusterResourcesSeen)
 		if !ok {
 			haveAllResources = false
@@ -459,13 +473,6 @@ func (m *DependencyManager) applyRouteConfigUpdateLocked(update *xdsresource.Rou
 				newClusters[cluster.Name] = true
 			}
 		}
-		// Cancel watch for clusters not seen in route config
-		for name := range m.clustersFromRouteConfig {
-			if _, ok := newClusters[name]; !ok {
-				m.clusterWatchers[name].stop()
-				delete(m.clusterWatchers, name)
-			}
-		}
 
 		// Watch for new clusters is started in populateClusterConfigLocked to
 		// avoid repeating the code.
@@ -845,15 +852,19 @@ func (m *DependencyManager) newDNSResolver(target string) *xdsResourceState[xdsr
 		err := fmt.Errorf("failed to parse DNS target %q: %v", target, m.annotateErrorWithNodeID(err))
 		m.logger.Warningf("%v", err)
 		rcc.ReportError(err)
-		return &xdsResourceState[xdsresource.DNSUpdate, dnsExtras]{}
+		return &xdsResourceState[xdsresource.DNSUpdate, dnsExtras]{
+			stop: func() {},
+		}
 	}
 
 	r, err := resolver.Get("dns").Build(resolver.Target{URL: *u}, rcc, resolver.BuildOptions{})
 	if err != nil {
 		rcc.ReportError(err)
 		err := fmt.Errorf("failed to build DNS resolver for target %q: %v", target, m.annotateErrorWithNodeID(err))
 		m.logger.Warningf("%v", err)
-		return nil
+		return &xdsResourceState[xdsresource.DNSUpdate, dnsExtras]{
+			stop: func() {},
+		}
 	}
 
 	return &xdsResourceState[xdsresource.DNSUpdate, dnsExtras]{
@@ -881,3 +892,74 @@ func (x *xdsResourceWatcher[T]) ResourceError(err error, onDone func()) {
 func (x *xdsResourceWatcher[T]) AmbientError(err error, onDone func()) {
 	x.onAmbientError(err, onDone)
 }
+
+// dependencyManagerKey is the type used as the key to store DependencyManager
+// in the Attributes field of resolver.states.
+type dependencyManagerKey struct{}
+
+// SetDependencyManager returns a copy of state in which the Attributes field is
+// updated with the DependencyManager.
+func SetDependencyManager(state resolver.State, depmngr *DependencyManager) resolver.State {
+	state.Attributes = state.Attributes.WithValue(dependencyManagerKey{}, depmngr)
+	return state
+}
+
+// DependencyManagerFromResolverState returns DependencyManager stored as an
+// attribute in the resolver state.
+func DependencyManagerFromResolverState(state resolver.State) *DependencyManager {
+	if v := state.Attributes.Value(dependencyManagerKey{}); v != nil {
+		return v.(*DependencyManager)
+	}
+	return nil
+}
+
+// ClusterRef represents a reference to a cluster being used by some component.
+// It maintains a reference count of the number of users of the cluster. It has
+// a method to unsubscribe from the cluster, which decrements the reference
+// count and removes the cluster from the clusterSubscriptions map in the
+// DependencyManager if the reference count reaches zero.
+type ClusterRef struct {
+	name     string
+	refCount int32
+	m        *DependencyManager
+}
+
+// ClusterSubscription increments the reference count for the cluster and
+// returns the ClusterRef. If the cluster is not already being tracked, it adds
+// it to the ClusterSubs map.
+func (m *DependencyManager) ClusterSubscription(name string) *ClusterRef {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	subs, ok := m.clusterSubscriptions[name]
+	if ok {
+		ref := &subs.refCount
+		atomic.AddInt32(ref, 1)
+		return subs
+	}
+	m.clusterSubscriptions[name] = &ClusterRef{
+		name:     name,
+		refCount: 1,
+		m:        m,
+	}
+	if _, ok := m.clustersFromRouteConfig[name]; !ok {
+		m.maybeSendUpdateLocked()
+	}
+	return m.clusterSubscriptions[name]
+}
+
+// Unsubscribe decrements the reference count for the cluster. If the reference
+// count reaches zero, it removes the cluster from the clusterSubscriptions map
+// in the DependencyManager.
+func (c *ClusterRef) Unsubscribe() {
+	c.m.mu.Lock()
+	defer c.m.mu.Unlock()
+	ref := atomic.AddInt32(&c.refCount, -1)
+	if ref <= 0 {
+		delete(c.m.clusterSubscriptions, c.name)
+		// This cluster is no longer in the route config, and it has no more
+		// references. Now is the time to cancel the watch.
+		if _, ok := c.m.clustersFromRouteConfig[c.name]; !ok {
+			c.m.maybeSendUpdateLocked()
+		}
+	}
+}