package uploader import ( "context" "fmt" "io" "path" "sync" "time" "gitlink.org.cn/cloudream/common/pkgs/ioswitch/exec" cdssdk "gitlink.org.cn/cloudream/common/sdks/storage" "gitlink.org.cn/cloudream/storage2/client/types" stgglb "gitlink.org.cn/cloudream/storage2/common/globals" "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 CreateLoadUploader struct { pkg cdssdk.Package userID cdssdk.UserID targetSpaces []types.UserSpaceDetail loadRoots []string uploader *Uploader distlock *distlock.Mutex successes []coormq.AddObjectEntry lock sync.Mutex commited bool } type CreateLoadResult struct { Package cdssdk.Package Objects map[string]cdssdk.Object } func (u *CreateLoadUploader) Upload(pa string, size int64, stream io.Reader) error { uploadTime := time.Now() spaceIDs := make([]cdssdk.StorageID, 0, len(u.targetSpaces)) ft := ioswitch2.FromTo{} fromExec, hd := ioswitch2.NewFromDriver(ioswitch2.RawStream()) ft.AddFrom(fromExec) for i, space := range u.targetSpaces { ft.AddTo(ioswitch2.NewToShardStore(*space.MasterHub, space, ioswitch2.RawStream(), "fileHash")) ft.AddTo(ioswitch2.NewLoadToPublic(*space.MasterHub, space, path.Join(u.loadRoots[i], pa))) spaceIDs = append(spaceIDs, space.Storage.StorageID) } plans := exec.NewPlanBuilder() err := parser.Parse(ft, plans) if err != nil { return fmt.Errorf("parsing plan: %w", err) } exeCtx := exec.NewExecContext() exec.SetValueByType(exeCtx, u.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) } u.lock.Lock() defer u.lock.Unlock() // 记录上传结果 fileHash := ret["fileHash"].(*ops2.ShardInfoValue).Hash u.successes = append(u.successes, coormq.AddObjectEntry{ Path: pa, Size: size, FileHash: fileHash, UploadTime: uploadTime, StorageIDs: spaceIDs, }) return nil } func (u *CreateLoadUploader) Commit() (CreateLoadResult, error) { u.lock.Lock() defer u.lock.Unlock() if u.commited { return CreateLoadResult{}, fmt.Errorf("package already commited") } u.commited = true defer u.distlock.Unlock() coorCli, err := stgglb.CoordinatorMQPool.Acquire() if err != nil { return CreateLoadResult{}, fmt.Errorf("new coordinator client: %w", err) } defer stgglb.CoordinatorMQPool.Release(coorCli) updateResp, err := coorCli.UpdatePackage(coormq.NewUpdatePackage(u.pkg.PackageID, u.successes)) if err != nil { return CreateLoadResult{}, fmt.Errorf("updating package: %w", err) } ret := CreateLoadResult{ Package: u.pkg, Objects: make(map[string]cdssdk.Object), } for _, entry := range updateResp.Added { ret.Objects[entry.Path] = entry } for i, stg := range u.targetSpaces { // 不关注是否成功 coorCli.StoragePackageLoaded(coormq.ReqStoragePackageLoaded(u.userID, stg.Storage.StorageID, u.pkg.PackageID, u.loadRoots[i], nil)) } return ret, nil } func (u *CreateLoadUploader) Abort() { u.lock.Lock() defer u.lock.Unlock() if u.commited { return } u.commited = true u.distlock.Unlock() // TODO 可以考虑删除PackageID }