Emit per workload labels for existing per table vttablet metrics

go/flags/endtoend/vttablet.txt

@@ -94,12 +94,14 @@ 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                                Synonym to -enable_per_workload_table_metrics
       --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.
       --enable_hot_row_protection                                        If true, incoming transactions for the same row (range) will be queued and cannot consume all txpool slots.
       --enable_hot_row_protection_dry_run                                If true, hot row protection is not enforced but logs if transactions would have been queued.
       --enable_lag_throttler                                             If true, vttablet will run a throttler service, and will implicitly enable heartbeats
+      --enable_per_workload_table_metrics                                If true, query counts and query error metrics include a label that identifies the workload
       --enable_replication_reporter                                      Use polling to track replication lag.
       --enable_transaction_limit                                         If true, limit on number of transactions open at the same time will be enforced for all users. User trying to open a new transaction after exhausting their limit will receive an error immediately, regardless of whether there are available slots or not.
       --enable_transaction_limit_dry_run                                 If true, limit on number of transactions open at the same time will be tracked for all users, but not enforced.
@@ -373,6 +375,8 @@ Usage of vttablet:
       --vttablet_skip_buildinfo_tags string                              comma-separated list of buildinfo tags to skip from merging with --init_tags. each tag is either an exact match or a regular expression of the form '/regexp/'. (default "/.*/")
       --wait_for_backup_interval duration                                (init restore parameter) if this is greater than 0, instead of starting up empty when no backups are found, keep checking at this interval for a backup to appear
       --watch_replication_stream                                         When enabled, vttablet will stream the MySQL replication stream from the local server, and use it to update schema when it sees a DDL.
+      --workload-label string                                            Synonym to -workload_label (default "WORKLOAD")
+      --workload_label string                                            The label looked for in query comments to identify the workload. Used in conjunction with enable-per-workload-table-metrics (default "WORKLOAD")
       --xbstream_restore_flags string                                    Flags to pass to xbstream command during restore. These should be space separated and will be added to the end of the command. These need to match the ones used for backup e.g. --compress / --decompress, --encrypt / --decrypt
       --xtrabackup_backup_flags string                                   Flags to pass to backup command. These should be space separated and will be added to the end of the command
       --xtrabackup_prepare_flags string                                  Flags to pass to prepare command. These should be space separated and will be added to the end of the command

go/vt/sqlparser/comments.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 	"unicode"
 
+	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
 
@@ -396,3 +397,49 @@ func Consolidator(stmt Statement) querypb.ExecuteOptions_Consolidator {
 	}
 	return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
 }
+
+// GetWorkloadFromComments returns the value of the workload as specified in the query comments. It assumes the presence
+// of a comment that includes a string in the form `workloadLabel=workloadname;` (including the semicolon). The function
+// returns `unspecified` if it fails to obtain the workload from the query comments. If there are multiple comments that
+// specify the workload, or multiple workload specifications in the same comment, it returns the first one. It trims
+// whitespaces both at the beginning and end of the workload name.
+func GetWorkloadFromComments(query string, workloadLabel string) string {
+	const unspecified = "unspecified"
+
+	statement, _, err := Parse2(query)
+	if err != nil {
+		log.Warningf("Failed to parse statement while attempting to retrieve workload from SQL string")
+
+		return unspecified
+	}
+
+	commentedStatement, ok := statement.(Commented)
+	// This would mean that the statement lacks comments, so we can't obtain the workload from it. Hence default to
+	// unspecified workload
+	if !ok {
+		return unspecified
+	}
+
+	parsedComments := commentedStatement.GetParsedComments()
+
+	if parsedComments != nil {
+		for _, comment := range parsedComments.comments {
+
+			workloadLabelStart := strings.Index(comment, workloadLabel)
+			if workloadLabelStart == -1 {
+				continue
+			}
+
+			workloadLabelEnd := strings.Index(comment[workloadLabelStart:], ";")
+			if workloadLabelEnd == -1 {
+				continue
+			}
+
+			workloadLabelEnd += workloadLabelStart
+			return strings.TrimSpace(comment[workloadLabelStart+len(workloadLabel)+1 : workloadLabelEnd])
+
+		}
+	}
+
+	return unspecified
+}

