tidb mockstore 源码

tidb mockstore 代码

文件路径：/store/mockstore/mockstore.go

// Copyright 2020 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mockstore

import (
	"net/url"
	"strings"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/store/mockstore/unistore"
	"github.com/tikv/client-go/v2/testutils"
	"github.com/tikv/client-go/v2/tikv"
	pd "github.com/tikv/pd/client"
)

// MockTiKVDriver is in memory mock TiKV driver.
type MockTiKVDriver struct{}

// Open creates a MockTiKV storage.
func (d MockTiKVDriver) Open(path string) (kv.Storage, error) {
	u, err := url.Parse(path)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if !strings.EqualFold(u.Scheme, "mocktikv") {
		return nil, errors.Errorf("Uri scheme expected(mocktikv) but found (%s)", u.Scheme)
	}

	opts := []MockTiKVStoreOption{WithPath(u.Path), WithStoreType(MockTiKV)}
	txnLocalLatches := config.GetGlobalConfig().TxnLocalLatches
	if txnLocalLatches.Enabled {
		opts = append(opts, WithTxnLocalLatches(txnLocalLatches.Capacity))
	}

	return NewMockStore(opts...)
}

// EmbedUnistoreDriver is in embedded unistore driver.
type EmbedUnistoreDriver struct{}

// Open creates a EmbedUnistore storage.
func (d EmbedUnistoreDriver) Open(path string) (kv.Storage, error) {
	u, err := url.Parse(path)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if !strings.EqualFold(u.Scheme, "unistore") {
		return nil, errors.Errorf("Uri scheme expected(unistore) but found (%s)", u.Scheme)
	}

	opts := []MockTiKVStoreOption{WithPath(u.Path), WithStoreType(EmbedUnistore)}
	txnLocalLatches := config.GetGlobalConfig().TxnLocalLatches
	if txnLocalLatches.Enabled {
		opts = append(opts, WithTxnLocalLatches(txnLocalLatches.Capacity))
	}

	return NewMockStore(opts...)
}

// StoreType is the type of backend mock storage.
type StoreType uint8

const (
	// MockTiKV is the mock storage based on goleveldb.
	MockTiKV StoreType = iota
	// EmbedUnistore is the mock storage based on unistore.
	EmbedUnistore

	defaultStoreType = EmbedUnistore
)

type mockOptions struct {
	clusterInspector func(testutils.Cluster)
	clientHijacker   func(tikv.Client) tikv.Client
	pdClientHijacker func(pd.Client) pd.Client
	path             string
	txnLocalLatches  uint
	storeType        StoreType
	ddlCheckerHijack bool
}

// MockTiKVStoreOption is used to control some behavior of mock tikv.
type MockTiKVStoreOption func(*mockOptions)

// WithMultipleOptions merges multiple options into one option.
func WithMultipleOptions(opts ...MockTiKVStoreOption) MockTiKVStoreOption {
	return func(args *mockOptions) {
		for _, opt := range opts {
			opt(args)
		}
	}
}

// WithClientHijacker hijacks KV client's behavior, makes it easy to simulate the network
// problem between TiDB and TiKV.
func WithClientHijacker(hijacker func(tikv.Client) tikv.Client) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.clientHijacker = hijacker
	}
}

// WithPDClientHijacker hijacks PD client's behavior, makes it easy to simulate the network
// problem between TiDB and PD.
func WithPDClientHijacker(hijacker func(pd.Client) pd.Client) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.pdClientHijacker = hijacker
	}
}

// WithClusterInspector lets user to inspect the mock cluster handler.
func WithClusterInspector(inspector func(testutils.Cluster)) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.clusterInspector = inspector
	}
}

// WithStoreType lets user choose the backend storage's type.
func WithStoreType(tp StoreType) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.storeType = tp
	}
}

// WithPath specifies the mocktikv path.
func WithPath(path string) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.path = path
	}
}

// WithTxnLocalLatches enable txnLocalLatches, when capacity > 0.
func WithTxnLocalLatches(capacity uint) MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.txnLocalLatches = capacity
	}
}

// WithDDLChecker prepare injected DDL implementation for the domain of this store. It must be done before bootstrap to
// avoid data race with dom.ddl.
func WithDDLChecker() MockTiKVStoreOption {
	return func(c *mockOptions) {
		c.ddlCheckerHijack = true
	}
}

// DDLCheckerInjector is used to break import cycle.
var DDLCheckerInjector func(kv.Storage) kv.Storage

// NewMockStore creates a mocked tikv store, the path is the file path to store the data.
// If path is an empty string, a memory storage will be created.
func NewMockStore(options ...MockTiKVStoreOption) (kv.Storage, error) {
	opt := mockOptions{
		clusterInspector: func(c testutils.Cluster) {
			BootstrapWithSingleStore(c)
		},
		storeType: defaultStoreType,
	}
	for _, f := range options {
		f(&opt)
	}

	var (
		store kv.Storage
		err   error
	)

	switch opt.storeType {
	case MockTiKV:
		store, err = newMockTikvStore(&opt)
	case EmbedUnistore:
		store, err = newUnistore(&opt)
	default:
		panic("unsupported mockstore")
	}
	if err != nil {
		return nil, errors.Trace(err)
	}

	if opt.ddlCheckerHijack {
		store = DDLCheckerInjector(store)
	}
	return store, nil
}

// BootstrapWithSingleStore initializes a Cluster with 1 Region and 1 Store.
func BootstrapWithSingleStore(cluster testutils.Cluster) (storeID, peerID, regionID uint64) {
	switch x := cluster.(type) {
	case *testutils.MockCluster:
		return testutils.BootstrapWithSingleStore(x)
	case *unistore.Cluster:
		return unistore.BootstrapWithSingleStore(x)
	default:
		panic("unsupported cluster type")
	}
}

// BootstrapWithMultiStores initializes a Cluster with 1 Region and n Stores.
func BootstrapWithMultiStores(cluster testutils.Cluster, n int) (storeIDs, peerIDs []uint64, regionID uint64, leaderPeer uint64) {
	switch x := cluster.(type) {
	case *testutils.MockCluster:
		return testutils.BootstrapWithMultiStores(x, n)
	case *unistore.Cluster:
		return unistore.BootstrapWithMultiStores(x, n)
	default:
		panic("unsupported cluster type")
	}
}

// BootstrapWithMultiRegions initializes a Cluster with multiple Regions and 1
// Store. The number of Regions will be len(splitKeys) + 1.
func BootstrapWithMultiRegions(cluster testutils.Cluster, splitKeys ...[]byte) (storeID uint64, regionIDs, peerIDs []uint64) {
	switch x := cluster.(type) {
	case *testutils.MockCluster:
		return testutils.BootstrapWithMultiRegions(x, splitKeys...)
	case *unistore.Cluster:
		return unistore.BootstrapWithMultiRegions(x, splitKeys...)
	default:
		panic("unsupported cluster type")
	}
}

相关信息

tidb 源码目录

相关文章

tidb redirector 源码

tidb tikv 源码

tidb unistore 源码