metadata: add ValueFromIncomingContext to more efficiently retrieve a single value

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,19 +182,46 @@ 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
 		// created using our helper functions.
 		key := strings.ToLower(k)
-		s := make([]string, len(v))
-		copy(s, v)
-		out[key] = s
+		out[key] = copyOf(v)
 	}
 	return out, true
 }
 
+// ValueFromIncomingContext returns the metadata value corresponding to the metadata
+// key from the incoming metadata if it exists. Key must be lower-case.
+func ValueFromIncomingContext(ctx context.Context, key string) []string {
+	md, ok := ctx.Value(mdIncomingKey{}).(MD)
+	if !ok {
+		return nil
+	}
+
+	if v, ok := md[key]; ok {
+		return copyOf(v)
+	}
+	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 {
+			return copyOf(v)
+		}
+	}
+	return nil
+}
+
+// the returned slice must not be modified in place
+func copyOf(v []string) (vals []string) {
+	vals = make([]string, len(v))
+	copy(vals, v)
+	return
+}
+
 // FromOutgoingContextRaw returns the un-merged, intermediary contents of rawMD.
 //
 // Remember to perform strings.ToLower on the keys, for both the returned MD (MD
@@ -222,15 +249,18 @@ func FromOutgoingContext(ctx context.Context) (MD, bool) {
 		return nil, false
 	}
 
-	out := MD{}
+	mdSize := len(raw.md)
+	for i := range raw.added {
+		mdSize += len(raw.added[i]) / 2
+	}
+
+	out := make(MD, mdSize)
 	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
 		// created using our helper functions.
 		key := strings.ToLower(k)
-		s := make([]string, len(v))
-		copy(s, v)
-		out[key] = s
+		out[key] = copyOf(v)
 	}
 	for _, added := range raw.added {
 		if len(added)%2 == 1 {

metadata/metadata_test.go

@@ -198,6 +198,43 @@ 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",
+		"x-my-header-3", "44",
+	)
+	ctx := NewIncomingContext(context.Background(), md)
+
+	for _, test := range []struct {
+		key  string
+		want []string
+	}{
+		{
+			key:  "x-my-header-1",
+			want: []string{"42"},
+		},
+		{
+			key:  "x-my-header-2",
+			want: []string{"43-1", "43-2"},
+		},
+		{
+			key:  "x-my-header-3",
+			want: []string{"44"},
+		},
+		{
+			key:  "x-unknown",
+			want: nil,
+		},
+	} {
+		v := ValueFromIncomingContext(ctx, test.key)
+		if !reflect.DeepEqual(v, test.want) {
+			t.Errorf("value of metadata is %v, want %v", v, test.want)
+		}
+	}
+}
+
 func (s) TestAppendToOutgoingContext(t *testing.T) {
 	// Pre-existing metadata
 	tCtx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)