Skip to content

Commit d29021d

Browse files
liubog2008liubog2008
authored
tests: add availability test for pd and tikv (#6454)
1 parent 936e413 commit d29021d

10 files changed

Lines changed: 385 additions & 53 deletions

File tree

‎cmd/testing-workload/main.go‎

Lines changed: 57 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,14 @@ var (
5454
tlsMountPath string
5555
tlsFromEnv bool
5656
insecureSkipVerify bool
57+
58+
// Flags for PD region API access
59+
pdEndpointsStr string
5760
)
5861

5962
//nolint:mnd,errcheck
6063
func main() {
61-
flag.StringVar(&action, "action", "ping", "ping, workload, import")
64+
flag.StringVar(&action, "action", "ping", "ping, workload, import, pd-region")
6265
flag.StringVar(&host, "host", "", "host")
6366
flag.StringVar(&port, "port", "4000", "port")
6467
flag.StringVar(&user, "user", "root", "db user")
@@ -85,53 +88,71 @@ func main() {
8588
flag.BoolVar(&tlsFromEnv, "tls-from-env", false, "load TLS certificates from environment variables")
8689
flag.BoolVar(&insecureSkipVerify, "tls-insecure-skip-verify", false, "skip TLS certificate verification")
8790

88-
flag.Parse()
91+
// Flags for PD region API access
92+
flag.StringVar(&pdEndpointsStr, "pd-endpoints", "", "comma-separated PD endpoints for pd-region action")
8993

90-
// enable "cleartext client side plugin" for `tidb_auth_token`.
91-
// ref: https://github.com/go-sql-driver/mysql?tab=readme-ov-file#allowcleartextpasswords
92-
params := []string{
93-
"charset=utf8mb4",
94-
"allowCleartextPasswords=true",
95-
"timeout=5s",
96-
}
94+
flag.Parse()
9795

98-
// Setup TLS if enabled
99-
if enableTLS {
100-
tlsConfigName, err := setupTLSConfig()
101-
if err != nil {
102-
panic(fmt.Errorf("failed to setup TLS config: %w", err))
96+
// Parse PD endpoints for pd-region action
97+
if action == "pd-region" && pdEndpointsStr != "" {
98+
pdEndpoints = strings.Split(pdEndpointsStr, ",")
99+
for i, endpoint := range pdEndpoints {
100+
pdEndpoints[i] = strings.TrimSpace(endpoint)
103101
}
104-
params = append(params, fmt.Sprintf("tls=%s", tlsConfigName))
105102
}
106103

107-
db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@(%s:%s)/test?%s", user, password, host, port, strings.Join(params, "&")))
108-
if err != nil {
109-
panic(err)
110-
}
111-
defer db.Close()
112-
113104
switch action {
114-
case "ping":
115-
if err := Ping(db); err != nil {
105+
case "pd-region":
106+
if err := PDRegionAccess(); err != nil {
116107
panic(err)
117108
}
118-
case "workload":
119-
if err := Workload(db); err != nil {
120-
panic(err)
109+
default:
110+
// enable "cleartext client side plugin" for `tidb_auth_token`.
111+
// ref: https://github.com/go-sql-driver/mysql?tab=readme-ov-file#allowcleartextpasswords
112+
params := []string{
113+
"charset=utf8mb4",
114+
"allowCleartextPasswords=true",
115+
"timeout=5s",
121116
}
122-
case "import":
123-
importCfg := ImportDataConfig{
124-
DB: db,
125-
BatchSize: batchSize,
126-
TotalRows: totalRows,
127-
TableName: importTable,
128-
SplitRegionCount: splitRegionCount,
117+
118+
// Setup TLS if enabled
119+
if enableTLS {
120+
tlsConfigName, err := setupTLSConfig()
121+
if err != nil {
122+
panic(fmt.Errorf("failed to setup TLS config: %w", err))
123+
}
124+
params = append(params, fmt.Sprintf("tls=%s", tlsConfigName))
129125
}
130-
if err := ImportData(importCfg); err != nil {
126+
127+
db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@(%s:%s)/test?%s", user, password, host, port, strings.Join(params, "&")))
128+
if err != nil {
131129
panic(err)
132130
}
133-
default:
134-
panic("unknown action: " + action)
131+
defer db.Close()
132+
133+
switch action {
134+
case "ping":
135+
if err := Ping(db); err != nil {
136+
panic(err)
137+
}
138+
case "workload":
139+
if err := Workload(db); err != nil {
140+
panic(err)
141+
}
142+
case "import":
143+
importCfg := ImportDataConfig{
144+
DB: db,
145+
BatchSize: batchSize,
146+
TotalRows: totalRows,
147+
TableName: importTable,
148+
SplitRegionCount: splitRegionCount,
149+
}
150+
if err := ImportData(importCfg); err != nil {
151+
panic(err)
152+
}
153+
default:
154+
panic("unknown action: " + action)
155+
}
135156
}
136157

137158
fmt.Println("workload is done")

‎cmd/testing-workload/pd_region.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
// Copyright 2024 PingCAP, Inc.
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+
package main
16+
17+
import (
18+
"context"
19+
"fmt"
20+
"strings"
21+
"sync"
22+
"sync/atomic"
23+
"time"
24+
25+
pd "github.com/tikv/pd/client"
26+
"github.com/tikv/pd/client/clients/router"
27+
"github.com/tikv/pd/client/opt"
28+
"github.com/tikv/pd/client/pkg/caller"
29+
"google.golang.org/grpc"
30+
"google.golang.org/grpc/keepalive"
31+
)
32+
33+
const (
34+
defaultGRPCKeepAliveTime = 10 * time.Second
35+
defaultGRPCKeepAliveTimeout = 3 * time.Second
36+
defaultPDServerTimeout = 3 * time.Second
37+
38+
defaultLeaderTransferTime = 100 * time.Millisecond
39+
)
40+
41+
func NewPDClient(pdAddrs []string) (pd.Client, error) {
42+
fmt.Println("set backoff with 3s")
43+
// init pd-client
44+
pdCli, err := pd.NewClient(
45+
caller.Component("tidb-operator"),
46+
pdAddrs,
47+
pd.SecurityOption{
48+
// CAPath: cfg.Security.ClusterSSLCA,
49+
// CertPath: cfg.Security.ClusterSSLCert,
50+
// KeyPath: cfg.Security.ClusterSSLKey,
51+
},
52+
opt.WithGRPCDialOptions(
53+
grpc.WithKeepaliveParams(
54+
keepalive.ClientParameters{
55+
Time: defaultGRPCKeepAliveTime,
56+
Timeout: defaultGRPCKeepAliveTimeout,
57+
},
58+
),
59+
),
60+
opt.WithCustomTimeoutOption(defaultPDServerTimeout),
61+
)
62+
if err != nil {
63+
return nil, err
64+
}
65+
return pdCli, nil
66+
}
67+
68+
var pdEndpoints []string
69+
70+
func PDRegionAccess() error {
71+
if len(pdEndpoints) == 0 {
72+
return fmt.Errorf("pd endpoints not provided")
73+
}
74+
75+
var totalCount, failCount atomic.Uint64
76+
var wg sync.WaitGroup
77+
78+
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(durationInMinutes)*time.Minute)
79+
defer cancel()
80+
client, err := NewPDClient(pdEndpoints)
81+
if err != nil {
82+
return fmt.Errorf("failed to create PD client: %w", err)
83+
}
84+
defer client.Close()
85+
86+
for i := 1; i <= maxConnections; i++ {
87+
wg.Add(1)
88+
go func(id int) {
89+
defer wg.Done()
90+
for {
91+
select {
92+
case <-ctx.Done():
93+
return
94+
default:
95+
err := accessPDRegionAPI(ctx, client)
96+
totalCount.Add(1)
97+
if err != nil {
98+
fmt.Printf("[%d-%s] failed to access PD region API: %v\n",
99+
id, time.Now().String(), err,
100+
)
101+
failCount.Add(1)
102+
}
103+
time.Sleep(time.Duration(sleepInterval) * time.Millisecond)
104+
}
105+
}
106+
}(i)
107+
}
108+
wg.Wait()
109+
fmt.Printf("total count: %d, fail count: %d\n", totalCount.Load(), failCount.Load())
110+
if failCount.Load() > 0 {
111+
return fmt.Errorf("there are failed PD region API accesses")
112+
}
113+
114+
return nil
115+
}
116+
117+
func accessPDRegionAPI(ctx context.Context, client pd.Client) error {
118+
// retry once
119+
var lastErr error
120+
for range 2 {
121+
if lastErr != nil {
122+
fmt.Println("retry because of ", lastErr)
123+
}
124+
_, err := client.BatchScanRegions(ctx, []router.KeyRange{
125+
{
126+
StartKey: []byte(""),
127+
EndKey: []byte(""),
128+
},
129+
}, 1)
130+
if err == nil {
131+
return nil
132+
}
133+
lastErr = err
134+
// only retry for not leader err
135+
if !strings.Contains(err.Error(), "not leader") {
136+
break
137+
}
138+
// wait 100 ms
139+
time.Sleep(defaultLeaderTransferTime)
140+
}
141+
if lastErr != nil {
142+
return fmt.Errorf("failed to scan regions from PD: %w", lastErr)
143+
}
144+
145+
return nil
146+
}

‎go.mod‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,8 @@ require (
5252
github.com/onsi/ginkgo/v2 v2.22.1
5353
github.com/onsi/gomega v1.36.2
5454
github.com/pelletier/go-toml/v2 v2.2.2
55-
github.com/pingcap/errors v0.11.4
56-
github.com/pingcap/kvproto v0.0.0-20240403065636-c699538f7aa1
55+
github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
56+
github.com/pingcap/kvproto v0.0.0-20250616075548-d951fb623bb3
5757
github.com/pingcap/tidb-operator/api/v2 v2.0.0-00000000000000-000000000000
5858
github.com/prometheus/client_golang v1.22.0
5959
github.com/prometheus/client_model v0.6.1
@@ -62,6 +62,7 @@ require (
6262
github.com/spf13/cobra v1.8.1
6363
github.com/spf13/pflag v1.0.5
6464
github.com/stretchr/testify v1.10.0
65+
github.com/tikv/pd/client v0.0.0-20250917050259-056e04c3e3a3
6566
go.etcd.io/etcd/client/v3 v3.5.21
6667
go.uber.org/mock v0.4.0
6768
go.uber.org/zap v1.27.0
@@ -70,6 +71,14 @@ require (
7071
gopkg.in/yaml.v3 v3.0.1
7172
)
7273

74+
require (
75+
github.com/opentracing/opentracing-go v1.2.0 // indirect
76+
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
77+
github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
78+
go.uber.org/atomic v1.10.0 // indirect
79+
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
80+
)
81+
7382
require (
7483
cel.dev/expr v0.19.1 // indirect
7584
cloud.google.com/go v0.116.0 // indirect
@@ -212,7 +221,7 @@ require (
212221
google.golang.org/genproto v0.0.0-20241118233622-e639e219e697 // indirect
213222
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 // indirect
214223
google.golang.org/genproto/googleapis/rpc v0.0.0-20241209162323-e6fa225c2576 // indirect
215-
google.golang.org/grpc v1.70.0 // indirect
224+
google.golang.org/grpc v1.70.0
216225
gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
217226
gopkg.in/inf.v0 v0.9.1 // indirect
218227
gopkg.in/ini.v1 v1.42.0 // indirect

0 commit comments

Comments
 (0)