JCS-pub/common/pkgs/ioswitch2/ops2/multipart.go

package ops2

import (
	"fmt"
	"time"

	log "gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/dag"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/pool"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

func init() {
	exec.UseOp[*MultipartInitiator]()
	exec.UseOp[*MultipartUpload]()
	exec.UseVarValue[*MultipartUploadArgsValue]()
	exec.UseVarValue[*UploadedPartInfoValue]()
}

type MultipartUploadArgsValue struct {
	InitState stgtypes.MultipartInitState
}

func (v *MultipartUploadArgsValue) Clone() exec.VarValue {
	return &MultipartUploadArgsValue{
		InitState: v.InitState,
	}
}

type UploadedPartInfoValue struct {
	stgtypes.UploadedPartInfo
}

func (v *UploadedPartInfoValue) Clone() exec.VarValue {
	return &UploadedPartInfoValue{
		UploadedPartInfo: v.UploadedPartInfo,
	}
}

type MultipartInitiator struct {
	UserSpace     jcstypes.UserSpaceDetail
	UploadArgs    exec.VarID
	UploadedParts []exec.VarID
	FileOutput    exec.VarID // 分片上传之后的临时文件的路径
}

func (o *MultipartInitiator) Execute(ctx *exec.ExecContext, e *exec.Executor) error {
	stgPool, err := exec.GetValueByType[*pool.Pool](ctx)
	if err != nil {
		return fmt.Errorf("getting storage pool: %w", err)
	}

	multi, err := stgPool.GetMultiparter(&o.UserSpace)
	if err != nil {
		return err
	}

	// 启动一个新的上传任务W
	multiTask, err := multi.Initiate(ctx.Context)
	if err != nil {
		return err
	}
	defer multiTask.Close()

	// 分发上传参数
	e.PutVar(o.UploadArgs, &MultipartUploadArgsValue{
		InitState: multiTask.InitState(),
	})

	// 收集分片上传结果
	partInfoValues, err := exec.BindArray[*UploadedPartInfoValue](e, ctx.Context, o.UploadedParts)
	if err != nil {
		return fmt.Errorf("getting uploaded parts: %v", err)
	}

	partInfos := make([]stgtypes.UploadedPartInfo, len(partInfoValues))
	for i, v := range partInfoValues {
		partInfos[i] = v.UploadedPartInfo
	}

	// 合并分片
	fileInfo, err := multiTask.JoinParts(ctx.Context, partInfos)
	if err != nil {
		return fmt.Errorf("completing multipart upload: %v", err)
	}

	// 告知后续Op临时文件的路径
	e.PutVar(o.FileOutput, &FileInfoValue{
		FileInfo: fileInfo,
	})

	return nil
}

func (o *MultipartInitiator) String() string {
	return fmt.Sprintf("MultipartInitiator Args: %v, Parts: %v, BypassFileOutput: %v", o.UploadArgs, o.UploadedParts, o.FileOutput)
}

type MultipartUpload struct {
	UserSpace    jcstypes.UserSpaceDetail
	UploadArgs   exec.VarID
	UploadResult exec.VarID
	PartStream   exec.VarID
	PartNumber   int
	PartSize     int64
}

func (o *MultipartUpload) Execute(ctx *exec.ExecContext, e *exec.Executor) error {
	stgPool, err := exec.GetValueByType[*pool.Pool](ctx)
	if err != nil {
		return fmt.Errorf("getting storage pool: %w", err)
	}

	uploadArgs, err := exec.BindVar[*MultipartUploadArgsValue](e, ctx.Context, o.UploadArgs)
	if err != nil {
		return err
	}

	partStr, err := exec.BindVar[*exec.StreamValue](e, ctx.Context, o.PartStream)
	if err != nil {
		return err
	}
	defer partStr.Stream.Close()

	multi, err := stgPool.GetMultiparter(&o.UserSpace)
	if err != nil {
		return err
	}

	startTime := time.Now()
	uploadedInfo, err := multi.UploadPart(ctx.Context, uploadArgs.InitState, o.PartSize, o.PartNumber, partStr.Stream)
	if err != nil {
		return err
	}
	log.Debugf("upload finished in %v", time.Since(startTime))

	e.PutVar(o.UploadResult, &UploadedPartInfoValue{
		uploadedInfo,
	})

	return nil
}

func (o *MultipartUpload) String() string {
	return fmt.Sprintf("MultipartUpload[PartNumber=%v,PartSize=%v] Args: %v, Result: %v, Stream: %v", o.PartNumber, o.PartSize, o.UploadArgs, o.UploadResult, o.PartStream)
}

type MultipartInitiatorNode struct {
	dag.NodeBase
	UserSpace jcstypes.UserSpaceDetail
}

func (b *GraphNodeBuilder) NewMultipartInitiator(userSpace jcstypes.UserSpaceDetail) *MultipartInitiatorNode {
	node := &MultipartInitiatorNode{
		UserSpace: userSpace,
	}
	b.AddNode(node)

	node.OutputValues().Init(node, 2)
	return node
}

func (n *MultipartInitiatorNode) UploadArgsVar() dag.ValueOutputSlot {
	return dag.ValueOutputSlot{
		Node:  n,
		Index: 0,
	}
}

func (n *MultipartInitiatorNode) FileInfoVar() dag.ValueOutputSlot {
	return dag.ValueOutputSlot{
		Node:  n,
		Index: 1,
	}
}

func (n *MultipartInitiatorNode) PartInfoSlot() dag.ValueInputSlot {
	return dag.ValueInputSlot{
		Node:  n,
		Index: n.InputValues().EnlargeOne(),
	}
}

func (n *MultipartInitiatorNode) GenerateOp() (exec.Op, error) {
	return &MultipartInitiator{
		UserSpace:     n.UserSpace,
		UploadArgs:    n.UploadArgsVar().Var().VarID,
		UploadedParts: n.InputValues().GetVarIDs(),
		FileOutput:    n.FileInfoVar().Var().VarID,
	}, nil
}

type MultipartUploadNode struct {
	dag.NodeBase
	UserSpace  jcstypes.UserSpaceDetail
	PartNumber int
	PartSize   int64
}

func (b *GraphNodeBuilder) NewMultipartUpload(userSpace jcstypes.UserSpaceDetail, partNumber int, partSize int64) *MultipartUploadNode {
	node := &MultipartUploadNode{
		UserSpace:  userSpace,
		PartNumber: partNumber,
		PartSize:   partSize,
	}
	b.AddNode(node)

	node.InputValues().Init(1)
	node.OutputValues().Init(node, 1)
	node.InputStreams().Init(1)
	return node
}

func (n *MultipartUploadNode) UploadArgsSlot() dag.ValueInputSlot {
	return dag.ValueInputSlot{
		Node:  n,
		Index: 0,
	}
}

func (n *MultipartUploadNode) UploadResultVar() dag.ValueOutputSlot {
	return dag.ValueOutputSlot{
		Node:  n,
		Index: 0,
	}
}

func (n *MultipartUploadNode) PartStreamSlot() dag.StreamInputSlot {
	return dag.StreamInputSlot{
		Node:  n,
		Index: 0,
	}
}

func (n *MultipartUploadNode) GenerateOp() (exec.Op, error) {
	return &MultipartUpload{
		UserSpace:    n.UserSpace,
		UploadArgs:   n.UploadArgsSlot().Var().VarID,
		UploadResult: n.UploadResultVar().Var().VarID,
		PartStream:   n.PartStreamSlot().Var().VarID,
		PartNumber:   n.PartNumber,
		PartSize:     n.PartSize,
	}, nil
}