From 10d4525862d996f0b398bbfac99b335b8659db93 Mon Sep 17 00:00:00 2001
From: Michael Demmer <mdemmer@slack-corp.com>
Date: Mon, 6 May 2024 14:44:03 -0700
Subject: [PATCH] add execute timing stats (#329)

---
 go/vt/vtgateproxy/vtgateproxy.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/go/vt/vtgateproxy/vtgateproxy.go b/go/vt/vtgateproxy/vtgateproxy.go
index 86ebb6350ff..1911f016a88 100644
--- a/go/vt/vtgateproxy/vtgateproxy.go
+++ b/go/vt/vtgateproxy/vtgateproxy.go
@@ -25,11 +25,13 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 
 	"google.golang.org/grpc"
 
 	"vitess.io/vitess/go/sqlescape"
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/grpcclient"
 	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -39,6 +41,12 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
 )
 
+const (
+	// timing metric keys
+	executeTimingKey       = "Execute"
+	streamExecuteTimingKey = "StreamExecute"
+)
+
 var (
 	vtgateHostsFile = flag.String("vtgate_hosts_file", "", "json file describing the host list to use for vtgate:// resolution")
 	numConnections  = flag.Int("num_connections", 4, "number of outbound GPRC connections to maintain")
@@ -48,6 +56,8 @@ var (
 	addressField    = flag.String("address_field", "address", "field name in the json file containing the address")
 	portField       = flag.String("port_field", "port", "field name in the json file containing the port")
 
+	timings = stats.NewTimings("Timings", "proxy timings by operation", "operation")
+
 	vtGateProxy *VTGateProxy = &VTGateProxy{
 		targetConns: map[string]*vtgateconn.VTGateConn{},
 		mu:          sync.RWMutex{},
@@ -160,10 +170,17 @@ func (proxy *VTGateProxy) Execute(ctx context.Context, session *vtgateconn.VTGat
 		return &sqltypes.Result{}, nil
 	}
 
+	startTime := time.Now()
+	defer timings.Record(executeTimingKey, startTime)
+
 	return session.Execute(ctx, sql, bindVariables)
+
 }
 
 func (proxy *VTGateProxy) StreamExecute(ctx context.Context, session *vtgateconn.VTGateSession, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) error {
+	startTime := time.Now()
+	defer timings.Record(streamExecuteTimingKey, startTime)
+
 	stream, err := session.StreamExecute(ctx, sql, bindVariables)
 	if err != nil {
 		return err