From e73bfc234642e8d3e6c5a9303febb284a2cfbd42 Mon Sep 17 00:00:00 2001
From: horpto <__singleton__@hackerdom.ru>
Date: Fri, 19 Aug 2022 20:01:22 +0300
Subject: [PATCH] Add helper for getting value from incoming context

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.
---
 metadata/metadata.go      | 38 ++++++++++++++++++++++++++++++++++----
 metadata/metadata_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/metadata/metadata.go b/metadata/metadata.go
index 8e0f6abe89d7..346e4a22b351 100644
--- a/metadata/metadata.go
+++ b/metadata/metadata.go
@@ -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)
@@ -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])
@@ -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
@@ -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
@@ -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
diff --git a/metadata/metadata_test.go b/metadata/metadata_test.go
index 89be06eaada0..3980e0d40658 100644
--- a/metadata/metadata_test.go
+++ b/metadata/metadata_test.go
@@ -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)