go/vt/vttablet/tabletserver/query_engine.go

@@ -173,6 +173,9 @@ type QueryEngine struct {
 	// stats
 	queryCounts, queryTimes, queryErrorCounts, queryRowsAffected, queryRowsReturned *stats.CountersWithMultiLabels
 
+	// stats flags
+	enablePerWorkloadTableMetrics bool
+
 	// Loggers
 	accessCheckerLogger *logutil.ThrottledLogger
 }
@@ -189,11 +192,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)
@@ -246,11 +250,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)
 
 	env.Exporter().HandleFunc("/debug/hotrows", qe.txSerializer.ServeHTTP)
 	env.Exporter().HandleFunc("/debug/tablet_plans", qe.handleHTTPQueryPlans)
@@ -478,9 +488,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) {
+func (qe *QueryEngine) AddStats(planType planbuilder.PlanType, tableName, workload string, queryCount int64, duration, mysqlTime time.Duration, rowsAffected, rowsReturned, errorCount int64) {
 	// 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

@@ -574,77 +574,151 @@ func TestPlanCachePollution(t *testing.T) {
 
 func TestAddQueryStats(t *testing.T) {
 	testcases := []struct {
-		name                      string
-		planType                  planbuilder.PlanType
-		tableName                 string
-		queryCount                int64
-		duration                  time.Duration
-		mysqlTime                 time.Duration
-		rowsAffected              int64
-		rowsReturned              int64
-		errorCount                int64
-		expectedQueryCounts       string
-		expectedQueryTimes        string
-		expectedQueryRowsAffected string
-		expectedQueryRowsReturned string
-		expectedQueryErrorCounts  string
+		name                          string
+		planType                      planbuilder.PlanType
+		tableName                     string
+		queryCount                    int64
+		duration                      time.Duration
+		mysqlTime                     time.Duration
+		rowsAffected                  int64
+		rowsReturned                  int64
+		errorCount                    int64
+		enablePerWorkloadTableMetrics bool
+		workload                      string
+		expectedQueryCounts           string
+		expectedQueryTimes            string
+		expectedQueryRowsAffected     string
+		expectedQueryRowsReturned     string
+		expectedQueryErrorCounts      string
 	}{
 		{
-			name:                      "select query",
-			planType:                  planbuilder.PlanSelect,
-			tableName:                 "A",
-			queryCount:                1,
-			duration:                  10,
-			rowsAffected:              0,
-			rowsReturned:              15,
-			errorCount:                0,
-			expectedQueryCounts:       `{"A.Select": 1}`,
-			expectedQueryTimes:        `{"A.Select": 10}`,
-			expectedQueryRowsAffected: `{}`,
-			expectedQueryRowsReturned: `{"A.Select": 15}`,
-			expectedQueryErrorCounts:  `{"A.Select": 0}`,
+			name:                          "select query",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  0,
+			rowsReturned:                  15,
+			errorCount:                    0,
+			enablePerWorkloadTableMetrics: false,
+			workload:                      "some-workload",
+			expectedQueryCounts:           `{"A.Select": 1}`,
+			expectedQueryTimes:            `{"A.Select": 10}`,
+			expectedQueryRowsAffected:     `{}`,
+			expectedQueryRowsReturned:     `{"A.Select": 15}`,
+			expectedQueryErrorCounts:      `{"A.Select": 0}`,
 		}, {
-			name:                      "select into query",
-			planType:                  planbuilder.PlanSelect,
-			tableName:                 "A",
-			queryCount:                1,
-			duration:                  10,
-			rowsAffected:              15,
-			rowsReturned:              0,
-			errorCount:                0,
-			expectedQueryCounts:       `{"A.Select": 1}`,
-			expectedQueryTimes:        `{"A.Select": 10}`,
-			expectedQueryRowsAffected: `{"A.Select": 15}`,
-			expectedQueryRowsReturned: `{"A.Select": 0}`,
-			expectedQueryErrorCounts:  `{"A.Select": 0}`,
+			name:                          "select into query",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  15,
+			rowsReturned:                  0,
+			errorCount:                    0,
+			enablePerWorkloadTableMetrics: false,
+			workload:                      "some-workload",
+			expectedQueryCounts:           `{"A.Select": 1}`,
+			expectedQueryTimes:            `{"A.Select": 10}`,
+			expectedQueryRowsAffected:     `{"A.Select": 15}`,
+			expectedQueryRowsReturned:     `{"A.Select": 0}`,
+			expectedQueryErrorCounts:      `{"A.Select": 0}`,
 		}, {
-			name:                      "error",
-			planType:                  planbuilder.PlanSelect,
-			tableName:                 "A",
-			queryCount:                1,
-			duration:                  10,
-			rowsAffected:              0,
-			rowsReturned:              0,
-			errorCount:                1,
-			expectedQueryCounts:       `{"A.Select": 1}`,
-			expectedQueryTimes:        `{"A.Select": 10}`,
-			expectedQueryRowsAffected: `{}`,
-			expectedQueryRowsReturned: `{"A.Select": 0}`,
-			expectedQueryErrorCounts:  `{"A.Select": 1}`,
+			name:                          "error",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  0,
+			rowsReturned:                  0,
+			errorCount:                    1,
+			enablePerWorkloadTableMetrics: false,
+			workload:                      "some-workload",
+			expectedQueryCounts:           `{"A.Select": 1}`,
+			expectedQueryTimes:            `{"A.Select": 10}`,
+			expectedQueryRowsAffected:     `{}`,
+			expectedQueryRowsReturned:     `{"A.Select": 0}`,
+			expectedQueryErrorCounts:      `{"A.Select": 1}`,
 		}, {
-			name:                      "insert query",
-			planType:                  planbuilder.PlanInsert,
-			tableName:                 "A",
-			queryCount:                1,
-			duration:                  10,
-			rowsAffected:              15,
-			rowsReturned:              0,
-			errorCount:                0,
-			expectedQueryCounts:       `{"A.Insert": 1}`,
-			expectedQueryTimes:        `{"A.Insert": 10}`,
-			expectedQueryRowsAffected: `{"A.Insert": 15}`,
-			expectedQueryRowsReturned: `{}`,
-			expectedQueryErrorCounts:  `{"A.Insert": 0}`,
+			name:                          "insert query",
+			planType:                      planbuilder.PlanInsert,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  15,
+			rowsReturned:                  0,
+			errorCount:                    0,
+			enablePerWorkloadTableMetrics: false,
+			workload:                      "some-workload",
+			expectedQueryCounts:           `{"A.Insert": 1}`,
+			expectedQueryTimes:            `{"A.Insert": 10}`,
+			expectedQueryRowsAffected:     `{"A.Insert": 15}`,
+			expectedQueryRowsReturned:     `{}`,
+			expectedQueryErrorCounts:      `{"A.Insert": 0}`,
+		}, {
+			name:                          "select query with per workload metrics",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  0,
+			rowsReturned:                  15,
+			errorCount:                    0,
+			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}`,
+		}, {
+			name:                          "select into query with per workload metrics",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  15,
+			rowsReturned:                  0,
+			errorCount:                    0,
+			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}`,
+		}, {
+			name:                          "error with per workload metrics",
+			planType:                      planbuilder.PlanSelect,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  0,
+			rowsReturned:                  0,
+			errorCount:                    1,
+			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}`,
+		}, {
+			name:                          "insert query with per workload metrics",
+			planType:                      planbuilder.PlanInsert,
+			tableName:                     "A",
+			queryCount:                    1,
+			duration:                      10,
+			rowsAffected:                  15,
+			rowsReturned:                  0,
+			errorCount:                    0,
+			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}`,
 		},
 	}
 
@@ -653,10 +727,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)
+			qe.AddStats(testcase.planType, testcase.tableName, testcase.workload, testcase.queryCount, testcase.duration, testcase.mysqlTime, testcase.rowsAffected, testcase.rowsReturned, testcase.errorCount)
 			assert.Equal(t, testcase.expectedQueryCounts, qe.queryCounts.String())
 			assert.Equal(t, testcase.expectedQueryTimes, qe.queryTimes.String())
 			assert.Equal(t, testcase.expectedQueryRowsAffected, qe.queryRowsAffected.String())