JCS-pub/client/internal/uploader/update.go

package uploader

import (
	"context"
	"fmt"
	"io"
	"path"
	"sync"
	"time"

	"github.com/samber/lo"
	"gitlink.org.cn/cloudream/common/pkgs/ioswitch/exec"
	cdssdk "gitlink.org.cn/cloudream/common/sdks/storage"
	stgglb "gitlink.org.cn/cloudream/storage2/common/globals"
	stgmod "gitlink.org.cn/cloudream/storage2/common/models"
	"gitlink.org.cn/cloudream/storage2/common/pkgs/distlock"
	"gitlink.org.cn/cloudream/storage2/common/pkgs/ioswitch2"
	"gitlink.org.cn/cloudream/storage2/common/pkgs/ioswitch2/ops2"
	"gitlink.org.cn/cloudream/storage2/common/pkgs/ioswitch2/parser"
	coormq "gitlink.org.cn/cloudream/storage2/common/pkgs/mq/coordinator"
)

type UpdateUploader struct {
	uploader   *Uploader
	pkgID      cdssdk.PackageID
	targetStg  stgmod.StorageDetail
	distMutex  *distlock.Mutex
	loadToStgs []stgmod.StorageDetail
	loadToPath []string
	successes  []coormq.AddObjectEntry
	lock       sync.Mutex
	commited   bool
}

type UploadStorageInfo struct {
	Storage        stgmod.StorageDetail
	Delay          time.Duration
	IsSameLocation bool
}

type UpdateResult struct {
	// 上传成功的文件列表，Key为Path
	Objects map[string]cdssdk.Object
}

func (w *UpdateUploader) Upload(pat string, stream io.Reader) error {
	uploadTime := time.Now()

	ft := ioswitch2.NewFromTo()
	fromExec, hd := ioswitch2.NewFromDriver(ioswitch2.RawStream())
	ft.AddFrom(fromExec).
		AddTo(ioswitch2.NewToShardStore(*w.targetStg.MasterHub, w.targetStg, ioswitch2.RawStream(), "shardInfo"))

	for i, stg := range w.loadToStgs {
		ft.AddTo(ioswitch2.NewLoadToPublic(*stg.MasterHub, stg, path.Join(w.loadToPath[i], pat)))
	}

	plans := exec.NewPlanBuilder()
	err := parser.Parse(ft, plans)
	if err != nil {
		return fmt.Errorf("parsing plan: %w", err)
	}

	exeCtx := exec.NewExecContext()
	exec.SetValueByType(exeCtx, w.uploader.stgAgts)
	exec := plans.Execute(exeCtx)
	exec.BeginWrite(io.NopCloser(stream), hd)
	ret, err := exec.Wait(context.TODO())
	if err != nil {
		return fmt.Errorf("executing plan: %w", err)
	}

	w.lock.Lock()
	defer w.lock.Unlock()

	// 记录上传结果
	shardInfo := ret["shardInfo"].(*ops2.ShardInfoValue)
	w.successes = append(w.successes, coormq.AddObjectEntry{
		Path:       pat,
		Size:       shardInfo.Size,
		FileHash:   shardInfo.Hash,
		UploadTime: uploadTime,
		StorageIDs: []cdssdk.StorageID{w.targetStg.Storage.StorageID},
	})
	return nil
}

// 取消上传对象。必须在对象调用了Upload之后调用。
func (w *UpdateUploader) CancelObject(path string) {
	w.lock.Lock()
	defer w.lock.Unlock()

	w.successes = lo.Reject(w.successes, func(e coormq.AddObjectEntry, i int) bool {
		return e.Path == path
	})
}

// 重命名对象。必须在对象调用了Upload之后调用。不会检查新路径是否已经存在，需由调用方去做保证。
func (w *UpdateUploader) RenameObject(path string, newPath string) {
	w.lock.Lock()
	defer w.lock.Unlock()

	for i := range w.successes {
		if w.successes[i].Path == path {
			w.successes[i].Path = newPath
			break
		}
	}
}

func (w *UpdateUploader) Commit() (UpdateResult, error) {
	w.lock.Lock()
	defer w.lock.Unlock()

	if w.commited {
		return UpdateResult{}, fmt.Errorf("package already commited")
	}
	w.commited = true

	defer w.distMutex.Unlock()

	coorCli, err := stgglb.CoordinatorMQPool.Acquire()
	if err != nil {
		return UpdateResult{}, fmt.Errorf("new coordinator client: %w", err)
	}
	defer stgglb.CoordinatorMQPool.Release(coorCli)

	updateResp, err := coorCli.UpdatePackage(coormq.NewUpdatePackage(w.pkgID, w.successes))
	if err != nil {
		return UpdateResult{}, fmt.Errorf("updating package: %w", err)
	}

	ret := UpdateResult{
		Objects: make(map[string]cdssdk.Object),
	}

	for _, entry := range updateResp.Added {
		ret.Objects[entry.Path] = entry
	}

	return ret, nil
}

func (w *UpdateUploader) Abort() {
	w.lock.Lock()
	defer w.lock.Unlock()

	if w.commited {
		return
	}

	w.commited = true
	w.distMutex.Unlock()
}