package export_task import ( "91porn-server/common" "91porn-server/common/log" "91porn-server/models/l/playlgmod" "91porn-server/models/v/export_task_mod" "91porn-server/models/v/usermod" "encoding/json" "errors" "fmt" "go.mongodb.org/mongo-driver/bson" "sync" ) // ExecExportUserTask 执行用户导出任务 func ExecExportUserTask(task *export_task_mod.ExportTask) (total int64, err error) { type userListReq struct { usermod.UserListSelector common.StandQuery IsPretendAcc *int `form:"isPretendAcc" json:"isPretendAcc"` //是否马甲账号 } var arg userListReq err = json.Unmarshal([]byte(task.Param), &arg) if err != nil { return } if arg.StartTime == nil || arg.EndTime == nil { return 0, errors.New("没有选择时间") } cond, _ := common.StandQueryMap(arg.StandQuery, arg.UserListSelector) //马甲账户 if arg.IsPretendAcc != nil && *arg.IsPretendAcc == 1 { cond["devID"] = bson.M{"$regex": usermod.SystemDevIDPrex, "$options": "i"} } var skip int64 var size int64 = 500 records := []usermod.ExportUser{} zipName := fmt.Sprintf("%v_%v_用户列表_%v.zip", task.Admin, GetProName(), task.CreatedAt.Format("2006-01-02_15:04:05")) index := 1 path := fmt.Sprintf("./temp/%v%v", "用户列表-", task.ID.Hex()) excelList := []string{} // 最后一条记录的时间 for { userList, err := usermod.ExportFindMany(cond, skip, size) if err != nil { log.Error("ExecExportUserTask fail", log.E(err)) return total, err } videoCntMap := make(map[uint64]int) unDealVideoCntMap := make(map[uint64]int) // 针对91PORN做的优化,不然数据太多,导出数据文本太大,速度慢 //if arg.IsPretendAcc != nil && *arg.IsPretendAcc == 1 { // videoCntMap, unDealVideoCntMap, _ = vidmod.GetVideosByUIDs(getUIDs(userList)) //} eUsers := encodeUsers(userList, videoCntMap, unDealVideoCntMap) records = append(records, eUsers...) var fileName = fmt.Sprintf("用户列表_%v.xlsx", index) // 最多容纳5000条 if len(records) >= maxDataNum { // 直接写入excel filePath, err := SaveExcel(records, path, fileName) if err != nil { log.Error("ExecExportVideoTask SaveExcel fail", log.E(err)) return total, err } // 重置 records = []usermod.ExportUser{} index = index + 1 excelList = append(excelList, filePath) total = total + int64(len(records)) } else if len(userList) < int(size) { if len(records) > 0 { // 直接写入excel filePath, err := SaveExcel(records, path, fileName) if err != nil { log.Error("ExecExportVideoTask SaveExcel fail", log.E(err)) return total, err } excelList = append(excelList, filePath) total = total + int64(len(records)) } // 已经没有更多数据了,直接返回 break } skip = skip + size } // 发送到tg err = SendTg(task.Admin, excelList, zipName) if err != nil { log.Error("ExecExportUserTask SendTg fail", log.E(err)) return } return total, nil } func encodeUsers(users []*usermod.User, videoCntMap map[uint64]int, unDealVideoCntMap map[uint64]int) []usermod.ExportUser { usersLen := len(users) uids := make([]uint64, 0, usersLen) uInfos := make([]usermod.ExportUser, 0, usersLen) for _, u := range users { if u != nil { user := usermod.ExportUser{ UID: u.UID, //DevID: u.DevID, DevType: u.DevType, RegisterIP: u.RegisterIP, Mobile: u.Mobile, //Gender: u.Gender, Channel: u.DistrictCode, //Name: u.Name, PromotionCode: u.PromCode, //Summary: u.Summary, //Region: u.Region, //Birthday: u.Birthday, //VipLevel: u.VipLevel, VipExpireDate: u.VipExpireDate, CreatedAt: u.CreatedAt, //MobileBindAt: u.MobileBindAt, //HasLocked: u.HasLocked, //HasBanned: u.HasBanned, //TotalVideoCnt: videoCntMap[u.UID], //UndealVideoCnt: unDealVideoCntMap[u.UID], LastVisitAt: u.LastVisitAt, } uInfos = append(uInfos, user) uids = append(uids, u.UID) } } wCnt := setArray2Map(calWatchedVideoCnt(uids)) for i, u := range uInfos { uInfos[i].WatchCount = wCnt[u.UID] } return uInfos } func getUIDs(users []*usermod.User) []uint64 { uids := make([]uint64, 0, len(users)) for _, v := range users { if v == nil { continue } uids = append(uids, v.UID) } return uids } func setArray2Map(datas []playlgmod.UIDCount) map[uint64]int { m := make(map[uint64]int) for _, d := range datas { m[d.UID] = d.Count } return m } func calWatchedVideoCnt(uids []uint64) []playlgmod.UIDCount { le := len(uids) if le < 1000 { datas, _ := playlgmod.HasWatchedVideoCnt(uids) return datas } //创建多个协程去拉取 const goRoutineCnt = 10 var wg sync.WaitGroup var mData [goRoutineCnt][]playlgmod.UIDCount size := le / goRoutineCnt wg.Add(goRoutineCnt) for i := 0; i < goRoutineCnt-1; i++ { common.GoParam(i, func(i int) { defer wg.Done() begin := i * size end := begin + size - 1 mData[i], _ = playlgmod.HasWatchedVideoCnt(uids[begin:end]) }) } common.Go(func() { defer wg.Done() begin := (goRoutineCnt - 1) * size end := le - 1 mData[goRoutineCnt-1], _ = playlgmod.HasWatchedVideoCnt(uids[begin:end]) }) wg.Wait() datas := []playlgmod.UIDCount{} for i := 0; i < goRoutineCnt; i++ { datas = append(datas, mData[i]...) } return datas }