Emit per workload labels for existing per table vttablet metrics

go/flags/endtoend/vttablet.txt

@@ -94,6 +94,7 @@ Usage of vttablet:
       --enable-consolidator                                              Synonym to -enable_consolidator (default true)
       --enable-consolidator-replicas                                     Synonym to -enable_consolidator_replicas
       --enable-lag-throttler                                             Synonym to -enable_lag_throttler
+      --enable-per-workload-table-metrics                                If true, query counts and query error metrics include a label that identifies the workload
       --enable-tx-throttler                                              Synonym to -enable_tx_throttler
       --enable_consolidator                                              This option enables the query consolidator. (default true)
       --enable_consolidator_replicas                                     This option enables the query consolidator only on replicas.

go/vt/sqlparser/comments.go

@@ -48,6 +48,8 @@ const (
 	DirectiveVExplainRunDMLQueries = "EXECUTE_DML_QUERIES"
 	// DirectiveConsolidator enables the query consolidator.
 	DirectiveConsolidator = "CONSOLIDATOR"
+	// DirectiveWorkloadName specifies the name of the client application workload issuing the query.
+	DirectiveWorkloadName = "WORKLOAD_NAME"
 )
 
 func isNonSpace(r rune) bool {
@@ -396,3 +398,19 @@ func Consolidator(stmt Statement) querypb.ExecuteOptions_Consolidator {
 	}
 	return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
 }
+
+// GetWorkloadNameFromStatement gets the workload name from the provided Statement, using workloadLabel as the name of
+// the query directive that specifies it.
+func GetWorkloadNameFromStatement(statement Statement) string {
+	commentedStatement, ok := statement.(Commented)
+	// This would mean that the statement lacks comments, so we can't obtain the workload from it. Hence default to
+	// empty workload name
+	if !ok {
+		return ""
+	}
+
+	directives := commentedStatement.GetParsedComments().Directives()
+	workloadName, _ := directives.GetString(DirectiveWorkloadName, "")
+
+	return workloadName
+}

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -255,6 +255,10 @@ func (t *noopVCursor) SetWorkload(querypb.ExecuteOptions_Workload) {
 	panic("implement me")
 }
 
+func (t *noopVCursor) SetWorkloadName(string) {
+	panic("implement me")
+}
+
 func (t *noopVCursor) SetPlannerVersion(querypb.ExecuteOptions_PlannerVersion) {
 	panic("implement me")
 }
@@ -682,6 +686,10 @@ func (f *loggingVCursor) SetWorkload(querypb.ExecuteOptions_Workload) {
 	panic("implement me")
 }
 
+func (f *loggingVCursor) SetWorkloadName(string) {
+	panic("implement me")
+}
+
 func (f *loggingVCursor) SetPlannerVersion(querypb.ExecuteOptions_PlannerVersion) {
 	panic("implement me")
 }

go/vt/vtgate/engine/primitive.go

