add retry for put object

jojoliang · jojoliang · commit b037b1d3efc0 · 2023-12-20T17:00:35.000+08:00
diff --git a/ci_test.go b/ci_test.go
@@ -612,6 +612,7 @@ func TestCIService_Put(t *testing.T) {
 		&ObjectPutHeaderOptions{
 			XOptionHeader: &http.Header{},
 		},
+		nil,
 	}
 	opt.XOptionHeader.Add("Pic-Operations", EncodePicOperations(pic))
 	res, _, err := client.CI.Put(context.Background(), name, f, opt)
@@ -734,6 +735,7 @@ func TestCIService_PutFromFile(t *testing.T) {
 		&ObjectPutHeaderOptions{
 			XOptionHeader: &http.Header{},
 		},
+		nil,
 	}
 	opt.XOptionHeader.Add("Pic-Operations", EncodePicOperations(pic))
 	res, _, err := client.CI.PutFromFile(context.Background(), name, filePath, opt)
diff --git a/cos.go b/cos.go
@@ -373,6 +373,24 @@ func toSwitchHost(oldURL *url.URL) *url.URL {
 	return newURL
 }
 
+func (c *Client) CheckRetrieable(u *url.URL, resp *Response, err error) (*url.URL, bool) {
+	res := u
+	if err != nil && err != invalidBucketErr {
+		// 不重试
+		if resp != nil && resp.StatusCode < 500 {
+			return res, false
+		}
+		if c.Conf.RetryOpt.AutoSwitchHost {
+			// 收不到报文 或者 不存在RequestId
+			if resp == nil || resp.Header.Get("X-Cos-Request-Id") == "" {
+				res = toSwitchHost(u)
+			}
+		}
+		return res, true
+	}
+	return res, false
+}
+
 func (c *Client) doRetry(ctx context.Context, opt *sendOptions) (resp *Response, err error) {
 	if opt.body != nil {
 		if _, ok := opt.body.(io.Reader); ok {
@@ -384,29 +402,21 @@ func (c *Client) doRetry(ctx context.Context, opt *sendOptions) (resp *Response,
 	if c.Conf.RetryOpt.Count > 0 {
 		count = c.Conf.RetryOpt.Count
 	}
-	interval := c.Conf.RetryOpt.Interval
+	var retrieable bool
 	for nr := 0; nr < count; nr++ {
 		resp, err = c.send(ctx, opt)
-		if err != nil && err != invalidBucketErr {
-			// 不重试
-			if resp != nil && resp.StatusCode < 500 {
-				break
-			}
-			if c.Conf.RetryOpt.AutoSwitchHost {
-				// 收不到报文 或者 不存在RequestId
-				if resp == nil || resp.Header.Get("X-Cos-Request-Id") == "" {
-					opt.baseURL = toSwitchHost(opt.baseURL)
-				}
-			}
-			if interval > 0 && nr+1 < count {
-				time.Sleep(interval)
+		opt.baseURL, retrieable = c.CheckRetrieable(opt.baseURL, resp, err)
+		if retrieable {
+			if c.Conf.RetryOpt.Interval > 0 && nr+1 < count {
+				time.Sleep(c.Conf.RetryOpt.Interval)
 			}
 			continue
 		}
 		break
 	}
 	return
 }
+
 func (c *Client) send(ctx context.Context, opt *sendOptions) (resp *Response, err error) {
 	req, err := c.newRequest(ctx, opt.baseURL, opt.uri, opt.method, opt.body, opt.optQuery, opt.optHeader)
 	if err != nil {
diff --git a/helper.go b/helper.go
@@ -233,6 +233,7 @@ func CloneObjectPutOptions(opt *ObjectPutOptions) *ObjectPutOptions {
 	res := &ObjectPutOptions{
 		&ACLHeaderOptions{},
 		&ObjectPutHeaderOptions{},
+		nil,
 	}
 	if opt != nil {
 		if opt.ACLHeaderOptions != nil {
@@ -243,6 +244,9 @@ func CloneObjectPutOptions(opt *ObjectPutOptions) *ObjectPutOptions {
 			res.XCosMetaXXX = cloneHeader(opt.XCosMetaXXX)
 			res.XOptionHeader = cloneHeader(opt.XOptionHeader)
 		}
+		if opt.innerSwitchURL != nil {
+			res.innerSwitchURL = opt.innerSwitchURL
+		}
 	}
 	return res
 }
diff --git a/object.go b/object.go
@@ -332,6 +332,9 @@ type ObjectPutHeaderOptions struct {
 type ObjectPutOptions struct {
 	*ACLHeaderOptions       `header:",omitempty" url:"-" xml:"-"`
 	*ObjectPutHeaderOptions `header:",omitempty" url:"-" xml:"-"`
+
+	// PutFromFile 使用
+	innerSwitchURL *url.URL `header:"-" url:"-" xml:"-"`
 }
 
 // Put Object请求可以将一个文件（Oject）上传至指定Bucket。
@@ -358,27 +361,58 @@ func (s *ObjectService) Put(ctx context.Context, name string, r io.Reader, uopt
 			opt.ContentLength = totalBytes
 		}
 	}
-	reader := TeeReader(r, nil, totalBytes, nil)
-	if s.client.Conf.EnableCRC {
-		reader.writer = crc64.New(crc64.MakeTable(crc64.ECMA))
-	}
-	if opt != nil && opt.Listener != nil {
-		reader.listener = opt.Listener
-	}
-	sendOpt := sendOptions{
-		baseURL:   s.client.BaseURL.BucketURL,
-		uri:       "/" + encodeURIComponent(name),
-		method:    http.MethodPut,
-		body:      reader,
-		optHeader: opt,
+	// 如果是io.Seeker，则重试
+	count := 1
+	var position int64
+	if seeker, ok := r.(io.Seeker); ok {
+		// 记录原始位置
+		position, err = seeker.Seek(0, io.SeekCurrent)
+		if err == nil && s.client.Conf.RetryOpt.Count > 0 {
+			count = s.client.Conf.RetryOpt.Count
+		}
+	}
+	var resp *Response
+	var retrieable bool
+	sUrl := s.client.BaseURL.BucketURL
+	if opt.innerSwitchURL != nil {
+		sUrl = opt.innerSwitchURL
+	}
+	for nr := 0; nr < count; nr++ {
+		reader := TeeReader(r, nil, totalBytes, nil)
+		if s.client.Conf.EnableCRC {
+			reader.writer = crc64.New(crc64.MakeTable(crc64.ECMA))
+		}
+		if opt != nil && opt.Listener != nil {
+			reader.listener = opt.Listener
+		}
+		sendOpt := sendOptions{
+			baseURL:   sUrl,
+			uri:       "/" + encodeURIComponent(name),
+			method:    http.MethodPut,
+			body:      reader,
+			optHeader: opt,
+		}
+		resp, err = s.client.send(ctx, &sendOpt)
+		sUrl, retrieable = s.client.CheckRetrieable(sUrl, resp, err)
+		if retrieable && nr+1 < count {
+			if seeker, ok := r.(io.Seeker); ok {
+				_, e := seeker.Seek(position, io.SeekStart)
+				if e != nil {
+					break
+				}
+				continue
+			}
+		}
+		break
 	}
-	resp, err := s.client.send(ctx, &sendOpt)
-
 	return resp, err
 }
 
 // PutFromFile put object from local file
 func (s *ObjectService) PutFromFile(ctx context.Context, name string, filePath string, opt *ObjectPutOptions) (resp *Response, err error) {
+	if opt == nil {
+		opt = &ObjectPutOptions{}
+	}
 	nr := 0
 	for nr < 3 {
 		fd, e := os.Open(filePath)
@@ -390,6 +424,12 @@ func (s *ObjectService) PutFromFile(ctx context.Context, name string, filePath s
 		if err != nil {
 			nr++
 			fd.Close()
+			if s.client.Conf.RetryOpt.AutoSwitchHost {
+				// 收不到报文 或者 不存在RequestId
+				if resp == nil || resp.Header.Get("X-Cos-Request-Id") == "" {
+					opt.innerSwitchURL = toSwitchHost(s.client.BaseURL.BucketURL)
+				}
+			}
 			continue
 		}
 		fd.Close()
@@ -900,6 +940,12 @@ func worker(ctx context.Context, s *ObjectService, jobs <-chan *Jobs, results ch
 					results <- &res
 					break
 				}
+				if s.client.Conf.RetryOpt.AutoSwitchHost {
+					// 收不到报文 或者 不存在RequestId
+					if resp == nil || resp.Header.Get("X-Cos-Request-Id") == "" {
+						j.Opt.innerSwitchURL = toSwitchHost(s.client.BaseURL.BucketURL)
+					}
+				}
 				time.Sleep(time.Millisecond)
 				continue
 			}
@@ -1128,6 +1174,7 @@ func (s *ObjectService) Upload(ctx context.Context, name string, filepath string
 			opt0 = &ObjectPutOptions{
 				opt.OptIni.ACLHeaderOptions,
 				opt.OptIni.ObjectPutHeaderOptions,
+				nil,
 			}
 		}
 		rsp, err := s.PutFromFile(ctx, name, filepath, opt0)
diff --git a/object_part.go b/object_part.go
@@ -60,6 +60,9 @@ type ObjectUploadPartOptions struct {
 	XOptionHeader *http.Header `header:"-,omitempty" url:"-" xml:"-"`
 	// 上传进度, ProgressCompleteEvent不能表示对应API调用成功，API是否调用成功的判断标准为返回err==nil
 	Listener ProgressListener `header:"-" url:"-" xml:"-"`
+
+	// Upload方法使用
+	innerSwitchURL *url.URL `header:"-" url:"-" xml:"-"`
 }
 
 // UploadPart 请求实现在初始化以后的分块上传，支持的块的数量为1到10000，块的大小为1 MB 到5 GB。
@@ -93,22 +96,51 @@ func (s *ObjectService) UploadPart(ctx context.Context, name, uploadID string, p
 			opt.ContentLength = totalBytes
 		}
 	}
-	reader := TeeReader(r, nil, totalBytes, nil)
-	if s.client.Conf.EnableCRC {
-		reader.writer = crc64.New(crc64.MakeTable(crc64.ECMA))
+	// 如果是io.Seeker，则重试
+	count := 1
+	var position int64
+	if seeker, ok := r.(io.Seeker); ok {
+		// 记录原始位置
+		position, err = seeker.Seek(0, io.SeekCurrent)
+		if err == nil && s.client.Conf.RetryOpt.Count > 0 {
+			count = s.client.Conf.RetryOpt.Count
+		}
 	}
-	if opt != nil && opt.Listener != nil {
-		reader.listener = opt.Listener
+	var resp *Response
+	var retrieable bool
+	sUrl := s.client.BaseURL.BucketURL
+	if opt.innerSwitchURL != nil {
+		sUrl = opt.innerSwitchURL
 	}
-	u := fmt.Sprintf("/%s?partNumber=%d&uploadId=%s", encodeURIComponent(name), partNumber, uploadID)
-	sendOpt := sendOptions{
-		baseURL:   s.client.BaseURL.BucketURL,
-		uri:       u,
-		method:    http.MethodPut,
-		optHeader: opt,
-		body:      reader,
+	for nr := 0; nr < count; nr++ {
+		reader := TeeReader(r, nil, totalBytes, nil)
+		if s.client.Conf.EnableCRC {
+			reader.writer = crc64.New(crc64.MakeTable(crc64.ECMA))
+		}
+		if opt != nil && opt.Listener != nil {
+			reader.listener = opt.Listener
+		}
+		u := fmt.Sprintf("/%s?partNumber=%d&uploadId=%s", encodeURIComponent(name), partNumber, uploadID)
+		sendOpt := sendOptions{
+			baseURL:   sUrl,
+			uri:       u,
+			method:    http.MethodPut,
+			optHeader: opt,
+			body:      reader,
+		}
+		resp, err = s.client.send(ctx, &sendOpt)
+		sUrl, retrieable = s.client.CheckRetrieable(sUrl, resp, err)
+		if retrieable && nr+1 < count {
+			if seeker, ok := r.(io.Seeker); ok {
+				_, e := seeker.Seek(position, io.SeekStart)
+				if e != nil {
+					break
+				}
+				continue
+			}
+		}
+		break
 	}
-	resp, err := s.client.send(ctx, &sendOpt)
 	return resp, err
 }
 
diff --git a/object_part_test.go b/object_part_test.go
@@ -10,9 +10,12 @@ import (
 	"hash/crc64"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"reflect"
 	"strconv"
+	"strings"
 	"testing"
+	"time"
 )
 
 func TestObjectService_AbortMultipartUpload(t *testing.T) {
@@ -540,3 +543,79 @@ func TestObjectService_MultiCopy(t *testing.T) {
 	}
 
 }
+
+func TestObjectService_UploadPartRetry(t *testing.T) {
+	setup()
+	defer teardown()
+
+	opt := &ObjectUploadPartOptions{
+		Listener: &DefaultProgressListener{},
+	}
+	name := "test/hello.txt"
+	data := make([]byte, 1024*1024*3)
+	_, err := rand.Read(data)
+	tb := crc64.MakeTable(crc64.ECMA)
+	realcrc := crc64.Update(0, tb, data)
+	uploadID := "xxxxx"
+	partNumber := 1
+
+	nr, count := 0, 3
+	mux.HandleFunc("/test/hello.txt", func(w http.ResponseWriter, r *http.Request) {
+		testMethod(t, r, http.MethodPut)
+		vs := values{
+			"uploadId":   uploadID,
+			"partNumber": "1",
+		}
+		testFormValues(t, r, vs)
+
+		b, _ := ioutil.ReadAll(r.Body)
+		crc := crc64.Update(0, tb, b)
+		if !reflect.DeepEqual(crc, realcrc) {
+			t.Errorf("Object.Put crc: %v, want: %v", crc, realcrc)
+		}
+		w.Header().Add("x-cos-hash-crc64ecma", strconv.FormatUint(crc, 10))
+		nr++
+		if nr < count {
+			w.WriteHeader(http.StatusInternalServerError)
+		}
+	})
+
+	_, err = client.Object.UploadPart(context.Background(), name, uploadID, partNumber, bytes.NewReader(data), opt)
+	if err != nil || nr != count {
+		t.Errorf("Object.UploadPart failed: %v", err)
+	}
+
+	nr, count = 0, 3
+	_, err = client.Object.UploadPart(context.Background(), name, uploadID, partNumber, strings.NewReader(string(data)), opt)
+	if err != nil || nr != count {
+		t.Errorf("Object.UploadPart failed: %v", err)
+	}
+
+	// 非io.Seeker不做重试
+	nr, count = 0, 3
+	_, err = client.Object.UploadPart(context.Background(), name, uploadID, partNumber, bytes.NewBuffer(data), opt)
+	if err == nil || nr != 1 {
+		t.Errorf("Object.UploadPart failed: %v", err)
+	}
+
+	filePath := "tmpfile" + time.Now().Format(time.RFC3339)
+	newfile, err := os.Create(filePath)
+	if err != nil {
+		t.Fatalf("create tmp file failed")
+	}
+	defer os.Remove(filePath)
+	// 源文件内容
+	newfile.Write(data)
+	newfile.Close()
+
+	fd, err := os.Open(filePath)
+	if err != nil {
+		t.Fatalf("open file failed: %v", err)
+	}
+	// 文件关闭, 不做重试
+	nr, count = 0, 3
+	_, err = client.Object.UploadPart(context.Background(), name, uploadID, partNumber, fd, opt)
+	if err == nil || nr != 1 {
+		t.Errorf("Object.UploadPart failed: %v", err)
+	}
+}
diff --git a/object_test.go b/object_test.go

Original file line number	Diff line number	Diff line change
`@@ -233,6 +233,7 @@ func CloneObjectPutOptions(opt ObjectPutOptions) ObjectPutOptions {`
`233`	`233`	`res := &ObjectPutOptions{`
`234`	`234`	`&ACLHeaderOptions{},`
`235`	`235`	`&ObjectPutHeaderOptions{},`
	`236`	`+ nil,`
`236`	`237`	`}`
`237`	`238`	`if opt != nil {`
`238`	`239`	`if opt.ACLHeaderOptions != nil {`
`@@ -243,6 +244,9 @@ func CloneObjectPutOptions(opt ObjectPutOptions) ObjectPutOptions {`
`243`	`244`	`res.XCosMetaXXX = cloneHeader(opt.XCosMetaXXX)`
`244`	`245`	`res.XOptionHeader = cloneHeader(opt.XOptionHeader)`
`245`	`246`	`}`
	`247`	`+ if opt.innerSwitchURL != nil {`
	`248`	`+ res.innerSwitchURL = opt.innerSwitchURL`
	`249`	`+ }`
`246`	`250`	`}`
`247`	`251`	`return res`
`248`	`252`	`}`