*: add keyspace observability labels

cmd/tidb-server/BUILD.bazel

@@ -107,7 +107,7 @@ go_test(
     srcs = ["main_test.go"],
     embed = [":tidb-server_lib"],
     flaky = True,
-    shard_count = 6,
+    shard_count = 8,
     deps = [
         "//pkg/config",
         "//pkg/config/deploymode",

cmd/tidb-server/main.go

@@ -19,6 +19,7 @@ import (
 	"flag"
 	"fmt"
 	"io/fs"
+	"maps"
 	"os"
 	"runtime"
 	"strconv"
@@ -314,6 +315,8 @@ func main() {
 	}
 
 	var standbyController server.StandbyController
+	var activationKeyspaceID uint32
+	var activationMetadata map[string]string
 	if config.GetGlobalConfig().Standby.StandByMode {
 		standbyController = standby.NewLoadKeyspaceController()
 	}
@@ -330,11 +333,17 @@ func main() {
 		defer standbyController.EndStandby(err)
 		// need to validate config again in case of config change via standby
 		terror.MustNil(config.GetGlobalConfig().Valid())
+		if c, ok := standbyController.(*standby.LoadKeyspaceController); ok {
+			activationKeyspaceID = c.ActivationKeyspaceID()
+			activationMetadata = c.ActivationMetadata()
+		}
 	}
 
 	signal.SetupUSR1Handler()
 	err = registerStores()
 	terror.MustNil(err)
+	err = prepareKeyspaceObservability(activationKeyspaceID, activationMetadata)
+	terror.MustNil(err)
 	err = metricsutil.RegisterMetrics()
 	terror.MustNil(err)
 
@@ -1146,6 +1155,50 @@ func closeStmtSummary() {
 	}
 }
 
+const (
+	keyspaceIDMetricLabel   = "keyspace_id"
+	keyspaceNameMetricLabel = "keyspace_name"
+)
+
+func prepareKeyspaceObservability(keyspaceID uint32, metadata map[string]string) error {
+	cfg := config.GetGlobalConfig()
+
+	if kerneltype.IsClassic() || cfg.Store != config.StoreTypeTiKV {
+		return nil
+	}
+	resolvedValues := &config.KeyspaceObservabilityValues{
+		MetricLabels: map[string]string{
+			keyspaceNameMetricLabel: cfg.KeyspaceName,
+		},
+	}
+
+	if deploymode.IsStarter() {
+		resolvedValues.MetricLabels[keyspaceIDMetricLabel] = fmt.Sprint(keyspaceID)
+		err := prepareKeyspaceObservabilityForStarter(metadata, resolvedValues)
+		if err != nil {
+			return err
+		}
+	}
+
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.KeyspaceObservabilityValues = *resolvedValues
+	})
+
+	return nil
+}
+
+func prepareKeyspaceObservabilityForStarter(metadata map[string]string, resolvedValues *config.KeyspaceObservabilityValues) error {
+	copiedConfig := *config.GetGlobalConfig()
+	if err := copiedConfig.ResolveKeyspaceObservability(metadata); err != nil {
+		return err
+	}
+	configuredValues := copiedConfig.KeyspaceObservabilityValues.Clone()
+	maps.Copy(resolvedValues.MetricLabels, configuredValues.MetricLabels)
+	resolvedValues.SlowLogFields = configuredValues.SlowLogFields
+	resolvedValues.StmtLogFields = configuredValues.StmtLogFields
+	return nil
+}
+
 func enablePyroscope() {
 	if os.Getenv("PYROSCOPE_SERVER_ADDRESS") != "" {
 		runtime.SetMutexProfileFraction(5)

cmd/tidb-server/main_test.go

@@ -154,3 +154,71 @@ func TestSetVersionByConfigNormalizeLegacyPlaceholderForNextGen(t *testing.T) {
 	require.Equal(t, "v26.3.0", mysql.TiDBReleaseVersion)
 	require.Equal(t, "8.0.11-TiDB-CLOUD.202603.0", mysql.ServerVersion)
 }
+
+func TestSetupKeyspaceObservabilityForStarter(t *testing.T) {
+	if kerneltype.IsClassic() {
+		t.Skip("only for nextgen kernel")
+	}
+	restore := config.RestoreFunc()
+	defer restore()
+
+	originalMode := deploymode.Get()
+	t.Cleanup(func() {
+		require.NoError(t, deploymode.Set(originalMode))
+	})
+	require.NoError(t, deploymode.Set(deploymode.Starter))
+
+	keyspaceID := uint32(42)
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.Store = config.StoreTypeTiKV
+		conf.KeyspaceName = "ks"
+	})
+	err := prepareKeyspaceObservability(keyspaceID, nil)
+	require.NoError(t, err)
+	require.Equal(t, map[string]string{"keyspace_id": "42", "keyspace_name": "ks"}, config.GetGlobalConfig().GetKeyspaceObservabilityMetricLabels())
+
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.KeyspaceObservability = config.KeyspaceObservability{
+			Fields: []config.KeyspaceObservabilityField{{
+				Source:       "meta_a",
+				MetricLabel:  "keyspace_meta_label_a",
+				SlowLogField: "Slow_meta_a",
+				StmtLogField: "stmt_meta_a",
+				Required:     true,
+			}},
+		}
+	})
+
+	err = prepareKeyspaceObservability(keyspaceID, map[string]string{
+		"meta_a": "value_a",
+	})
+	require.NoError(t, err)
+
+	cfg := config.GetGlobalConfig()
+	require.Equal(t, map[string]string{"keyspace_id": "42", "keyspace_name": "ks", "keyspace_meta_label_a": "value_a"}, cfg.GetKeyspaceObservabilityMetricLabels())
+	require.Equal(t, map[string]string{"Slow_meta_a": "value_a"}, cfg.GetKeyspaceObservabilitySlowLogFields())
+	require.Equal(t, map[string]string{"stmt_meta_a": "value_a"}, cfg.GetKeyspaceObservabilityStmtLogFields())
+}
+
+func TestSetupKeyspaceObservabilityForStarterSkipsNonTiKV(t *testing.T) {
+	if kerneltype.IsClassic() {
+		t.Skip("only for nextgen kernel")
+	}
+	restore := config.RestoreFunc()
+	defer restore()
+	originalMode := deploymode.Get()
+	t.Cleanup(func() {
+		require.NoError(t, deploymode.Set(originalMode))
+	})
+	require.NoError(t, deploymode.Set(deploymode.Starter))
+
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.Store = config.StoreTypeUniStore
+		conf.Path = "invalid-pd-path"
+		conf.KeyspaceName = "test_keyspace"
+	})
+
+	keyspaceID := uint32(42)
+	require.NoError(t, prepareKeyspaceObservability(keyspaceID, nil))
+	require.Empty(t, config.GetGlobalConfig().GetKeyspaceObservabilityMetricLabels())
+}