@@ -145,6 +145,7 @@ type (
 		SetWorkload(querypb.ExecuteOptions_Workload)
 		SetPlannerVersion(querypb.ExecuteOptions_PlannerVersion)
 		SetConsolidator(querypb.ExecuteOptions_Consolidator)
+		SetWorkloadName(string)
 		SetFoundRows(uint64)
 
 		SetDDLStrategy(string)

go/vt/vtgate/executor.go

@@ -984,6 +984,7 @@ func (e *Executor) getPlan(ctx context.Context, vcursor *vcursorImpl, sql string
 	vcursor.SetIgnoreMaxMemoryRows(ignoreMaxMemoryRows)
 	consolidator := sqlparser.Consolidator(stmt)
 	vcursor.SetConsolidator(consolidator)
+	vcursor.SetWorkloadName(sqlparser.GetWorkloadNameFromStatement(stmt))
 
 	setVarComment, err := prepareSetVarComment(vcursor, stmt)
 	if err != nil {

go/vt/vtgate/vcursor_impl.go

@@ -851,6 +851,12 @@ func (vc *vcursorImpl) SetConsolidator(consolidator querypb.ExecuteOptions_Conso
 	vc.safeSession.GetOrCreateOptions().Consolidator = consolidator
 }
 
+func (vc *vcursorImpl) SetWorkloadName(workloadName string) {
+	if workloadName != "" {
+		vc.safeSession.GetOrCreateOptions().WorkloadName = workloadName
+	}
+}
+
 // SetFoundRows implements the SessionActions interface
 func (vc *vcursorImpl) SetFoundRows(foundRows uint64) {
 	vc.safeSession.FoundRows = foundRows

go/vt/vttablet/tabletserver/query_engine.go

@@ -174,6 +174,9 @@ type QueryEngine struct {
 	// Note: queryErrorCountsWithCode is similar to queryErrorCounts except it contains error code as an additional dimension
 	queryCounts, queryTimes, queryErrorCounts, queryErrorCountsWithCode, queryRowsAffected, queryRowsReturned *stats.CountersWithMultiLabels
 
+	// stats flags
+	enablePerWorkloadTableMetrics bool
+
 	// Loggers
 	accessCheckerLogger *logutil.ThrottledLogger
 }
@@ -190,11 +193,12 @@ func NewQueryEngine(env tabletenv.Env, se *schema.Engine) *QueryEngine {
 	}
 
 	qe := &QueryEngine{
-		env:              env,
-		se:               se,
-		tables:           make(map[string]*schema.Table),
-		plans:            cache.NewDefaultCacheImpl(cacheCfg),
-		queryRuleSources: rules.NewMap(),
+		env:                           env,
+		se:                            se,
+		tables:                        make(map[string]*schema.Table),
+		plans:                         cache.NewDefaultCacheImpl(cacheCfg),
+		queryRuleSources:              rules.NewMap(),
+		enablePerWorkloadTableMetrics: config.EnablePerWorkloadTableMetrics,
 	}
 
 	qe.conns = connpool.NewPool(env, "ConnPool", config.OltpReadPool)
@@ -247,11 +251,17 @@ func NewQueryEngine(env tabletenv.Env, se *schema.Engine) *QueryEngine {
 	env.Exporter().NewGaugeFunc("QueryCacheSize", "Query engine query cache size", qe.plans.UsedCapacity)
 	env.Exporter().NewGaugeFunc("QueryCacheCapacity", "Query engine query cache capacity", qe.plans.MaxCapacity)
 	env.Exporter().NewCounterFunc("QueryCacheEvictions", "Query engine query cache evictions", qe.plans.Evictions)
-	qe.queryCounts = env.Exporter().NewCountersWithMultiLabels("QueryCounts", "query counts", []string{"Table", "Plan"})
-	qe.queryTimes = env.Exporter().NewCountersWithMultiLabels("QueryTimesNs", "query times in ns", []string{"Table", "Plan"})
-	qe.queryRowsAffected = env.Exporter().NewCountersWithMultiLabels("QueryRowsAffected", "query rows affected", []string{"Table", "Plan"})
-	qe.queryRowsReturned = env.Exporter().NewCountersWithMultiLabels("QueryRowsReturned", "query rows returned", []string{"Table", "Plan"})
-	qe.queryErrorCounts = env.Exporter().NewCountersWithMultiLabels("QueryErrorCounts", "query error counts", []string{"Table", "Plan"})
+
+	labels := []string{"Table", "Plan"}
+	if config.EnablePerWorkloadTableMetrics {
+		labels = []string{"Table", "Plan", "Workload"}
+	}
+
+	qe.queryCounts = env.Exporter().NewCountersWithMultiLabels("QueryCounts", "query counts", labels)
+	qe.queryTimes = env.Exporter().NewCountersWithMultiLabels("QueryTimesNs", "query times in ns", labels)
+	qe.queryRowsAffected = env.Exporter().NewCountersWithMultiLabels("QueryRowsAffected", "query rows affected", labels)
+	qe.queryRowsReturned = env.Exporter().NewCountersWithMultiLabels("QueryRowsReturned", "query rows returned", labels)
+	qe.queryErrorCounts = env.Exporter().NewCountersWithMultiLabels("QueryErrorCounts", "query error counts", labels)
 	qe.queryErrorCountsWithCode = env.Exporter().NewCountersWithMultiLabels("QueryErrorCountsWithCode", "query error counts with error code", []string{"Table", "Plan", "Code"})
 
 	env.Exporter().HandleFunc("/debug/hotrows", qe.txSerializer.ServeHTTP)
@@ -480,9 +490,13 @@ func (qe *QueryEngine) QueryPlanCacheLen() int {
 }
 
 // AddStats adds the given stats for the planName.tableName
-func (qe *QueryEngine) AddStats(planType planbuilder.PlanType, tableName string, queryCount int64, duration, mysqlTime time.Duration, rowsAffected, rowsReturned, errorCount int64, errorCode string) {
+func (qe *QueryEngine) AddStats(planType planbuilder.PlanType, tableName, workload string, queryCount int64, duration, mysqlTime time.Duration, rowsAffected, rowsReturned, errorCount int64, errorCode string) {
 	// table names can contain "." characters, replace them!
 	keys := []string{tableName, planType.String()}
+	// Only use the workload as a label if that's enabled in the configuration.
+	if qe.enablePerWorkloadTableMetrics {
+		keys = append(keys, workload)
+	}
 	qe.queryCounts.Add(keys, queryCount)
 	qe.queryTimes.Add(keys, int64(duration))
 	qe.queryErrorCounts.Add(keys, errorCount)

go/vt/vttablet/tabletserver/query_engine_test.go

@@ -584,6 +584,8 @@ func TestAddQueryStats(t *testing.T) {
 		rowsReturned                     int64
 		errorCount                       int64
 		errorCode                        string
+		enablePerWorkloadTableMetrics    bool
+		workload                         string
 		expectedQueryCounts              string
 		expectedQueryTimes               string
 		expectedQueryRowsAffected        string
@@ -601,6 +603,8 @@ func TestAddQueryStats(t *testing.T) {
 			rowsReturned:                     15,
 			errorCount:                       0,
 			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    false,
+			workload:                         "some-workload",
 			expectedQueryCounts:              `{"A.Select": 1}`,
 			expectedQueryTimes:               `{"A.Select": 10}`,
 			expectedQueryRowsAffected:        `{}`,
@@ -617,6 +621,8 @@ func TestAddQueryStats(t *testing.T) {
 			rowsReturned:                     0,
 			errorCount:                       0,
 			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    false,
+			workload:                         "some-workload",
 			expectedQueryCounts:              `{"A.Select": 1}`,
 			expectedQueryTimes:               `{"A.Select": 10}`,
 			expectedQueryRowsAffected:        `{"A.Select": 15}`,
@@ -633,6 +639,8 @@ func TestAddQueryStats(t *testing.T) {
 			rowsReturned:                     0,
 			errorCount:                       1,
 			errorCode:                        "RESOURCE_EXHAUSTED",
+			enablePerWorkloadTableMetrics:    false,
+			workload:                         "some-workload",
 			expectedQueryCounts:              `{"A.Select": 1}`,
 			expectedQueryTimes:               `{"A.Select": 10}`,
 			expectedQueryRowsAffected:        `{}`,
@@ -649,12 +657,86 @@ func TestAddQueryStats(t *testing.T) {
 			rowsReturned:                     0,
 			errorCount:                       0,
 			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    false,
+			workload:                         "some-workload",
 			expectedQueryCounts:              `{"A.Insert": 1}`,
 			expectedQueryTimes:               `{"A.Insert": 10}`,
 			expectedQueryRowsAffected:        `{"A.Insert": 15}`,
 			expectedQueryRowsReturned:        `{}`,
 			expectedQueryErrorCounts:         `{"A.Insert": 0}`,
 			expectedQueryErrorCountsWithCode: `{}`,
+		}, {
+			name:                             "select query with per workload metrics",
+			planType:                         planbuilder.PlanSelect,
+			tableName:                        "A",
+			queryCount:                       1,
+			duration:                         10,
+			rowsAffected:                     0,
+			rowsReturned:                     15,
+			errorCount:                       0,
+			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    true,
+			workload:                         "some-workload",
+			expectedQueryCounts:              `{"A.Select.some-workload": 1}`,
+			expectedQueryTimes:               `{"A.Select.some-workload": 10}`,
+			expectedQueryRowsAffected:        `{}`,
+			expectedQueryRowsReturned:        `{"A.Select.some-workload": 15}`,
+			expectedQueryErrorCounts:         `{"A.Select.some-workload": 0}`,
+			expectedQueryErrorCountsWithCode: `{}`,
+		}, {
+			name:                             "select into query with per workload metrics",
+			planType:                         planbuilder.PlanSelect,
+			tableName:                        "A",
+			queryCount:                       1,
+			duration:                         10,
+			rowsAffected:                     15,
+			rowsReturned:                     0,
+			errorCount:                       0,
+			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    true,
+			workload:                         "some-workload",
+			expectedQueryCounts:              `{"A.Select.some-workload": 1}`,
+			expectedQueryTimes:               `{"A.Select.some-workload": 10}`,
+			expectedQueryRowsAffected:        `{"A.Select.some-workload": 15}`,
+			expectedQueryRowsReturned:        `{"A.Select.some-workload": 0}`,
+			expectedQueryErrorCounts:         `{"A.Select.some-workload": 0}`,
+			expectedQueryErrorCountsWithCode: `{}`,
+		}, {
+			name:                             "error with per workload metrics",
+			planType:                         planbuilder.PlanSelect,
+			tableName:                        "A",
+			queryCount:                       1,
+			duration:                         10,
+			rowsAffected:                     0,
+			rowsReturned:                     0,
+			errorCount:                       1,
+			errorCode:                        "RESOURCE_EXHAUSTED",
+			enablePerWorkloadTableMetrics:    true,
+			workload:                         "some-workload",
+			expectedQueryCounts:              `{"A.Select.some-workload": 1}`,
+			expectedQueryTimes:               `{"A.Select.some-workload": 10}`,
+			expectedQueryRowsAffected:        `{}`,
+			expectedQueryRowsReturned:        `{"A.Select.some-workload": 0}`,
+			expectedQueryErrorCounts:         `{"A.Select.some-workload": 1}`,
+			expectedQueryErrorCountsWithCode: `{"A.Select.RESOURCE_EXHAUSTED": 1}`,
+		}, {
+			name:                             "insert query with per workload metrics",
+			planType:                         planbuilder.PlanInsert,
+			tableName:                        "A",
+			queryCount:                       1,
+			duration:                         10,
+			rowsAffected:                     15,
+			rowsReturned:                     0,
+			errorCount:                       0,
+			errorCode:                        "OK",
+			enablePerWorkloadTableMetrics:    true,
+			workload:                         "some-workload",
+			expectedQueryCounts:              `{"A.Insert.some-workload": 1}`,
+			expectedQueryTimes:               `{"A.Insert.some-workload": 10}`,
+			expectedQueryRowsAffected:        `{"A.Insert.some-workload": 15}`,
+			expectedQueryRowsReturned:        `{}`,
+			expectedQueryErrorCounts:         `{"A.Insert.some-workload": 0}`,
+			expectedQueryErrorCountsWithCode: `{}`,
 		},
 	}
 
@@ -663,10 +745,11 @@ func TestAddQueryStats(t *testing.T) {
 		t.Run(testcase.name, func(t *testing.T) {
 			config := tabletenv.NewDefaultConfig()
 			config.DB = newDBConfigs(fakesqldb.New(t))
+			config.EnablePerWorkloadTableMetrics = testcase.enablePerWorkloadTableMetrics
 			env := tabletenv.NewEnv(config, "TestAddQueryStats_"+testcase.name)
 			se := schema.NewEngine(env)
 			qe := NewQueryEngine(env, se)
-			qe.AddStats(testcase.planType, testcase.tableName, testcase.queryCount, testcase.duration, testcase.mysqlTime, testcase.rowsAffected, testcase.rowsReturned, testcase.errorCount, testcase.errorCode)
+			qe.AddStats(testcase.planType, testcase.tableName, testcase.workload, testcase.queryCount, testcase.duration, testcase.mysqlTime, testcase.rowsAffected, testcase.rowsReturned, testcase.errorCount, testcase.errorCode)
 			assert.Equal(t, testcase.expectedQueryCounts, qe.queryCounts.String())
 			assert.Equal(t, testcase.expectedQueryTimes, qe.queryTimes.String())
 			assert.Equal(t, testcase.expectedQueryRowsAffected, qe.queryRowsAffected.String())

go/vt/vttablet/tabletserver/query_executor.go

@@ -62,6 +62,7 @@ type QueryExecutor struct {
 	tsv            *TabletServer
 	tabletType     topodatapb.TabletType
 	setting        *pools.Setting
+	workload       string
 }
 
 const (
@@ -132,11 +133,12 @@ func (qre *QueryExecutor) Execute() (reply *sqltypes.Result, err error) {
 		vtErrorCode := vterrors.Code(err)
 		errCode = vtErrorCode.String()
 		if reply == nil {
-			qre.tsv.qe.AddStats(qre.plan.PlanID, tableName, 1, duration, mysqlTime, 0, 0, 1, errCode)
+			qre.tsv.qe.AddStats(qre.plan.PlanID, tableName, qre.options.GetWorkloadName(), 1, duration, mysqlTime, 0, 0, 1, errCode)
 			qre.plan.AddStats(1, duration, mysqlTime, 0, 0, 1)
 			return
 		}
-		qre.tsv.qe.AddStats(qre.plan.PlanID, tableName, 1, duration, mysqlTime, int64(reply.RowsAffected), int64(len(reply.Rows)), 0, errCode)
+
+		qre.tsv.qe.AddStats(qre.plan.PlanID, tableName, qre.options.GetWorkloadName(), 1, duration, mysqlTime, int64(reply.RowsAffected), int64(len(reply.Rows)), 0, errCode)
 		qre.plan.AddStats(1, duration, mysqlTime, reply.RowsAffected, uint64(len(reply.Rows)), 0)
 		qre.logStats.RowsAffected = int(reply.RowsAffected)
 		qre.logStats.Rows = reply.Rows