keyspace: introduce keyspace conf and etcd path

config/config.go

@@ -89,6 +89,8 @@ const (
 	DefTempDir = "/tmp/tidb"
 	// DefAuthTokenRefreshInterval is the default time interval to refresh tidb auth token.
 	DefAuthTokenRefreshInterval = time.Hour
+	// EnvVarKeyspaceName is the system env name for keyspace name.
+	EnvVarKeyspaceName = "KEYSPACE_NAME"
 )
 
 // Valid config maps
@@ -183,6 +185,7 @@ type Config struct {
 	VersionComment             string                  `toml:"version-comment" json:"version-comment"`
 	TiDBEdition                string                  `toml:"tidb-edition" json:"tidb-edition"`
 	TiDBReleaseVersion         string                  `toml:"tidb-release-version" json:"tidb-release-version"`
+	KeyspaceName               string                  `toml:"keyspace-name" json:"keyspace-name"`
 	Log                        Log                     `toml:"log" json:"log"`
 	Instance                   Instance                `toml:"instance" json:"instance"`
 	Security                   Security                `toml:"security" json:"security"`
@@ -1457,3 +1460,15 @@ func ContainHiddenConfig(s string) bool {
 	}
 	return false
 }
+
+// GetGlobalKeyspaceName is used to get global keyspace name
+// from config file or environment variable.
+func GetGlobalKeyspaceName() string {
+	// The cfg.keyspaceName get higher priority than KEYSPACE_NAME in system env.
+	keyspaceName := GetGlobalConfig().KeyspaceName
+	if keyspaceName != "" {
+		return keyspaceName
+	}
+
+	return os.Getenv(EnvVarKeyspaceName)
+}

config/config_test.go

@@ -1288,3 +1288,20 @@ func TestStatsLoadLimit(t *testing.T) {
 	checkQueueSizeValid(DefMaxOfStatsLoadQueueSizeLimit, true)
 	checkQueueSizeValid(DefMaxOfStatsLoadQueueSizeLimit+1, false)
 }
+
+func TestGetGlobalKeyspaceName(t *testing.T) {
+	conf := NewConfig()
+	require.Empty(t, conf.KeyspaceName)
+
+	UpdateGlobal(func(conf *Config) {
+		conf.KeyspaceName = "test"
+	})
+
+	require.Equal(t, "test", GetGlobalKeyspaceName())
+
+	UpdateGlobal(func(conf *Config) {
+		conf.KeyspaceName = ""
+	})
+	require.Nil(t, os.Setenv(EnvVarKeyspaceName, "test2"))
+	require.Equal(t, "test2", GetGlobalKeyspaceName())
+}

domain/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "//errno",
         "//infoschema",
         "//infoschema/perfschema",
+        "//keyspace",
         "//kv",
         "//meta",
         "//metrics",
