feat(OneDrive): support Retry-After throttling control from Graph API (#280)

pkg/filesystem/chunk/backoff/backoff.go

@@ -1,14 +1,22 @@
 package backoff
 
-import "time"
+import (
+	"errors"
+	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"net/http"
+	"strconv"
+	"time"
+)
 
 // Backoff used for retry sleep backoff
 type Backoff interface {
-	Next() bool
+	Next(err error) bool
 	Reset()
 }
 
-// ConstantBackoff implements Backoff interface with constant sleep time
+// ConstantBackoff implements Backoff interface with constant sleep time. If the error
+// is retryable and with `RetryAfter` defined, the `RetryAfter` will be used as sleep duration.
 type ConstantBackoff struct {
 	Sleep time.Duration
 	Max   int
@@ -16,16 +24,51 @@ type ConstantBackoff struct {
 	tried int
 }
 
-func (c *ConstantBackoff) Next() bool {
+func (c *ConstantBackoff) Next(err error) bool {
 	c.tried++
 	if c.tried > c.Max {
 		return false
 	}
 
-	time.Sleep(c.Sleep)
+	var e *RetryableError
+	if errors.As(err, &e) && e.RetryAfter > 0 {
+		util.Log().Warning("Retryable error %q occurs in backoff, will sleep after %s.", e, e.RetryAfter)
+		time.Sleep(e.RetryAfter)
+	} else {
+		time.Sleep(c.Sleep)
+	}
+
 	return true
 }
 
 func (c *ConstantBackoff) Reset() {
 	c.tried = 0
 }
+
+type RetryableError struct {
+	Err        error
+	RetryAfter time.Duration
+}
+
+// NewRetryableErrorFromHeader constructs a new RetryableError from http response header
+// and existing error.
+func NewRetryableErrorFromHeader(err error, header http.Header) *RetryableError {
+	retryAfter := header.Get("retry-after")
+	if retryAfter == "" {
+		retryAfter = "0"
+	}
+
+	res := &RetryableError{
+		Err: err,
+	}
+
+	if retryAfterSecond, err := strconv.ParseInt(retryAfter, 10, 64); err == nil {
+		res.RetryAfter = time.Duration(retryAfterSecond) * time.Second
+	}
+
+	return res
+}
+
+func (e *RetryableError) Error() string {
+	return fmt.Sprintf("retryable error with retry-after=%s: %s", e.RetryAfter, e.Err)
+}

pkg/filesystem/chunk/backoff/backoff_test.go

@@ -1,7 +1,9 @@
 package backoff
 
 import (
+	"errors"
 	"github.com/stretchr/testify/assert"
+	"net/http"
 	"testing"
 	"time"
 )
@@ -9,14 +11,51 @@ import (
 func TestConstantBackoff_Next(t *testing.T) {
 	a := assert.New(t)
 
-	b := &ConstantBackoff{Sleep: time.Duration(0), Max: 3}
-	a.True(b.Next())
-	a.True(b.Next())
-	a.True(b.Next())
-	a.False(b.Next())
-	b.Reset()
-	a.True(b.Next())
-	a.True(b.Next())
-	a.True(b.Next())
-	a.False(b.Next())
+	// General error
+	{
+		err := errors.New("error")
+		b := &ConstantBackoff{Sleep: time.Duration(0), Max: 3}
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.False(b.Next(err))
+		b.Reset()
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.False(b.Next(err))
+	}
+
+	// Retryable error
+	{
+		err := &RetryableError{RetryAfter: time.Duration(1)}
+		b := &ConstantBackoff{Sleep: time.Duration(0), Max: 3}
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.False(b.Next(err))
+		b.Reset()
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.True(b.Next(err))
+		a.False(b.Next(err))
+	}
+
+}
+
+func TestNewRetryableErrorFromHeader(t *testing.T) {
+	a := assert.New(t)
+	// no retry-after header
+	{
+		err := NewRetryableErrorFromHeader(nil, http.Header{})
+		a.Empty(err.RetryAfter)
+	}
+
+	// with retry-after header
+	{
+		header := http.Header{}
+		header.Add("retry-after", "120")
+		err := NewRetryableErrorFromHeader(nil, header)
+		a.EqualValues(time.Duration(120)*time.Second, err.RetryAfter)
+	}
 }

pkg/filesystem/chunk/chunk.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/chunk/backoff"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
+	"github.com/cloudreve/Cloudreve/v3/pkg/request"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 	"io"
 	"os"
@@ -66,7 +67,7 @@ func (c *ChunkGroup) TempAvailable() bool {
 
 // Process a chunk with retry logic
 func (c *ChunkGroup) Process(processor ChunkProcessFunc) error {
-	reader := io.LimitReader(c.file, int64(c.chunkSize))
+	reader := io.LimitReader(c.file, c.Length())
 
 	// If useBuffer is enabled, tee the reader to a temp file
 	if c.enableRetryBuffer && c.bufferTemp == nil && !c.file.Seekable() {
@@ -90,13 +91,17 @@ func (c *ChunkGroup) Process(processor ChunkProcessFunc) error {
 			}
 
 			util.Log().Debug("Chunk %d will be read from temp file %q.", c.Index(), c.bufferTemp.Name())
-			reader = c.bufferTemp
+			reader = io.NopCloser(c.bufferTemp)
 		}
 	}
 
 	err := processor(c, reader)
 	if err != nil {
-		if err != context.Canceled && (c.file.Seekable() || c.TempAvailable()) && c.backoff.Next() {
+		if c.enableRetryBuffer {
+			request.BlackHole(reader)
+		}
+
+		if err != context.Canceled && (c.file.Seekable() || c.TempAvailable()) && c.backoff.Next(err) {
 			if c.file.Seekable() {
 				if _, seekErr := c.file.Seek(c.Start(), io.SeekStart); seekErr != nil {
 					return fmt.Errorf("failed to seek back to chunk start: %w, last error: %s", seekErr, err)