mirror of https://github.com/tikv/client-go.git
216 lines
6.5 KiB
Go
216 lines
6.5 KiB
Go
// Copyright 2023 TiKV Authors
|
|
//
|
|
// 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.
|
|
|
|
// NOTE: The code in this file is based on code from the
|
|
// TiDB project, licensed under the Apache License v 2.0
|
|
//
|
|
// https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/tests/prewrite_test.go
|
|
//
|
|
|
|
// Copyright 2023 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 tikv_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/pingcap/failpoint"
|
|
"github.com/pingcap/kvproto/pkg/kvrpcpb"
|
|
"github.com/pingcap/kvproto/pkg/metapb"
|
|
"github.com/stretchr/testify/suite"
|
|
"github.com/tikv/client-go/v2/config"
|
|
"github.com/tikv/client-go/v2/oracle"
|
|
"github.com/tikv/client-go/v2/tikv"
|
|
"github.com/tikv/client-go/v2/tikvrpc"
|
|
"github.com/tikv/client-go/v2/util"
|
|
pd "github.com/tikv/pd/client"
|
|
)
|
|
|
|
func TestPDAPI(t *testing.T) {
|
|
if !*withTiKV {
|
|
t.Skip("skipping TestPDAPI because with-tikv is not enabled")
|
|
}
|
|
suite.Run(t, new(apiTestSuite))
|
|
}
|
|
|
|
type apiTestSuite struct {
|
|
suite.Suite
|
|
store *tikv.KVStore
|
|
}
|
|
|
|
func (s *apiTestSuite) SetupTest() {
|
|
addrs := strings.Split(*pdAddrs, ",")
|
|
pdClient, err := pd.NewClient(addrs, pd.SecurityOption{})
|
|
s.Require().NoError(err)
|
|
rpcClient := tikv.NewRPCClient()
|
|
// Set PD HTTP client.
|
|
store, err := tikv.NewTestTiKVStore(rpcClient, pdClient, nil, nil, 0, tikv.WithPDHTTPClient(nil, addrs))
|
|
s.store = store
|
|
storeID := uint64(1)
|
|
s.store.GetRegionCache().SetRegionCacheStore(storeID, s.storeAddr(storeID), s.storeAddr(storeID), tikvrpc.TiKV, 1, nil)
|
|
}
|
|
|
|
func (s *apiTestSuite) storeAddr(id uint64) string {
|
|
return fmt.Sprintf("store%d", id)
|
|
}
|
|
|
|
type storeSafeTsMockClient struct {
|
|
tikv.Client
|
|
requestCount int32
|
|
}
|
|
|
|
func (c *storeSafeTsMockClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
|
|
if req.Type != tikvrpc.CmdStoreSafeTS {
|
|
return c.Client.SendRequest(ctx, addr, req, timeout)
|
|
}
|
|
atomic.AddInt32(&c.requestCount, 1)
|
|
resp := &tikvrpc.Response{}
|
|
resp.Resp = &kvrpcpb.StoreSafeTSResponse{SafeTs: 150}
|
|
return resp, nil
|
|
}
|
|
|
|
func (c *storeSafeTsMockClient) Close() error {
|
|
return c.Client.Close()
|
|
}
|
|
|
|
func (c *storeSafeTsMockClient) CloseAddr(addr string) error {
|
|
return c.Client.CloseAddr(addr)
|
|
}
|
|
|
|
func (s *apiTestSuite) TestGetStoresMinResolvedTS() {
|
|
util.EnableFailpoints()
|
|
require := s.Require()
|
|
require.NoError(failpoint.Enable("tikvclient/mockFastSafeTSUpdater", `return()`))
|
|
defer func() {
|
|
require.NoError(failpoint.Disable("tikvclient/mockFastSafeTSUpdater"))
|
|
}()
|
|
|
|
// Set DC label for store 1.
|
|
// Mock Cluster-level min resolved ts failed.
|
|
dcLabel := "testDC"
|
|
restore := config.UpdateGlobal(func(conf *config.Config) {
|
|
conf.TxnScope = dcLabel
|
|
})
|
|
defer restore()
|
|
|
|
labels := []*metapb.StoreLabel{
|
|
{
|
|
Key: tikv.DCLabelKey,
|
|
Value: dcLabel,
|
|
},
|
|
}
|
|
storeID := uint64(1)
|
|
s.store.GetRegionCache().SetRegionCacheStore(storeID, s.storeAddr(storeID), s.storeAddr(storeID), tikvrpc.TiKV, 1, labels)
|
|
// Try to get the minimum resolved timestamp of the stores from PD.
|
|
require.NoError(failpoint.Enable("tikvclient/InjectMinResolvedTS", `return(100)`))
|
|
mockClient := storeSafeTsMockClient{
|
|
Client: s.store.GetTiKVClient(),
|
|
}
|
|
s.store.SetTiKVClient(&mockClient)
|
|
var retryCount int
|
|
for s.store.GetMinSafeTS(dcLabel) != 100 {
|
|
time.Sleep(100 * time.Millisecond)
|
|
if retryCount > 5 {
|
|
break
|
|
}
|
|
retryCount++
|
|
}
|
|
require.Equal(int32(0), atomic.LoadInt32(&mockClient.requestCount))
|
|
require.Equal(uint64(100), s.store.GetMinSafeTS(dcLabel))
|
|
require.Nil(failpoint.Disable("tikvclient/InjectMinResolvedTS"))
|
|
}
|
|
|
|
func (s *apiTestSuite) TestDCLabelClusterMinResolvedTS() {
|
|
util.EnableFailpoints()
|
|
// Try to get the minimum resolved timestamp of the cluster from PD.
|
|
require := s.Require()
|
|
require.NoError(failpoint.Enable("tikvclient/mockFastSafeTSUpdater", `return()`))
|
|
defer func() {
|
|
require.NoError(failpoint.Disable("tikvclient/mockFastSafeTSUpdater"))
|
|
}()
|
|
require.NoError(failpoint.Enable("tikvclient/InjectMinResolvedTS", `return(100)`))
|
|
mockClient := storeSafeTsMockClient{
|
|
Client: s.store.GetTiKVClient(),
|
|
}
|
|
s.store.SetTiKVClient(&mockClient)
|
|
var retryCount int
|
|
for s.store.GetMinSafeTS(oracle.GlobalTxnScope) != 100 {
|
|
time.Sleep(100 * time.Millisecond)
|
|
if retryCount > 5 {
|
|
break
|
|
}
|
|
retryCount++
|
|
}
|
|
require.Equal(atomic.LoadInt32(&mockClient.requestCount), int32(0))
|
|
require.Equal(uint64(100), s.store.GetMinSafeTS(oracle.GlobalTxnScope))
|
|
require.NoError(failpoint.Disable("tikvclient/InjectMinResolvedTS"))
|
|
|
|
// Set DC label for store 1.
|
|
// Mock PD server not support get min resolved ts by stores.
|
|
require.NoError(failpoint.Enable("tikvclient/InjectMinResolvedTS", `return(0)`))
|
|
defer func() {
|
|
require.NoError(failpoint.Disable("tikvclient/InjectMinResolvedTS"))
|
|
}()
|
|
dcLabel := "testDC"
|
|
restore := config.UpdateGlobal(func(conf *config.Config) {
|
|
conf.TxnScope = dcLabel
|
|
})
|
|
defer restore()
|
|
|
|
labels := []*metapb.StoreLabel{
|
|
{
|
|
Key: tikv.DCLabelKey,
|
|
Value: dcLabel,
|
|
},
|
|
}
|
|
storeID := uint64(1)
|
|
s.store.GetRegionCache().SetRegionCacheStore(storeID, s.storeAddr(storeID), s.storeAddr(storeID), tikvrpc.TiKV, 1, labels)
|
|
|
|
// Try to get the minimum resolved timestamp of the store from TiKV.
|
|
retryCount = 0
|
|
for s.store.GetMinSafeTS(dcLabel) != 150 {
|
|
time.Sleep(100 * time.Millisecond)
|
|
if retryCount > 5 {
|
|
break
|
|
}
|
|
retryCount++
|
|
}
|
|
|
|
require.GreaterOrEqual(atomic.LoadInt32(&mockClient.requestCount), int32(1))
|
|
require.Equal(uint64(150), s.store.GetMinSafeTS(dcLabel))
|
|
}
|
|
|
|
func (s *apiTestSuite) TearDownTest() {
|
|
if s.store != nil {
|
|
s.Require().Nil(s.store.Close())
|
|
}
|
|
}
|