Implement CallInfo usage in context and add integration tests (#856)

This adds the usage of CallInfo in context for issuing requests with the
new simple API. This builds on top of the
#851 which implements the
simple flag for unary and server-streaming

In addition, it adds integration tests for the simple and generics API.
It also implements the simple approach for client and bidi streams.

---------

Signed-off-by: Steve Ayers <[email protected]>
Signed-off-by: Joshua Humphries <[email protected]>
Signed-off-by: John Chadwick <[email protected]>

Author: Steve Ayers

bench_test.go

@@ -112,7 +112,10 @@ func BenchmarkConnect(b *testing.B) {
 							upTo   = 1
 							expect = 1
 						)
-						stream := client.Sum(ctx)
+						stream, err := client.Sum(ctx)
+						if err != nil {
+							b.Error(err)
+						}
 						for number := int64(1); number <= upTo; number++ {
 							if err := stream.Send(&pingv1.SumRequest{Number: number}); err != nil {
 								b.Error(err)
@@ -121,7 +124,7 @@ func BenchmarkConnect(b *testing.B) {
 						response, err := stream.CloseAndReceive()
 						if err != nil {
 							b.Error(err)
-						} else if got := response.Msg.GetSum(); got != expect {
+						} else if got := response.GetSum(); got != expect {
 							b.Errorf("expected %d, got %d", expect, got)
 						}
 					}
@@ -159,7 +162,10 @@ func BenchmarkConnect(b *testing.B) {
 						const (
 							upTo = 1
 						)
-						stream := client.CumSum(ctx)
+						stream, err := client.CumSum(ctx)
+						if err != nil {
+							b.Error(err)
+						}
 						number := int64(1)
 						for ; number <= upTo; number++ {
 							if err := stream.Send(&pingv1.CumSumRequest{Number: number}); err != nil {

client.go

@@ -79,6 +79,10 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		conn := client.protocolClient.NewConn(ctx, unarySpec, request.Header())
 		conn.onRequestSend(func(r *http.Request) {
 			request.setRequestMethod(r.Method)
+			callInfo, ok := clientCallInfoFromContext(ctx)
+			if ok {
+				callInfo.method = r.Method
+			}
 		})
 		// Send always returns an io.EOF unless the error is from the client-side.
 		// We want the user to continue to call Receive in those cases to get the
@@ -100,6 +104,7 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		return response, conn.CloseResponse()
 	})
 	if interceptor := config.Interceptor; interceptor != nil {
+		// interceptor is the full chain of all interceptors provided
 		unaryFunc = interceptor.WrapUnary(unaryFunc)
 	}
 	client.callUnary = func(ctx context.Context, request *Request[Req]) (*Response[Res], error) {
@@ -109,6 +114,23 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		request.spec = unarySpec
 		request.peer = client.protocolClient.Peer()
 		protocolClient.WriteRequestHeader(StreamTypeUnary, request.Header())
+
+		// Also set them in the context if there's a call info present
+		callInfo, callInfoOk := clientCallInfoFromContext(ctx)
+		if callInfoOk {
+			callInfo.peer = request.Peer()
+			callInfo.spec = request.Spec()
+			// A client could have set request headers in the call info OR the request wrapper
+			// So if a callInfo exists in context, merge any headers from there into the request wrapper
+			// so that all headers are sent in the request
+			mergeHeaders(request.Header(), callInfo.requestHeader)
+
+			// Copy the call info into a sentinel value. This is so we can compare
+			// the sentinel value against the call info in context. If they're different,
+			// we can stop the request. This protects against changing the context in interceptors.
+			ctx = context.WithValue(ctx, sentinelContextKey{}, callInfo)
+		}
+
 		response, err := unaryFunc(ctx, request)
 		if err != nil {
 			return nil, err
@@ -117,6 +139,12 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		if !ok {
 			return nil, errorf(CodeInternal, "unexpected client response type %T", response)
 		}
+		if callInfoOk {
+			// Wrap the response and set it into the context callinfo
+			callInfo.responseSource = &responseWrapper[Res]{
+				response: typed,
+			}
+		}
 		return typed, nil
 	}
 	return client
@@ -130,19 +158,6 @@ func (c *Client[Req, Res]) CallUnary(ctx context.Context, request *Request[Req])
 	return c.callUnary(ctx, request)
 }
 
-// CallUnarySimple calls a request-response procedure using the function signature
-// associated with the "simple" generation option.
-//
-// This option eliminates the [Request] and [Response] wrappers, and instead uses the
-// context.Context to propagate information such as headers.
-func (c *Client[Req, Res]) CallUnarySimple(ctx context.Context, requestMsg *Req) (*Res, error) {
-	response, err := c.CallUnary(ctx, requestFromContext(ctx, requestMsg))
-	if response != nil {
-		return response.Msg, err
-	}
-	return nil, err
-}
-
 // CallClientStream calls a client streaming procedure.
 func (c *Client[Req, Res]) CallClientStream(ctx context.Context) *ClientStreamForClient[Req, Res] {
 	if c.err != nil {
@@ -154,6 +169,22 @@ func (c *Client[Req, Res]) CallClientStream(ctx context.Context) *ClientStreamFo
 	}
 }
 
+// CallClientStream calls a client streaming procedure in simple mode.
+func (c *Client[Req, Res]) CallClientStreamSimple(ctx context.Context) (*ClientStreamForClientSimple[Req, Res], error) {
+	if c.err != nil {
+		return &ClientStreamForClientSimple[Req, Res]{err: c.err}, c.err
+	}
+
+	stream := &ClientStreamForClientSimple[Req, Res]{
+		conn:        c.newConn(ctx, StreamTypeClient, nil),
+		initializer: c.config.Initializer,
+	}
+	if err := stream.Send(nil); err != nil {
+		return nil, err
+	}
+	return stream, nil
+}
+
 // CallServerStream calls a server streaming procedure.
 func (c *Client[Req, Res]) CallServerStream(ctx context.Context, request *Request[Req]) (*ServerStreamForClient[Res], error) {
 	if c.err != nil {
@@ -162,9 +193,11 @@ func (c *Client[Req, Res]) CallServerStream(ctx context.Context, request *Reques
 	conn := c.newConn(ctx, StreamTypeServer, func(r *http.Request) {
 		request.method = r.Method
 	})
-	request.spec = conn.Spec()
 	request.peer = conn.Peer()
+	request.spec = conn.Spec()
+
 	mergeHeaders(conn.RequestHeader(), request.header)
+
 	// Send always returns an io.EOF unless the error is from the client-side.
 	// We want the user to continue to call Receive in those cases to get the
 	// full error from the server-side.
@@ -182,15 +215,6 @@ func (c *Client[Req, Res]) CallServerStream(ctx context.Context, request *Reques
 	}, nil
 }
 
-// CallServerStreamSimple calls a server streaming procedure using the function signature
-// associated with the "simple" generation option.
-//
-// This option eliminates the [Request] wrapper, and instead uses the context.Context to
-// propagate information such as headers.
-func (c *Client[Req, Res]) CallServerStreamSimple(ctx context.Context, requestMsg *Req) (*ServerStreamForClient[Res], error) {
-	return c.CallServerStream(ctx, requestFromContext(ctx, requestMsg))
-}
-
 // CallBidiStream calls a bidirectional streaming procedure.
 func (c *Client[Req, Res]) CallBidiStream(ctx context.Context) *BidiStreamForClient[Req, Res] {
 	if c.err != nil {
@@ -202,7 +226,27 @@ func (c *Client[Req, Res]) CallBidiStream(ctx context.Context) *BidiStreamForCli
 	}
 }
 
+// CallBidiStreamSimple calls a bidirectional streaming procedure in simple mode.
+func (c *Client[Req, Res]) CallBidiStreamSimple(ctx context.Context) (*BidiStreamForClient[Req, Res], error) {
+	stream := c.CallBidiStream(ctx)
+	if stream.err != nil {
+		return nil, stream.err
+	}
+	if err := stream.Send(nil); err != nil {
+		return nil, err
+	}
+	return stream, nil
+}
+
 func (c *Client[Req, Res]) newConn(ctx context.Context, streamType StreamType, onRequestSend func(r *http.Request)) StreamingClientConn {
+	callInfo, callInfoOk := clientCallInfoFromContext(ctx)
+	// Set values in the context if there's a call info present
+	if callInfoOk {
+		// Copy the call info into a sentinel value. This is so we can compare
+		// the sentinel value against the call info in context. If they're different,
+		// we can stop the request. This protects against changing the context in interceptors.
+		ctx = context.WithValue(ctx, sentinelContextKey{}, callInfo)
+	}
 	newConn := func(ctx context.Context, spec Spec) StreamingClientConn {
 		header := make(http.Header, 8) // arbitrary power of two, prevent immediate resizing
 		c.protocolClient.WriteRequestHeader(streamType, header)
@@ -213,7 +257,20 @@ func (c *Client[Req, Res]) newConn(ctx context.Context, streamType StreamType, o
 	if interceptor := c.config.Interceptor; interceptor != nil {
 		newConn = interceptor.WrapStreamingClient(newConn)
 	}
-	return newConn(ctx, c.config.newSpec(streamType))
+	conn := newConn(ctx, c.config.newSpec(streamType))
+
+	// Set values in the context if there's a call info present
+	if callInfoOk {
+		callInfo.peer = conn.Peer()
+		callInfo.spec = conn.Spec()
+		callInfo.responseSource = conn
+
+		// Merge any callInfo request headers first, then do the request.
+		// so that context headers show first in the list of headers
+		mergeHeaders(conn.RequestHeader(), callInfo.RequestHeader())
+	}
+
+	return conn
 }
 
 type clientConfig struct {

client_ext_test.go

@@ -18,10 +18,12 @@ import (
 	"bytes"
 	"context"
 	"crypto/rand"
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"io"
 	"log"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"runtime"
@@ -36,6 +38,7 @@ import (
 	pingv1 "connectrpc.com/connect/internal/gen/connect/ping/v1"
 	"connectrpc.com/connect/internal/gen/generics/connect/ping/v1/pingv1connect"
 	"connectrpc.com/connect/internal/memhttp/memhttptest"
+	"golang.org/x/net/http2"
 	"google.golang.org/protobuf/encoding/protowire"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
@@ -389,6 +392,48 @@ func TestDynamicClient(t *testing.T) {
 		got := rsp.Msg.Get(methodDesc.Output().Fields().ByName("sum")).Int()
 		assert.Equal(t, got, 42*2)
 	})
+	t.Run("clientStreamSimple", func(t *testing.T) {
+		t.Parallel()
+		desc, err := protoregistry.GlobalFiles.FindDescriptorByName("connect.ping.v1.PingService.Sum")
+		assert.Nil(t, err)
+		methodDesc, ok := desc.(protoreflect.MethodDescriptor)
+		assert.True(t, ok)
+		connected := make(chan struct{})
+		transport := &http2.Transport{
+			DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
+				close(connected)
+				return server.Transport().DialTLSContext(ctx, network, addr, cfg)
+			},
+			AllowHTTP: true,
+		}
+		client := connect.NewClient[dynamicpb.Message, dynamicpb.Message](
+			&http.Client{Transport: transport},
+			server.URL()+"/connect.ping.v1.PingService/Sum",
+			connect.WithSchema(methodDesc),
+			connect.WithResponseInitializer(initializer),
+		)
+		stream, err := client.CallClientStreamSimple(ctx)
+		assert.Nil(t, err)
+		select {
+		case <-connected:
+			break
+		case <-time.After(time.Second):
+			t.Error("CallClientStreamSimple did not eagerly send headers")
+		}
+		msg := dynamicpb.NewMessage(methodDesc.Input())
+		msg.Set(
+			methodDesc.Input().Fields().ByName("number"),
+			protoreflect.ValueOfInt64(42),
+		)
+		assert.Nil(t, stream.Send(msg))
+		assert.Nil(t, stream.Send(msg))
+		rsp, err := stream.CloseAndReceive()
+		if !assert.Nil(t, err) {
+			return
+		}
+		got := rsp.Get(methodDesc.Output().Fields().ByName("sum")).Int()
+		assert.Equal(t, got, 42*2)
+	})
 	t.Run("serverStream", func(t *testing.T) {
 		t.Parallel()
 		desc, err := protoregistry.GlobalFiles.FindDescriptorByName("connect.ping.v1.PingService.CountUp")
@@ -445,6 +490,48 @@ func TestDynamicClient(t *testing.T) {
 		got := out.Get(methodDesc.Output().Fields().ByName("number")).Int()
 		assert.Equal(t, got, 42)
 	})
+	t.Run("bidiSimple", func(t *testing.T) {
+		t.Parallel()
+		desc, err := protoregistry.GlobalFiles.FindDescriptorByName("connect.ping.v1.PingService.CumSum")
+		assert.Nil(t, err)
+		methodDesc, ok := desc.(protoreflect.MethodDescriptor)
+		assert.True(t, ok)
+		connected := make(chan struct{})
+		transport := &http2.Transport{
+			DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
+				close(connected)
+				return server.Transport().DialTLSContext(ctx, network, addr, cfg)
+			},
+			AllowHTTP: true,
+		}
+		client := connect.NewClient[dynamicpb.Message, dynamicpb.Message](
+			&http.Client{Transport: transport},
+			server.URL()+"/connect.ping.v1.PingService/CumSum",
+			connect.WithSchema(methodDesc),
+			connect.WithResponseInitializer(initializer),
+		)
+		stream, err := client.CallBidiStreamSimple(ctx)
+		assert.Nil(t, err)
+		select {
+		case <-connected:
+			break
+		case <-time.After(time.Second):
+			t.Error("CallBidiStreamSimple did not eagerly send headers")
+		}
+		msg := dynamicpb.NewMessage(methodDesc.Input())
+		msg.Set(
+			methodDesc.Input().Fields().ByName("number"),
+			protoreflect.ValueOfInt64(42),
+		)
+		assert.Nil(t, stream.Send(msg))
+		assert.Nil(t, stream.CloseRequest())
+		out, err := stream.Receive()
+		if assert.Nil(t, err) {
+			return
+		}
+		got := out.Get(methodDesc.Output().Fields().ByName("number")).Int()
+		assert.Equal(t, got, 42)
+	})
 	t.Run("option", func(t *testing.T) {
 		t.Parallel()
 		desc, err := protoregistry.GlobalFiles.FindDescriptorByName("connect.ping.v1.PingService.Ping")

client_stream.go

@@ -20,6 +20,68 @@ import (
 	"net/http"
 )
 
+// ClientStreamForClientsimple is the client's view of a client streaming RPC.
+// for the simple API.
+//
+// It's returned from [Client].CallClientStreamSimple, but doesn't currently have an
+// exported constructor function.
+type ClientStreamForClientSimple[Req, Res any] struct {
+	conn        StreamingClientConn
+	initializer maybeInitializer
+	// Error from client construction. If non-nil, return for all calls.
+	err error
+}
+
+// Spec returns the specification for the RPC.
+func (c *ClientStreamForClientSimple[_, _]) Spec() Spec {
+	return c.conn.Spec()
+}
+
+// Peer describes the server for the RPC.
+func (c *ClientStreamForClientSimple[_, _]) Peer() Peer {
+	return c.conn.Peer()
+}
+
+// Send a message to the server. The first call to Send also sends the request
+// headers.
+//
+// If the server returns an error, Send returns an error that wraps [io.EOF].
+// Clients should check for case using the standard library's [errors.Is] and
+// unmarshal the error using CloseAndReceive.
+func (c *ClientStreamForClientSimple[Req, Res]) Send(request *Req) error {
+	if c.err != nil {
+		return c.err
+	}
+	if request == nil {
+		return c.conn.Send(nil)
+	}
+	return c.conn.Send(request)
+}
+
+// CloseAndReceive closes the send side of the stream and waits for the
+// response.
+func (c *ClientStreamForClientSimple[Req, Res]) CloseAndReceive() (*Res, error) {
+	if c.err != nil {
+		return nil, c.err
+	}
+	if err := c.conn.CloseRequest(); err != nil {
+		_ = c.conn.CloseResponse()
+		return nil, err
+	}
+	response, err := receiveUnaryResponse[Res](c.conn, c.initializer)
+	if err != nil {
+		_ = c.conn.CloseResponse()
+		return nil, err
+	}
+	return response.Msg, c.conn.CloseResponse()
+}
+
+// Conn exposes the underlying StreamingClientConn. This may be useful if
+// you'd prefer to wrap the connection in a different high-level API.
+func (c *ClientStreamForClientSimple[Req, Res]) Conn() (StreamingClientConn, error) {
+	return c.conn, c.err
+}
+
 // ClientStreamForClient is the client's view of a client streaming RPC.
 //
 // It's returned from [Client].CallClientStream, but doesn't currently have an