Merge branch 'master' into outer2inner

.github/workflows/integration-test-compile-br.yml

@@ -51,10 +51,10 @@ jobs:
         os: [macos-latest,  ubuntu-latest, windows-latest]
 
     steps:
-      - uses: actions/checkout@v2.1.0
+      - uses: actions/checkout@v4
 
       - name: Set up Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v5
         with:
           go-version-file: 'go.mod'
 
@@ -66,10 +66,10 @@ jobs:
     name: Compile for FreeBSD job
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2.1.0
+      - uses: actions/checkout@v4
 
       - name: Set up Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v5
         with:
           go-version-file: 'go.mod'
 

.github/workflows/integration-test-dumpling.yml

@@ -52,6 +52,7 @@ jobs:
           - 5.7.35
           - 8.0.22
           - 8.0.26
+          - 8.0.37
     runs-on: ubuntu-latest
     timeout-minutes: 15
     services:
@@ -64,13 +65,13 @@ jobs:
         options: --health-cmd="mysqladmin ping" --health-interval=10s --health-timeout=5s --health-retries=3
     steps:
       - name: "checkout repository"
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: "set up golang"
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v5
         with:
           go-version-file: 'go.mod'
       - name: "try to use build cache"
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: |
             ~/.cache/go-build

WORKSPACE

@@ -64,7 +64,7 @@ go_download_sdk(
         "https://mirrors.aliyun.com/golang/{}",
         "https://dl.google.com/go/{}",
     ],
