add primitive stats to LogStats

Signed-off-by: Andres Taylor <[email protected]>
planetscale · Sep 11, 2024 · 9de346e · 9de346e
1 parent d955e51
commit 9de346e
Show file tree

Hide file tree

Showing 16 changed files with 117 additions and 43 deletions.
diff --git a/go/logstats/logger.go b/go/logstats/logger.go
@@ -87,6 +87,35 @@ func (log *Logger) appendBVarsJSON(b []byte, bvars map[string]*querypb.BindVaria
 	return append(b, '}')
 }
 
+func appendJsonInt(b []byte, i int) []byte {
+	b = append(b, '"')
+	b = strconv.AppendInt(b, int64(i), 10)
+	return append(b, '"')
+}
+
+func (log *Logger) appendPrimitiveStats(b []byte, m map[int]PrimitiveStats) []byte {
+	b = append(b, '{')
+	for i := range len(m) {
+		if i > 0 {
+			b = append(b, ',')
+		}
+		op := i + 1
+		b = appendJsonInt(b, op)
+		b = append(b, `:{"Calls":`...)
+		stats := m[op]
+		b = strconv.AppendInt(b, int64(stats.NoOfCalls), 10)
+		b = append(b, `,"Rows":[`...)
+		for rowIdx, rows := range stats.Rows {
+			if rowIdx > 0 {
+				b = append(b, ',')
+			}
+			b = appendJsonInt(b, rows)
+		}
+		b = append(b, `]}`...)
+	}
+	return append(b, '}')
+}
+
 func (log *Logger) Init(json bool) {
 	log.n = 0
 	log.json = json
@@ -165,6 +194,10 @@ func (log *Logger) BindVariables(bvars map[string]*querypb.BindVariable, full bo
 	log.b = log.appendBVarsJSON(log.b, bvars, full)
 }
 
+func (log *Logger) OpStats(stats map[int]PrimitiveStats) {
+	log.b = log.appendPrimitiveStats(log.b, stats)
+}
+
 func (log *Logger) Int(i int64) {
 	log.b = strconv.AppendInt(log.b, i, 10)
 }

diff --git a/go/logstats/logger_test.go b/go/logstats/logger_test.go
@@ -244,6 +244,29 @@ func TestFlush(t *testing.T) {
 	assert.Equal(t, []byte("{\"testKey\": \"testValue\"}\n"), calledValue)
 }
 
+func TestOpStats(t *testing.T) {
+	tl := NewLogger()
+	tl.Init(false)
+	stats := map[int]PrimitiveStats{
+		1: {
+			NoOfCalls: 1,
+			Rows:      []int{100},
+		},
+		2: {
+			NoOfCalls: 2,
+			Rows:      []int{100, 200},
+		},
+		3: {
+			NoOfCalls: 10,
+			Rows:      []int{100, 200, 100, 200, 300, 100, 200, 100, 200, 300},
+		},
+	}
+	tl.OpStats(stats)
+
+	want := []byte(`{"1":{"Calls":1,"Rows":["100"]},"2":{"Calls":2,"Rows":["100","200"]},"3":{"Calls":10,"Rows":["100","200","100","200","300","100","200","100","200","300"]}}`)
+	assert.Equal(t, string(want), string(tl.b))
+}
+
 func TestBindVariables(t *testing.T) {
 	tcases := []struct {
 		name  string

diff --git a/go/vt/vtgate/logstats/logstats.go → go/logstats/logstats.go b/go/vt/vtgate/logstats/logstats.go → go/logstats/logstats.go
@@ -24,35 +24,42 @@ import (
 
 	"github.com/google/safehtml"
 
-	"vitess.io/vitess/go/logstats"
 	"vitess.io/vitess/go/streamlog"
 	"vitess.io/vitess/go/vt/callerid"
 	"vitess.io/vitess/go/vt/callinfo"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
 
-// LogStats records the stats for a single vtgate query
-type LogStats struct {
-	Ctx            context.Context
-	Method         string
-	TabletType     string
-	StmtType       string
-	SQL            string
-	BindVariables  map[string]*querypb.BindVariable
-	StartTime      time.Time
-	EndTime        time.Time
-	ShardQueries   uint64
-	RowsAffected   uint64
-	RowsReturned   uint64
-	PlanTime       time.Duration
-	ExecuteTime    time.Duration
-	CommitTime     time.Duration
-	Error          error
-	TablesUsed     []string
-	SessionUUID    string
-	CachedPlan     bool
-	ActiveKeyspace string // ActiveKeyspace is the selected keyspace `use ks`
-}
+type (
+	// LogStats records the stats for a single vtgate query
+	LogStats struct {
+		Ctx            context.Context
+		Method         string
+		TabletType     string
+		StmtType       string
+		SQL            string
+		BindVariables  map[string]*querypb.BindVariable
+		StartTime      time.Time
+		EndTime        time.Time
+		ShardQueries   uint64
+		RowsAffected   uint64
+		RowsReturned   uint64
+		PlanTime       time.Duration
+		ExecuteTime    time.Duration
+		CommitTime     time.Duration
+		Error          error
+		TablesUsed     []string
+		SessionUUID    string
+		CachedPlan     bool
+		ActiveKeyspace string // ActiveKeyspace is the selected keyspace `use ks`
+		PrimitiveStats map[int]PrimitiveStats
+	}
+
+	PrimitiveStats struct {
+		NoOfCalls int
+		Rows      []int
+	}
+)
 
 // NewLogStats constructs a new LogStats with supplied Method and ctx
 // field values, and the StartTime field set to the present time.
@@ -127,7 +134,7 @@ func (stats *LogStats) Logf(w io.Writer, params url.Values) error {
 	_, fullBindParams := params["full"]
 	remoteAddr, username := stats.RemoteAddrUsername()
 
-	log := logstats.NewLogger()
+	log := NewLogger()
 	log.Init(streamlog.GetQueryLogFormat() == streamlog.QueryLogFormatJSON)
 	log.Key("Method")
 	log.StringUnquoted(stats.Method)
@@ -177,6 +184,9 @@ func (stats *LogStats) Logf(w io.Writer, params url.Values) error {
 	log.Strings(stats.TablesUsed)
 	log.Key("ActiveKeyspace")
 	log.String(stats.ActiveKeyspace)
+	if len(stats.PrimitiveStats) > 0 {
+		log.Key("PrimitiveStats")
+	}
 
 	return log.Flush(w)
 }
diff --git a/go/vt/vtgate/logstats/logstats_test.go → go/logstats/logstats_test.go b/go/vt/vtgate/logstats/logstats_test.go → go/logstats/logstats_test.go
diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go
@@ -26,6 +26,8 @@ import (
 	"sort"
 	"strings"
 
+	"vitess.io/vitess/go/logstats"
+
 	"vitess.io/vitess/go/cache/theine"
 	"vitess.io/vitess/go/json2"
 	"vitess.io/vitess/go/stats"
@@ -38,7 +40,6 @@ import (
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate"
 	"vitess.io/vitess/go/vt/vtgate/engine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vttablet/queryservice"
 

diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
@@ -28,6 +28,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/spf13/pflag"
 
 	"vitess.io/vitess/go/acl"
@@ -57,7 +59,6 @@ import (
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"

diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go
@@ -25,6 +25,8 @@ import (
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -43,7 +45,6 @@ import (
 	"vitess.io/vitess/go/vt/srvtopo"
 	"vitess.io/vitess/go/vt/vtenv"
 	"vitess.io/vitess/go/vt/vtgate/engine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
 )

diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go
@@ -26,6 +26,8 @@ import (
 	"testing"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -43,7 +45,6 @@ import (
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtenv"
 	"vitess.io/vitess/go/vt/vterrors"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	_ "vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
 )

diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go
@@ -21,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -31,7 +33,6 @@ import (
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 	"vitess.io/vitess/go/vt/vtenv"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	_ "vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
 )

diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go
@@ -30,6 +30,8 @@ import (
 	"testing"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/safehtml/template"
 	"github.com/stretchr/testify/assert"
@@ -46,7 +48,6 @@ import (
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/vtgate/buffer"
 	"vitess.io/vitess/go/vt/vtgate/engine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vtgate/vschemaacl"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"

diff --git a/go/vt/vtgate/plan_execute.go b/go/vt/vtgate/plan_execute.go
@@ -22,14 +22,15 @@ import (
 	"strings"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"
 )
 

diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json
@@ -1,8 +1 @@
-[
-  {
-    "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
-    "plan": {
-    }
-  }
-]
+{"1":{"Calls":1,"Rows":["100"]},"2":{"Calls":2,"Rows":["100","200"]},"3":{"Calls":10,"Rows":["100","200","100","200","300","100","200","100","200","300"]}}
diff --git a/go/vt/vtgate/querylog.go b/go/vt/vtgate/querylog.go
@@ -19,9 +19,10 @@ package vtgate
 import (
 	"net/http"
 
+	"vitess.io/vitess/go/logstats"
+
 	"vitess.io/vitess/go/streamlog"
 	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 )
 
 var (

diff --git a/go/vt/vtgate/querylogz.go b/go/vt/vtgate/querylogz.go
@@ -23,7 +23,7 @@ import (
 	"text/template"
 	"time"
 
-	"vitess.io/vitess/go/vt/vtgate/logstats"
+	"vitess.io/vitess/go/logstats"
 
 	"vitess.io/vitess/go/acl"
 	"vitess.io/vitess/go/vt/log"

diff --git a/go/vt/vtgate/querylogz_test.go b/go/vt/vtgate/querylogz_test.go
@@ -26,8 +26,9 @@ import (
 	"testing"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"vitess.io/vitess/go/vt/sqlparser"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 
 	"vitess.io/vitess/go/streamlog"
 	"vitess.io/vitess/go/vt/callerid"

diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go
@@ -25,6 +25,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"vitess.io/vitess/go/logstats"
+
 	"github.com/google/uuid"
 
 	"vitess.io/vitess/go/mysql/collations"
@@ -50,7 +52,6 @@ import (
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/buffer"
 	"vitess.io/vitess/go/vt/vtgate/engine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
@@ -515,6 +516,11 @@ func (vc *vcursorImpl) ExecutePrimitive(ctx context.Context, primitive engine.Pr
 		if err != nil && vterrors.RootCause(err) == buffer.ShardMissingError {
 			continue
 		}
+		if vc.logOperatorTraffic {
+			stats := vc.logStats.PrimitiveStats[int(primitive.GetID())]
+			stats.NoOfCalls++
+			stats.Rows = append(stats.Rows, len(res.Rows))
+		}
 		return res, err
 	}
 	return nil, vterrors.New(vtrpcpb.Code_UNAVAILABLE, "upstream shards are not available")