xds/resolver: ensure OnCommitted is invoked on early stream creation failure to prune unreferenced clusters.

internal/xds/resolver/serviceconfig.go

@@ -24,6 +24,7 @@ import (
 	"math/bits"
 	rand "math/rand/v2"
 	"strings"
+	"sync"
 	"time"
 
 	xxhash "github.com/cespare/xxhash/v2"
@@ -196,14 +197,6 @@ func (cs *configSelector) SelectConfig(rpcInfo iresolver.RPCInfo) (*iresolver.RP
 		return nil, annotateErrorWithNodeID(status.Errorf(codes.Internal, "error retrieving cluster for match: %v (%T)", cluster, cluster), cs.xdsNodeID)
 	}
 
-	// Add a ref to the selected cluster/plugin, as this RPC needs this
-	// cluster/plugin until it is committed.
-	if info, ok := cs.clusters[cluster.name]; ok {
-		info.refCount.Add(1)
-	} else if info, ok := cs.plugins[cluster.name]; ok {
-		info.refCount.Add(1)
-	}
-
 	lbCtx := clustermanager.SetPickedCluster(rpcInfo.Context, cluster.name)
 	lbCtx = xdsresource.NewContextWithXDSConfig(lbCtx, cs.xdsConfig)
 	lbCtx = iringhash.SetXDSRequestHash(lbCtx, cs.generateHash(rpcInfo, rt.hashPolicies))
@@ -212,12 +205,17 @@ func (cs *configSelector) SelectConfig(rpcInfo iresolver.RPCInfo) (*iresolver.RP
 	}
 
 	config := &iresolver.RPCConfig{
-		// Communicate to the LB policy the chosen cluster and request hash, if Ring Hash LB policy.
-		Context: lbCtx,
-		OnCommitted: func() {
-			// When the RPC is committed, the cluster is no longer required.
-			// Decrease its ref.
-			if info, ok := cs.clusters[cluster.name]; ok {
+		Context:     lbCtx,
+		Interceptor: cluster.interceptor,
+	}
+
+	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)
+		var once sync.Once
+		config.OnCommitted = func() {
+			once.Do(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
@@ -227,16 +225,22 @@ func (cs *configSelector) SelectConfig(rpcInfo iresolver.RPCInfo) (*iresolver.RP
 					// manager to handle the update flow once and for all.
 					info.unsubscribe()
 				}
-			}
-			if info, ok := cs.plugins[cluster.name]; ok {
+			})
+		}
+	} 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)
+		var once sync.Once
+		config.OnCommitted = func() {
+			once.Do(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()
 				}
-			}
-		},
-		Interceptor: cluster.interceptor,
+			})
+		}
 	}
 
 	if rt.maxStreamDuration != 0 {

internal/xds/resolver/xds_resolver_test.go

@@ -935,128 +935,6 @@ func (s) TestResolverMaxStreamDuration(t *testing.T) {
 	}
 }
 
