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.

Sign up for GitHub

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

Jump to bottom

rls: support extra_keys and constant_keys #4995

Merged
merged 3 commits into from Dec 3, 2021
Merged

rls: support extra_keys and constant_keys #4995

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6a6155d
rls: support extra_keys and constant_keys
easwars Nov 19, 2021
a4ec23f
review comments
easwars Nov 29, 2021
bdccb6b
use the constant_keys map from the proto
easwars Dec 3, 2021
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
116 changes: 70 additions & 46 deletions balancer/rls/internal/keys/builder.go
View file
Open in desktop
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,41 @@ 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 := make(map[string]string)
for k, v := range kb.GetConstantKeys() {
seenKeys[k] = true
constantKeys[k] = v
Copy link
Member

Choose a reason for hiding this comment

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

constantKeys becomes a copy of kb.GetConstantKeys(). Do we need to make a copy value by value instead of using it directly?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

If we use it directly, then the proto cannot be garbage collected, right. Not sure if this is a big deal either way. Wdyt?

Copy link
Member

Choose a reason for hiding this comment

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

Why? The proto that references the map can still be GC'd can't it?

To be clear, the only thing I'm suggesting is:

constantKeys := kb.GetConstantKeys()

Copy link
Contributor Author

Choose a reason for hiding this comment

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

You are right. Done.

}
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 across headers, constant_keys and extra_keys {%+v}", kbs)
Copy link
Member

Choose a reason for hiding this comment

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

Please include key which was repeated. And this can leave out extra_keys since that hasn't been checked yet.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Done.

}
seenKeys[key] = true
matchers = append(matchers, matcher{key: h.GetKey(), names: h.GetNames()})
}
b := builder{matchers: matchers}
if seenKeys[kb.GetExtraKeys().GetHost()] || seenKeys[kb.GetExtraKeys().GetService()] || seenKeys[kb.GetExtraKeys().GetMethod()] {
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key across headers, constant_keys and extra_keys {%+v}", kbs)
Copy link
Member

Choose a reason for hiding this comment

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

Same: include keys and we know it's extra_keys that has the dupe. So "extra_keys contains duplicate key from constant_keys or headers" or something.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Done.

}
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 +114,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 +162,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 field in `extra_keys`.
Copy link
Member

Choose a reason for hiding this comment

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

s/field/fields/

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Done.

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 +209,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 +225,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 md == nil {
Copy link
Member

Choose a reason for hiding this comment

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

if len(md) == 0 is probably even better.

Also it seems this is an allocation in the RPC path, so I'd return nil and move the make down below this if to avoid any unnecessary allocations.

And if it's possible to avoid this allocation altogether, that would be even better.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Switched to checking for non-zero length.

We cannot avoid this allocation as the caller of this method, RLSKey(), will try to add extra_keys and constant_keys to the returned map. We could add the following block before every possible assignment in RLSKey

if kvMap == nil {
  kvMap = make(map[string]string)
}

but that would be a lot of boilerplate. Wdyt?

Copy link
Member

Choose a reason for hiding this comment

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

Oops, missed this. No, what you have is fine in that case.

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