*: add keyspace observability labels

cmd/tidb-server/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
         "//pkg/keyspace",
         "//pkg/kv",
         "//pkg/metrics",
+        "//pkg/metrics/common",
         "//pkg/parser/mysql",
         "//pkg/parser/terror",
         "//pkg/parser/types",
@@ -67,11 +68,16 @@ go_library(
         "@com_github_opentracing_opentracing_go//:opentracing-go",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
+        "@com_github_pingcap_kvproto//pkg/keyspacepb",
         "@com_github_pingcap_log//:log",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_prometheus_client_golang//prometheus/push",
+        "@com_github_tikv_client_go_v2//config",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//txnkv/transaction",
+        "@com_github_tikv_pd_client//:client",
+        "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@org_uber_go_automaxprocs//maxprocs",
         "@org_uber_go_zap//:zap",
     ],
@@ -107,7 +113,7 @@ go_test(
     srcs = ["main_test.go"],
     embed = [":tidb-server_lib"],
     flaky = True,
-    shard_count = 6,
+    shard_count = 9,
     deps = [
         "//pkg/config",
         "//pkg/config/deploymode",
@@ -116,6 +122,7 @@ go_test(
         "//pkg/sessionctx/vardef",
         "//pkg/sessionctx/variable",
         "//pkg/testkit/testsetup",
+        "@com_github_pingcap_kvproto//pkg/keyspacepb",
         "@com_github_stretchr_testify//require",
         "@io_opencensus_go//stats/view",
         "@org_uber_go_goleak//:goleak",

cmd/tidb-server/main.go

@@ -19,6 +19,7 @@ import (
 	"flag"
 	"fmt"
 	"io/fs"
+	"maps"
 	"os"
 	"runtime"
 	"strconv"
@@ -30,6 +31,7 @@ import (
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/keyspacepb"
 	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/pkg/bindinfo"
 	"github.com/pingcap/tidb/pkg/config"
@@ -44,6 +46,7 @@ import (
 	"github.com/pingcap/tidb/pkg/keyspace"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/metrics"
+	metricscommon "github.com/pingcap/tidb/pkg/metrics/common"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
 	parsertypes "github.com/pingcap/tidb/pkg/parser/types"
@@ -90,8 +93,12 @@ import (
 	repository "github.com/pingcap/tidb/pkg/util/workloadrepo"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/push"
+	tikvconfig "github.com/tikv/client-go/v2/config"
 	"github.com/tikv/client-go/v2/tikv"
 	"github.com/tikv/client-go/v2/txnkv/transaction"
+	pd "github.com/tikv/pd/client"
+	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/automaxprocs/maxprocs"
 	"go.uber.org/zap"
 )
@@ -335,6 +342,8 @@ func main() {
 	signal.SetupUSR1Handler()
 	err = registerStores()
 	terror.MustNil(err)
+	err = prepareKeyspaceObservability()
+	terror.MustNil(err)
 	err = metricsutil.RegisterMetrics()
 	terror.MustNil(err)
 
@@ -1146,6 +1155,87 @@ func closeStmtSummary() {
 	}
 }
 
+var keyspaceMetaComponentName = caller.Component("tidb-keyspace-meta")
+
+const (
+	keyspaceIDMetricLabel   = "keyspace_id"
+	keyspaceNameMetricLabel = "keyspace_name"
+)
+
+func prepareKeyspaceObservability() error {
+	cfg := config.GetGlobalConfig()
+	if !kerneltype.IsNextGen() {
+		return nil
+	}
+	if keyspace.IsKeyspaceNameEmpty(cfg.KeyspaceName) || cfg.Store != config.StoreTypeTiKV {
+		return nil
+	}
+	metricscommon.SetConstLabels(keyspaceNameMetricLabel, cfg.KeyspaceName)
+	pdAddrs, _, _, err := tikvconfig.ParsePath("tikv://" + cfg.Path)
+	if err != nil {
+		return err
+	}
+	timeoutSec := time.Duration(cfg.PDClient.PDServerTimeout) * time.Second
+	pdCli, err := pd.NewClient(keyspaceMetaComponentName, pdAddrs, pd.SecurityOption{
+		CAPath:   cfg.Security.ClusterSSLCA,
+		CertPath: cfg.Security.ClusterSSLCert,
+		KeyPath:  cfg.Security.ClusterSSLKey,
+	}, opt.WithCustomTimeoutOption(timeoutSec), opt.WithInitMetricsOption(false))
+	if err != nil {
+		return err
+	}
+	defer pdCli.Close()
+
+	keyspaceMeta, err := getKeyspaceMeta(pdCli, cfg.KeyspaceName)
+	if err != nil {
+		return err
+	}
+	keyspace.SetKeyspaceMeta(keyspaceMeta)
+	return prepareKeyspaceObservabilityWithKeyspaceMeta(keyspaceMeta, cfg.KeyspaceName, deploymode.IsStarter())
+}
+
+func getKeyspaceMeta(pdCli pd.Client, keyspaceName string) (*keyspacepb.KeyspaceMeta, error) {
+	var keyspaceMeta *keyspacepb.KeyspaceMeta
+	err := util.RunWithRetry(util.DefaultMaxRetries, util.RetryInterval, func() (bool, error) {
+		var errInner error
+		keyspaceMeta, errInner = pdCli.LoadKeyspace(context.TODO(), keyspaceName)
+		if kvstore.IsNotBootstrappedError(errInner) || kvstore.IsKeyspaceNotExistError(errInner) {
+			return true, errInner
+		}
+		return false, errInner
+	})
+	if err != nil {
+		return nil, err
+	}
+	return keyspaceMeta, nil
+}
+
+func prepareKeyspaceObservabilityWithKeyspaceMeta(keyspaceMeta *keyspacepb.KeyspaceMeta, keyspaceName string, includeConfiguredFields bool) error {
+	if keyspaceMeta == nil {
+		return nil
+	}
+	resolvedValues := config.KeyspaceObservabilityValues{
+		MetricLabels: map[string]string{
+			keyspaceIDMetricLabel:   fmt.Sprint(keyspaceMeta.GetId()),
+			keyspaceNameMetricLabel: keyspaceName,
+		},
+	}
+	if includeConfiguredFields {
+		copiedConfig := *config.GetGlobalConfig()
+		if err := copiedConfig.ResolveKeyspaceObservability(keyspaceMeta.GetConfig()); err != nil {
+			return err
+		}
+		configuredValues := copiedConfig.KeyspaceObservabilityValues.Clone()
+		maps.Copy(resolvedValues.MetricLabels, configuredValues.MetricLabels)
+		resolvedValues.SlowLogFields = configuredValues.SlowLogFields
+		resolvedValues.StmtLogFields = configuredValues.StmtLogFields
+	}
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.KeyspaceObservabilityValues = resolvedValues
+	})
+	return nil
+}
+
 func enablePyroscope() {
 	if os.Getenv("PYROSCOPE_SERVER_ADDRESS") != "" {
 		runtime.SetMutexProfileFraction(5)

cmd/tidb-server/main_test.go

@@ -18,6 +18,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/pingcap/kvproto/pkg/keyspacepb"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/config/deploymode"
 	"github.com/pingcap/tidb/pkg/config/kerneltype"
@@ -154,3 +155,63 @@ 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) {
+	restore := config.RestoreFunc()
+	defer restore()
+	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 := prepareKeyspaceObservabilityWithKeyspaceMeta(&keyspacepb.KeyspaceMeta{
+		Id:     42,
+		Config: map[string]string{"meta_a": "value_a"},
+	}, "ks", true)
+	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, []config.KeyspaceObservabilityFieldPair{{Key: "Slow_meta_a", Value: "value_a"}}, cfg.GetKeyspaceObservabilitySlowLogFields())
+	require.Equal(t, []config.KeyspaceObservabilityFieldPair{{Key: "stmt_meta_a", Value: "value_a"}}, cfg.GetKeyspaceObservabilityStmtLogFields())
+}
+
+func TestSetupKeyspaceObservabilityForNonStarter(t *testing.T) {
+	restore := config.RestoreFunc()
+	defer restore()
+
+	err := prepareKeyspaceObservabilityWithKeyspaceMeta(&keyspacepb.KeyspaceMeta{
+		Id:     42,
+		Config: map[string]string{"meta_a": "value_a"},
+	}, "ks", false)
+	require.NoError(t, err)
+
+	cfg := config.GetGlobalConfig()
+	require.Equal(t, map[string]string{"keyspace_id": "42", "keyspace_name": "ks"}, cfg.GetKeyspaceObservabilityMetricLabels())
+	require.Empty(t, cfg.GetKeyspaceObservabilitySlowLogFields())
+	require.Empty(t, cfg.GetKeyspaceObservabilityStmtLogFields())
+}
+
+func TestSetupKeyspaceObservabilityForStartSkipsClassic(t *testing.T) {
+	if !kerneltype.IsClassic() {
+		t.Skip("only verifies the classic-mode short-circuit path")
+	}
+
+	restore := config.RestoreFunc()
+	defer restore()
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.Store = config.StoreTypeTiKV
+		conf.Path = "invalid-pd-path"
+		conf.KeyspaceName = "test_keyspace"
+	})
+
+	require.NoError(t, prepareKeyspaceObservability())
+	require.Empty(t, config.GetGlobalConfig().GetKeyspaceObservabilityMetricLabels())
+}

pkg/config/BUILD.bazel

@@ -42,7 +42,7 @@ go_test(
     data = glob(["**"]),
     embed = [":config"],
     flaky = True,
-    shard_count = 32,
+    shard_count = 34,
     deps = [
         "//pkg/config/deploymode",
         "//pkg/config/kerneltype",