From 093b38e2fbaa3a13e6d14df04e162c7431c4529c Mon Sep 17 00:00:00 2001
From: Leavrth <jianjun.liao@outlook.com>
Date: Mon, 24 Apr 2023 11:34:38 +0800
Subject: [PATCH] add unit test

Signed-off-by: Leavrth <jianjun.liao@outlook.com>
---
 br/pkg/storage/BUILD.bazel |  2 ++
 br/pkg/storage/s3.go       |  8 +++++++
 br/pkg/storage/s3_test.go  | 49 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/br/pkg/storage/BUILD.bazel b/br/pkg/storage/BUILD.bazel
index a21eb4cb03e9a..e8b9522015676 100644
--- a/br/pkg/storage/BUILD.bazel
+++ b/br/pkg/storage/BUILD.bazel
@@ -42,6 +42,7 @@ go_library(
         "@com_github_google_uuid//:uuid",
         "@com_github_klauspost_compress//zstd",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_pingcap_log//:log",
         "@com_github_spf13_pflag//:pflag",
@@ -80,6 +81,7 @@ go_test(
         "@com_github_fsouza_fake_gcs_server//fakestorage",
         "@com_github_golang_mock//gomock",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_stretchr_testify//require",
     ],
diff --git a/br/pkg/storage/s3.go b/br/pkg/storage/s3.go
index 460266236fedf..d6d29706825bc 100644
--- a/br/pkg/storage/s3.go
+++ b/br/pkg/storage/s3.go
@@ -27,6 +27,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/s3/s3iface"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/log"
 	berrors "github.com/pingcap/tidb/br/pkg/errors"
@@ -961,6 +962,13 @@ func isConnectionResetError(err error) bool {
 }
 
 func (rl retryerWithLog) ShouldRetry(r *request.Request) bool {
+	// for unit test
+	failpoint.Inject("replace-error-to-connection-reset-by-peer", func(_ failpoint.Value) {
+		log.Info("original error", zap.Error(r.Error))
+		if r.Error != nil {
+			r.Error = errors.New("read tcp *.*.*.*:*->*.*.*.*:*: read: connection reset by peer")
+		}
+	})
 	if isConnectionResetError(r.Error) {
 		return true
 	}
diff --git a/br/pkg/storage/s3_test.go b/br/pkg/storage/s3_test.go
index 0541e66b8617f..265733f61dfdb 100644
--- a/br/pkg/storage/s3_test.go
+++ b/br/pkg/storage/s3_test.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -20,6 +21,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/golang/mock/gomock"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/tidb/br/pkg/mock"
 	. "github.com/pingcap/tidb/br/pkg/storage"
@@ -1292,3 +1294,50 @@ func TestS3StorageBucketRegion(t *testing.T) {
 		}(ca.name, ca.expectRegion, ca.s3)
 	}
 }
+
+func TestRetryError(t *testing.T) {
+	var count int32 = 0
+	var errString string = "read tcp *.*.*.*:*->*.*.*.*:*: read: connection reset by peer"
+	var lock sync.Mutex
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == "PUT" {
+			var curCnt int32
+			t.Log(r.URL)
+			lock.Lock()
+			count += 1
+			curCnt = count
+			lock.Unlock()
+			if curCnt < 2 {
+				// write an cannot-retry error, but we modify the error to specific error, so client would retry.
+				w.WriteHeader(403)
+				return
+			}
+		}
+
+		w.WriteHeader(200)
+
+	}))
+
+	defer server.Close()
+	t.Log(server.URL)
+
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/storage/replace-error-to-connection-reset-by-peer", "return(true)"))
+	defer func() {
+		failpoint.Disable("github.com/pingcap/tidb/br/pkg/storage/replace-error-to-connection-reset-by-peer")
+	}()
+
+	ctx := context.Background()
+	s, err := NewS3Storage(ctx, &backuppb.S3{
+		Endpoint:        server.URL,
+		Bucket:          "test",
+		Prefix:          "retry",
+		AccessKey:       "none",
+		SecretAccessKey: "none",
+		Provider:        "skip check region",
+		ForcePathStyle:  true,
+	}, &ExternalStorageOptions{})
+	require.NoError(t, err)
+	err = s.WriteFile(ctx, "reset", []byte(errString))
+	require.NoError(t, err)
+	require.Equal(t, count, int32(2))
+}