vttablet: implement conn benchmark

Signed-off-by: Vicent Marti <vmg@strn.cat>

go/vt/vttablet/grpctabletconn/conn_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package grpctabletconn
 
 import (
+	"context"
 	"io"
 	"io/ioutil"
 	"net"
@@ -25,13 +26,64 @@ import (
 
 	"google.golang.org/grpc"
 
+	"vitess.io/vitess/go/vt/callerid"
 	"vitess.io/vitess/go/vt/servenv"
 	"vitess.io/vitess/go/vt/vttablet/grpcqueryservice"
+	"vitess.io/vitess/go/vt/vttablet/tabletconn"
 	"vitess.io/vitess/go/vt/vttablet/tabletconntest"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
+func BenchmarkGRPCTabletConn(b *testing.B) {
+	// fake service
+	service := tabletconntest.CreateFakeServer(b)
+
+	// listen on a random port
+	listener, err := net.Listen("tcp", ":0")
+	if err != nil {
+		b.Fatalf("Cannot listen: %v", err)
+	}
+	host := listener.Addr().(*net.TCPAddr).IP.String()
+	port := listener.Addr().(*net.TCPAddr).Port
+
+	// Create a gRPC server and listen on the port
+	server := grpc.NewServer()
+	grpcqueryservice.Register(server, service)
+	go server.Serve(listener)
+
+	tablet := &topodatapb.Tablet{
+		Keyspace: tabletconntest.TestTarget.Keyspace,
+		Shard:    tabletconntest.TestTarget.Shard,
+		Type:     tabletconntest.TestTarget.TabletType,
+		Alias:    tabletconntest.TestAlias,
+		Hostname: host,
+		PortMap: map[string]int32{
+			"grpc": int32(port),
+		},
+	}
+
+	conn, err := tabletconn.GetDialer()(tablet, false)
+	if err != nil {
+		b.Fatalf("dial failed: %v", err)
+	}
+	defer conn.Close(context.Background())
+
+	service.ExpectedTransactionID = tabletconntest.ExecuteBatchTransactionID
+
+	b.SetParallelism(4)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			ctx := context.Background()
+			ctx = callerid.NewContext(ctx, tabletconntest.TestCallerID, tabletconntest.TestVTGateCallerID)
+			_, err := conn.ExecuteBatch(ctx, tabletconntest.TestTarget, tabletconntest.ExecuteBatchQueries, tabletconntest.TestAsTransaction, tabletconntest.ExecuteBatchTransactionID, tabletconntest.TestExecuteOptions)
+			if err != nil {
+				b.Fatalf("ExecuteBatch failed: %v", err)
+			}
+		}
+	})
+}
+
 // This test makes sure the go rpc service works
 func TestGRPCTabletConn(t *testing.T) {
 	// fake service

go/vt/vttablet/tabletconntest/fakequeryservice.go

@@ -39,7 +39,7 @@ import (
 // FakeQueryService implements a programmable fake for the query service
 // server side.
 type FakeQueryService struct {
-	t              *testing.T
+	t              testing.TB
 	TestingGateway bool
 
 	// these fields are used to simulate and synchronize on errors
@@ -740,7 +740,7 @@ func (f *FakeQueryService) Release(ctx context.Context, target *querypb.Target,
 }
 
 // CreateFakeServer returns the fake server for the tests
-func CreateFakeServer(t *testing.T) *FakeQueryService {
+func CreateFakeServer(t testing.TB) *FakeQueryService {
 	return &FakeQueryService{
 		t: t,
 	}