fix full duplex streaming

Author: Michal Witkowski

proxy/handler.go

@@ -9,6 +9,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/transport"
+	"golang.org/x/net/context"
 )
 
 var (
@@ -64,21 +65,28 @@ func (s *handler) handler(srv interface{}, serverStream grpc.ServerStream) error
 		return grpc.Errorf(codes.Internal, "lowLevelServerStream not exists in context")
 	}
 	fullMethodName := lowLevelServerStream.Method()
+	clientCtx, clientCancel := context.WithCancel(serverStream.Context())
 	backendConn, err := s.director(serverStream.Context(), fullMethodName)
 	if err != nil {
 		return err
 	}
 	// TODO(mwitkow): Add a `forwarded` header to metadata, https://en.wikipedia.org/wiki/X-Forwarded-For.
-	clientStream, err := grpc.NewClientStream(serverStream.Context(), clientStreamDescForProxying, backendConn, fullMethodName)
+	clientStream, err := grpc.NewClientStream(clientCtx, clientStreamDescForProxying, backendConn, fullMethodName)
 	if err != nil {
 		return err
 	}
-	defer clientStream.CloseSend() // always close this!
-	s2cErr := <-s.forwardServerToClient(serverStream, clientStream)
-	c2sErr := <-s.forwardClientToServer(clientStream, serverStream)
+
+	s2cErrChan := s.forwardServerToClient(serverStream, clientStream)
+	c2sErrChan := s.forwardClientToServer(clientStream, serverStream)
+	s2cErr := <-s2cErrChan
 	if s2cErr != io.EOF {
+		clientCancel()
 		return grpc.Errorf(codes.Internal, "failed proxying s2c: %v", s2cErr)
+	} else {
+		clientStream.CloseSend()
 	}
+	c2sErr := <-c2sErrChan
+
 	serverStream.SetTrailer(clientStream.Trailer())
 	// c2sErr will contain RPC error from client code. If not io.EOF return the RPC error as server stream error.
 	if c2sErr != io.EOF {

proxy/handler_test.go

@@ -22,6 +22,8 @@ import (
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/metadata"
 
+	"fmt"
+
 	pb "github.com/mwitkow/grpc-proxy/testservice"
 )
 
@@ -72,6 +74,27 @@ func (s *assertingService) PingList(ping *pb.PingRequest, stream pb.TestService_
 	return nil
 }
 
+func (s *assertingService) PingStream(stream pb.TestService_PingStreamServer) error {
+	stream.SendHeader(metadata.Pairs(serverHeaderMdKey, "I like turtles."))
+	counter := int32(0)
+	for {
+		ping, err := stream.Recv()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			require.NoError(s.t, err, "can't fail reading stream")
+			return err
+		}
+		pong := &pb.PingResponse{Value: ping.Value, Counter: counter}
+		if err := stream.Send(pong); err != nil {
+			require.NoError(s.t, err, "can't fail sending back a pong")
+		}
+		counter += 1
+	}
+	stream.SetTrailer(metadata.Pairs(serverTrailerMdKey, "I like ending turtles."))
+	return nil
+}
+
 // ProxyHappySuite tests the "happy" path of handling: that everything works in absence of connection issues.
 type ProxyHappySuite struct {
 	suite.Suite
@@ -125,24 +148,28 @@ func (s *ProxyHappySuite) TestDirectorErrorIsPropagated() {
 	assert.Equal(s.T(), "testing rejection", grpc.ErrorDesc(err))
 }
 
-func (s *ProxyHappySuite) TestPingListStreamsAll() {
-	stream, err := s.testClient.PingList(s.ctx(), &pb.PingRequest{Value: "foo"})
-	require.NoError(s.T(), err, "PingList request should be successful.")
-	// Check that the header arrives before all entries.
-	headerMd, err := stream.Header()
-	require.NoError(s.T(), err, "PingList headers should not error.")
-	assert.Len(s.T(), headerMd, 1, "PingList response headers user contain metadata")
-	count := 0
-	for {
+func (s *ProxyHappySuite) TestPingStream_FullDuplexWorks() {
+	stream, err := s.testClient.PingStream(s.ctx())
+	require.NoError(s.T(), err, "PingStream request should be successful.")
+
+	for i := 0; i < countListResponses; i++ {
+		ping := &pb.PingRequest{Value: fmt.Sprintf("foo:%d", i)}
+		require.NoError(s.T(), stream.Send(ping), "sending to PingStream must not fail")
 		resp, err := stream.Recv()
 		if err == io.EOF {
 			break
 		}
-		require.NoError(s.T(), err, "PingList stream should not be interrupted.")
-		require.Equal(s.T(), "foo", resp.Value)
-		count = count + 1
+		if i == 0 {
+			// Check that the header arrives before all entries.
+			headerMd, err := stream.Header()
+			require.NoError(s.T(), err, "PingStream headers should not error.")
+			assert.Len(s.T(), headerMd, 1, "PingStream response headers user contain metadata")
+		}
+		assert.EqualValues(s.T(), i, resp.Counter, "ping roundtrip must succeed with the correct id")
 	}
-	assert.Equal(s.T(), countListResponses, count, "PingList must successfully return all outputs")
+	require.NoError(s.T(), stream.CloseSend(), "no error on close send")
+	_, err = stream.Recv()
+	require.Equal(s.T(), io.EOF, err, "stream should close with io.EOF, meaining OK")
 	// Check that the trailer headers are here.
 	trailerMd := stream.Trailer()
 	assert.Len(s.T(), trailerMd, 1, "PingList trailer headers user contain metadata")

testservice/test.pb.go

@@ -22,5 +22,8 @@ service TestService {
   rpc PingError(PingRequest) returns (Empty) {}
 
   rpc PingList(PingRequest) returns (stream PingResponse) {}
+
+  rpc PingStream(stream PingRequest) returns (stream PingResponse) {}
+
 }