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

package obs

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

	"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"
	clitypes "gitlink.org.cn/cloudream/jcs-pub/client/types"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/s3"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
)

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

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

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

// 执行数据直传。返回传输后的文件路径
func (s *S2STransfer) Transfer(ctx context.Context, src *clitypes.UserSpaceDetail, srcPath string, opt types.S2SOption) (string, error) {
	req := s.makeRequest(src, srcPath)
	if req == nil {
		return "", 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 "", err
	}

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

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

	tempPrefix := s3.JoinKey(s.feat.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 "", fmt.Errorf("create task: %w", err)
	}

	s.taskID = resp.Id

	err = s.waitTask(ctx, *resp.Id)
	if err != nil {
		return "", fmt.Errorf("wait task: %w", err)
	}

	return s3.JoinKey(tempPrefix, srcPath), nil
}

func (s *S2STransfer) makeRequest(srcStg *clitypes.UserSpaceDetail, srcPath string) *model.SrcNodeReq {
	switch srcType := srcStg.UserSpace.Storage.(type) {
	case *cortypes.OBSType:
		cloudType := "HuaweiCloud"

		cred, ok := srcStg.UserSpace.Credential.(*cortypes.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},
		}

	default:
		return nil
	}
}

func (s *S2STransfer) waitTask(ctx context.Context, taskId 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 fmt.Errorf("show task failed too many times: %w", err)
		}
		failures = 0

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

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

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

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

// 完成传输
func (s *S2STransfer) Complete() {

}

// 取消传输。如果已经调用了Complete，则这个方法应该无效果
func (s *S2STransfer) Abort() {
	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),
		})

		// TODO 清理临时文件
	}
}