Skip to content

Commit

Permalink
Add helper for getting value from incoming context
Browse files Browse the repository at this point in the history
FromIncomingContext creates a copy of md every time.
So even if you need a single header value you will get a value and an overhead for this.
Many interceptors and user code doesn't need all values at a time.
  • Loading branch information
horpto committed Aug 19, 2022
1 parent 92cee34 commit e73bfc2
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 4 deletions.
38 changes: 34 additions & 4 deletions metadata/metadata.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ type MD map[string][]string
// Keys beginning with "grpc-" are reserved for grpc-internal use only and may
// result in errors if set in metadata.
func New(m map[string]string) MD {
md := MD{}
md := make(MD, len(m))
for k, val := range m {
key := strings.ToLower(k)
md[key] = append(md[key], val)
Expand All @@ -74,7 +74,7 @@ func Pairs(kv ...string) MD {
if len(kv)%2 == 1 {
panic(fmt.Sprintf("metadata: Pairs got the odd number of input pairs for metadata: %d", len(kv)))
}
md := MD{}
md := make(MD, len(kv)/2)
for i := 0; i < len(kv); i += 2 {
key := strings.ToLower(kv[i])
md[key] = append(md[key], kv[i+1])
Expand Down Expand Up @@ -182,7 +182,7 @@ func FromIncomingContext(ctx context.Context) (MD, bool) {
if !ok {
return nil, false
}
out := MD{}
out := make(MD, len(md))
for k, v := range md {
// We need to manually convert all keys to lower case, because MD is a
// map, and there's no guarantee that the MD attached to the context is
Expand All @@ -195,6 +195,36 @@ func FromIncomingContext(ctx context.Context) (MD, bool) {
return out, true
}

// ValueFromIncomingContext returns value from the incoming metadata if exists.
//
// This is intended for using as fast access to context value with string constant.
func ValueFromIncomingContext(ctx context.Context, key string) []string {
md, ok := ctx.Value(mdIncomingKey{}).(MD)
if !ok {
return nil
}
// fastpath
if v, ok := md[key]; ok {
res := make([]string, len(v))
copy(res, v)
return res
}

// slowpath
key = strings.ToLower(key)
for k, v := range md {
// We need to manually convert all keys to lower case, because MD is a
// map, and there's no guarantee that the MD attached to the context is
// created using our helper functions.
if strings.ToLower(k) == key {
s := make([]string, len(v))
copy(s, v)
return s
}
}
return nil
}

// FromOutgoingContextRaw returns the un-merged, intermediary contents of rawMD.
//
// Remember to perform strings.ToLower on the keys, for both the returned MD (MD
Expand Down Expand Up @@ -222,7 +252,7 @@ func FromOutgoingContext(ctx context.Context) (MD, bool) {
return nil, false
}

out := MD{}
out := make(MD, len(raw.md)+len(raw.added))
for k, v := range raw.md {
// We need to manually convert all keys to lower case, because MD is a
// map, and there's no guarantee that the MD attached to the context is
Expand Down
25 changes: 25 additions & 0 deletions metadata/metadata_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,31 @@ func (s) TestDelete(t *testing.T) {
}
}

func (s) TestValueFromIncomingContext(t *testing.T) {
md := Pairs(
"X-My-Header-1", "42",
"X-My-Header-2", "43-1",
"X-My-Header-2", "43-2",
)
ctx := NewIncomingContext(context.Background(), md)

var v []string
v = ValueFromIncomingContext(ctx, "X-My-Header-1")
if !reflect.DeepEqual(v, []string{"42"}) {
t.Errorf("value from context is %v", v)
}

v = ValueFromIncomingContext(ctx, "x-my-header-1")
if !reflect.DeepEqual(v, []string{"42"}) {
t.Errorf("value from context is %v", v)
}

v = ValueFromIncomingContext(ctx, "x-my-header-2")
if !reflect.DeepEqual(v, []string{"43-1", "43-2"}) {
t.Errorf("value from context is %v", v)
}
}

func (s) TestAppendToOutgoingContext(t *testing.T) {
// Pre-existing metadata
tCtx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
Expand Down

0 comments on commit e73bfc2

Please sign in to comment.