github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/filesystem/chunk/chunk.go

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)
}