@@ -55,6 +56,7 @@ go_library(
         "//util/dbterror",
         "//util/domainutil",
         "//util/engine",
+        "//util/etcd",
         "//util/execdetails",
         "//util/expensivequery",
         "//util/logutil",

domain/domain.go

@@ -45,6 +45,7 @@ import (
 	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/infoschema/perfschema"
+	"github.com/pingcap/tidb/keyspace"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta"
 	"github.com/pingcap/tidb/metrics"
@@ -65,6 +66,7 @@ import (
 	"github.com/pingcap/tidb/util/dbterror"
 	"github.com/pingcap/tidb/util/domainutil"
 	"github.com/pingcap/tidb/util/engine"
+	"github.com/pingcap/tidb/util/etcd"
 	"github.com/pingcap/tidb/util/expensivequery"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/memory"
@@ -966,6 +968,9 @@ func (do *Domain) Init(
 			if err != nil {
 				return errors.Trace(err)
 			}
+
+			etcd.SetEtcdCliByNamespace(cli, keyspace.MakeKeyspaceEtcdNamespace(do.store.GetCodec()))
+
 			do.etcdClient = cli
 		}
 	}

go.mod

@@ -260,5 +260,6 @@ replace (
 	// fix potential security issue(CVE-2020-26160) introduced by indirect dependency.
 	github.com/dgrijalva/jwt-go => github.com/form3tech-oss/jwt-go v3.2.6-0.20210809144907-32ab6a8243d7+incompatible
 	github.com/pingcap/tidb/parser => ./parser
+	github.com/tikv/client-go/v2 => github.com/AmoebaProtozoa/client-go/v2 v2.0.1-0.20230106075830-5c53ec2b712a
 	go.opencensus.io => go.opencensus.io v0.23.1-0.20220331163232-052120675fac
 )

go.sum

@@ -58,6 +58,8 @@ cloud.google.com/go/storage v1.21.0 h1:HwnT2u2D309SFDHQII6m18HlrCi3jAXhUMTLOWXYH
 cloud.google.com/go/storage v1.21.0/go.mod h1:XmRlxkgPjlBONznT2dDUU/5XlpU2OjMnKuqnZI01LAA=
 contrib.go.opencensus.io/exporter/ocagent v0.4.12/go.mod h1:450APlNTSR6FrvC3CTRqYosuDstRB9un7SOx2k/9ckA=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
+github.com/AmoebaProtozoa/client-go/v2 v2.0.1-0.20230106075830-5c53ec2b712a h1:zvh9psq64H7+MRJASAXNVUyudNft81a8k9nJRw+QEpI=
+github.com/AmoebaProtozoa/client-go/v2 v2.0.1-0.20230106075830-5c53ec2b712a/go.mod h1:ntjnKL/fFGCmKcpw90CSySxJPsGaMr39ZPOtSFK5zqs=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/Azure/azure-sdk-for-go v23.2.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v0.20.0 h1:KQgdWmEOmaJKxaUUZwHAYh12t+b+ZJf8q3friycK1kA=
@@ -936,8 +938,6 @@ github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJf
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW93SG+q0F8KI+yFrcIDT4c/RNoc4=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
-github.com/tikv/client-go/v2 v2.0.4 h1:cPtMXTExqjzk8L40qhrgB/mXiBXKP5LRU0vwjtI2Xxo=
-github.com/tikv/client-go/v2 v2.0.4/go.mod h1:v52O5zDtv2BBus4lm5yrSQhxGW4Z4RaXWfg0U1Kuyqo=
 github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07 h1:ckPpxKcl75mO2N6a4cJXiZH43hvcHPpqc9dh1TmH1nc=
 github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07/go.mod h1:CipBxPfxPUME+BImx9MUYXCnAVLS3VJUr3mnSJwh40A=
 github.com/timakin/bodyclose v0.0.0-20210704033933-f49887972144 h1:kl4KhGNsJIbDHS9/4U9yQo1UcPQM0kOMJHn29EoH/Ro=

keyspace/BUILD.bazel

@@ -0,0 +1,12 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "keyspace",
+    srcs = ["keyspace.go"],
+    importpath = "github.com/pingcap/tidb/keyspace",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_github_pingcap_kvproto//pkg/kvrpcpb",
+        "@com_github_tikv_client_go_v2//tikv",
+    ],
+)

keyspace/keyspace.go

@@ -0,0 +1,24 @@
+package keyspace
+
+import (
+	"fmt"
+
+	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/tikv/client-go/v2/tikv"
+)
+
+const (
+	// tidbKeyspaceEtcdPathPrefix is the keyspace prefix for etcd namespace
+	tidbKeyspaceEtcdPathPrefix = "/keyspaces/tidb/"
+)
+
+// CodecV1 represents api v1 codec.
+var CodecV1 = tikv.NewCodecV1(tikv.ModeTxn)
+
+// MakeKeyspaceEtcdNamespace return the keyspace prefix path for etcd namespace
+func MakeKeyspaceEtcdNamespace(c tikv.Codec) string {
+	if c.GetAPIVersion() == kvrpcpb.APIVersion_V1 {
+		return ""
+	}
+	return fmt.Sprintf(tidbKeyspaceEtcdPathPrefix+"%d", c.GetKeyspaceID())
+}

kv/kv.go

@@ -651,6 +651,8 @@ type Storage interface {
 	GetMinSafeTS(txnScope string) uint64
 	// GetLockWaits return all lock wait information
 	GetLockWaits() ([]*deadlockpb.WaitForEntry, error)
+	// GetCodec gets the codec of the storage.
+	GetCodec() tikv.Codec
 }
 
 // EtcdBackend is used for judging a storage is a real TiKV.

