xinmi/internal/workbench/job_control.go

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
			}
		}
	}
}