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"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/db"
	"gitlink.org.cn/cloudream/jcs-pub/client/types"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/distlock"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/ops2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/parser"
)

type UpdateUploader struct {
	uploader     *Uploader
	pkgID        types.PackageID
	targetSpace  types.UserSpaceDetail
	pubLock      *distlock.Mutex
	loadToSpaces []types.UserSpaceDetail
	loadToPath   []string
	successes    []db.AddObjectEntry
	lock         sync.Mutex
	commited     bool
}

type UploadSpaceInfo struct {
	Space          types.UserSpaceDetail
	Delay          time.Duration
	IsSameLocation bool
}

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

type UploadOption struct {
	CreateTime time.Time // 设置文件的上传时间，如果为0值，则使用开始上传时的时间。
}

func (w *UpdateUploader) Upload(pat string, stream io.Reader, opts ...UploadOption) error {
	opt := UploadOption{}
	if len(opts) > 0 {
		opt = opts[0]
	}

	if opt.CreateTime.IsZero() {
		opt.CreateTime = time.Now()
	}

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

	for i, space := range w.loadToSpaces {
		ft.AddTo(ioswitch2.NewToPublicStore(*space.MasterHub, space, 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.stgPool)
	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, db.AddObjectEntry{
		Path:         pat,
		Size:         shardInfo.Size,
		FileHash:     shardInfo.Hash,
		CreateTime:   opt.CreateTime,
		UserSpaceIDs: []types.UserSpaceID{w.targetSpace.UserSpace.UserSpaceID},
	})
	return nil
}

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

	w.successes = lo.Reject(w.successes, func(e db.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.pubLock.Unlock()

	var addedObjs []types.Object
	err := w.uploader.db.DoTx(func(tx db.SQLContext) error {
		var err error
		addedObjs, err = w.uploader.db.Object().BatchAdd(tx, w.pkgID, w.successes)
		return err
	})
	if err != nil {
		return UpdateResult{}, fmt.Errorf("adding objects: %w", err)
	}

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

	for _, entry := range addedObjs {
		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.pubLock.Unlock()
}