xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-12-03 06:46:07 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

重构异步任务更新方式，使用 Http 替代 websocket

Browse Source
This commit is contained in:
RockYang
2025-03-03 19:00:10 +08:00
parent 8369e18bf0
commit ed063a1d9d
22 changed files with 308 additions and 541 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
api/service/dalle/service.go
View File

@@ -34,10 +34,8 @@ type Service struct {
db *gorm.DB
uploadManager *oss.UploaderManager
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
userService *service.UserService
wsService *service.WebsocketService
clientIds map[uint]string
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService, wsService *service.WebsocketService) *Service {
@@ -45,11 +43,9 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
taskQueue: store.NewRedisQueue("DallE_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("DallE_Notify_Queue", redisCli),
wsService: wsService,
uploadManager: manager,
userService: userService,
clientIds: map[uint]string{},
}
}
@@ -60,7 +56,7 @@ func (s *Service) PushTask(task types.DallTask) {
}
func (s *Service) Run() {
// 将数据库中未提交的人物加载到队列
// 将数据库中未提交的任务加载到队列
var jobs []model.DallJob
s.db.Where("progress", 0).Find(&jobs)
for _, v := range jobs {
@@ -84,16 +80,16 @@ func (s *Service) Run() {
continue
}
logger.Infof("handle a new DALL-E task: %+v", task)
s.clientIds[task.Id] = task.ClientId
_, err = s.Image(task, false)
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
}
go func() {
_, err = s.Image(task, false)
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
}
}()
}
}()
}
@@ -212,10 +208,9 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return "", fmt.Errorf("err with update database: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
var content string
if sync {
imgURL, err := s.downloadImage(task.Id, int(task.UserId), res.Data[0].Url)
imgURL, err := s.downloadImage(task.Id, res.Data[0].Url)
if err != nil {
return "", fmt.Errorf("error with download image: %v", err)
}
@@ -225,26 +220,6 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return content, nil
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running DALL-E task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChDall, message.Message)
}
}()
}
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running DALL-E task status checking ...")
@@ -254,7 +229,7 @@ func (s *Service) CheckTaskStatus() {
s.db.Where("progress < ?", 100).Find(&jobs)
for _, job := range jobs {
// 超时的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
if time.Since(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
@@ -301,7 +276,7 @@ func (s *Service) DownloadImages() {
}
logger.Infof("try to download image: %s", v.OrgURL)
imgURL, err := s.downloadImage(v.Id, int(v.UserId), v.OrgURL)
imgURL, err := s.downloadImage(v.Id, v.OrgURL)
if err != nil {
logger.Error("error with download image: %s, error: %v", imgURL, err)
continue
@@ -316,7 +291,7 @@ func (s *Service) DownloadImages() {
}()
}
func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string, error) {
func (s *Service) downloadImage(jobId uint, orgURL string) (string, error) {
// sava image
imgURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(orgURL, false)
if err != nil {
@@ -328,6 +303,5 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
if res.Error != nil {
return "", err
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[jobId], UserId: userId, JobId: int(jobId), Message: service.TaskStatusFinished})
return imgURL, nil
}

55
api/service/mj/service.go
View File

@@ -15,10 +15,11 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"strings"
"time"
"github.com/go-redis/redis/v8"
"gorm.io/gorm"
)
@@ -26,23 +27,19 @@ import (
type Service struct {
client *Client // MJ Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
wsService *service.WebsocketService
uploaderManager *oss.UploaderManager
userService *service.UserService
clientIds map[uint]string
}
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
db: db,
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("MidJourney_Notify_Queue", redisCli),
client: client,
wsService: wsService,
uploaderManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -59,7 +56,6 @@ func (s *Service) Run() {
continue
}
task.Id = v.Id
s.clientIds[task.Id] = task.ClientId
s.PushTask(task)
}
@@ -96,7 +92,6 @@ func (s *Service) Run() {
if task.Mode == "" {
task.Mode = "fast"
}
s.clientIds[task.Id] = task.ClientId
var job model.MidJourneyJob
tx := s.db.Where("id = ?", task.Id).First(&job)
@@ -139,7 +134,6 @@ func (s *Service) Run() {
// update the task progress
s.db.Updates(&job)
// 任务失败，通知前端
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
continue
}
logger.Infof("任务提交成功：%+v", res)
@@ -178,24 +172,6 @@ func GetImageHash(action string) string {
return split[len(split)-1]
}
func (s *Service) CheckTaskNotify() {
go func() {
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("receive a new mj notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChMj, message.Message)
}
}()
}
func (s *Service) DownloadImages() {
go func() {
var items []model.MidJourneyJob
@@ -228,12 +204,6 @@ func (s *Service) DownloadImages() {
v.ImgURL = imgURL
s.db.Updates(&v)
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[v.Id],
UserId: v.UserId,
JobId: int(v.Id),
Message: service.TaskStatusFinished})
}
time.Sleep(time.Second * 5)
@@ -259,7 +229,7 @@ func (s *Service) SyncTaskProgress() {
for _, job := range jobs {
// 10 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
if time.Since(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
@@ -279,18 +249,12 @@ func (s *Service) SyncTaskProgress() {
"err_msg": task.FailReason,
})
logger.Errorf("task failed: %v", task.FailReason)
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[job.Id],
UserId: job.UserId,
JobId: int(job.Id),
Message: service.TaskStatusFailed})
continue
}
if len(task.Buttons) > 0 {
job.Hash = GetImageHash(task.Buttons[0].CustomId)
}
oldProgress := job.Progress
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
if task.ImageUrl != "" {
job.OrgURL = task.ImageUrl
@@ -300,19 +264,6 @@ func (s *Service) SyncTaskProgress() {
logger.Errorf("error with update database: %v", err)
continue
}
// 通知前端更新任务进度
if oldProgress != job.Progress {
message := service.TaskStatusRunning
if job.Progress == 100 {
message = service.TaskStatusFinished
}
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[job.Id],
UserId: job.UserId,
JobId: int(job.Id),
Message: message})
}
}
// 找出失败的任务，并恢复其扣减算力

