restore use power when removed not finish jobs

api/service/dalle/service.go

@@ -70,7 +70,7 @@ func (s *Service) Run() {
 			if err != nil {
 				logger.Errorf("error with image task: %v", err)
 				s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
-					"progress": -1,
+					"progress": 101,
 					"err_msg":  err.Error(),
 				})
 				s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Failed})
@@ -148,7 +148,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
 		Where("enabled", true).
 		Order("last_used_at ASC").First(&apiKey)
 	if tx.Error != nil {
-		return "", fmt.Errorf("no available IMG api key: %v", tx.Error)
+		return "", fmt.Errorf("no available DALL-E api key: %v", tx.Error)
 	}
 
 	var res imgRes
@@ -225,6 +225,30 @@ func (s *Service) CheckTaskNotify() {
 	}()
 }
 
+func (s *Service) CheckTaskStatus() {
+	go func() {
+		logger.Info("Running DALL-E task status checking ...")
+		for {
+			var jobs []model.DallJob
+			res := s.db.Where("progress < ?", 100).Find(&jobs)
+			if res.Error != nil {
+				time.Sleep(5 * time.Second)
+				continue
+			}
+
+			for _, job := range jobs {
+				// 超时的任务标记为失败
+				if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
+					job.Progress = 101
+					job.ErrMsg = "任务超时"
+					s.db.Updates(&job)
+				}
+			}
+			time.Sleep(time.Second * 10)
+		}
+	}()
+}
+
 func (s *Service) DownloadImages() {
 	go func() {
 		var items []model.DallJob
@@ -271,44 +295,3 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
 	s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Finished})
 	return imgURL, nil
 }
-
-// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
-func (s *Service) CheckTaskStatus() {
-	go func() {
-		logger.Info("Running Stable-Diffusion task status checking ...")
-		for {
-			var jobs []model.DallJob
-			res := s.db.Where("progress < ?", 100).Find(&jobs)
-			if res.Error != nil {
-				time.Sleep(5 * time.Second)
-				continue
-			}
-
-			for _, job := range jobs {
-				// 5 分钟还没完成的任务直接删除
-				if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
-					s.db.Delete(&job)
-					var user model.User
-					s.db.Where("id = ?", job.UserId).First(&user)
-					// 退回绘图次数
-					res = s.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
-					if res.Error == nil && res.RowsAffected > 0 {
-						s.db.Create(&model.PowerLog{
-							UserId:    user.Id,
-							Username:  user.Username,
-							Type:      types.PowerConsume,
-							Amount:    job.Power,
-							Balance:   user.Power + job.Power,
-							Mark:      types.PowerAdd,
-							Model:     "dall-e-3",
-							Remark:    fmt.Sprintf("任务失败，退回算力。任务ID：%d", job.Id),
-							CreatedAt: time.Now(),
-						})
-					}
-					continue
-				}
-			}
-			time.Sleep(time.Second * 10)
-		}
-	}()
-}

api/service/mj/pool.go

@@ -187,6 +187,14 @@ func (p *ServicePool) SyncTaskProgress() {
 			}
 
 			for _, job := range jobs {
+				// 5 分钟还没完成的任务标记为失败
+				if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
+					job.Progress = 101
+					job.ErrMsg = "任务超时"
+					p.db.Updates(&job)
+					continue
+				}
+
 				if servicePlus := p.getService(job.ChannelId); servicePlus != nil {
 					_ = servicePlus.Notify(job)
 				}

api/service/sd/pool.go

@@ -109,27 +109,11 @@ func (p *ServicePool) CheckTaskStatus() {
 			}
 
 			for _, job := range jobs {
-				// 5 分钟还没完成的任务直接删除
-				if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
-					p.db.Delete(&job)
-					var user model.User
-					p.db.Where("id = ?", job.UserId).First(&user)
-					// 退回绘图次数
-					res = p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
-					if res.Error == nil && res.RowsAffected > 0 {
-						p.db.Create(&model.PowerLog{
-							UserId:    user.Id,
-							Username:  user.Username,
-							Type:      types.PowerConsume,
-							Amount:    job.Power,
-							Balance:   user.Power + job.Power,
-							Mark:      types.PowerAdd,
-							Model:     "stable-diffusion",
-							Remark:    fmt.Sprintf("任务失败，退回算力。任务ID：%s", job.TaskId),
-							CreatedAt: time.Now(),
-						})
-					}
-					continue
+				// 5 分钟还没完成的任务标记为失败
+				if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
+					job.Progress = 101
+					job.ErrMsg = "任务超时"
+					p.db.Updates(&job)
 				}
 			}
 			time.Sleep(time.Second * 5)

api/service/sd/service.go

@@ -87,7 +87,7 @@ func (s *Service) Run() {
 			logger.Error("绘画任务执行失败：", err.Error())
 			// update the task progress
 			s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
-				"progress": -1,
+				"progress": 101,
 				"err_msg":  err.Error(),
 			})
 			// 通知前端，任务失败