163 lines
3.9 KiB
Go
163 lines
3.9 KiB
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()
|
||
}
|