server/http_status.go

@@ -461,8 +461,9 @@ func (s *Server) startStatusServerAndRPCServer(serverMux *http.ServeMux) {
 	grpcServer := NewRPCServer(s.cfg, s.dom, s)
 	service.RegisterChannelzServiceToServer(grpcServer)
 	if s.cfg.Store == "tikv" {
+		keyspaceName := config.GetGlobalKeyspaceName()
 		for {
-			fullPath := fmt.Sprintf("tikv://%s", s.cfg.Path)
+			fullPath := fmt.Sprintf("%s://%s?keyspaceName=%s", s.cfg.Store, s.cfg.Path, keyspaceName)
 			store, err := store.New(fullPath)
 			if err != nil {
 				logutil.BgLogger().Error("new tikv store fail", zap.Error(err))

store/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//util",
         "//util/logutil",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_kvproto//pkg/pdpb",
         "@org_uber_go_zap//:zap",
     ],
 )

store/driver/tikv_driver.go

@@ -88,6 +88,7 @@ func WithPDClientConfig(client config.PDClient) Option {
 
 // TiKVDriver implements engine TiKV.
 type TiKVDriver struct {
+	keyspaceName    string
 	pdConfig        config.PDClient
 	security        config.Security
 	tikvConfig      config.TiKVClient
@@ -117,7 +118,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 	mc.Lock()
 	defer mc.Unlock()
 	d.setDefaultAndOptions(options...)
-	etcdAddrs, disableGC, err := config.ParsePath(path)
+	etcdAddrs, disableGC, keyspaceName, err := config.ParsePath(path)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -142,7 +143,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 	}
 
 	// FIXME: uuid will be a very long and ugly string, simplify it.
-	uuid := fmt.Sprintf("tikv-%v", pdCli.GetClusterID(context.TODO()))
+	uuid := fmt.Sprintf("tikv-%v-%s", pdCli.GetClusterID(context.TODO()), keyspaceName)
 	if store, ok := mc.cache[uuid]; ok {
 		return store, nil
 	}
@@ -157,11 +158,39 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		return nil, errors.Trace(err)
 	}
 
-	pdClient := tikv.CodecPDClient{Client: pdCli}
-	s, err := tikv.NewKVStore(uuid, &pdClient, spkv, tikv.NewRPCClient(tikv.WithSecurity(d.security)))
+	// ---------------- keyspace logic  ----------------
+	var (
+		pdClient *tikv.CodecPDClient
+	)
+
+	if keyspaceName == "" {
+		logutil.BgLogger().Info("using API V1.")
+		pdClient = tikv.NewCodecPDClient(tikv.ModeTxn, pdCli)
+	} else {
+		logutil.BgLogger().Info("using API V2.", zap.String("keyspaceName", keyspaceName))
+		pdClient, err = tikv.NewCodecPDClientWithKeyspace(tikv.ModeTxn, pdCli, keyspaceName)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		// If there's setting keyspace-name, then skipped GC worker logic.
+		// It needs a group of special tidb nodes to execute GC worker logic.
+		// TODO: remove this restriction while merged keyspace GC worker logic.
+		disableGC = true
+	}
+
+	codec := pdClient.GetCodec()
+
+	rpcClient := tikv.NewRPCClient(
+		tikv.WithSecurity(d.security),
+		tikv.WithCodec(codec),
+	)
+
+	s, err := tikv.NewKVStore(uuid, pdClient, spkv, rpcClient)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+
+	// ---------------- keyspace logic  ----------------
 	if d.txnLocalLatches.Enabled {
 		s.EnableTxnLocalLatches(d.txnLocalLatches.Capacity)
 	}
@@ -178,6 +207,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		memCache:  kv.NewCacheDB(),
 		enableGC:  !disableGC,
 		coprStore: coprStore,
+		codec:     codec,
 	}
 
 	mc.cache[uuid] = store
@@ -192,6 +222,7 @@ type tikvStore struct {
 	enableGC  bool
 	gcWorker  *gcworker.GCWorker
 	coprStore *copr.Store
+	codec     tikv.Codec
 }
 
 // Name gets the name of the storage engine
@@ -343,3 +374,7 @@ func (s *tikvStore) GetLockWaits() ([]*deadlockpb.WaitForEntry, error) {
 	}
 	return result, nil
 }
