Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

rls: support extra_keys and constant_keys #4995

Merged
merged 3 commits into from Dec 3, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
121 changes: 75 additions & 46 deletions balancer/rls/internal/keys/builder.go
Expand Up @@ -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 {
Expand All @@ -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)
Expand Down Expand Up @@ -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)}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to store Str in KeyMap? Or can we compute it when we add/retrieve from the map? (Basically: if we're only consuming Str once, we shouldn't store it.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RLSKey method will be invoked for every Pick, and the returned keys (string form) will be used to lookup the data cache. There is no sharing of the key map between RPCs. So, we can either compute it here or Pick will have to compute it right after it calls this method. I don't have much of a preference and I'm leaning towards leaving it as is, to keep Pick working the way it currently does. Do you have a preference?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For now it seems better to avoid storing it to me, but let's wait for the rest of the changes to decide if it's worth changing.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ack.

}

// Equal reports whether bm and am represent equivalent BuilderMaps.
Expand All @@ -141,40 +167,43 @@ 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
// are specified as a repeated field inside the KeyBuilder proto. If the
// 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.
Expand All @@ -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
}
Expand All @@ -204,17 +230,20 @@ 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, ",")
break
}
}
}
return KeyMap{Map: kvMap, Str: mapToString(kvMap)}
return kvMap
}

func mapToString(kv map[string]string) string {
Expand Down