package workbench import ( "context" "fmt" "time" "github.com/Tencent/WeKnora/internal/accesspolicy" "github.com/gin-gonic/gin" "github.com/google/uuid" "gorm.io/gorm" "gorm.io/gorm/clause" ) // Lock the resource before its job, matching the generation commit order. // Revoking a lease prevents a late model response from replacing newer content. func (s *Service) controlJob(ctx context.Context, id, action string) (*Job, error) { tenant, user, _, err := identity(ctx) if err != nil { return nil, err } if action != "pause" && action != "cancel" && action != "resume" && action != "retry" { return nil, ErrInvalid } var job Job if err = s.db.WithContext(ctx).Where("id=? AND tenant_id=? AND owner_id=?", id, tenant, user).Take(&job).Error; err != nil { return nil, err } err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { resource, err := lockResource(ctx, tx, job.Kind, job.ResourceID) if err != nil { return err } if err = tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND owner_id=?", id, tenant, user).Take(&job).Error; err != nil { return err } if !accesspolicy.Allowed(ctx, accesspolicy.WorkbenchActionPermission(job.Kind, job.Action)) { return ErrForbidden } request, err := decode[ActionRequest](job.Request) if err != nil { return ErrInvalid } if resource.Version != request.Version { return fmt.Errorf("%w: 内容已有新版本,请从当前内容重新发起操作", ErrConflict) } if action == "retry" || action == "resume" { if (action == "resume" && job.Status != "paused") || (action == "retry" && job.Status != "failed" && job.Status != "cancelled") { return ErrConflict } if err = requireExecutable(resource); err != nil { return err } if err = requireIdle(tx, tenant, resource.ID); err != nil { return err } job = Job{ID: uuid.NewString(), TenantID: tenant, OwnerID: user, ResourceID: resource.ID, Kind: resource.Kind, Action: job.Action, Status: "queued", Stage: "等待执行", Request: job.Request, CreatedAt: time.Now(), UpdatedAt: time.Now()} return tx.Create(&job).Error } if job.Status != "queued" && job.Status != "running" && !(action == "cancel" && job.Status == "paused") { return ErrConflict } status, stage := "cancelled", "已取消" if action == "pause" { status, stage = "paused", "已暂停,继续时重新执行本步" } if err = tx.Model(&job).Updates(map[string]any{"status": status, "stage": stage, "lease_token": "", "lease_until": nil, "updated_at": time.Now()}).Error; err != nil { return err } return finishConsumption(tx, &job, false) }) if err != nil { return nil, err } _ = s.event(ctx, job.ResourceID, action, actionName(action), "info") return &job, nil } func (h *Handler) ControlJob(c *gin.Context) { job, err := h.service.controlJob(c.Request.Context(), c.Param("id"), c.Param("action")) if err != nil { respondError(c, err) return } result(c, job) } func (s *Service) watchJobLease(ctx context.Context, cancel context.CancelFunc, job *Job) { ticker := time.NewTicker(time.Second) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: var count int64 err := s.db.WithContext(ctx).Model(&Job{}).Where("id=? AND lease_token=? AND status=? AND lease_until>?", job.ID, job.LeaseToken, "running", time.Now()).Count(&count).Error if err == nil && count == 0 { cancel() return } } } }