Fix gRPC over cmux and add unit tests

cmd/query/app/grpc_handler_test.go

@@ -40,7 +40,6 @@ import (
 )
 
 var (
-	grpcServerPort       = ":0"
 	errStorageMsgGRPC    = "Storage error"
 	errStorageGRPC       = errors.New(errStorageMsgGRPC)
 	errStatusStorageGRPC = status.Error(2, errStorageMsgGRPC)
@@ -138,7 +137,7 @@ type grpcClient struct {
 }
 
 func newGRPCServer(t *testing.T, q *querysvc.QueryService, logger *zap.Logger, tracer opentracing.Tracer) (*grpc.Server, net.Addr) {
-	lis, _ := net.Listen("tcp", grpcServerPort)
+	lis, _ := net.Listen("tcp", ":0")
 	grpcServer := grpc.NewServer()
 	grpcHandler := NewGRPCHandler(q, logger, tracer)
 	api_v2.RegisterQueryServiceServer(grpcServer, grpcHandler)
@@ -151,8 +150,8 @@ func newGRPCServer(t *testing.T, q *querysvc.QueryService, logger *zap.Logger, t
 	return grpcServer, lis.Addr()
 }
 
-func newGRPCClient(t *testing.T, addr net.Addr) *grpcClient {
-	conn, err := grpc.Dial(addr.String(), grpc.WithInsecure())
+func newGRPCClient(t *testing.T, addr string) *grpcClient {
+	conn, err := grpc.Dial(addr, grpc.WithInsecure())
 	require.NoError(t, err)
 
 	return &grpcClient{
@@ -192,7 +191,7 @@ func initializeTestServerGRPCWithOptions(t *testing.T) *grpcServer {
 
 func withServerAndClient(t *testing.T, actualTest func(server *grpcServer, client *grpcClient)) {
 	server := initializeTestServerGRPCWithOptions(t)
-	client := newGRPCClient(t, server.lisAddr)
+	client := newGRPCClient(t, server.lisAddr.String())
 	defer server.server.Stop()
 	defer client.conn.Close()
 

cmd/query/app/server.go

@@ -104,7 +104,12 @@ func (s *Server) Start() error {
 
 	grpcListener := cmuxServer.Match(
 		cmux.HTTP2HeaderField("content-type", "application/grpc"),
-		cmux.HTTP2HeaderField("content-type", "application/grpc+proto"))
+		cmux.HTTP2HeaderField("content-type", "application/grpc+proto"),
+	)
+	// grpcListener := cmuxServer.MatchWithWriters(
+	// 	cmux.HTTP2MatchHeaderFieldSendSettings("content-type", "application/grpc"),
+	// 	cmux.HTTP2MatchHeaderFieldSendSettings("content-type", "application/grpc+proto"),
+	// )
 	httpListener := cmuxServer.Match(cmux.Any())
 
 	go func() {

cmd/query/app/server_test.go

@@ -15,6 +15,7 @@
 package app
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
@@ -28,6 +29,9 @@ import (
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
 	"github.com/jaegertracing/jaeger/pkg/healthcheck"
 	"github.com/jaegertracing/jaeger/ports"
+	"github.com/jaegertracing/jaeger/proto-gen/api_v2"
+	depsmocks "github.com/jaegertracing/jaeger/storage/dependencystore/mocks"
+	spanstoremocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks"
 )
 
 func TestServerError(t *testing.T) {
@@ -43,13 +47,40 @@ func TestServer(t *testing.T) {
 	flagsSvc := flags.NewService(ports.AgentAdminHTTP)
 	flagsSvc.Logger = zap.NewNop()
 
-	querySvc := &querysvc.QueryService{}
+	spanReader := &spanstoremocks.Reader{}
+	dependencyReader := &depsmocks.Reader{}
+
+	querySvc := querysvc.NewQueryService(spanReader, dependencyReader, querysvc.QueryServiceOptions{})
+
 	tracer := opentracing.NoopTracer{}
 
-	server := NewServer(flagsSvc, querySvc, &QueryOptions{Port: ports.QueryAdminHTTP,
-		BearerTokenPropagation: true}, tracer)
+	server := NewServer(flagsSvc, querySvc,
+		&QueryOptions{Port: ports.QueryAdminHTTP, BearerTokenPropagation: true},
+		tracer)
 	assert.NoError(t, server.Start())
 
+	client := newGRPCClient(t, fmt.Sprintf(":%d", ports.QueryHTTP))
+	defer client.conn.Close()
+
+	var queryErr error
+	for i := 0; i < 10; i++ {
+		queryErr = func() error {
+			ctx, cancel := context.WithTimeout(context.Background(), 2000*time.Millisecond)
+			defer cancel()
+
+			_, err := client.GetTrace(ctx, &api_v2.GetTraceRequest{})
+			if err != nil {
+				t.Log("cannot GetTrace", err)
+			}
+			return err
+		}()
+		if queryErr == nil {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	assert.NoError(t, queryErr, "Connection test did not succeed")
+
 	// TODO wait for servers to come up and test http and grpc endpoints
 	time.Sleep(1 * time.Second)
 
@@ -82,6 +113,6 @@ func TestServerGracefulExit(t *testing.T) {
 
 	for _, logEntry := range logs.All() {
 		assert.True(t, logEntry.Level != zap.ErrorLevel,
-			fmt.Sprintf("Error log found on server exit: %v", logEntry))
+			"Error log found on server exit: %v", logEntry)
 	}
 }