metadata: provide AppendToOutgoingContext interface (#1794)

- Provide AppendToOutgoingContext interface for imrproved performance
over manually creating md and joining with existing md
- Add benchmarks for old/new approaches

Fixes #1390

Documentation/grpc-metadata.md

@@ -82,13 +82,16 @@ func (s *server) SomeRPC(ctx context.Context, in *pb.SomeRequest) (*pb.SomeRespo
 
 ### Sending metadata
 
-To send metadata to server, the client can wrap the metadata into a context using `NewOutgoingContext`, and make the RPC with this context:
+There are two ways to send metadata to the server. The recommended way is to append kv pairs to the context using
+`AppendToOutgoingContext`. This can be used with or without existing metadata on the context. When there is no prior
+metadata, metadata is added; when metadata already exists on the context, kv pairs are merged in.
 
 ```go
-md := metadata.Pairs("key", "val")
+// create a new context with some metadata
+ctx := metadata.AppendToOutgoingContext(ctx, "k1", "v1", "k1", "v2", "k2", "v3")
 
-// create a new context with this metadata
-ctx := metadata.NewOutgoingContext(context.Background(), md)
+// later, add some more metadata to the context (e.g. in an interceptor)
+ctx := metadata.AppendToOutgoingContext(ctx, "k3", "v4")
 
 // make unary RPC
 response, err := client.SomeRPC(ctx, someRequest)
@@ -97,7 +100,27 @@ response, err := client.SomeRPC(ctx, someRequest)
 stream, err := client.SomeStreamingRPC(ctx)
 ```
 
-To read this back from the context on the client (e.g. in an interceptor) before the RPC is sent, use `FromOutgoingContext`.
+Alternatively, metadata may be attached to the context using `NewOutgoingContext`. However, this
+replaces any existing metadata in the context, so care must be taken to preserve the existing
+metadata if desired. This is slower than using `AppendToOutgoingContext`. An example of this
+is below:
+
+```go
+// create a new context with some metadata
+md := metadata.Pairs("k1", "v1", "k1", "v2", "k2", "v3")
+ctx := metadata.NewOutgoingContext(context.Background(), md)
+
+// later, add some more metadata to the context (e.g. in an interceptor)
+md, _ := metadata.FromOutgoingContext(ctx)
+newMD := metadata.Pairs("k3", "v3")
+ctx = metadata.NewContext(ctx, metadata.Join(metadata.New(send), newMD))
+
+// make unary RPC
+response, err := client.SomeRPC(ctx, someRequest)
+
+// or make streaming RPC
+stream, err := client.SomeStreamingRPC(ctx)
+```
 
 ### Receiving metadata
 

metadata/metadata.go

@@ -116,9 +116,22 @@ func NewIncomingContext(ctx context.Context, md MD) context.Context {
 	return context.WithValue(ctx, mdIncomingKey{}, md)
 }
 
-// NewOutgoingContext creates a new context with outgoing md attached.
+// NewOutgoingContext creates a new context with outgoing md attached. If used
+// in conjunction with AppendToOutgoingContext, NewOutgoingContext will
+// overwrite any previously-appended metadata.
 func NewOutgoingContext(ctx context.Context, md MD) context.Context {
-	return context.WithValue(ctx, mdOutgoingKey{}, md)
+	return context.WithValue(ctx, mdOutgoingKey{}, rawMD{md: md})
+}
+
+// AppendToOutgoingContext returns a new context with the provided kv merged
+// with any existing metadata in the context. Please refer to the
+// documentation of Pairs for a description of kv.
+func AppendToOutgoingContext(ctx context.Context, kv ...string) context.Context {
+	if len(kv)%2 == 1 {
+		panic(fmt.Sprintf("metadata: AppendToOutgoingContext got an odd number of input pairs for metadata: %d", len(kv)))
+	}
+	md, _ := ctx.Value(mdOutgoingKey{}).(rawMD)
+	return context.WithValue(ctx, mdOutgoingKey{}, rawMD{md: md.md, added: append(md.added, kv)})
 }
 
 // FromIncomingContext returns the incoming metadata in ctx if it exists.  The
@@ -129,10 +142,39 @@ func FromIncomingContext(ctx context.Context) (md MD, ok bool) {
 	return
 }
 
+// FromOutgoingContextRaw returns the un-merged, intermediary contents
+// of rawMD. Remember to perform strings.ToLower on the keys. The returned
+// MD should not be modified. Writing to it may cause races. Modification
+// should be made to copies of the returned MD.
+//
+// This is intended for gRPC-internal use ONLY.
+func FromOutgoingContextRaw(ctx context.Context) (MD, [][]string, bool) {
+	raw, ok := ctx.Value(mdOutgoingKey{}).(rawMD)
+	if !ok {
+		return nil, nil, false
+	}
+
+	return raw.md, raw.added, true
+}
+
 // FromOutgoingContext returns the outgoing metadata in ctx if it exists.  The
 // returned MD should not be modified. Writing to it may cause races.
 // Modification should be made to the copies of the returned MD.
-func FromOutgoingContext(ctx context.Context) (md MD, ok bool) {
-	md, ok = ctx.Value(mdOutgoingKey{}).(MD)
-	return
+func FromOutgoingContext(ctx context.Context) (MD, bool) {
+	raw, ok := ctx.Value(mdOutgoingKey{}).(rawMD)
+	if !ok {
+		return nil, false
+	}
+
+	mds := make([]MD, 0, len(raw.added)+1)
+	mds = append(mds, raw.md)
+	for _, vv := range raw.added {
+		mds = append(mds, Pairs(vv...))
+	}
+	return Join(mds...), ok
+}
+
+type rawMD struct {
+	md    MD
+	added [][]string
 }

metadata/metadata_test.go

@@ -21,6 +21,8 @@ package metadata
 import (
 	"reflect"
 	"testing"
+
+	"golang.org/x/net/context"
 )
 
 func TestPairsMD(t *testing.T) {
@@ -69,3 +71,55 @@ func TestJoin(t *testing.T) {
 		}
 	}
 }
+
+func TestAppendToOutgoingContext(t *testing.T) {
+	// Pre-existing metadata
+	ctx := NewOutgoingContext(context.Background(), Pairs("k1", "v1", "k2", "v2"))
+	ctx = AppendToOutgoingContext(ctx, "k1", "v3")
+	ctx = AppendToOutgoingContext(ctx, "k1", "v4")
+	md, ok := FromOutgoingContext(ctx)
+	if !ok {
+		t.Errorf("Expected MD to exist in ctx, but got none")
+	}
+	want := Pairs("k1", "v1", "k1", "v3", "k1", "v4", "k2", "v2")
+	if !reflect.DeepEqual(md, want) {
+		t.Errorf("context's metadata is %v, want %v", md, want)
+	}
+
+	// No existing metadata
+	ctx = AppendToOutgoingContext(context.Background(), "k1", "v1")
+	md, ok = FromOutgoingContext(ctx)
+	if !ok {
+		t.Errorf("Expected MD to exist in ctx, but got none")
+	}
+	want = Pairs("k1", "v1")
+	if !reflect.DeepEqual(md, want) {
+		t.Errorf("context's metadata is %v, want %v", md, want)
+	}
+}
+
+// Old/slow approach to adding metadata to context
+func Benchmark_AddingMetadata_ContextManipulationApproach(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		ctx := context.Background()
+		md, _ := FromOutgoingContext(ctx)
+		NewOutgoingContext(ctx, Join(Pairs("k1", "v1", "k2", "v2"), md))
+	}
+}
+
+// Newer/faster approach to adding metadata to context
+func BenchmarkAppendToOutgoingContext(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		AppendToOutgoingContext(context.Background(), "k1", "v1", "k2", "v2")
+	}
+}
+
+func BenchmarkFromOutgoingContext(b *testing.B) {
+	ctx := context.Background()
+	ctx = NewOutgoingContext(ctx, MD{"k3": {"v3", "v4"}})
+	ctx = AppendToOutgoingContext(ctx, "k1", "v1", "k2", "v2")
+
+	for n := 0; n < b.N; n++ {
+		FromOutgoingContext(ctx)
+	}
+}

transport/http2_client.go

@@ -464,7 +464,22 @@ func (t *http2Client) NewStream(ctx context.Context, callHdr *CallHdr) (_ *Strea
 	if b := stats.OutgoingTrace(ctx); b != nil {
 		headerFields = append(headerFields, hpack.HeaderField{Name: "grpc-trace-bin", Value: encodeBinHeader(b)})
 	}
-	if md, ok := metadata.FromOutgoingContext(ctx); ok {
+
+	if md, added, ok := metadata.FromOutgoingContextRaw(ctx); ok {
+		var k string
+		for _, vv := range added {
+			for i, v := range vv {
+				if i%2 == 0 {
+					k = v
+					continue
+				}
+				// HTTP doesn't allow you to set pseudoheaders after non pseudoheaders were set.
+				if isReservedHeader(k) {
+					continue
+				}
+				headerFields = append(headerFields, hpack.HeaderField{Name: strings.ToLower(k), Value: encodeMetadataHeader(k, v)})
+			}
+		}
 		for k, vv := range md {
 			// HTTP doesn't allow you to set pseudoheaders after non pseudoheaders were set.
 			if isReservedHeader(k) {