-    version = "1.21.9",
+    version = "1.21.10",
 )
 
 go_register_toolchains(

br/cmd/br/BUILD.bazel

@@ -21,8 +21,9 @@ go_library(
         "//br/pkg/logutil",
         "//br/pkg/metautil",
         "//br/pkg/mock/mockid",
-        "//br/pkg/restore",
+        "//br/pkg/restore/utils",
         "//br/pkg/rtree",
+        "//br/pkg/stream",
         "//br/pkg/streamhelper/config",
         "//br/pkg/summary",
         "//br/pkg/task",
@@ -62,4 +63,5 @@ go_test(
     srcs = ["main_test.go"],
     embed = [":br_lib"],
     flaky = True,
+    deps = ["@org_uber_go_goleak//:goleak"],
 )

br/cmd/br/debug.go

@@ -21,8 +21,9 @@ import (
 	"github.com/pingcap/tidb/br/pkg/logutil"
 	"github.com/pingcap/tidb/br/pkg/metautil"
 	"github.com/pingcap/tidb/br/pkg/mock/mockid"
-	"github.com/pingcap/tidb/br/pkg/restore"
+	restoreutils "github.com/pingcap/tidb/br/pkg/restore/utils"
 	"github.com/pingcap/tidb/br/pkg/rtree"
+	"github.com/pingcap/tidb/br/pkg/stream"
 	"github.com/pingcap/tidb/br/pkg/task"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/pingcap/tidb/br/pkg/version/build"
@@ -207,7 +208,7 @@ func newBackupMetaValidateCommand() *cobra.Command {
 			for offset := uint64(0); offset < tableIDOffset; offset++ {
 				_, _ = tableIDAllocator.Alloc() // Ignore error
 			}
-			rewriteRules := &restore.RewriteRules{
+			rewriteRules := &restoreutils.RewriteRules{
 				Data: make([]*import_sstpb.RewriteRule, 0),
 			}
 			tableIDMap := make(map[int64]int64)
@@ -245,13 +246,13 @@ func newBackupMetaValidateCommand() *cobra.Command {
 					}
 				}
 
-				rules := restore.GetRewriteRules(newTable, table.Info, 0, true)
+				rules := restoreutils.GetRewriteRules(newTable, table.Info, 0, true)
 				rewriteRules.Data = append(rewriteRules.Data, rules.Data...)
 				tableIDMap[table.Info.ID] = int64(tableID)
 			}
 			// Validate rewrite rules
 			for _, file := range files {
-				err = restore.ValidateFileRewriteRule(file, rewriteRules)
+				err = restoreutils.ValidateFileRewriteRule(file, rewriteRules)
 				if err != nil {
 					return errors.Trace(err)
 				}
@@ -447,8 +448,8 @@ func searchStreamBackupCommand() *cobra.Command {
 			if err != nil {
 				return errors.Trace(err)
 			}
-			comparator := restore.NewStartWithComparator()
-			bs := restore.NewStreamBackupSearch(s, comparator, keyBytes)
+			comparator := stream.NewStartWithComparator()
+			bs := stream.NewStreamBackupSearch(s, comparator, keyBytes)
 			bs.SetStartTS(startTs)
 			bs.SetEndTs(endTs)
 

br/cmd/br/main_test.go

@@ -20,8 +20,40 @@ import (
 	"os"
 	"strings"
 	"testing"
+
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	skipLeakTest := false
+	var newArgs []string
+	for _, arg := range os.Args {
+		if arg == "--skip-goleak" {
+			skipLeakTest = true
+		} else {
+			newArgs = append(newArgs, arg)
+		}
+	}
+	os.Args = newArgs
+
+	if !skipLeakTest {
+		goleak.VerifyTestMain(m,
+			goleak.IgnoreCurrent(),
+			goleak.IgnoreTopFunction("github.com/pingcap/tidb/br/pkg/utils.StartExitSingleListener.func1"),
+			goleak.IgnoreTopFunction("github.com/pingcap/tidb/br/pkg/utils.StartDynamicPProfListener.func1"),
+			goleak.IgnoreTopFunction("gopkg.in/natefinch/lumberjack%2ev2.(*Logger).millRun"),
+			goleak.IgnoreTopFunction("go.etcd.io/etcd/client/v3.waitRetryBackoff"),
+			goleak.IgnoreTopFunction("net/http.(*persistConn).writeLoop"),
+			goleak.IgnoreTopFunction("internal/poll.runtime_pollWait"),
+			goleak.IgnoreTopFunction("google.golang.org/grpc.(*ClientConn).WaitForStateChange"),
+			goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/config/retry.(*Config).createBackoffFn.newBackoffFn.func2"),
+			goleak.IgnoreTopFunction("google.golang.org/grpc/internal/grpcsync.NewCallbackSerializer"),
+		)
+	}
+
+	os.Exit(m.Run())
+}
+
 func TestRunMain(*testing.T) {
 	var args []string
 	for _, arg := range os.Args {

br/pkg/aws/BUILD.bazel

@@ -29,9 +29,11 @@ go_test(
     srcs = ["ebs_test.go"],
     embed = [":aws"],
     flaky = True,
+    shard_count = 3,
     deps = [
         "@com_github_aws_aws_sdk_go//aws",
         "@com_github_aws_aws_sdk_go//service/ec2",
+        "@com_github_aws_aws_sdk_go//service/ec2/ec2iface",
         "@com_github_stretchr_testify//require",
     ],
 )

br/pkg/aws/ebs.go

@@ -251,6 +251,9 @@ func (e *EC2Session) WaitSnapshotsCreated(snapIDMap map[string]string, progress
 			if *s.State == ec2.SnapshotStateCompleted {
 				log.Info("snapshot completed", zap.String("id", *s.SnapshotId))
 				totalVolumeSize += *s.VolumeSize
+			} else if *s.State == ec2.SnapshotStateError {
+				log.Error("snapshot failed", zap.String("id", *s.SnapshotId), zap.String("error", (*s.StateMessage)))
+				return 0, errors.Errorf("snapshot %s failed", *s.SnapshotId)
 			} else {
 				log.Debug("snapshot creating...", zap.Stringer("snap", s))
 				uncompletedSnapshots = append(uncompletedSnapshots, s.SnapshotId)

br/pkg/aws/ebs_test.go

@@ -14,10 +14,12 @@
 package aws
 
 import (
+	"context"
 	"testing"
 
 	awsapi "github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/stretchr/testify/require"
 )
 
@@ -76,3 +78,127 @@ func TestHandleDescribeVolumesResponse(t *testing.T) {
 	require.Equal(t, int64(4), createdVolumeSize)
 	require.Equal(t, 1, len(unfinishedVolumes))
 }
+
+type mockEC2 struct {
+	ec2iface.EC2API
+	output ec2.DescribeSnapshotsOutput
+}
+
+func (m mockEC2) DescribeSnapshots(*ec2.DescribeSnapshotsInput) (*ec2.DescribeSnapshotsOutput, error) {
+	return &m.output, nil
+}
+
+func NewMockEc2Session(mock mockEC2) *EC2Session {
+	return &EC2Session{
+		ec2: mock,
+	}
+}
+
+func TestWaitSnapshotsCreated(t *testing.T) {
+	snapIdMap := map[string]string{
+		"vol-1": "snap-1",
+		"vol-2": "snap-2",
+	}
+
+	cases := []struct {
+		desc            string
+		snapshotsOutput ec2.DescribeSnapshotsOutput
+		expectedSize    int64
+		expectErr       bool
+		expectTimeout   bool
+	}{
+		{
+			desc: "snapshots are all completed",
+			snapshotsOutput: ec2.DescribeSnapshotsOutput{
+				Snapshots: []*ec2.Snapshot{
+					{
+						SnapshotId: awsapi.String("snap-1"),
+						VolumeSize: awsapi.Int64(1),
+						State:      awsapi.String(ec2.SnapshotStateCompleted),
+					},
+					{
+						SnapshotId: awsapi.String("snap-2"),
+						VolumeSize: awsapi.Int64(2),
+						State:      awsapi.String(ec2.SnapshotStateCompleted),
+					},
+				},
+			},
+			expectedSize: 3,
+			expectErr:    false,
+		},
+		{
+			desc: "snapshot failed",
+			snapshotsOutput: ec2.DescribeSnapshotsOutput{
+				Snapshots: []*ec2.Snapshot{
+					{
+						SnapshotId: awsapi.String("snap-1"),
+						VolumeSize: awsapi.Int64(1),
+						State:      awsapi.String(ec2.SnapshotStateCompleted),
+					},
+					{
+						SnapshotId:   awsapi.String("snap-2"),
+						State:        awsapi.String(ec2.SnapshotStateError),
+						StateMessage: awsapi.String("snapshot failed"),
+					},
+				},
+			},
+			expectedSize: 0,
+			expectErr:    true,
+		},
+		{
+			desc: "snapshots pending",
+			snapshotsOutput: ec2.DescribeSnapshotsOutput{
+				Snapshots: []*ec2.Snapshot{
+					{
+						SnapshotId: awsapi.String("snap-1"),
+						VolumeSize: awsapi.Int64(1),
+						State:      awsapi.String(ec2.SnapshotStateCompleted),
+					},
+					{
+						SnapshotId: awsapi.String("snap-2"),
+						State:      awsapi.String(ec2.SnapshotStatePending),
+					},
+				},
+			},
+			expectTimeout: true,
+		},
+	}
+
+	for _, c := range cases {
+		e := NewMockEc2Session(mockEC2{
+			output: c.snapshotsOutput,
+		})
+
+		if c.expectTimeout {
+			func() {
+				// We wait 5s before checking snapshots
+				ctx, cancel := context.WithTimeout(context.Background(), 6)
+				defer cancel()
+
+				done := make(chan struct{})
+				go func() {
+					_, _ = e.WaitSnapshotsCreated(snapIdMap, nil)
+					done <- struct{}{}
+				}()
+
+				select {
+				case <-done:
+					t.Fatal("WaitSnapshotsCreated should not return before timeout")
+				case <-ctx.Done():
+					require.True(t, true)
+				}
+			}()
+
+			continue
+		}
+
+		size, err := e.WaitSnapshotsCreated(snapIdMap, nil)
+		if c.expectErr {
+			require.Error(t, err)
+		} else {
+			require.NoError(t, err)
+		}
+
+		require.Equal(t, c.expectedSize, size)
+	}
+}