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 copyToSpaces []types.UserSpaceDetail copyToPath []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, ioswitch2.RawStream(), "shardInfo")) for i, space := range w.copyToSpaces { ft.AddTo(ioswitch2.NewToBaseStore(space, path.Join(w.copyToPath[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.FileInfoValue) 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() }