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

add sync lock for sub or add user's power

Browse Source
This commit is contained in:
RockYang
2024-09-03 12:09:36 +08:00
parent 2e5bd238b7
commit 65dde9e69d
16 changed files with 341 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
api/handler/dalle_handler.go
View File

@@ -11,32 +11,33 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/dalle"
"geekai/service/oss"
"geekai/store/model"
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gorilla/websocket"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/gorilla/websocket"
"gorm.io/gorm"
"net/http"
)
type DallJobHandler struct {
BaseHandler
redis *redis.Client
service *dalle.Service
uploader *oss.UploaderManager
redis *redis.Client
dallService *dalle.Service
uploader *oss.UploaderManager
userService *service.UserService
}
func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service, manager *oss.UploaderManager) *DallJobHandler {
func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service, manager *oss.UploaderManager, userService *service.UserService) *DallJobHandler {
return &DallJobHandler{
service: service,
uploader: manager,
dallService: service,
uploader: manager,
userService: userService,
BaseHandler: BaseHandler{
App: app,
DB: db,
@@ -61,14 +62,14 @@ func (h *DallJobHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.dallService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
go func() {
for {
_, msg, err := client.Receive()
if err != nil {
client.Close()
h.service.Clients.Delete(uint(userId))
h.dallService.Clients.Delete(uint(userId))
return
}
@@ -127,7 +128,7 @@ func (h *DallJobHandler) Image(c *gin.Context) {
return
}
h.service.PushTask(types.DallTask{
h.dallService.PushTask(types.DallTask{
JobId: job.Id,
UserId: uint(userId),
Prompt: data.Prompt,
@@ -137,7 +138,7 @@ func (h *DallJobHandler) Image(c *gin.Context) {
Power: job.Power,
})
client := h.service.Clients.Get(job.UserId)
client := h.dallService.Clients.Get(job.UserId)
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -233,26 +234,11 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err := h.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())