fix the "i don't know who finished" case

Author: Michal Witkowski

proxy/handler.go

@@ -6,10 +6,10 @@ package proxy
 import (
 	"io"
 
+	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/transport"
-	"golang.org/x/net/context"
 )
 
 var (
@@ -75,24 +75,39 @@ func (s *handler) handler(srv interface{}, serverStream grpc.ServerStream) error
 	if err != nil {
 		return err
 	}
-
 	s2cErrChan := s.forwardServerToClient(serverStream, clientStream)
+	defer close(s2cErrChan)
 	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 {
-		return c2sErr
+	defer close(c2sErrChan)
+	// We don't know which side is going to stop sending first, so we need a select between the two.
+	for i := 0; i < 2; i++ {
+		select {
+		case s2cErr := <-s2cErrChan:
+			if s2cErr == io.EOF {
+				// this is the happy case where the sender has encountered io.EOF, and won't be sending anymore./
+				// the clientStream>serverStream may continue pumping though.
+				clientStream.CloseSend()
+				break
+			} else {
+				// however, we may have gotten a receive error (stream disconnected, a read error etc) in which case we need
+				// to cancel the clientStream to the backend, let all of its goroutines be freed up by the CancelFunc and
+				// exit with an error to the stack
+				clientCancel()
+				return grpc.Errorf(codes.Internal, "failed proxying s2c: %v", s2cErr)
+			}
+		case c2sErr := <-c2sErrChan:
+			// This happens when the clientStream has nothing else to offer (io.EOF), returned a gRPC error. In those two
+			// cases we may have received Trailers as part of the call. In case of other errors (stream closed) the trailers
+			// will be nil.
+			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 {
+				return c2sErr
+			}
+			return nil
+		}
 	}
-	return nil
+	return grpc.Errorf(codes.Internal, "gRPC proxying should never reach this stage.")
 }
 
 func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerStream) chan error {
@@ -123,7 +138,6 @@ func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerSt
 				break
 			}
 		}
-		close(ret)
 	}()
 	return ret
 }
@@ -134,15 +148,16 @@ func (s *handler) forwardServerToClient(src grpc.ServerStream, dst grpc.ClientSt
 		f := &frame{}
 		for i := 0; ; i++ {
 			if err := src.RecvMsg(f); err != nil {
+				//grpclog.Printf("s2c err: %v", err)
 				ret <- err // this can be io.EOF which is happy case
 				break
 			}
 			if err := dst.SendMsg(f); err != nil {
+				//grpclog.Printf("s2c err: %v", err)
 				ret <- err
 				break
 			}
 		}
-		close(ret)
 	}()
 	return ret
 }

proxy/handler_test.go

@@ -210,12 +210,12 @@ func (s *ProxyHappySuite) SetupSuite() {
 		"Ping")
 
 	// Start the serving loops.
+	s.T().Logf("starting grpc.Server at: %v", s.serverListener.Addr().String())
 	go func() {
-		s.T().Logf("starting grpc.Server at: %v", s.serverListener.Addr().String())
 		s.server.Serve(s.serverListener)
 	}()
+	s.T().Logf("starting grpc.Proxy at: %v", s.proxyListener.Addr().String())
 	go func() {
-		s.T().Logf("starting grpc.Proxy at: %v", s.proxyListener.Addr().String())
 		s.proxy.Serve(s.proxyListener)
 	}()