xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-09 10:43:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add websocket heartbeat message for mj page

Browse Source
This commit is contained in:
RockYang
2024-01-24 09:33:04 +08:00
parent 855a953010
commit 04e990aeb3
4 changed files with 49 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
api/service/sd/pool.go
View File

@@ -11,13 +11,16 @@ import (
)
type ServicePool struct {
services []*Service
taskQueue *store.RedisQueue
services []*Service
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
}
func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig) *ServicePool {
services := make([]*Service, 0)
queue := store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli)
taskQueue := store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli)
notifyQueue := store.NewRedisQueue("StableDiffusion_Queue", redisCli)
// create mj client and service
for k, config := range appConfig.SdConfigs {
if config.Enabled == false {
@@ -26,7 +29,7 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
// create sd service
name := fmt.Sprintf("StableDifffusion Service-%d", k)
service := NewService(name, 1, 300, config, queue, db, manager)
service := NewService(name, 1, 300, config, taskQueue, notifyQueue, db, manager)
// run sd service
go func() {
service.Run()
@@ -36,8 +39,10 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
}
return &ServicePool{
taskQueue: queue,
services: services,
taskQueue: taskQueue,
notifyQueue: notifyQueue,
services: services,
Clients: types.NewLMap[uint, *types.WsClient](),
}
}

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

@@ -24,6 +24,7 @@ type Service struct {
httpClient *req.Client
config types.StableDiffusionConfig
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
uploadManager *oss.UploaderManager
name string // service name
@@ -33,12 +34,13 @@ type Service struct {
taskTimeout int64
}
func NewService(name string, maxTaskNum int32, timeout int64, config types.StableDiffusionConfig, queue *store.RedisQueue, db *gorm.DB, manager *oss.UploaderManager) *Service {
func NewService(name string, maxTaskNum int32, timeout int64, config types.StableDiffusionConfig, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, manager *oss.UploaderManager) *Service {
return &Service{
name: name,
config: config,
httpClient: req.C(),
taskQueue: queue,
taskQueue: taskQueue,
notifyQueue: notifyQueue,
db: db,
uploadManager: manager,
taskTimeout: timeout,
@@ -73,6 +75,8 @@ func (s *Service) Run() {
s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1))
// release task num
atomic.AddInt32(&s.handledTaskNum, -1)
// 通知前端，任务失败
s.notifyQueue.RPush(task.UserId)
continue
}