github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/aria2/monitor/monitor.go (about) 1 package monitor 2 3 import ( 4 "context" 5 "encoding/json" 6 "errors" 7 "path/filepath" 8 "strconv" 9 "time" 10 11 model "github.com/cloudreve/Cloudreve/v3/models" 12 "github.com/cloudreve/Cloudreve/v3/pkg/aria2/common" 13 "github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc" 14 "github.com/cloudreve/Cloudreve/v3/pkg/cluster" 15 "github.com/cloudreve/Cloudreve/v3/pkg/filesystem" 16 "github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx" 17 "github.com/cloudreve/Cloudreve/v3/pkg/mq" 18 "github.com/cloudreve/Cloudreve/v3/pkg/task" 19 "github.com/cloudreve/Cloudreve/v3/pkg/util" 20 ) 21 22 // Monitor 离线下载状态监控 23 type Monitor struct { 24 Task *model.Download 25 Interval time.Duration 26 27 notifier <-chan mq.Message 28 node cluster.Node 29 retried int 30 } 31 32 var MAX_RETRY = 10 33 34 // NewMonitor 新建离线下载状态监控 35 func NewMonitor(task *model.Download, pool cluster.Pool, mqClient mq.MQ) { 36 monitor := &Monitor{ 37 Task: task, 38 notifier: make(chan mq.Message), 39 node: pool.GetNodeByID(task.GetNodeID()), 40 } 41 42 if monitor.node != nil { 43 monitor.Interval = time.Duration(monitor.node.GetAria2Instance().GetConfig().Interval) * time.Second 44 go monitor.Loop(mqClient) 45 46 monitor.notifier = mqClient.Subscribe(monitor.Task.GID, 0) 47 } else { 48 monitor.setErrorStatus(errors.New("node not avaliable")) 49 } 50 } 51 52 // Loop 开启监控循环 53 func (monitor *Monitor) Loop(mqClient mq.MQ) { 54 defer mqClient.Unsubscribe(monitor.Task.GID, monitor.notifier) 55 56 // 首次循环立即更新 57 interval := 50 * time.Millisecond 58 59 for { 60 select { 61 case <-monitor.notifier: 62 if monitor.Update() { 63 return 64 } 65 case <-time.After(interval): 66 interval = monitor.Interval 67 if monitor.Update() { 68 return 69 } 70 } 71 } 72 } 73 74 // Update 更新状态,返回值表示是否退出监控 75 func (monitor *Monitor) Update() bool { 76 status, err := monitor.node.GetAria2Instance().Status(monitor.Task) 77 78 if err != nil { 79 monitor.retried++ 80 util.Log().Warning("Cannot get status of download task %q: %s", monitor.Task.GID, err) 81 82 // 十次重试后认定为任务失败 83 if monitor.retried > MAX_RETRY { 84 util.Log().Warning("Cannot get status of download task %q,exceed maximum retry threshold: %s", 85 monitor.Task.GID, err) 86 monitor.setErrorStatus(err) 87 monitor.RemoveTempFolder() 88 return true 89 } 90 91 return false 92 } 93 monitor.retried = 0 94 95 // 磁力链下载需要跟随 96 if len(status.FollowedBy) > 0 { 97 util.Log().Debug("Redirected download task from %q to %q.", monitor.Task.GID, status.FollowedBy[0]) 98 monitor.Task.GID = status.FollowedBy[0] 99 monitor.Task.Save() 100 return false 101 } 102 103 // 更新任务信息 104 if err := monitor.UpdateTaskInfo(status); err != nil { 105 util.Log().Warning("Failed to update status of download task %q: %s", monitor.Task.GID, err) 106 monitor.setErrorStatus(err) 107 monitor.RemoveTempFolder() 108 return true 109 } 110 111 util.Log().Debug("Remote download %q status updated to %q.", status.Gid, status.Status) 112 113 switch common.GetStatus(status) { 114 case common.Complete, common.Seeding: 115 return monitor.Complete(task.TaskPoll) 116 case common.Error: 117 return monitor.Error(status) 118 case common.Downloading, common.Ready, common.Paused: 119 return false 120 case common.Canceled: 121 monitor.Task.Status = common.Canceled 122 monitor.Task.Save() 123 monitor.RemoveTempFolder() 124 return true 125 default: 126 util.Log().Warning("Download task %q returns unknown status %q.", monitor.Task.GID, status.Status) 127 return true 128 } 129 } 130 131 // UpdateTaskInfo 更新数据库中的任务信息 132 func (monitor *Monitor) UpdateTaskInfo(status rpc.StatusInfo) error { 133 originSize := monitor.Task.TotalSize 134 135 monitor.Task.GID = status.Gid 136 monitor.Task.Status = common.GetStatus(status) 137 138 // 文件大小、已下载大小 139 total, err := strconv.ParseUint(status.TotalLength, 10, 64) 140 if err != nil { 141 total = 0 142 } 143 downloaded, err := strconv.ParseUint(status.CompletedLength, 10, 64) 144 if err != nil { 145 downloaded = 0 146 } 147 monitor.Task.TotalSize = total 148 monitor.Task.DownloadedSize = downloaded 149 monitor.Task.GID = status.Gid 150 monitor.Task.Parent = status.Dir 151 152 // 下载速度 153 speed, err := strconv.Atoi(status.DownloadSpeed) 154 if err != nil { 155 speed = 0 156 } 157 158 monitor.Task.Speed = speed 159 attrs, _ := json.Marshal(status) 160 monitor.Task.Attrs = string(attrs) 161 162 if err := monitor.Task.Save(); err != nil { 163 return err 164 } 165 166 if originSize != monitor.Task.TotalSize { 167 // 文件大小更新后,对文件限制等进行校验 168 if err := monitor.ValidateFile(); err != nil { 169 // 验证失败时取消任务 170 monitor.node.GetAria2Instance().Cancel(monitor.Task) 171 return err 172 } 173 } 174 175 return nil 176 } 177 178 // ValidateFile 上传过程中校验文件大小、文件名 179 func (monitor *Monitor) ValidateFile() error { 180 // 找到任务创建者 181 user := monitor.Task.GetOwner() 182 if user == nil { 183 return common.ErrUserNotFound 184 } 185 186 // 创建文件系统 187 fs, err := filesystem.NewFileSystem(user) 188 if err != nil { 189 return err 190 } 191 defer fs.Recycle() 192 193 // 创建上下文环境 194 file := &fsctx.FileStream{ 195 Size: monitor.Task.TotalSize, 196 } 197 198 // 验证用户容量 199 if err := filesystem.HookValidateCapacity(context.Background(), fs, file); err != nil { 200 return err 201 } 202 203 // 验证每个文件 204 for _, fileInfo := range monitor.Task.StatusInfo.Files { 205 if fileInfo.Selected == "true" { 206 // 创建上下文环境 207 fileSize, _ := strconv.ParseUint(fileInfo.Length, 10, 64) 208 file := &fsctx.FileStream{ 209 Size: fileSize, 210 Name: filepath.Base(fileInfo.Path), 211 } 212 if err := filesystem.HookValidateFile(context.Background(), fs, file); err != nil { 213 return err 214 } 215 } 216 217 } 218 219 return nil 220 } 221 222 // Error 任务下载出错处理,返回是否中断监控 223 func (monitor *Monitor) Error(status rpc.StatusInfo) bool { 224 monitor.setErrorStatus(errors.New(status.ErrorMessage)) 225 226 // 清理临时文件 227 monitor.RemoveTempFolder() 228 229 return true 230 } 231 232 // RemoveTempFolder 清理下载临时目录 233 func (monitor *Monitor) RemoveTempFolder() { 234 monitor.node.GetAria2Instance().DeleteTempFile(monitor.Task) 235 } 236 237 // Complete 完成下载,返回是否中断监控 238 func (monitor *Monitor) Complete(pool task.Pool) bool { 239 // 未开始转存,提交转存任务 240 if monitor.Task.TaskID == 0 { 241 return monitor.transfer(pool) 242 } 243 244 // 做种完成 245 if common.GetStatus(monitor.Task.StatusInfo) == common.Complete { 246 transferTask, err := model.GetTasksByID(monitor.Task.TaskID) 247 if err != nil { 248 monitor.setErrorStatus(err) 249 monitor.RemoveTempFolder() 250 return true 251 } 252 253 // 转存完成,回收下载目录 254 if transferTask.Type == task.TransferTaskType && transferTask.Status >= task.Error { 255 job, err := task.NewRecycleTask(monitor.Task) 256 if err != nil { 257 monitor.setErrorStatus(err) 258 monitor.RemoveTempFolder() 259 return true 260 } 261 262 // 提交回收任务 263 pool.Submit(job) 264 265 return true 266 } 267 } 268 269 return false 270 } 271 272 func (monitor *Monitor) transfer(pool task.Pool) bool { 273 // 创建中转任务 274 file := make([]string, 0, len(monitor.Task.StatusInfo.Files)) 275 sizes := make(map[string]uint64, len(monitor.Task.StatusInfo.Files)) 276 for i := 0; i < len(monitor.Task.StatusInfo.Files); i++ { 277 fileInfo := monitor.Task.StatusInfo.Files[i] 278 if fileInfo.Selected == "true" { 279 file = append(file, fileInfo.Path) 280 size, _ := strconv.ParseUint(fileInfo.Length, 10, 64) 281 sizes[fileInfo.Path] = size 282 } 283 } 284 285 job, err := task.NewTransferTask( 286 monitor.Task.UserID, 287 file, 288 monitor.Task.Dst, 289 monitor.Task.Parent, 290 true, 291 monitor.node.ID(), 292 sizes, 293 ) 294 if err != nil { 295 monitor.setErrorStatus(err) 296 monitor.RemoveTempFolder() 297 return true 298 } 299 300 // 提交中转任务 301 pool.Submit(job) 302 303 // 更新任务ID 304 monitor.Task.TaskID = job.Model().ID 305 monitor.Task.Save() 306 307 return false 308 } 309 310 func (monitor *Monitor) setErrorStatus(err error) { 311 monitor.Task.Status = common.Error 312 monitor.Task.Error = err.Error() 313 monitor.Task.Save() 314 }