41
api/service/sd/service.go
View File

@@ -16,9 +16,10 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"time"
"github.com/go-redis/redis/v8"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
@@ -30,7 +31,6 @@ var logger = logger2.GetLogger()
type Service struct {
httpClient *req.Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
uploadManager *oss.UploaderManager
wsService *service.WebsocketService
@@ -41,7 +41,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelD
return &Service{
httpClient: req.C(),
taskQueue: store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("StableDiffusion_Queue", redisCli),
db: db,
wsService: wsService,
uploadManager: manager,
@@ -102,8 +101,6 @@ func (s *Service) Run() {
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
// 通知前端，任务失败
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFailed})
continue
}
}
@@ -225,15 +222,12 @@ func (s *Service) Txt2Img(task types.SdTask) error {
// task finished
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFinished})
return nil
default:
err, resp := s.checkTaskProgress(apiKey)
resp, err := s.checkTaskProgress(apiKey)
// 更新任务进度
if err == nil && resp.Progress > 0 {
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
// 发送更新状态信号
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusRunning})
}
time.Sleep(time.Second)
}
@@ -242,7 +236,7 @@ func (s *Service) Txt2Img(task types.SdTask) error {
}
// 执行任务
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressResp) {
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (*TaskProgressResp, error) {
apiURL := fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", apiKey.ApiURL)
var res TaskProgressResp
response, err := s.httpClient.R().
@@ -250,13 +244,13 @@ func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressRe
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
return err, nil
return nil, err
}
if response.IsErrorState() {
return fmt.Errorf("error http code status: %v", response.Status), nil
return nil, fmt.Errorf("error http code status: %v", response.Status)
}
return nil, &res
return &res, nil
}
func (s *Service) PushTask(task types.SdTask) {
@@ -264,25 +258,6 @@ func (s *Service) PushTask(task types.SdTask) {
s.taskQueue.RPush(task)
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Stable-Diffusion task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChSd, message.Message)
}
}()
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
go func() {
@@ -297,7 +272,7 @@ func (s *Service) CheckTaskStatus() {
for _, job := range jobs {
// 5 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
if time.Since(job.CreatedAt) > time.Minute*5 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)

32
api/service/suno/service.go
View File

@@ -18,10 +18,11 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"io"
"time"
"github.com/go-redis/redis/v8"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
@@ -35,7 +36,6 @@ type Service struct {
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[string]string
userService *service.UserService
}
@@ -47,7 +47,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
notifyQueue: store.NewRedisQueue("Suno_Notify_Queue", redisCli),
uploadManager: manager,
wsService: wsService,
clientIds: map[string]string{},
userService: userService,
}
}
@@ -70,7 +69,6 @@ func (s *Service) Run() {
}
task.Id = v.Id
s.PushTask(task)
s.clientIds[v.TaskId] = task.ClientId
}
logger.Info("Starting Suno job consumer...")
go func() {
@@ -95,7 +93,6 @@ func (s *Service) Run() {
"err_msg": err.Error(),
"progress": service.FailTaskProgress,
})
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed})
continue
}
@@ -104,7 +101,6 @@ func (s *Service) Run() {
"task_id": r.Data,
"channel": r.Channel,
})
s.clientIds[r.Data] = task.ClientId
}
}()
}
@@ -262,27 +258,6 @@ func (s *Service) Upload(task types.SunoTask) (RespVo, error) {
return res, nil
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Suno task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
logger.Debugf("client id: %+v", s.wsService.Clients)
client := s.wsService.Clients.Get(message.ClientId)
logger.Debugf("%+v", client)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChSuno, message.Message)
}
}()
}
func (s *Service) DownloadFiles() {
go func() {
var items []model.SunoJob
@@ -311,7 +286,6 @@ func (s *Service) DownloadFiles() {
v.AudioURL = audioURL
v.Progress = 100
s.db.Updates(&v)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.TaskId], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished})
}
time.Sleep(time.Second * 10)
@@ -377,12 +351,10 @@ func (s *Service) SyncTaskProgress() {
}
}
tx.Commit()
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFinished})
} else if task.Data.FailReason != "" {
job.Progress = service.FailTaskProgress
job.ErrMsg = task.Data.FailReason
s.db.Updates(&job)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
}
}

