package mq import ( "io/fs" "os" "path/filepath" "time" "github.com/samber/lo" "gitlink.org.cn/cloudream/common/consts/errorcode" "gitlink.org.cn/cloudream/common/pkgs/logger" "gitlink.org.cn/cloudream/common/pkgs/mq" "gitlink.org.cn/cloudream/storage/agent/internal/config" mytask "gitlink.org.cn/cloudream/storage/agent/internal/task" "gitlink.org.cn/cloudream/storage/common/consts" "gitlink.org.cn/cloudream/storage/common/globals" "gitlink.org.cn/cloudream/storage/common/pkgs/iterator" agtmq "gitlink.org.cn/cloudream/storage/common/pkgs/mq/agent" coormq "gitlink.org.cn/cloudream/storage/common/pkgs/mq/coordinator" "gitlink.org.cn/cloudream/storage/common/utils" ) func (svc *Service) StartStorageLoadPackage(msg *agtmq.StartStorageLoadPackage) (*agtmq.StartStorageLoadPackageResp, *mq.CodeMessage) { coorCli, err := globals.CoordinatorMQPool.Acquire() if err != nil { logger.Warnf("new coordinator client: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "new coordinator client failed") } defer coorCli.Close() getStgResp, err := coorCli.GetStorageInfo(coormq.NewGetStorageInfo(msg.UserID, msg.StorageID)) if err != nil { logger.WithField("StorageID", msg.StorageID). Warnf("getting storage info: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "get storage info failed") } outputDirPath := filepath.Join(config.Cfg().StorageBaseDir, getStgResp.Directory, utils.MakeStorageLoadPackageDirName(msg.PackageID, msg.UserID)) if err = os.MkdirAll(outputDirPath, 0755); err != nil { logger.WithField("StorageID", msg.StorageID). Warnf("creating output directory: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "create output directory failed") } tsk := svc.taskManager.StartNew(mytask.NewDownloadPackage(msg.UserID, msg.PackageID, outputDirPath)) return mq.ReplyOK(agtmq.NewStartStorageLoadPackageResp(tsk.ID())) } func (svc *Service) WaitStorageLoadPackage(msg *agtmq.WaitStorageLoadPackage) (*agtmq.WaitStorageLoadPackageResp, *mq.CodeMessage) { logger.WithField("TaskID", msg.TaskID).Debugf("wait loading package") tsk := svc.taskManager.FindByID(msg.TaskID) if tsk == nil { return mq.ReplyFailed[agtmq.WaitStorageLoadPackageResp](errorcode.TaskNotFound, "task not found") } if msg.WaitTimeoutMs == 0 { tsk.Wait() errMsg := "" if tsk.Error() != nil { errMsg = tsk.Error().Error() } return mq.ReplyOK(agtmq.NewWaitStorageLoadPackageResp(true, errMsg)) } else { if tsk.WaitTimeout(time.Duration(msg.WaitTimeoutMs)) { errMsg := "" if tsk.Error() != nil { errMsg = tsk.Error().Error() } return mq.ReplyOK(agtmq.NewWaitStorageLoadPackageResp(true, errMsg)) } return mq.ReplyOK(agtmq.NewWaitStorageLoadPackageResp(false, "")) } } func (svc *Service) StorageCheck(msg *agtmq.StorageCheck) (*agtmq.StorageCheckResp, *mq.CodeMessage) { dirFullPath := filepath.Join(config.Cfg().StorageBaseDir, msg.Directory) infos, err := os.ReadDir(dirFullPath) if err != nil { logger.Warnf("list storage directory failed, err: %s", err.Error()) return mq.ReplyOK(agtmq.NewStorageCheckResp( err.Error(), nil, )) } dirInfos := lo.Filter(infos, func(info fs.DirEntry, index int) bool { return info.IsDir() }) if msg.IsComplete { return svc.checkStorageComplete(msg, dirInfos) } else { return svc.checkStorageIncrement(msg, dirInfos) } } func (svc *Service) checkStorageIncrement(msg *agtmq.StorageCheck, dirInfos []fs.DirEntry) (*agtmq.StorageCheckResp, *mq.CodeMessage) { infosMap := make(map[string]fs.DirEntry) for _, info := range dirInfos { infosMap[info.Name()] = info } var entries []agtmq.StorageCheckRespEntry for _, obj := range msg.Packages { dirName := utils.MakeStorageLoadPackageDirName(obj.PackageID, obj.UserID) _, ok := infosMap[dirName] if ok { // 不需要做处理 // 删除map中的记录,表示此记录已被检查过 delete(infosMap, dirName) } else { // 只要文件不存在,就删除StoragePackage表中的记录 entries = append(entries, agtmq.NewStorageCheckRespEntry(obj.PackageID, obj.UserID, agtmq.CHECK_STORAGE_RESP_OP_DELETE)) } } // 增量情况下,不需要对infosMap中没检查的记录进行处理 return mq.ReplyOK(agtmq.NewStorageCheckResp(consts.StorageDirectoryStateOK, entries)) } func (svc *Service) checkStorageComplete(msg *agtmq.StorageCheck, dirInfos []fs.DirEntry) (*agtmq.StorageCheckResp, *mq.CodeMessage) { infosMap := make(map[string]fs.DirEntry) for _, info := range dirInfos { infosMap[info.Name()] = info } var entries []agtmq.StorageCheckRespEntry for _, obj := range msg.Packages { dirName := utils.MakeStorageLoadPackageDirName(obj.PackageID, obj.UserID) _, ok := infosMap[dirName] if ok { // 不需要做处理 // 删除map中的记录,表示此记录已被检查过 delete(infosMap, dirName) } else { // 只要文件不存在,就删除StoragePackage表中的记录 entries = append(entries, agtmq.NewStorageCheckRespEntry(obj.PackageID, obj.UserID, agtmq.CHECK_STORAGE_RESP_OP_DELETE)) } } return mq.ReplyOK(agtmq.NewStorageCheckResp(consts.StorageDirectoryStateOK, entries)) } func (svc *Service) StartStorageCreatePackage(msg *agtmq.StartStorageCreatePackage) (*agtmq.StartStorageCreatePackageResp, *mq.CodeMessage) { coorCli, err := globals.CoordinatorMQPool.Acquire() if err != nil { logger.Warnf("new coordinator client: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "new coordinator client failed") } defer coorCli.Close() getStgResp, err := coorCli.GetStorageInfo(coormq.NewGetStorageInfo(msg.UserID, msg.StorageID)) if err != nil { logger.WithField("StorageID", msg.StorageID). Warnf("getting storage info: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "get storage info failed") } fullPath := filepath.Clean(filepath.Join(config.Cfg().StorageBaseDir, getStgResp.Directory, msg.Path)) var uploadFilePathes []string err = filepath.WalkDir(fullPath, func(fname string, fi os.DirEntry, err error) error { if err != nil { return nil } if !fi.IsDir() { uploadFilePathes = append(uploadFilePathes, fname) } return nil }) if err != nil { logger.Warnf("opening directory %s: %s", fullPath, err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "read directory failed") } objIter := iterator.NewUploadingObjectIterator(fullPath, uploadFilePathes) if msg.Redundancy.IsRepInfo() { repInfo, err := msg.Redundancy.ToRepInfo() if err != nil { logger.Warnf("getting rep redundancy info: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "get rep redundancy info failed") } tsk := svc.taskManager.StartNew(mytask.NewCreateRepPackage(msg.UserID, msg.BucketID, msg.Name, objIter, repInfo)) return mq.ReplyOK(agtmq.NewStartStorageCreatePackageResp(tsk.ID())) } ecInfo, err := msg.Redundancy.ToECInfo() if err != nil { logger.Warnf("getting ec redundancy info: %s", err.Error()) return nil, mq.Failed(errorcode.OperationFailed, "get ec redundancy info failed") } tsk := svc.taskManager.StartNew(mytask.NewCreateECPackage(msg.UserID, msg.BucketID, msg.Name, objIter, ecInfo)) return mq.ReplyOK(agtmq.NewStartStorageCreatePackageResp(tsk.ID())) } func (svc *Service) WaitStorageCreatePackage(msg *agtmq.WaitStorageCreatePackage) (*agtmq.WaitStorageCreatePackageResp, *mq.CodeMessage) { tsk := svc.taskManager.FindByID(msg.TaskID) if tsk == nil { return nil, mq.Failed(errorcode.TaskNotFound, "task not found") } if msg.WaitTimeoutMs == 0 { tsk.Wait() } else if !tsk.WaitTimeout(time.Duration(msg.WaitTimeoutMs)) { return mq.ReplyOK(agtmq.NewWaitStorageCreatePackageResp(false, "", 0)) } if tsk.Error() != nil { return mq.ReplyOK(agtmq.NewWaitStorageCreatePackageResp(true, tsk.Error().Error(), 0)) } // TODO 避免判断类型 if repTask, ok := tsk.Body().(*mytask.CreateRepPackage); ok { return mq.ReplyOK(agtmq.NewWaitStorageCreatePackageResp(true, "", repTask.Result.PackageID)) } if ecTask, ok := tsk.Body().(*mytask.CreateECPackage); ok { return mq.ReplyOK(agtmq.NewWaitStorageCreatePackageResp(true, "", ecTask.Result.PackageID)) } return nil, mq.Failed(errorcode.TaskNotFound, "task not found") }