From dacfb020193e53ea49cf1fef9767813690f796c2 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Fri, 16 Dec 2022 22:20:53 +0800
Subject: [PATCH] fix: ebs volume restore failure since incorrect unfinished
 volumes (#39975) (#40002)

close pingcap/tidb#39974
---
 br/pkg/aws/BUILD.bazel |  6 ++++-
 br/pkg/aws/ebs.go      | 31 +++++++++++++++---------
 br/pkg/aws/ebs_test.go | 53 +++++++++++++++++++++++++++++++++++-------
 3 files changed, 70 insertions(+), 20 deletions(-)

diff --git a/br/pkg/aws/BUILD.bazel b/br/pkg/aws/BUILD.bazel
index 6f33637abc5f3..28f58d2a1e1ad 100644
--- a/br/pkg/aws/BUILD.bazel
+++ b/br/pkg/aws/BUILD.bazel
@@ -25,5 +25,9 @@ go_test(
     name = "aws_test",
     srcs = ["ebs_test.go"],
     embed = [":aws"],
-    deps = ["@com_github_stretchr_testify//require"],
+    deps = [
+        "@com_github_aws_aws_sdk_go//aws",
+        "@com_github_aws_aws_sdk_go//service/ec2",
+        "@com_github_stretchr_testify//require",
+    ],
 )
diff --git a/br/pkg/aws/ebs.go b/br/pkg/aws/ebs.go
index 9ded291e1f9b9..fb96b95578ffb 100644
--- a/br/pkg/aws/ebs.go
+++ b/br/pkg/aws/ebs.go
@@ -307,17 +307,9 @@ func (e *EC2Session) WaitVolumesCreated(volumeIDMap map[string]string, progress
 			return 0, errors.Trace(err)
 		}
 
-		var unfinishedVolumes []*string
-		for _, volume := range resp.Volumes {
-			if *volume.State == ec2.VolumeStateAvailable {
-				log.Info("volume is available", zap.String("id", *volume.SnapshotId))
-				totalVolumeSize += *volume.Size
-				progress.Inc()
-			} else {
-				log.Debug("volume creating...", zap.Stringer("volume", volume))
-				unfinishedVolumes = append(unfinishedVolumes, volume.SnapshotId)
-			}
-		}
+		createdVolumeSize, unfinishedVolumes := e.HandleDescribeVolumesResponse(resp)
+		progress.IncBy(int64(len(pendingVolumes) - len(unfinishedVolumes)))
+		totalVolumeSize += createdVolumeSize
 		pendingVolumes = unfinishedVolumes
 	}
 	log.Info("all pending volume are created.")
@@ -357,3 +349,20 @@ func (e *EC2Session) DeleteVolumes(volumeIDMap map[string]string) {
 func ec2Tag(key, val string) *ec2.Tag {
 	return &ec2.Tag{Key: &key, Value: &val}
 }
+
+func (e *EC2Session) HandleDescribeVolumesResponse(resp *ec2.DescribeVolumesOutput) (int64, []*string) {
+	totalVolumeSize := int64(0)
+
+	var unfinishedVolumes []*string
+	for _, volume := range resp.Volumes {
+		if *volume.State == ec2.VolumeStateAvailable {
+			log.Info("volume is available", zap.String("id", *volume.VolumeId))
+			totalVolumeSize += *volume.Size
+		} else {
+			log.Debug("volume creating...", zap.Stringer("volume", volume))
+			unfinishedVolumes = append(unfinishedVolumes, volume.VolumeId)
+		}
+	}
+
+	return totalVolumeSize, unfinishedVolumes
+}
diff --git a/br/pkg/aws/ebs_test.go b/br/pkg/aws/ebs_test.go
index 695f31a73c477..d7f3be2a4a4a1 100644
--- a/br/pkg/aws/ebs_test.go
+++ b/br/pkg/aws/ebs_test.go
@@ -16,26 +16,63 @@ package aws
 import (
 	"testing"
 
+	awsapi "github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/stretchr/testify/require"
 )
 
 func TestEC2SessionExtractSnapProgress(t *testing.T) {
-	strPtr := func(s string) *string {
-		return &s
-	}
 	tests := []struct {
 		str  *string
 		want int64
 	}{
 		{nil, 0},
-		{strPtr("12.12%"), 12},
-		{strPtr("44.99%"), 44},
-		{strPtr("  89.89%  "), 89},
-		{strPtr("100%"), 100},
-		{strPtr("111111%"), 100},
+		{awsapi.String("12.12%"), 12},
+		{awsapi.String("44.99%"), 44},
+		{awsapi.String("  89.89%  "), 89},
+		{awsapi.String("100%"), 100},
+		{awsapi.String("111111%"), 100},
 	}
 	e := &EC2Session{}
 	for _, tt := range tests {
 		require.Equal(t, tt.want, e.extractSnapProgress(tt.str))
 	}
 }
+
+func createVolume(snapshotId string, volumeId string, state string) *ec2.Volume {
+	return &ec2.Volume{
+		Attachments:        nil,
+		AvailabilityZone:   awsapi.String("us-west-2"),
+		CreateTime:         nil,
+		Encrypted:          awsapi.Bool(true),
+		FastRestored:       awsapi.Bool(true),
+		Iops:               awsapi.Int64(3000),
+		KmsKeyId:           nil,
+		MultiAttachEnabled: awsapi.Bool(true),
+		OutpostArn:         awsapi.String("arn:12342"),
+		Size:               awsapi.Int64(1),
+		SnapshotId:         awsapi.String(snapshotId),
+		State:              awsapi.String(state),
+		Tags:               nil,
+		Throughput:         nil,
+		VolumeId:           awsapi.String(volumeId),
+		VolumeType:         awsapi.String("gp3"),
+	}
+}
+func TestHandleDescribeVolumesResponse(t *testing.T) {
+	curentVolumesStates := &ec2.DescribeVolumesOutput{
+		NextToken: awsapi.String("fake token"),
+		Volumes: []*ec2.Volume{
+			createVolume("snap-0873674883", "vol-98768979", "available"),
+			createVolume("snap-0873674883", "vol-98768979", "creating"),
+			createVolume("snap-0873674883", "vol-98768979", "available"),
+			createVolume("snap-0873674883", "vol-98768979", "available"),
+			createVolume("snap-0873674883", "vol-98768979", "available"),
+		},
+	}
+
+	e := &EC2Session{}
+	createdVolumeSize, unfinishedVolumes := e.HandleDescribeVolumesResponse(curentVolumesStates)
+	require.Equal(t, int64(4), createdVolumeSize)
+	require.Equal(t, 1, len(unfinishedVolumes))
+}