34
api/service/video/video.go
View File

@@ -20,7 +20,6 @@ import (
"geekai/store/model"
"geekai/utils"
"io"
"io/ioutil"
"net/http"
"time"
@@ -37,9 +36,7 @@ type Service struct {
db *gorm.DB
uploadManager *oss.UploaderManager
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[uint]string
userService *service.UserService
}
@@ -48,10 +45,8 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
taskQueue: store.NewRedisQueue("Video_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("Video_Notify_Queue", redisCli),
wsService: wsService,
uploadManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -74,7 +69,6 @@ func (s *Service) Run() {
}
task.Id = v.Id
s.PushTask(task)
s.clientIds[v.Id] = task.ClientId
}
logger.Info("Starting Video job consumer...")
go func() {
@@ -86,10 +80,6 @@ func (s *Service) Run() {
continue
}
if task.ClientId != "" {
s.clientIds[task.Id] = task.ClientId
}
if task.Type == types.VideoLuma {
// translate prompt
if utils.HasChinese(task.Prompt) {
@@ -112,7 +102,6 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("update task with error: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed, Type: types.VideoLuma})
continue
}
@@ -150,7 +139,6 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("update task with error: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed, Type: types.VideoKeLing})
continue
}
@@ -170,25 +158,6 @@ func (s *Service) Run() {
}()
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Suno task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("Receive notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChLuma, message.Message)
}
}()
}
func (s *Service) DownloadFiles() {
go func() {
var items []model.VideoJob
@@ -232,7 +201,6 @@ func (s *Service) DownloadFiles() {
continue
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.Id], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished, Type: videoTask.Type})
}
time.Sleep(time.Second * 10)
@@ -672,7 +640,7 @@ func (s *Service) QueryKeLingTask(taskId string, channel string, action string)
return VideoCallbackData{}, fmt.Errorf("unexpected status code: %d", res.StatusCode)
}
body, err := ioutil.ReadAll(res.Body)
body, err := io.ReadAll(res.Body)
if err != nil {
return VideoCallbackData{}, fmt.Errorf("failed to read response body: %w", err)
}