JCS-pub/common/pkgs/storage/obs/s2s.go

package obs

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/huaweicloud/huaweicloud-sdk-go-v3/core/auth/basic"
	oms "github.com/huaweicloud/huaweicloud-sdk-go-v3/services/oms/v2"
	"github.com/huaweicloud/huaweicloud-sdk-go-v3/services/oms/v2/model"
	omsregion "github.com/huaweicloud/huaweicloud-sdk-go-v3/services/oms/v2/region"
	"gitlink.org.cn/cloudream/common/utils/os2"
	stgs3 "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/s3"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

type S2STransfer struct {
	detail  *jcstypes.UserSpaceDetail
	stgType *jcstypes.OBSType
	cred    *jcstypes.OBSCred
	feat    *jcstypes.S2STransferFeature
	taskID  *int64
	omsCli  *oms.OmsClient
}

func NewS2STransfer(detail *jcstypes.UserSpaceDetail, stgType *jcstypes.OBSType, cred *jcstypes.OBSCred, feat *jcstypes.S2STransferFeature) *S2STransfer {
	return &S2STransfer{
		detail:  detail,
		stgType: stgType,
		cred:    cred,
		feat:    feat,
	}
}

// 判断是否能从指定的源存储中直传到当前存储的目的路径
func (*S2STransfer) CanTransfer(src, dst *jcstypes.UserSpaceDetail) bool {
	req := makeRequest(src, jcstypes.JPath{})
	return req != nil
}

// 执行数据直传。返回传输后的文件路径
func (s *S2STransfer) Transfer(ctx context.Context, src *jcstypes.UserSpaceDetail, srcPath jcstypes.JPath, dstPath jcstypes.JPath) (stgtypes.FileInfo, error) {
	req := makeRequest(src, srcPath)
	if req == nil {
		return stgtypes.FileInfo{}, fmt.Errorf("unsupported source storage type: %T", src.UserSpace.Storage)
	}

	auth, err := basic.NewCredentialsBuilder().
		WithAk(s.cred.AK).
		WithSk(s.cred.SK).
		WithProjectId(s.stgType.ProjectID).
		SafeBuild()
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	region, err := omsregion.SafeValueOf(s.stgType.Region)
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	cli, err := oms.OmsClientBuilder().
		WithRegion(region).
		WithCredential(auth).
		SafeBuild()
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	// 先上传成一个临时文件
	tempDir := stgs3.JoinKey(s.detail.UserSpace.WorkingDir.String(), stgtypes.TempWorkingDir)
	tempPrefix := stgs3.JoinKey(tempDir, os2.GenerateRandomFileName(10)) + "/"

	taskType := model.GetCreateTaskReqTaskTypeEnum().OBJECT
	s.omsCli = oms.NewOmsClient(cli)
	resp, err := s.omsCli.CreateTask(&model.CreateTaskRequest{
		Body: &model.CreateTaskReq{
			TaskType: &taskType,
			SrcNode:  req,
			DstNode: &model.DstNodeReq{
				Region:     s.stgType.Region,
				Ak:         s.cred.AK,
				Sk:         s.cred.SK,
				Bucket:     s.stgType.Bucket,
				SavePrefix: &tempPrefix,
			},
		},
	})
	if err != nil {
		return stgtypes.FileInfo{}, fmt.Errorf("create task: %w", err)
	}

	s.taskID = resp.Id

	// 轮询任务状态，直到完成
	size, err := s.waitTask(ctx, *resp.Id)
	if err != nil {
		return stgtypes.FileInfo{}, fmt.Errorf("wait task: %w", err)
	}

	// 传输完成后，将文件名改成目标路径
	obsCli, bkt, err := createClient(s.stgType, s.cred)
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	_, err = obsCli.CopyObject(ctx, &awss3.CopyObjectInput{
		Bucket:     aws.String(bkt),
		CopySource: aws.String(stgs3.JoinKey(bkt, tempPrefix, srcPath.String())),
		Key:        aws.String(dstPath.String()),
	})
	if err != nil {
		return stgtypes.FileInfo{}, fmt.Errorf("copy object: %w", err)
	}

	return stgtypes.FileInfo{
		Path: dstPath,
		Size: size,
		Hash: "",
	}, nil
}

func (s *S2STransfer) waitTask(ctx context.Context, taskId int64) (int64, error) {
	ticker := time.NewTicker(time.Second * 5)
	defer ticker.Stop()

	failures := 0

	for {
		resp, err := s.omsCli.ShowTask(&model.ShowTaskRequest{
			TaskId: fmt.Sprintf("%v", taskId),
		})
		if err != nil {
			if failures < 3 {
				failures++
				continue
			}

			return 0, fmt.Errorf("show task failed too many times: %w", err)
		}
		failures = 0

		if *resp.Status == 3 {
			return 0, fmt.Errorf("task stopped")
		}

		if *resp.Status == 4 {
			return 0, errors.New(resp.ErrorReason.String())
		}

		if *resp.Status == 5 {
			return *resp.CompleteSize, nil
		}

		select {
		case <-ticker.C:
			continue
		case <-ctx.Done():
			return 0, ctx.Err()
		}
	}
}

func (s *S2STransfer) Close() {
	if s.taskID != nil {
		s.omsCli.StopTask(&model.StopTaskRequest{
			TaskId: fmt.Sprintf("%v", *s.taskID),
		})

		s.omsCli.DeleteTask(&model.DeleteTaskRequest{
			TaskId: fmt.Sprintf("%v", *s.taskID),
		})
	}
}

func makeRequest(srcStg *jcstypes.UserSpaceDetail, srcPath jcstypes.JPath) *model.SrcNodeReq {
	switch srcType := srcStg.UserSpace.Storage.(type) {
	case *jcstypes.OBSType:
		cloudType := "HuaweiCloud"

		cred, ok := srcStg.UserSpace.Credential.(*jcstypes.OBSCred)
		if !ok {
			return nil
		}

		return &model.SrcNodeReq{
			CloudType: &cloudType,
			Region:    &srcType.Region,
			Ak:        &cred.AK,
			Sk:        &cred.SK,
			Bucket:    &srcType.Bucket,
			ObjectKey: &[]string{srcPath.String()},
		}

	default:
		return nil
	}
}