github.com/KinWaiYuen/client-go/v2@v2.5.4/rawkv/rawkv_test.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/rawkv_test.go
    19  //
    20  
    21  // Copyright 2021 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 rawkv
    36  
    37  import (
    38  	"context"
    39  	"fmt"
    40  	"testing"
    41  
    42  	"github.com/KinWaiYuen/client-go/v2/internal/locate"
    43  	"github.com/KinWaiYuen/client-go/v2/internal/mockstore/mocktikv"
    44  	"github.com/KinWaiYuen/client-go/v2/internal/retry"
    45  	"github.com/KinWaiYuen/client-go/v2/kv"
    46  	"github.com/stretchr/testify/suite"
    47  )
    48  
    49  func TestRawKV(t *testing.T) {
    50  	suite.Run(t, new(testRawkvSuite))
    51  }
    52  
    53  type testRawkvSuite struct {
    54  	suite.Suite
    55  	mvccStore mocktikv.MVCCStore
    56  	cluster   *mocktikv.Cluster
    57  	store1    uint64 // store1 is leader
    58  	store2    uint64 // store2 is follower
    59  	peer1     uint64 // peer1 is leader
    60  	peer2     uint64 // peer2 is follower
    61  	region1   uint64
    62  	bo        *retry.Backoffer
    63  }
    64  
    65  func (s *testRawkvSuite) SetupTest() {
    66  	s.mvccStore = mocktikv.MustNewMVCCStore()
    67  	s.cluster = mocktikv.NewCluster(s.mvccStore)
    68  	storeIDs, peerIDs, regionID, _ := mocktikv.BootstrapWithMultiStores(s.cluster, 2)
    69  	s.region1 = regionID
    70  	s.store1 = storeIDs[0]
    71  	s.store2 = storeIDs[1]
    72  	s.peer1 = peerIDs[0]
    73  	s.peer2 = peerIDs[1]
    74  	s.bo = retry.NewBackofferWithVars(context.Background(), 5000, nil)
    75  }
    76  
    77  func (s *testRawkvSuite) TearDownTest() {
    78  	s.mvccStore.Close()
    79  }
    80  
    81  func (s *testRawkvSuite) storeAddr(id uint64) string {
    82  	return fmt.Sprintf("store%d", id)
    83  }
    84  
    85  func (s *testRawkvSuite) TestReplaceAddrWithNewStore() {
    86  	mvccStore := mocktikv.MustNewMVCCStore()
    87  	defer mvccStore.Close()
    88  
    89  	client := &Client{
    90  		clusterID:   0,
    91  		regionCache: locate.NewRegionCache(mocktikv.NewPDClient(s.cluster)),
    92  		rpcClient:   mocktikv.NewRPCClient(s.cluster, mvccStore, nil),
    93  	}
    94  	defer client.Close()
    95  	testKey := []byte("test_key")
    96  	testValue := []byte("test_value")
    97  	err := client.Put(context.Background(), testKey, testValue)
    98  	s.Nil(err)
    99  
   100  	// make store2 using store1's addr and store1 offline
   101  	store1Addr := s.storeAddr(s.store1)
   102  	s.cluster.UpdateStoreAddr(s.store1, s.storeAddr(s.store2))
   103  	s.cluster.UpdateStoreAddr(s.store2, store1Addr)
   104  	s.cluster.RemoveStore(s.store1)
   105  	s.cluster.ChangeLeader(s.region1, s.peer2)
   106  	s.cluster.RemovePeer(s.region1, s.peer1)
   107  
   108  	getVal, err := client.Get(context.Background(), testKey)
   109  
   110  	s.Nil(err)
   111  	s.Equal(getVal, testValue)
   112  }
   113  
   114  func (s *testRawkvSuite) TestUpdateStoreAddr() {
   115  	mvccStore := mocktikv.MustNewMVCCStore()
   116  	defer mvccStore.Close()
   117  
   118  	client := &Client{
   119  		clusterID:   0,
   120  		regionCache: locate.NewRegionCache(mocktikv.NewPDClient(s.cluster)),
   121  		rpcClient:   mocktikv.NewRPCClient(s.cluster, mvccStore, nil),
   122  	}
   123  	defer client.Close()
   124  	testKey := []byte("test_key")
   125  	testValue := []byte("test_value")
   126  	err := client.Put(context.Background(), testKey, testValue)
   127  	s.Nil(err)
   128  	// tikv-server reports `StoreNotMatch` And retry
   129  	store1Addr := s.storeAddr(s.store1)
   130  	s.cluster.UpdateStoreAddr(s.store1, s.storeAddr(s.store2))
   131  	s.cluster.UpdateStoreAddr(s.store2, store1Addr)
   132  
   133  	getVal, err := client.Get(context.Background(), testKey)
   134  
   135  	s.Nil(err)
   136  	s.Equal(getVal, testValue)
   137  }
   138  
   139  func (s *testRawkvSuite) TestReplaceNewAddrAndOldOfflineImmediately() {
   140  	mvccStore := mocktikv.MustNewMVCCStore()
   141  	defer mvccStore.Close()
   142  
   143  	client := &Client{
   144  		clusterID:   0,
   145  		regionCache: locate.NewRegionCache(mocktikv.NewPDClient(s.cluster)),
   146  		rpcClient:   mocktikv.NewRPCClient(s.cluster, mvccStore, nil),
   147  	}
   148  	defer client.Close()
   149  	testKey := []byte("test_key")
   150  	testValue := []byte("test_value")
   151  	err := client.Put(context.Background(), testKey, testValue)
   152  	s.Nil(err)
   153  
   154  	// pre-load store2's address into cache via follower-read.
   155  	loc, err := client.regionCache.LocateKey(s.bo, testKey)
   156  	s.Nil(err)
   157  	fctx, err := client.regionCache.GetTiKVRPCContext(s.bo, loc.Region, kv.ReplicaReadFollower, 0)
   158  	s.Nil(err)
   159  	s.Equal(fctx.Store.StoreID(), s.store2)
   160  	s.Equal(fctx.Addr, "store2")
   161  
   162  	// make store2 using store1's addr and store1 offline
   163  	store1Addr := s.storeAddr(s.store1)
   164  	s.cluster.UpdateStoreAddr(s.store1, s.storeAddr(s.store2))
   165  	s.cluster.UpdateStoreAddr(s.store2, store1Addr)
   166  	s.cluster.RemoveStore(s.store1)
   167  	s.cluster.ChangeLeader(s.region1, s.peer2)
   168  	s.cluster.RemovePeer(s.region1, s.peer1)
   169  
   170  	getVal, err := client.Get(context.Background(), testKey)
   171  	s.Nil(err)
   172  	s.Equal(getVal, testValue)
   173  }
   174  
   175  func (s *testRawkvSuite) TestReplaceStore() {
   176  	mvccStore := mocktikv.MustNewMVCCStore()
   177  	defer mvccStore.Close()
   178  
   179  	client := &Client{
   180  		clusterID:   0,
   181  		regionCache: locate.NewRegionCache(mocktikv.NewPDClient(s.cluster)),
   182  		rpcClient:   mocktikv.NewRPCClient(s.cluster, mvccStore, nil),
   183  	}
   184  	defer client.Close()
   185  	testKey := []byte("test_key")
   186  	testValue := []byte("test_value")
   187  	err := client.Put(context.Background(), testKey, testValue)
   188  	s.Nil(err)
   189  
   190  	s.cluster.MarkTombstone(s.store1)
   191  	store3 := s.cluster.AllocID()
   192  	peer3 := s.cluster.AllocID()
   193  	s.cluster.AddStore(store3, s.storeAddr(s.store1))
   194  	s.cluster.AddPeer(s.region1, store3, peer3)
   195  	s.cluster.RemovePeer(s.region1, s.peer1)
   196  	s.cluster.ChangeLeader(s.region1, peer3)
   197  
   198  	err = client.Put(context.Background(), testKey, testValue)
   199  	s.Nil(err)
   200  }