pkg/config/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "config.go",
         "config_util.go",
         "const.go",
+        "keyspace_observability.go",
         "store.go",
         "tiflash.go",
     ],
@@ -23,6 +24,7 @@ go_library(
         "@com_github_burntsushi_toml//:toml",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_log//:log",
+        "@com_github_prometheus_common//model",
         "@com_github_tikv_client_go_v2//config",
         "@com_github_uber_jaeger_client_go//config",
         "@org_uber_go_atomic//:atomic",
@@ -42,7 +44,7 @@ go_test(
     data = glob(["**"]),
     embed = [":config"],
     flaky = True,
-    shard_count = 32,
+    shard_count = 34,
     deps = [
         "//pkg/config/deploymode",
         "//pkg/config/kerneltype",

pkg/config/config.go

@@ -262,6 +262,9 @@ type Config struct {
 	// key will be the default value of the session variable `txn_scope` for this tidb-server.
 	Labels map[string]string `toml:"labels" json:"labels"`
 
+	KeyspaceObservability       KeyspaceObservability       `toml:"keyspace-observability" json:"keyspace-observability"`
+	KeyspaceObservabilityValues KeyspaceObservabilityValues `toml:"-" json:"-"`
+
 	// EnableGlobalIndex is deprecated.
 	EnableGlobalIndex bool `toml:"enable-global-index" json:"enable-global-index"`
 
@@ -1473,12 +1476,18 @@ func (c *Config) Valid() error {
 	if !kerneltype.IsNextGen() && c.DeployMode != deploymode.Premium {
 		return fmt.Errorf("deploy-mode can only be configured for nextgen TiDB")
 	}
+	if len(c.KeyspaceObservability.Fields) > 0 && c.DeployMode != deploymode.Starter {
+		return fmt.Errorf("keyspace-observability.fields can only be configured when deploy-mode is starter")
+	}
 	if c.DXFResourceLimit < MinDXFResourceLimit || c.DXFResourceLimit > MaxDXFResourceLimit {
 		return fmt.Errorf("dxf-resource-limit should be between %d and %d", MinDXFResourceLimit, MaxDXFResourceLimit)
 	}
 	if c.DXFResourceLimit != DefDXFResourceLimit && c.DeployMode != deploymode.PremiumReserved {
 		return fmt.Errorf("dxf-resource-limit can only be configured when deploy-mode is premium_reserved")
 	}
+	if err := c.KeyspaceObservability.Valid(); err != nil {
+		return err
+	}
 	if c.Store == StoreTypeMockTiKV && !c.Instance.TiDBEnableDDL.Load() {
 		return fmt.Errorf("can't disable DDL on mocktikv")
 	}

pkg/config/config.toml.nextgen.example

@@ -446,6 +446,15 @@ allow-expression-index = false
 # engines means allow the tidb server read data from which types of engines. options: "tikv", "tiflash", "tidb".
 engines = ["tikv", "tiflash", "tidb"]
 
+# Map selected keyspace metadata entries to observability outputs.
+# Only valid when deploy-mode is starter.
+# [[keyspace-observability.fields]]
+# source = "meta_key"
+# metric-label = "metric_label"
+# slow-log-field = "Slow_log_field"
+# stmt-log-field = "stmt_log_field"
+# required = false
+
 # instance scope variables
 # These options are also available as a system variable for online configuration
 # changes to the system variable do not persist to the cluster. You must make changes