init

pkg/filesystem/chunk/backoff/backoff.go

@@ -0,0 +1,74 @@
+package backoff
+
+import (
+	"errors"
+	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"net/http"
+	"strconv"
+	"time"
+)
+
+// Backoff used for retry sleep backoff
+type Backoff interface {
+	Next(err error) bool
+	Reset()
+}
+
+// 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
+
+	tried int
+}
+
+func (c *ConstantBackoff) Next(err error) bool {
+	c.tried++
+	if c.tried > c.Max {
+		return false
+	}
+
+	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/chunk.go

@@ -0,0 +1,167 @@
+package chunk
+
+import (
+	"context"
+	"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"
+)
+
+const bufferTempPattern = "cdChunk.*.tmp"
+
+// ChunkProcessFunc callback function for processing a chunk
+type ChunkProcessFunc func(c *ChunkGroup, chunk io.Reader) error
+
+// ChunkGroup manage groups of chunks
+type ChunkGroup struct {
+	file              fsctx.FileHeader
+	chunkSize         uint64
+	backoff           backoff.Backoff
+	enableRetryBuffer bool
+
+	fileInfo     *fsctx.UploadTaskInfo
+	currentIndex int
+	chunkNum     uint64
+	bufferTemp   *os.File
+}
+
+func NewChunkGroup(file fsctx.FileHeader, chunkSize uint64, backoff backoff.Backoff, useBuffer bool) *ChunkGroup {
+	c := &ChunkGroup{
+		file:              file,
+		chunkSize:         chunkSize,
+		backoff:           backoff,
+		fileInfo:          file.Info(),
+		currentIndex:      -1,
+		enableRetryBuffer: useBuffer,
+	}
+
+	if c.chunkSize == 0 {
+		c.chunkSize = c.fileInfo.Size
+	}
+
+	if c.fileInfo.Size == 0 {
+		c.chunkNum = 1
+	} else {
+		c.chunkNum = c.fileInfo.Size / c.chunkSize
+		if c.fileInfo.Size%c.chunkSize != 0 {
+			c.chunkNum++
+		}
+	}
+
+	return c
+}
+
+// TempAvailable returns if current chunk temp file is available to be read
+func (c *ChunkGroup) TempAvailable() bool {
+	if c.bufferTemp != nil {
+		state, _ := c.bufferTemp.Stat()
+		return state != nil && state.Size() == c.Length()
+	}
+
+	return false
+}
+
+// Process a chunk with retry logic
+func (c *ChunkGroup) Process(processor ChunkProcessFunc) error {
+	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() {
+		c.bufferTemp, _ = os.CreateTemp("", bufferTempPattern)
+		reader = io.TeeReader(reader, c.bufferTemp)
+	}
+
+	if c.bufferTemp != nil {
+		defer func() {
+			if c.bufferTemp != nil {
+				c.bufferTemp.Close()
+				os.Remove(c.bufferTemp.Name())
+				c.bufferTemp = nil
+			}
+		}()
+
+		// if temp buffer file is available, use it
+		if c.TempAvailable() {
+			if _, err := c.bufferTemp.Seek(0, io.SeekStart); err != nil {
+				return fmt.Errorf("failed to seek temp file back to chunk start: %w", err)
+			}
+
+			util.Log().Debug("Chunk %d will be read from temp file %q.", c.Index(), c.bufferTemp.Name())
+			reader = io.NopCloser(c.bufferTemp)
+		}
+	}
+
+	err := processor(c, reader)
+	if err != nil {
+		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)
+				}
+			}
+
+			util.Log().Debug("Retrying chunk %d, last error: %s", c.currentIndex, err)
+			return c.Process(processor)
+		}
+
+		return err
+	}
+
+	util.Log().Debug("Chunk %d processed", c.currentIndex)
+	return nil
+}
+
+// Start returns the byte index of current chunk
+func (c *ChunkGroup) Start() int64 {
+	return int64(uint64(c.Index()) * c.chunkSize)
+}
+
+// Total returns the total length
+func (c *ChunkGroup) Total() int64 {
+	return int64(c.fileInfo.Size)
+}
+
+// Num returns the total chunk number
+func (c *ChunkGroup) Num() int {
+	return int(c.chunkNum)
+}
+
+// RangeHeader returns header value of Content-Range
+func (c *ChunkGroup) RangeHeader() string {
+	return fmt.Sprintf("bytes %d-%d/%d", c.Start(), c.Start()+c.Length()-1, c.Total())
+}
+
+// Index returns current chunk index, starts from 0
+func (c *ChunkGroup) Index() int {
+	return c.currentIndex
+}
+
+// Next switch to next chunk, returns whether all chunks are processed
+func (c *ChunkGroup) Next() bool {
+	c.currentIndex++
+	c.backoff.Reset()
+	return c.currentIndex < int(c.chunkNum)
+}
+
+// Length returns the length of current chunk
+func (c *ChunkGroup) Length() int64 {
+	contentLength := c.chunkSize
+	if c.Index() == int(c.chunkNum-1) {
+		contentLength = c.fileInfo.Size - c.chunkSize*(c.chunkNum-1)
+	}
+
+	return int64(contentLength)
+}
+
+// IsLast returns if current chunk is the last one
+func (c *ChunkGroup) IsLast() bool {
+	return c.Index() == int(c.chunkNum-1)
+}