github.com/KinWaiYuen/client-go/v2@v2.5.4/internal/mockstore/mocktikv/cluster_manipulate.go (about)

     1  // Copyright 2021 TiKV Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // NOTE: The code in this file is based on code from the
    16  // TiDB project, licensed under the Apache License v 2.0
    17  //
    18  // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/mockstore/mocktikv/cluster_manipulate.go
    19  //
    20  
    21  // Copyright 2016 PingCAP, Inc.
    22  //
    23  // Licensed under the Apache License, Version 2.0 (the "License");
    24  // you may not use this file except in compliance with the License.
    25  // You may obtain a copy of the License at
    26  //
    27  //     http://www.apache.org/licenses/LICENSE-2.0
    28  //
    29  // Unless required by applicable law or agreed to in writing, software
    30  // distributed under the License is distributed on an "AS IS" BASIS,
    31  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    32  // See the License for the specific language governing permissions and
    33  // limitations under the License.
    34  
    35  package mocktikv
    36  
    37  import (
    38  	"fmt"
    39  
    40  	"github.com/pingcap/kvproto/pkg/metapb"
    41  )
    42  
    43  // BootstrapWithSingleStore initializes a Cluster with 1 Region and 1 Store.
    44  func BootstrapWithSingleStore(cluster *Cluster) (storeID, peerID, regionID uint64) {
    45  	ids := cluster.AllocIDs(3)
    46  	storeID, peerID, regionID = ids[0], ids[1], ids[2]
    47  	cluster.AddStore(storeID, fmt.Sprintf("store%d", storeID))
    48  	cluster.Bootstrap(regionID, []uint64{storeID}, []uint64{peerID}, peerID)
    49  	return
    50  }
    51  
    52  // BootstrapWithMultiStores initializes a Cluster with 1 Region and n Stores.
    53  func BootstrapWithMultiStores(cluster *Cluster, n int) (storeIDs, peerIDs []uint64, regionID uint64, leaderPeer uint64) {
    54  	storeIDs = cluster.AllocIDs(n)
    55  	peerIDs = cluster.AllocIDs(n)
    56  	leaderPeer = peerIDs[0]
    57  	regionID = cluster.AllocID()
    58  	for _, storeID := range storeIDs {
    59  		labels := []*metapb.StoreLabel{
    60  			{
    61  				Key:   "id",
    62  				Value: fmt.Sprintf("%v", storeID),
    63  			},
    64  		}
    65  		cluster.AddStore(storeID, fmt.Sprintf("store%d", storeID), labels...)
    66  	}
    67  	cluster.Bootstrap(regionID, storeIDs, peerIDs, leaderPeer)
    68  	return
    69  }
    70  
    71  // BootstrapWithMultiRegions initializes a Cluster with multiple Regions and 1
    72  // Store. The number of Regions will be len(splitKeys) + 1.
    73  func BootstrapWithMultiRegions(cluster *Cluster, splitKeys ...[]byte) (storeID uint64, regionIDs, peerIDs []uint64) {
    74  	var firstRegionID, firstPeerID uint64
    75  	storeID, firstPeerID, firstRegionID = BootstrapWithSingleStore(cluster)
    76  	regionIDs = append([]uint64{firstRegionID}, cluster.AllocIDs(len(splitKeys))...)
    77  	peerIDs = append([]uint64{firstPeerID}, cluster.AllocIDs(len(splitKeys))...)
    78  	for i, k := range splitKeys {
    79  		cluster.Split(regionIDs[i], regionIDs[i+1], k, []uint64{peerIDs[i]}, peerIDs[i])
    80  	}
    81  	return
    82  }