From 512e89474bfbd567806e026e0afdec468b1ee1ea Mon Sep 17 00:00:00 2001 From: Easwar Swaminathan Date: Fri, 3 Dec 2021 15:25:06 -0800 Subject: [PATCH] rls: support extra_keys and constant_keys (#4995) * rls: support extra_keys and constant_keys * review comments * use the constant_keys map from the proto --- balancer/rls/internal/keys/builder.go | 121 +++++---- balancer/rls/internal/keys/builder_test.go | 276 +++++++++++++++------ balancer/rls/internal/picker.go | 14 +- 3 files changed, 280 insertions(+), 131 deletions(-) diff --git a/balancer/rls/internal/keys/builder.go b/balancer/rls/internal/keys/builder.go index 55cf7478c887..fbf45c668d7f 100644 --- a/balancer/rls/internal/keys/builder.go +++ b/balancer/rls/internal/keys/builder.go @@ -29,25 +29,11 @@ import ( "google.golang.org/grpc/metadata" ) -// BuilderMap provides a mapping from a request path to the key builder to be -// used for that path. -// The BuilderMap is constructed by parsing the RouteLookupConfig received by -// the RLS balancer as part of its ServiceConfig, and is used by the picker in -// the data path to build the RLS keys to be used for a given request. +// BuilderMap maps from request path to the key builder for that path. type BuilderMap map[string]builder // MakeBuilderMap parses the provided RouteLookupConfig proto and returns a map // from paths to key builders. -// -// The following conditions are validated, and an error is returned if any of -// them is not met: -// grpc_keybuilders field -// * must have at least one entry -// * must not have two entries with the same Name -// * must not have any entry with a Name with the service field unset or empty -// * must not have any entries without a Name -// * must not have a headers entry that has required_match set -// * must not have two headers entries with the same key within one entry func MakeBuilderMap(cfg *rlspb.RouteLookupConfig) (BuilderMap, error) { kbs := cfg.GetGrpcKeybuilders() if len(kbs) == 0 { @@ -56,21 +42,46 @@ func MakeBuilderMap(cfg *rlspb.RouteLookupConfig) (BuilderMap, error) { bm := make(map[string]builder) for _, kb := range kbs { + // Extract keys from `headers`, `constant_keys` and `extra_keys` fields + // and populate appropriate values in the builder struct. Also ensure + // that keys are not repeated. var matchers []matcher seenKeys := make(map[string]bool) + constantKeys := kb.GetConstantKeys() + for k := range kb.GetConstantKeys() { + seenKeys[k] = true + } for _, h := range kb.GetHeaders() { if h.GetRequiredMatch() { return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set {%+v}", kbs) } key := h.GetKey() if seenKeys[key] { - return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Key field in headers {%+v}", kbs) + return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q across headers, constant_keys and extra_keys {%+v}", key, kbs) } seenKeys[key] = true matchers = append(matchers, matcher{key: h.GetKey(), names: h.GetNames()}) } - b := builder{matchers: matchers} + if seenKeys[kb.GetExtraKeys().GetHost()] { + return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetHost(), kbs) + } + if seenKeys[kb.GetExtraKeys().GetService()] { + return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetService(), kbs) + } + if seenKeys[kb.GetExtraKeys().GetMethod()] { + return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetMethod(), kbs) + } + b := builder{ + headerKeys: matchers, + constantKeys: constantKeys, + hostKey: kb.GetExtraKeys().GetHost(), + serviceKey: kb.GetExtraKeys().GetService(), + methodKey: kb.GetExtraKeys().GetMethod(), + } + // Store the builder created above in the BuilderMap based on the value + // of the `Names` field, which wraps incoming request's service and + // method. Also, ensure that there are no repeated `Names` field. names := kb.GetNames() if len(names) == 0 { return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig does not contain any Name {%+v}", kbs) @@ -108,16 +119,31 @@ type KeyMap struct { // RLSKey builds the RLS keys to be used for the given request, identified by // the request path and the request headers stored in metadata. -func (bm BuilderMap) RLSKey(md metadata.MD, path string) KeyMap { +func (bm BuilderMap) RLSKey(md metadata.MD, host, path string) KeyMap { + i := strings.LastIndex(path, "/") + service, method := path[:i+1], path[i+1:] b, ok := bm[path] if !ok { - i := strings.LastIndex(path, "/") - b, ok = bm[path[:i+1]] + b, ok = bm[service] if !ok { return KeyMap{} } } - return b.keys(md) + + kvMap := b.buildHeaderKeys(md) + if b.hostKey != "" { + kvMap[b.hostKey] = host + } + if b.serviceKey != "" { + kvMap[b.serviceKey] = service + } + if b.methodKey != "" { + kvMap[b.methodKey] = method + } + for k, v := range b.constantKeys { + kvMap[k] = v + } + return KeyMap{Map: kvMap, Str: mapToString(kvMap)} } // Equal reports whether bm and am represent equivalent BuilderMaps. @@ -141,26 +167,19 @@ func (bm BuilderMap) Equal(am BuilderMap) bool { return true } -// builder provides the actual functionality of building RLS keys. These are -// stored in the BuilderMap. -// While processing a pick, the picker looks in the BuilderMap for the -// appropriate builder to be used for the given RPC. For each of the matchers -// in the found builder, we iterate over the list of request headers (available -// as metadata in the context). Once a header matches one of the names in the -// matcher, we set the value of the header in the keyMap (with the key being -// the one found in the matcher) and move on to the next matcher. If no -// KeyBuilder was found in the map, or no header match was found, an empty -// keyMap is returned. +// builder provides the actual functionality of building RLS keys. type builder struct { - matchers []matcher + headerKeys []matcher + constantKeys map[string]string + // The following keys mirror corresponding fields in `extra_keys`. + hostKey string + serviceKey string + methodKey string } // Equal reports whether b and a represent equivalent key builders. func (b builder) Equal(a builder) bool { - if (b.matchers == nil) != (a.matchers == nil) { - return false - } - if len(b.matchers) != len(a.matchers) { + if len(b.headerKeys) != len(a.headerKeys) { return false } // Protobuf serialization maintains the order of repeated fields. Matchers @@ -168,13 +187,23 @@ func (b builder) Equal(a builder) bool { // order changes, it means that the order in the protobuf changed. We report // this case as not being equal even though the builders could possible be // functionally equal. - for i, bMatcher := range b.matchers { - aMatcher := a.matchers[i] + for i, bMatcher := range b.headerKeys { + aMatcher := a.headerKeys[i] if !bMatcher.Equal(aMatcher) { return false } } - return true + + if len(b.constantKeys) != len(a.constantKeys) { + return false + } + for k, v := range b.constantKeys { + if a.constantKeys[k] != v { + return false + } + } + + return b.hostKey == a.hostKey && b.serviceKey == a.serviceKey && b.methodKey == a.methodKey } // matcher helps extract a key from request headers based on a given name. @@ -185,14 +214,11 @@ type matcher struct { names []string } -// Equal reports if m and are are equivalent matchers. +// Equal reports if m and are are equivalent headerKeys. func (m matcher) Equal(a matcher) bool { if m.key != a.key { return false } - if (m.names == nil) != (a.names == nil) { - return false - } if len(m.names) != len(a.names) { return false } @@ -204,9 +230,12 @@ func (m matcher) Equal(a matcher) bool { return true } -func (b builder) keys(md metadata.MD) KeyMap { +func (b builder) buildHeaderKeys(md metadata.MD) map[string]string { kvMap := make(map[string]string) - for _, m := range b.matchers { + if len(md) == 0 { + return kvMap + } + for _, m := range b.headerKeys { for _, name := range m.names { if vals := md.Get(name); vals != nil { kvMap[m.key] = strings.Join(vals, ",") @@ -214,7 +243,7 @@ func (b builder) keys(md metadata.MD) KeyMap { } } } - return KeyMap{Map: kvMap, Str: mapToString(kvMap)} + return kvMap } func mapToString(kv map[string]string) string { diff --git a/balancer/rls/internal/keys/builder_test.go b/balancer/rls/internal/keys/builder_test.go index d7b74b3c0dd4..64ace65bd9a0 100644 --- a/balancer/rls/internal/keys/builder_test.go +++ b/balancer/rls/internal/keys/builder_test.go @@ -37,6 +37,15 @@ var ( {Key: "k1", Names: []string{"n1"}}, {Key: "k2", Names: []string{"n1"}}, }, + ExtraKeys: &rlspb.GrpcKeyBuilder_ExtraKeys{ + Host: "host", + Service: "service", + Method: "method", + }, + ConstantKeys: map[string]string{ + "const-key-1": "const-val-1", + "const-key-2": "const-val-2", + }, } goodKeyBuilder2 = &rlspb.GrpcKeyBuilder{ Names: []*rlspb.GrpcKeyBuilder_Name{ @@ -50,13 +59,21 @@ var ( ) func TestMakeBuilderMap(t *testing.T) { - wantBuilderMap1 := map[string]builder{ - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}}}, + gFooBuilder := builder{ + headerKeys: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}}, + constantKeys: map[string]string{ + "const-key-1": "const-val-1", + "const-key-2": "const-val-2", + }, + hostKey: "host", + serviceKey: "service", + methodKey: "method", } + wantBuilderMap1 := map[string]builder{"/gFoo/": gFooBuilder} wantBuilderMap2 := map[string]builder{ - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}}}, - "/gBar/method1": {matchers: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, - "/gFoobar/": {matchers: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, + "/gFoo/": gFooBuilder, + "/gBar/method1": {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, + "/gFoobar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, } tests := []struct { @@ -91,33 +108,6 @@ func TestMakeBuilderMap(t *testing.T) { } func TestMakeBuilderMapErrors(t *testing.T) { - emptyServiceKeyBuilder := &rlspb.GrpcKeyBuilder{ - Names: []*rlspb.GrpcKeyBuilder_Name{ - {Service: "bFoo", Method: "method1"}, - {Service: "bBar"}, - {Method: "method1"}, - }, - Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, - } - requiredMatchKeyBuilder := &rlspb.GrpcKeyBuilder{ - Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "bFoo", Method: "method1"}}, - Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}, RequiredMatch: true}}, - } - repeatedHeadersKeyBuilder := &rlspb.GrpcKeyBuilder{ - Names: []*rlspb.GrpcKeyBuilder_Name{ - {Service: "gBar", Method: "method1"}, - {Service: "gFoobar"}, - }, - Headers: []*rlspb.NameMatcher{ - {Key: "k1", Names: []string{"n1", "n2"}}, - {Key: "k1", Names: []string{"n1", "n2"}}, - }, - } - methodNameWithSlashKeyBuilder := &rlspb.GrpcKeyBuilder{ - Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "gBar", Method: "method1/foo"}}, - Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, - } - tests := []struct { desc string cfg *rlspb.RouteLookupConfig @@ -138,7 +128,17 @@ func TestMakeBuilderMapErrors(t *testing.T) { { desc: "GrpcKeyBuilder with empty Service field", cfg: &rlspb.RouteLookupConfig{ - GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{emptyServiceKeyBuilder, goodKeyBuilder1}, + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{ + {Service: "bFoo", Method: "method1"}, + {Service: "bBar"}, + {Method: "method1"}, + }, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, + }, + goodKeyBuilder1, + }, }, wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a Name field with no Service", }, @@ -152,21 +152,96 @@ func TestMakeBuilderMapErrors(t *testing.T) { { desc: "GrpcKeyBuilder with requiredMatch field set", cfg: &rlspb.RouteLookupConfig{ - GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{requiredMatchKeyBuilder, goodKeyBuilder1}, + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "bFoo", Method: "method1"}}, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}, RequiredMatch: true}}, + }, + goodKeyBuilder1, + }, }, wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set", }, { desc: "GrpcKeyBuilder two headers with same key", cfg: &rlspb.RouteLookupConfig{ - GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{repeatedHeadersKeyBuilder, goodKeyBuilder1}, + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{ + {Service: "gBar", Method: "method1"}, + {Service: "gFoobar"}, + }, + Headers: []*rlspb.NameMatcher{ + {Key: "k1", Names: []string{"n1", "n2"}}, + {Key: "k1", Names: []string{"n1", "n2"}}, + }, + }, + goodKeyBuilder1, + }, + }, + wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" across headers, constant_keys and extra_keys", + }, + { + desc: "GrpcKeyBuilder repeated keys across headers and constant_keys", + cfg: &rlspb.RouteLookupConfig{ + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{ + {Service: "gBar", Method: "method1"}, + {Service: "gFoobar"}, + }, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, + ConstantKeys: map[string]string{"k1": "v1"}, + }, + goodKeyBuilder1, + }, + }, + wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" across headers, constant_keys and extra_keys", + }, + { + desc: "GrpcKeyBuilder repeated keys across headers and extra_keys", + cfg: &rlspb.RouteLookupConfig{ + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{ + {Service: "gBar", Method: "method1"}, + {Service: "gFoobar"}, + }, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, + ExtraKeys: &rlspb.GrpcKeyBuilder_ExtraKeys{Method: "k1"}, + }, + goodKeyBuilder1, + }, + }, + wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" in extra_keys from constant_keys or headers", + }, + { + desc: "GrpcKeyBuilder repeated keys across constant_keys and extra_keys", + cfg: &rlspb.RouteLookupConfig{ + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{ + {Service: "gBar", Method: "method1"}, + {Service: "gFoobar"}, + }, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, + ConstantKeys: map[string]string{"host": "v1"}, + ExtraKeys: &rlspb.GrpcKeyBuilder_ExtraKeys{Host: "host"}, + }, + goodKeyBuilder1, + }, }, - wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Key field in headers", + wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"host\" in extra_keys from constant_keys or headers", }, { desc: "GrpcKeyBuilder with slash in method name", cfg: &rlspb.RouteLookupConfig{ - GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{methodNameWithSlashKeyBuilder}, + GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{ + { + Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "gBar", Method: "method1/foo"}}, + Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}}, + }, + }, }, wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a method with a slash", }, @@ -257,11 +332,22 @@ func TestRLSKey(t *testing.T) { wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"}, }, { - // Multiple matchers find hits in the provided request headers. - desc: "multipleMatchers", - path: "/gFoo/method1", - md: metadata.Pairs("n2", "v2", "n1", "v1"), - wantKM: KeyMap{Map: map[string]string{"k1": "v1", "k2": "v1"}, Str: "k1=v1,k2=v1"}, + // Multiple headerKeys find hits in the provided request headers. + desc: "multipleMatchers", + path: "/gFoo/method1", + md: metadata.Pairs("n2", "v2", "n1", "v1"), + wantKM: KeyMap{ + Map: map[string]string{ + "const-key-1": "const-val-1", + "const-key-2": "const-val-2", + "host": "dummy-host", + "service": "/gFoo/", + "method": "method1", + "k1": "v1", + "k2": "v1", + }, + Str: "const-key-1=const-val-1,const-key-2=const-val-2,host=dummy-host,k1=v1,k2=v1,method=method1,service=/gFoo/", + }, }, { // A match is found for a header which is specified multiple times. @@ -275,7 +361,7 @@ func TestRLSKey(t *testing.T) { for _, test := range tests { t.Run(test.desc, func(t *testing.T) { - if gotKM := bm.RLSKey(test.md, test.path); !cmp.Equal(gotKM, test.wantKM) { + if gotKM := bm.RLSKey(test.md, "dummy-host", test.path); !cmp.Equal(gotKM, test.wantKM) { t.Errorf("RLSKey(%+v, %s) = %+v, want %+v", test.md, test.path, gotKM, test.wantKM) } }) @@ -351,57 +437,57 @@ func TestBuilderMapEqual(t *testing.T) { { desc: "nil and non-nil builder maps", a: nil, - b: map[string]builder{"/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}}, + b: map[string]builder{"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}}, wantEqual: false, }, { desc: "empty and non-empty builder maps", a: make(map[string]builder), - b: map[string]builder{"/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}}, + b: map[string]builder{"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}}, wantEqual: false, }, { desc: "different number of map keys", a: map[string]builder{ - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, b: map[string]builder{ - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, wantEqual: false, }, { desc: "different map keys", a: map[string]builder{ - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, b: map[string]builder{ - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, wantEqual: false, }, { desc: "equal keys different values", a: map[string]builder{ - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}}, }, b: map[string]builder{ - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, wantEqual: false, }, { desc: "good match", a: map[string]builder{ - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, b: map[string]builder{ - "/gBar/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, - "/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, + "/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}, }, wantEqual: true, }, @@ -425,44 +511,80 @@ func TestBuilderEqual(t *testing.T) { }{ { desc: "nil builders", - a: builder{matchers: nil}, - b: builder{matchers: nil}, + a: builder{headerKeys: nil}, + b: builder{headerKeys: nil}, wantEqual: true, }, { desc: "empty builders", - a: builder{matchers: []matcher{}}, - b: builder{matchers: []matcher{}}, + a: builder{headerKeys: []matcher{}}, + b: builder{headerKeys: []matcher{}}, wantEqual: true, }, { - desc: "nil and non-nil builders", - a: builder{matchers: nil}, - b: builder{matchers: []matcher{}}, + desc: "empty and non-empty builders", + a: builder{headerKeys: []matcher{}}, + b: builder{headerKeys: []matcher{{key: "foo"}}}, wantEqual: false, }, { - desc: "empty and non-empty builders", - a: builder{matchers: []matcher{}}, - b: builder{matchers: []matcher{{key: "foo"}}}, + desc: "different number of headerKeys", + a: builder{headerKeys: []matcher{{key: "foo"}, {key: "bar"}}}, + b: builder{headerKeys: []matcher{{key: "foo"}}}, wantEqual: false, }, { - desc: "different number of matchers", - a: builder{matchers: []matcher{{key: "foo"}, {key: "bar"}}}, - b: builder{matchers: []matcher{{key: "foo"}}}, + desc: "equal number but differing headerKeys", + a: builder{headerKeys: []matcher{{key: "bar"}}}, + b: builder{headerKeys: []matcher{{key: "foo"}}}, wantEqual: false, }, { - desc: "equal number but differing matchers", - a: builder{matchers: []matcher{{key: "bar"}}}, - b: builder{matchers: []matcher{{key: "foo"}}}, + desc: "different number of constantKeys", + a: builder{constantKeys: map[string]string{"k1": "v1"}}, + b: builder{constantKeys: map[string]string{"k1": "v1", "k2": "v2"}}, wantEqual: false, }, { - desc: "good match", - a: builder{matchers: []matcher{{key: "foo"}}}, - b: builder{matchers: []matcher{{key: "foo"}}}, + desc: "equal number but differing constantKeys", + a: builder{constantKeys: map[string]string{"k1": "v1"}}, + b: builder{constantKeys: map[string]string{"k2": "v2"}}, + wantEqual: false, + }, + { + desc: "different hostKey", + a: builder{hostKey: "host1"}, + b: builder{hostKey: "host2"}, + wantEqual: false, + }, + { + desc: "different serviceKey", + a: builder{hostKey: "service1"}, + b: builder{hostKey: "service2"}, + wantEqual: false, + }, + { + desc: "different methodKey", + a: builder{hostKey: "method1"}, + b: builder{hostKey: "method2"}, + wantEqual: false, + }, + { + desc: "equal", + a: builder{ + headerKeys: []matcher{{key: "foo"}}, + constantKeys: map[string]string{"k1": "v1"}, + hostKey: "host", + serviceKey: "/service/", + methodKey: "method", + }, + b: builder{ + headerKeys: []matcher{{key: "foo"}}, + constantKeys: map[string]string{"k1": "v1"}, + hostKey: "host", + serviceKey: "/service/", + methodKey: "method", + }, wantEqual: true, }, } diff --git a/balancer/rls/internal/picker.go b/balancer/rls/internal/picker.go index 738449446558..37e58759e256 100644 --- a/balancer/rls/internal/picker.go +++ b/balancer/rls/internal/picker.go @@ -41,6 +41,9 @@ type rlsPicker struct { // The keyBuilder map used to generate RLS keys for the RPC. This is built // by the LB policy based on the received ServiceConfig. kbm keys.BuilderMap + // Endpoint from the user's original dial target. Used to set the `host_key` + // field in `extra_keys`. + origEndpoint string // The following hooks are setup by the LB policy to enable the rlsPicker to // access state stored in the policy. This approach has the following @@ -76,14 +79,9 @@ type rlsPicker struct { // Pick makes the routing decision for every outbound RPC. func (p *rlsPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { - // For every incoming request, we first build the RLS keys using the - // keyBuilder we received from the LB policy. If no metadata is present in - // the context, we end up using an empty key. - km := keys.KeyMap{} - md, ok := metadata.FromOutgoingContext(info.Ctx) - if ok { - km = p.kbm.RLSKey(md, info.FullMethodName) - } + // Build the request's keys using the key builders from LB config. + md, _ := metadata.FromOutgoingContext(info.Ctx) + km := p.kbm.RLSKey(md, p.origEndpoint, info.FullMethodName) // We use the LB policy hook to read the data cache and the pending request // map (whether or not an entry exists) for the RPC path and the generated