+
+func (s *tikvStore) GetCodec() tikv.Codec {
+	return s.codec
+}

store/helper/helper.go

@@ -78,6 +78,7 @@ type Storage interface {
 	Closed() <-chan struct{}
 	GetMinSafeTS(txnScope string) uint64
 	GetLockWaits() ([]*deadlockpb.WaitForEntry, error)
+	GetCodec() tikv.Codec
 }
 
 // Helper is a middleware to get some information from tikv/pd. It can be used for TiDB's http api or mem table.

store/mockstore/mockstorage/storage.go

@@ -117,6 +117,12 @@ func (s *mockStorage) Close() error {
 	return s.KVStore.Close()
 }
 
+func (s *mockStorage) GetCodec() tikv.Codec {
+	pdClient := s.KVStore.GetPDClient()
+	pdCodecCli := tikv.NewCodecPDClient(tikv.ModeTxn, pdClient)
+	return pdCodecCli.GetCodec()
+}
+
 // MockLockWaitSetter is used to set the mocked lock wait information, which helps implementing tests that uses the
 // GetLockWaits function.
 type MockLockWaitSetter interface {

store/store.go

@@ -20,6 +20,7 @@ import (
 	"sync"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/logutil"
@@ -74,7 +75,7 @@ func newStoreWithRetry(path string, maxRetries int) (kv.Storage, error) {
 	err = util.RunWithRetry(maxRetries, util.RetryInterval, func() (bool, error) {
 		logutil.BgLogger().Info("new store", zap.String("path", path))
 		s, err = d.Open(path)
-		return kv.IsTxnRetryableError(err), err
+		return isNewStoreRetryableError(err), err
 	})
 
 	if err == nil {
@@ -91,3 +92,32 @@ func loadDriver(name string) (kv.Driver, bool) {
 	d, ok := stores[name]
 	return d, ok
 }
+
+// isOpenRetryableError check if the new store operation should be retried under given error
+// currently, it should be retried if:
+//
+//	Transaction conflict and is retryable (kv.IsTxnRetryableError)
+//	PD is not bootstrapped at the time of request
+//	Keyspace requested does not exist (request prior to PD keyspace pre-split)
+func isNewStoreRetryableError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return kv.IsTxnRetryableError(err) || IsNotBootstrappedError(err) || IsKeyspaceNotExistError(err)
+}
+
+// IsNotBootstrappedError returns true if the error is pd not bootstrapped error.
+func IsNotBootstrappedError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return strings.Contains(err.Error(), pdpb.ErrorType_NOT_BOOTSTRAPPED.String())
+}
+
+// IsKeyspaceNotExistError returns true the error is caused by keyspace not exists.
+func IsKeyspaceNotExistError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return strings.Contains(err.Error(), pdpb.ErrorType_ENTRY_NOT_FOUND.String())
+}

tidb-server/main.go

@@ -214,8 +214,11 @@ func main() {
 	printInfo()
 	setupBinlogClient()
 	setupMetrics()
+
+	keyspaceName := config.GetGlobalKeyspaceName()
+
 	resourcemanager.GlobalResourceManager.Start()
-	storage, dom := createStoreAndDomain()
+	storage, dom := createStoreAndDomain(keyspaceName)
 	svr := createServer(storage, dom)
 
 	// Register error API is not thread-safe, the caller MUST NOT register errors after initialization.
@@ -306,9 +309,9 @@ func registerMetrics() {
 	}
 }
 
-func createStoreAndDomain() (kv.Storage, *domain.Domain) {
+func createStoreAndDomain(keyspaceName string) (kv.Storage, *domain.Domain) {
 	cfg := config.GetGlobalConfig()
-	fullPath := fmt.Sprintf("%s://%s", cfg.Store, cfg.Path)
+	fullPath := fmt.Sprintf("%s://%s?keyspaceName=%s", cfg.Store, cfg.Path, keyspaceName)
 	var err error
 	storage, err := kvstore.New(fullPath)
 	terror.MustNil(err)

util/etcd/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
     deps = [
         "@com_github_pingcap_errors//:errors",
         "@io_etcd_go_etcd_client_v3//:client",
+        "@io_etcd_go_etcd_client_v3//namespace",
     ],
 )