client: fix wrong for check (tikv#7622)

close tikv#7613 fix wrong for loop check Signed-off-by: husharp <[email protected]>
HuSharp · Dec 27, 2023 · cee6e63 · cee6e63
1 parent ed9685a
commit cee6e63
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 2 deletions.
diff --git a/client/go.mod b/client/go.mod
@@ -13,6 +13,7 @@ require (
 	github.com/pingcap/kvproto v0.0.0-20231222062942-c0c73f41d0b2
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/prometheus/client_golang v1.11.1
+	github.com/prometheus/client_model v0.2.0
 	github.com/stretchr/testify v1.8.2
 	go.uber.org/atomic v1.10.0
 	go.uber.org/goleak v1.1.11
@@ -31,7 +32,6 @@ require (
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/common v0.26.0 // indirect
 	github.com/prometheus/procfs v0.6.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect

diff --git a/client/http/client.go b/client/http/client.go
@@ -161,7 +161,10 @@ func (ci *clientInner) requestWithRetry(
 			zap.String("source", ci.source), zap.Int("leader-idx", leaderAddrIdx), zap.String("addr", addr), zap.Error(err))
 	}
 	// Try to send the request to the other PD followers.
-	for idx := 0; idx < len(pdAddrs) && idx != leaderAddrIdx; idx++ {
+	for idx := 0; idx < len(pdAddrs); idx++ {
+		if idx == leaderAddrIdx {
+			continue
+		}
 		addr = ci.pdAddrs[idx]
 		err = ci.doRequest(ctx, addr, reqInfo, headerOpts...)
 		if err == nil {

diff --git a/client/http/client_test.go b/client/http/client_test.go
@@ -21,6 +21,9 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/pingcap/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 )
@@ -32,11 +35,13 @@ func TestPDAddrNormalization(t *testing.T) {
 	re.Len(pdAddrs, 1)
 	re.Equal(-1, leaderAddrIdx)
 	re.Contains(pdAddrs[0], httpScheme)
+	c.Close()
 	c = NewClient("test-https-pd-addr", []string{"127.0.0.1"}, WithTLSConfig(&tls.Config{}))
 	pdAddrs, leaderAddrIdx = c.(*client).inner.getPDAddrs()
 	re.Len(pdAddrs, 1)
 	re.Equal(-1, leaderAddrIdx)
 	re.Contains(pdAddrs[0], httpsScheme)
+	c.Close()
 }
 
 // requestChecker is used to check the HTTP request sent by the client.
@@ -93,3 +98,75 @@ func TestCallerID(t *testing.T) {
 	c.WithCallerID(expectedVal.Load()).GetRegions(context.Background())
 	c.Close()
 }
+
+func TestRedirectWithMetrics(t *testing.T) {
+	re := require.New(t)
+
+	pdAddrs := []string{"127.0.0.1", "172.0.0.1", "192.0.0.1"}
+	metricCnt := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "check",
+		}, []string{"name", ""})
+
+	// 1. Test all followers failed, need to send all followers.
+	httpClient := newHTTPClientWithRequestChecker(func(req *http.Request) error {
+		if req.URL.Path == Schedulers {
+			return errors.New("mock error")
+		}
+		return nil
+	})
+	c := NewClient("test-http-pd-redirect", pdAddrs, WithHTTPClient(httpClient), WithMetrics(metricCnt, nil))
+	pdAddrs, leaderAddrIdx := c.(*client).inner.getPDAddrs()
+	re.Equal(-1, leaderAddrIdx)
+	c.CreateScheduler(context.Background(), "test", 0)
+	var out dto.Metric
+	failureCnt, err := c.(*client).inner.requestCounter.GetMetricWithLabelValues([]string{createSchedulerName, networkErrorStatus}...)
+	re.NoError(err)
+	failureCnt.Write(&out)
+	re.Equal(float64(3), out.Counter.GetValue())
+
+	// 2. Test the Leader success, just need to send to leader.
+	httpClient = newHTTPClientWithRequestChecker(func(req *http.Request) error {
+		// mock leader success.
+		if !strings.Contains(pdAddrs[0], req.Host) {
+			return errors.New("mock error")
+		}
+		return nil
+	})
+	c.(*client).inner.cli = httpClient
+	// Force to update members info.
+	c.(*client).inner.leaderAddrIdx = 0
+	c.(*client).inner.pdAddrs = pdAddrs
+	c.CreateScheduler(context.Background(), "test", 0)
+	successCnt, err := c.(*client).inner.requestCounter.GetMetricWithLabelValues([]string{createSchedulerName, ""}...)
+	re.NoError(err)
+	successCnt.Write(&out)
+	re.Equal(float64(1), out.Counter.GetValue())
+
+	// 3. Test when the leader fails, needs to be sent to the follower in order,
+	// and returns directly if one follower succeeds
+	httpClient = newHTTPClientWithRequestChecker(func(req *http.Request) error {
+		// mock leader failure.
+		if strings.Contains(pdAddrs[0], req.Host) {
+			return errors.New("mock error")
+		}
+		return nil
+	})
+	c.(*client).inner.cli = httpClient
+	// Force to update members info.
+	c.(*client).inner.leaderAddrIdx = 0
+	c.(*client).inner.pdAddrs = pdAddrs
+	c.CreateScheduler(context.Background(), "test", 0)
+	successCnt, err = c.(*client).inner.requestCounter.GetMetricWithLabelValues([]string{createSchedulerName, ""}...)
+	re.NoError(err)
+	successCnt.Write(&out)
+	// only one follower success
+	re.Equal(float64(2), out.Counter.GetValue())
+	failureCnt, err = c.(*client).inner.requestCounter.GetMetricWithLabelValues([]string{createSchedulerName, networkErrorStatus}...)
+	re.NoError(err)
+	failureCnt.Write(&out)
+	// leader failure
+	re.Equal(float64(4), out.Counter.GetValue())
+
+	c.Close()
+}