-// Tests that clusters remain in service config if RPCs are in flight.
-func (s) TestResolverDelayedOnCommitted(t *testing.T) {
-	// Spin up an xDS management server for the test.
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
-	nodeID := uuid.New().String()
-	mgmtServer, _, _, bc := setupManagementServerForTest(t, nodeID)
-
-	// Configure resources on the management server.
-	resources := e2e.DefaultClientResources(e2e.ResourceParams{
-		DialTarget: defaultTestServiceName,
-		NodeID:     nodeID,
-		Host:       defaultTestHostname,
-		Port:       defaultTestPort[0],
-		SecLevel:   e2e.SecurityLevelNone,
-	})
-	if err := mgmtServer.Update(ctx, resources); err != nil {
-		t.Fatal(err)
-	}
-
-	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)
-
-	// Read the update pushed by the resolver to the ClientConn.
-	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantServiceConfig(resources.Clusters[0].Name))
-
-	// Make an RPC, but do not commit it yet.
-	resOld, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
-	if err != nil {
-		t.Fatalf("cs.SelectConfig(): %v", err)
-	}
-	wantClusterName := fmt.Sprintf("cluster:%s", resources.Clusters[0].Name)
-	if cluster := clustermanager.PickedCluster(resOld.Context); cluster != wantClusterName {
-		t.Fatalf("Picked cluster is %q, want %q", cluster, wantClusterName)
-	}
-
-	// Delay resOld.OnCommitted(). As long as there are pending RPCs to removed
-	// clusters, they still appear in the service config.
-	oldClusterName := resources.Clusters[0].Name
-	// Update the route configuration resource on the management server to
-	// return a new cluster.
-	newClusterName := "new-" + defaultTestClusterName
-	newEndpointName := "new-" + defaultTestEndpointName
-	resources.Routes = []*v3routepb.RouteConfiguration{e2e.DefaultRouteConfig(resources.Routes[0].Name, defaultTestServiceName, newClusterName)}
-	// Appending the new cluster and endpoint resources to avoid getting
-	// resource removed errors.
-	resources.Clusters = append(resources.Clusters, e2e.DefaultCluster(newClusterName, newEndpointName, e2e.SecurityLevelNone))
-	resources.Endpoints = append(resources.Endpoints, e2e.DefaultEndpoint(newEndpointName, defaultTestHostname, defaultTestPort))
-	if err := mgmtServer.Update(ctx, resources); err != nil {
-		t.Fatal(err)
-	}
-
-	// Read the update pushed by the resolver to the ClientConn and ensure the
-	// old cluster is present in the service config. Also ensure that the newly
-	// returned config selector does not hold a reference to the old cluster.
-	wantSC := fmt.Sprintf(`
-{
-	"loadBalancingConfig": [
-		{
-		  "xds_cluster_manager_experimental": {
-			"children": {
-			  "cluster:%s": {
-				"childPolicy": [
-				  {
-					"cds_experimental": {
-					  "cluster": "%s"
-					}
-				  }
-				]
-			  },
-			  "cluster:%s": {
-				"childPolicy": [
-				  {
-					"cds_experimental": {
-					  "cluster": "%s"
-					}
-				  }
-				]
-			  }
-			}
-		  }
-		}
-	  ]
-}`, oldClusterName, oldClusterName, newClusterName, newClusterName)
-	cs = verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
-
-	resNew, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
-	if err != nil {
-		t.Fatalf("cs.SelectConfig(): %v", err)
-	}
-	wantClusterName = fmt.Sprintf("cluster:%s", newClusterName)
-	if cluster := clustermanager.PickedCluster(resNew.Context); cluster != wantClusterName {
-		t.Fatalf("Picked cluster is %q, want %q", cluster, wantClusterName)
-	}
-
-	// Invoke OnCommitted on the old RPC; should lead to a service config update
-	// that deletes the old cluster, as the old cluster no longer has any
-	// pending RPCs.
-	resOld.OnCommitted()
-
-	wantSC = fmt.Sprintf(`
-{
-	"loadBalancingConfig": [
-		{
-		  "xds_cluster_manager_experimental": {
-			"children": {
-			  "cluster:%s": {
-				"childPolicy": [
-				  {
-					"cds_experimental": {
-					  "cluster": "%s"
-					}
-				  }
-				]
-			  }
-			}
-		  }
-		}
-	  ]
-}`, newClusterName, newClusterName)
-	verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
-}
-
 // Tests the case where two LDS updates with the same RDS name to watch are
 // received without an RDS in between. Those LDS updates shouldn't trigger a
 // service config update.

stream.go

@@ -253,10 +253,7 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 
 	mc := &emptyMethodConfig
 	var onCommit func()
-	newStream := func(ctx context.Context, filterOpts ...CallOption) (ClientStream, error) {
-		if filterOpts != nil {
-			opts = combine(opts, filterOpts)
-		}
+	newStream := func(ctx context.Context, opts ...CallOption) (ClientStream, error) {
 		return newClientStreamWithParams(ctx, desc, cc, method, mc, onCommit, nameResolutionDelayed, opts...)
 	}
 
@@ -278,13 +275,23 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 			ctx = rpcConfig.Context
 		}
 		mc = &rpcConfig.MethodConfig
-		onCommit = rpcConfig.OnCommitted
+
+		if rpcConfig.OnCommitted != nil {
+			onCommit = rpcConfig.OnCommitted
+			// Register an OnFinish CallOption with the OnCommitted callback to
+			// ensure it is invoked on stream termination, even if the stream
+			// fails early before committing. Implementations of OnCommitted are
+			// expected to be idempotent (e.g., guarded by sync.Once), since both
+			// onCommit and OnFinish may run for a single RPC.
+			opts = append(opts, OnFinish(func(error) { rpcConfig.OnCommitted() }))
+		}
+
 		if rpcConfig.Interceptor != nil {
 			rpcInfo.Context = nil
 			ns := newStream
 			if interceptor, ok := rpcConfig.Interceptor.(clientInterceptor); ok {
-				newStream = func(ctx context.Context, filterOpts ...CallOption) (ClientStream, error) {
-					cs, err := interceptor.NewStream(ctx, rpcInfo, ns, filterOpts...)
+				newStream = func(ctx context.Context, opts ...CallOption) (ClientStream, error) {
+					cs, err := interceptor.NewStream(ctx, rpcInfo, ns, opts...)
 					if err != nil {
 						return nil, toRPCErr(err)
 					}
@@ -296,7 +303,7 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 		}
 	}
 
-	return newStream(ctx)
+	return newStream(ctx, opts...)
 }
 
 func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, mc *serviceconfig.MethodConfig, onCommit func(), nameResolutionDelayed bool, opts ...CallOption) (_ ClientStream, err error) {