feat: replace http polling with webscoket notify in sd image page

api/core/app_server.go

@@ -155,6 +155,7 @@ func authorizeMiddleware(s *AppServer, client *redis.Client) gin.HandlerFunc {
 			c.Request.URL.Path == "/api/mj/notify" ||
 			c.Request.URL.Path == "/api/invite/hits" ||
 			c.Request.URL.Path == "/api/sd/jobs" ||
+			c.Request.URL.Path == "/api/sd/client" ||
 			strings.HasPrefix(c.Request.URL.Path, "/api/test") ||
 			strings.HasPrefix(c.Request.URL.Path, "/api/function/") ||
 			strings.HasPrefix(c.Request.URL.Path, "/api/sms/") ||

api/handler/sd_handler.go

@@ -156,6 +156,11 @@ func (h *SdJobHandler) Image(c *gin.Context) {
 		UserId:    userId,
 	})
 
+	client := h.pool.Clients.Get(uint(job.UserId))
+	if client != nil {
+		_ = client.Send([]byte("Task Updated"))
+	}
+
 	// update user's img calls
 	h.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
 
@@ -229,6 +234,7 @@ func (h *SdJobHandler) JobList(c *gin.Context) {
 func (h *SdJobHandler) Remove(c *gin.Context) {
 	var data struct {
 		Id     uint   `json:"id"`
+		UserId uint   `json:"user_id"`
 		ImgURL string `json:"img_url"`
 	}
 	if err := c.ShouldBindJSON(&data); err != nil {
@@ -249,6 +255,11 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
 		logger.Error("remove image failed: ", err)
 	}
 
+	client := h.pool.Clients.Get(data.UserId)
+	if client != nil {
+		_ = client.Send([]byte("Task Updated"))
+	}
+
 	resp.SUCCESS(c)
 }
 

api/main.go

@@ -247,6 +247,7 @@ func main() {
 		}),
 		fx.Invoke(func(s *core.AppServer, h *handler.SdJobHandler) {
 			group := s.Engine.Group("/api/sd")
+			group.Any("client", h.Client)
 			group.POST("image", h.Image)
 			group.GET("jobs", h.JobList)
 			group.POST("remove", h.Remove)

api/service/sd/pool.go

@@ -52,6 +52,26 @@ func (p *ServicePool) PushTask(task types.SdTask) {
 	p.taskQueue.RPush(task)
 }
 
+func (p *ServicePool) CheckTaskNotify() {
+	go func() {
+		for {
+			var userId uint
+			err := p.notifyQueue.LPop(&userId)
+			if err != nil {
+				continue
+			}
+			client := p.Clients.Get(userId)
+			if client == nil {
+				continue
+			}
+			err = client.Send([]byte("Task Updated"))
+			if err != nil {
+				continue
+			}
+		}
+	}()
+}
+
 // HasAvailableService check if it has available mj service in pool
 func (p *ServicePool) HasAvailableService() bool {
 	return len(p.services) > 0

api/service/sd/service.go

@@ -310,4 +310,7 @@ func (s *Service) callback(data CBReq) {
 		// restore img_calls
 		s.db.Model(&model.User{}).Where("id = ? AND img_calls > 0", data.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1))
 	}
+
+	// 发送更新状态信号
+	s.notifyQueue.RPush(data.UserId)
 }