cluster_specifier_plugin_test.go

/*
 *
 * Copyright 2021 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 resolver_test

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"testing"

	"github.com/google/uuid"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	iresolver "google.golang.org/grpc/internal/resolver"
	"google.golang.org/grpc/internal/testutils"
	"google.golang.org/grpc/internal/testutils/xds/e2e"
	"google.golang.org/grpc/internal/xds/balancer/clustermanager"
	"google.golang.org/grpc/internal/xds/clusterspecifier"
	"google.golang.org/grpc/internal/xds/httpfilter"
	"google.golang.org/grpc/resolver"
	"google.golang.org/grpc/serviceconfig"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
	"google.golang.org/protobuf/types/known/wrapperspb"

	v3corepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	v3listenerpb "github.com/envoyproxy/go-control-plane/envoy/config/listener/v3"
	v3routepb "github.com/envoyproxy/go-control-plane/envoy/config/route/v3"
	v3httppb "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/network/http_connection_manager/v3"
)

func init() {
	balancer.Register(cspBalancerBuilder{})
	clusterspecifier.Register(testClusterSpecifierPlugin{})
}

// cspBalancerBuilder is a no-op LB policy which is referenced by the
// testClusterSpecifierPlugin.
type cspBalancerBuilder struct{}

func (cspBalancerBuilder) Build(balancer.ClientConn, balancer.BuildOptions) balancer.Balancer {
	return nil
}

func (cspBalancerBuilder) Name() string {
	return "csp_experimental"
}

type cspBalancerConfig struct {
	serviceconfig.LoadBalancingConfig
	ArbitraryField string `json:"arbitrary_field"`
}

func (cspBalancerBuilder) ParseConfig(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
	cfg := &cspBalancerConfig{}
	if err := json.Unmarshal(lbCfg, cfg); err != nil {
		return nil, err
	}
	return cfg, nil

}

// testClusterSpecifierPlugin is a test cluster specifier plugin which returns
// an LB policy configuration specifying the cspBalancer.
type testClusterSpecifierPlugin struct {
}

func (testClusterSpecifierPlugin) TypeURLs() []string {
	// The config for this plugin contains a wrapperspb.StringValue, and since
	// we marshal that proto as an Any proto, the type URL on the latter gets
	// set to "type.googleapis.com/google.protobuf.StringValue". If we wanted a
	// more descriptive type URL for this test plugin, we would have to define a
	// proto package with a message for the configuration. That would be
	// overkill for a test. Therefore, this seems to be an acceptable tradeoff.
	return []string{"type.googleapis.com/google.protobuf.StringValue"}
}

func (testClusterSpecifierPlugin) ParseClusterSpecifierConfig(cfg proto.Message) (clusterspecifier.BalancerConfig, error) {
	if cfg == nil {
		return nil, fmt.Errorf("testClusterSpecifierPlugin: nil configuration message provided")
	}
	anyp, ok := cfg.(*anypb.Any)
	if !ok {
		return nil, fmt.Errorf("testClusterSpecifierPlugin: error parsing config %v: got type %T, want *anypb.Any", cfg, cfg)
	}
	lbCfg := new(wrapperspb.StringValue)
	if err := anypb.UnmarshalTo(anyp, lbCfg, proto.UnmarshalOptions{}); err != nil {
		return nil, fmt.Errorf("testClusterSpecifierPlugin: error parsing config %v: %v", cfg, err)
	}
	return []map[string]any{{"csp_experimental": cspBalancerConfig{ArbitraryField: lbCfg.GetValue()}}}, nil
}

// TestResolverClusterSpecifierPlugin tests the case where a route configuration
// containing cluster specifier plugins is sent by the management server. The
// test verifies that the service config output by the resolver contains the LB
// policy specified by the cluster specifier plugin, and the config selector
// returns the cluster associated with the cluster specifier plugin.
//
// The test also verifies that a change in the cluster specifier plugin config
// result in appropriate change in the service config pushed by the resolver.
func (s) TestResolverClusterSpecifierPlugin(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.
	listeners := []*v3listenerpb.Listener{e2e.DefaultClientListener(defaultTestServiceName, defaultTestRouteConfigName)}
	routes := []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
		RouteConfigName:              defaultTestRouteConfigName,
		ListenerName:                 defaultTestServiceName,
		ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
		ClusterSpecifierPluginName:   "cspA",
		ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: "anything"}),
	})}
	configureResources(ctx, t, mgmtServer, nodeID, listeners, routes, nil, nil)

	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)

	// Wait for an update from the resolver, and verify the service config.
	wantSC := `
 {
	 "loadBalancingConfig": [
		 {
		   "xds_cluster_manager_experimental": {
			 "children": {
			   "cluster_specifier_plugin:cspA": {
				 "childPolicy": [
				   {
					 "csp_experimental": {
					   "arbitrary_field": "anything"
					 }
				   }
				 ]
			   }
			 }
		   }
		 }
	   ]
 }`
	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
	res, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
	if err != nil {
		t.Fatalf("cs.SelectConfig(): %v", err)
	}

	gotCluster := clustermanager.PickedCluster(res.Context)
	wantCluster := "cluster_specifier_plugin:cspA"
	if gotCluster != wantCluster {
		t.Fatalf("config selector returned cluster: %v, want: %v", gotCluster, wantCluster)
	}

	// Change the cluster specifier plugin configuration.
	routes = []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
		RouteConfigName:              defaultTestRouteConfigName,
		ListenerName:                 defaultTestServiceName,
		ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
		ClusterSpecifierPluginName:   "cspA",
		ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: "changed"}),
	})}
	configureResources(ctx, t, mgmtServer, nodeID, listeners, routes, nil, nil)

	// Wait for an update from the resolver, and verify the service config.
	wantSC = `
 {
	 "loadBalancingConfig": [
		 {
		   "xds_cluster_manager_experimental": {
			 "children": {
			   "cluster_specifier_plugin:cspA": {
				 "childPolicy": [
				   {
					 "csp_experimental": {
					   "arbitrary_field": "changed"
					 }
				   }
				 ]
			   }
			 }
		   }
		 }
	   ]
 }`
	verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
}

// TestXDSResolverDelayedOnCommittedCSP tests that cluster specifier plugins and
// their corresponding configurations remain in service config if RPCs are in
// flight.
func (s) TestXDSResolverDelayedOnCommittedCSP(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.
	listeners := []*v3listenerpb.Listener{e2e.DefaultClientListener(defaultTestServiceName, defaultTestRouteConfigName)}
	routes := []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
		RouteConfigName:              defaultTestRouteConfigName,
		ListenerName:                 defaultTestServiceName,
		ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
		ClusterSpecifierPluginName:   "cspA",
		ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: "anythingA"}),
	})}
	configureResources(ctx, t, mgmtServer, nodeID, listeners, routes, nil, nil)

	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)

	// Wait for an update from the resolver, and verify the service config.
	wantSC := `
 {
	 "loadBalancingConfig": [
		 {
		   "xds_cluster_manager_experimental": {
			 "children": {
			   "cluster_specifier_plugin:cspA": {
				 "childPolicy": [
				   {
					 "csp_experimental": {
					   "arbitrary_field": "anythingA"
					 }
				   }
				 ]
			   }
			 }
		   }
		 }
	   ]
 }`
	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantSC)

	resOld, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
	if err != nil {
		t.Fatalf("cs.SelectConfig(): %v", err)
	}

	gotCluster := clustermanager.PickedCluster(resOld.Context)
	wantCluster := "cluster_specifier_plugin:cspA"
	if gotCluster != wantCluster {
		t.Fatalf("config selector returned cluster: %v, want: %v", gotCluster, wantCluster)
	}

	// Delay resOld.OnCommitted(). As long as there are pending RPCs to removed
	// clusters, they still appear in the service config.

	// Change the cluster specifier plugin configuration.
	routes = []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
		RouteConfigName:              defaultTestRouteConfigName,
		ListenerName:                 defaultTestServiceName,
		ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
		ClusterSpecifierPluginName:   "cspB",
		ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: "anythingB"}),
	})}
	configureResources(ctx, t, mgmtServer, nodeID, listeners, routes, nil, nil)

	// Wait for an update from the resolver, and verify the service config.
	wantSC = `
 {
	 "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"
					 }
				   }
				 ]
			   }
			 }
		   }
		 }
	   ]
 }`
	cs = verifyUpdateFromResolver(ctx, t, stateCh, wantSC)

	// Perform an RPC and ensure that it is routed to the new cluster.
	resNew, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
	if err != nil {
		t.Fatalf("cs.SelectConfig(): %v", err)
	}

	gotCluster = clustermanager.PickedCluster(resNew.Context)
	wantCluster = "cluster_specifier_plugin:cspB"
	if gotCluster != wantCluster {
		t.Fatalf("config selector returned cluster: %v, want: %v", gotCluster, wantCluster)
	}

	// Invoke resOld.OnCommitted; should lead to a service config update that deletes
	// cspA.
	resOld.OnCommitted()

	wantSC = `
 {
	 "loadBalancingConfig": [
		 {
		   "xds_cluster_manager_experimental": {
			 "children": {
			   "cluster_specifier_plugin:cspB": {
				 "childPolicy": [
				   {
					 "csp_experimental": {
					   "arbitrary_field": "anythingB"
					 }
				   }
				 ]
			   }
			 }
		   }
		 }
	   ]
 }`
	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
// interceptor chain is built for routes matching cluster specifier plugins.
func (s) TestResolverClusterSpecifierPlugin_WithFilters(t *testing.T) {
	// Register custom httpFilter builders for the test.
	testFilterTypeURL1 := "test-filter-type-url-1" + uuid.New().String()
	testFilterTypeURL2 := "test-filter-type-url-2" + uuid.New().String()
	newStreamChan := testutils.NewChannelWithSize(2)
	fb1 := &testHTTPFilterWithRPCMetadata{
		logger:        t,
		typeURL:       testFilterTypeURL1,
		newStreamChan: newStreamChan,
	}
	fb2 := &testHTTPFilterWithRPCMetadata{
		logger:        t,
		typeURL:       testFilterTypeURL2,
		newStreamChan: newStreamChan,
	}
	httpfilter.Register(fb1)
	httpfilter.Register(fb2)
	defer httpfilter.UnregisterForTesting(fb1.typeURL)
	defer httpfilter.UnregisterForTesting(fb2.typeURL)

	// Spin up an xDS management server.
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	nodeID := uuid.New().String()
	mgmtServer, _, _, bc := setupManagementServerForTest(t, nodeID)

	// Configure resources on the management server.
	// We need a listener with the filter, and a route with ClusterSpecifierPlugin.
	listeners := []*v3listenerpb.Listener{{
		Name: defaultTestServiceName,
		ApiListener: &v3listenerpb.ApiListener{
			ApiListener: testutils.MarshalAny(t, &v3httppb.HttpConnectionManager{
				RouteSpecifier: &v3httppb.HttpConnectionManager_Rds{Rds: &v3httppb.Rds{
					ConfigSource: &v3corepb.ConfigSource{
						ConfigSourceSpecifier: &v3corepb.ConfigSource_Ads{Ads: &v3corepb.AggregatedConfigSource{}},
					},
					RouteConfigName: defaultTestRouteConfigName,
				}},
				HttpFilters: []*v3httppb.HttpFilter{
					newHTTPFilter(t, "test-filter-1", testFilterTypeURL1, "filter-path-1", ""),
					newHTTPFilter(t, "test-filter-2", testFilterTypeURL2, "filter-path-2", ""),
					e2e.RouterHTTPFilter,
				},
			}),
		},
	}}

	routes := []*v3routepb.RouteConfiguration{e2e.RouteConfigResourceWithOptions(e2e.RouteConfigOptions{
		RouteConfigName:              defaultTestRouteConfigName,
		ListenerName:                 defaultTestServiceName,
		ClusterSpecifierType:         e2e.RouteConfigClusterSpecifierTypeClusterSpecifierPlugin,
		ClusterSpecifierPluginName:   "cspA",
		ClusterSpecifierPluginConfig: testutils.MarshalAny(t, &wrapperspb.StringValue{Value: "anything"}),
	})}
	// Override the configuration for "test-filter-1" in the route.
	routes[0].VirtualHosts[0].Routes[0].TypedPerFilterConfig = map[string]*anypb.Any{
		"test-filter-1": newHTTPFilter(t, "test-filter-1", testFilterTypeURL1, "override-path-1", "").GetTypedConfig(),
	}
	configureResources(ctx, t, mgmtServer, nodeID, listeners, routes, nil, nil)

	stateCh, _, _ := buildResolverForTarget(t, resolver.Target{URL: *testutils.MustParseURL("xds:///" + defaultTestServiceName)}, bc)

	// Wait for an update from the resolver, and verify the service config.
	wantSC := `
 {
	"loadBalancingConfig": [
		{
			"xds_cluster_manager_experimental": {
			"children": {
				"cluster_specifier_plugin:cspA": {
				"childPolicy": [
					{
					"csp_experimental": {
						"arbitrary_field": "anything"
					}
					}
				]
				}
			}
			}
		}
	]
 }`
	cs := verifyUpdateFromResolver(ctx, t, stateCh, wantSC)
	res, err := cs.SelectConfig(iresolver.RPCInfo{Context: ctx, Method: "/service/method"})
	if err != nil {
		t.Fatalf("cs.SelectConfig(): %v", err)
	}

	// Verify that the interceptor is not nil.
	if res.Interceptor == nil {
		t.Fatal("RPCInfo does not contain interceptors list")
	}
	newStream := func(context.Context, ...grpc.CallOption) (grpc.ClientStream, error) {
		return nil, nil
	}

	if interceptor, ok := res.Interceptor.(httpfilter.ClientInterceptor); ok {
		if _, err = interceptor.NewStream(ctx, iresolver.RPCInfo{Method: "/service/method", Context: ctx}, newStream); err != nil {
			t.Fatalf("NewStream() failed with error: %v", err)
		}
	} else {
		t.Fatalf("res.Interceptor is type %T, want httpfilter.ClientInterceptor", res.Interceptor)
	}

	// Verify that first filter receives the config.
	cfg, err := newStreamChan.Receive(ctx)
	if err != nil {
		t.Fatalf("Timeout waiting for first filter to receive config: %v", err)
	}
	ofc := cfg.(overallFilterConfig)
	if ofc.BasePath != "filter-path-1" {
		t.Fatalf("Unexpected base path for first filter, got: %q, want: %q", ofc.BasePath, "filter-path-1")
	}
	if ofc.OverridePath != "override-path-1" {
		t.Fatalf("Unexpected override path for first filter, got: %q, want: %q", ofc.OverridePath, "override-path-1")
	}

	// Verify that second filter receives the base path.
	cfg, err = newStreamChan.Receive(ctx)
	if err != nil {
		t.Fatalf("Timeout waiting for second filter to receive config: %v", err)
	}
	ofc = cfg.(overallFilterConfig)
	if ofc.BasePath != "filter-path-2" {
		t.Fatalf("Unexpected base path for second filter, got: %q, want: %q", ofc.BasePath, "filter-path-2")
	}
	if ofc.OverridePath != "" {
		t.Fatalf("Unexpected override path for second filter, got: %q, want: %q", ofc.OverridePath, "")
	}
}