JCS-pub/common/pkgs/storage/local/multipart_upload.go

package local

import (
	"context"
	"crypto/sha256"
	"fmt"
	"hash"
	"io"
	"os"
	"path/filepath"

	"gitlink.org.cn/cloudream/common/utils/io2"
	"gitlink.org.cn/cloudream/common/utils/os2"
	"gitlink.org.cn/cloudream/common/utils/sort2"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

type Multiparter struct {
	detail   *jcstypes.UserSpaceDetail
	localStg *jcstypes.LocalCred
	feat     *jcstypes.MultipartUploadFeature
}

func (*Multiparter) MinPartSize() int64 {
	return 1 * 1024 * 1024 // 1MB
}

func (*Multiparter) MaxPartSize() int64 {
	return 5 * 1024 * 1024 * 1024 // 5GB
}

func (m *Multiparter) Initiate(ctx context.Context) (stgtypes.MultipartTask, error) {
	tempDir := filepath.Join(m.localStg.RootDir, m.detail.UserSpace.WorkingDir.JoinOSPath(), stgtypes.TempWorkingDir)
	absTempDir, err := filepath.Abs(tempDir)
	if err != nil {
		return nil, fmt.Errorf("get abs temp dir %v: %v", tempDir, err)
	}

	tempFileName := os2.GenerateRandomFileName(10)
	tempPartsDir := filepath.Join(absTempDir, tempFileName)
	absJoinedFilePath := filepath.Join(absTempDir, tempFileName+".joined")

	err = os.MkdirAll(tempPartsDir, 0777)

	if err != nil {
		return nil, err
	}

	return &MultipartTask{
		absTempDir:        absTempDir,
		tempFileName:      tempFileName,
		tempPartsDir:      tempPartsDir,
		joinedFileJPath:   m.detail.UserSpace.WorkingDir.ConcatCompsNew(stgtypes.TempWorkingDir, tempFileName+".joined"),
		absJoinedFilePath: absJoinedFilePath,
		uploadID:          tempPartsDir,
	}, nil
}

func (m *Multiparter) UploadPart(ctx context.Context, init stgtypes.MultipartInitState, partSize int64, partNumber int, stream io.Reader) (stgtypes.UploadedPartInfo, error) {
	partFilePath := filepath.Join(init.UploadID, fmt.Sprintf("%v", partNumber))
	partFile, err := os.Create(partFilePath)
	if err != nil {
		return stgtypes.UploadedPartInfo{}, err
	}
	defer partFile.Close()

	_, err = io.Copy(partFile, stream)
	if err != nil {
		return stgtypes.UploadedPartInfo{}, err
	}
	return stgtypes.UploadedPartInfo{
		ETag:       partFilePath,
		PartNumber: partNumber,
	}, nil
}

type MultipartTask struct {
	absTempDir        string // 应该要是绝对路径
	tempFileName      string
	tempPartsDir      string
	joinedFileJPath   jcstypes.JPath
	absJoinedFilePath string
	uploadID          string
}

func (i *MultipartTask) InitState() stgtypes.MultipartInitState {
	return stgtypes.MultipartInitState{
		UploadID: i.uploadID,
	}
}

func (i *MultipartTask) JoinParts(ctx context.Context, parts []stgtypes.UploadedPartInfo) (stgtypes.FileInfo, error) {
	parts = sort2.Sort(parts, func(l, r stgtypes.UploadedPartInfo) int {
		return l.PartNumber - r.PartNumber
	})

	joined, err := os.Create(i.absJoinedFilePath)
	if err != nil {
		return stgtypes.FileInfo{}, err
	}
	defer joined.Close()

	size := int64(0)
	hasher := sha256.New()

	for _, part := range parts {
		partSize, err := i.writePart(part, joined, hasher)
		if err != nil {
			return stgtypes.FileInfo{}, err
		}

		size += partSize
	}

	h := hasher.Sum(nil)
	return stgtypes.FileInfo{
		Path: i.joinedFileJPath,
		Size: size,
		Hash: jcstypes.NewFullHash(h),
	}, nil
}

func (i *MultipartTask) writePart(partInfo stgtypes.UploadedPartInfo, joined *os.File, hasher hash.Hash) (int64, error) {
	part, err := os.Open(partInfo.ETag)
	if err != nil {
		return 0, err
	}
	defer part.Close()

	buf := make([]byte, 32*1024)
	size := int64(0)
	for {
		n, err := part.Read(buf)
		if n > 0 {
			size += int64(n)
			io2.WriteAll(hasher, buf[:n])
			err := io2.WriteAll(joined, buf[:n])
			if err != nil {
				return 0, err
			}
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return 0, err
		}
	}

	return size, nil
}

func (i *MultipartTask) Close() {
	os.RemoveAll(i.tempPartsDir)
}