xinmi/internal/workbench/jobs.go

package workbench

import (
	"context"
	"errors"
	"fmt"
	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"runtime/debug"
	"time"

	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

func requireIdle(tx *gorm.DB, tenant uint64, id string) error {
	var count int64
	e := tx.Model(&Job{}).Where("tenant_id = ? AND resource_id = ? AND status IN ?", tenant, id, []string{"queued", "running"}).Count(&count).Error
	if e != nil {
		return e
	}
	if count > 0 {
		return ErrConflict
	}
	return nil
}
func allowedAction(kind, action string) bool {
	actions := map[string][]string{"training-maps": {"generate", "revise", "generate-day", "confirm-day", "confirm"}, "content-tasks": {"decompose", "import-image", "generate", "revise", "accept", "skip", "review", "review-settings", "compose"}}
	for _, a := range actions[kind] {
		if a == action {
			return true
		}
	}
	return false
}
func (s *Service) submit(ctx context.Context, kind, id, action string, request ActionRequest) (*Job, error) {
	if !allowedAction(kind, action) {
		return nil, ErrInvalid
	}
	if !accesspolicy.Allowed(ctx, accesspolicy.WorkbenchActionPermission(kind, action)) {
		return nil, ErrForbidden
	}
	if len([]rune(request.Advice)) > 6000 || len([]rune(request.Quote)) > 6000 {
		return nil, ErrInvalid
	}
	if action == "revise" && request.Advice == "" {
		return nil, fmt.Errorf("%w: 请填写修改建议", ErrInvalid)
	}
	if action == "review-settings" {
		if request.QualitySettings == nil {
			return nil, fmt.Errorf("%w: 请填写质检设置", ErrInvalid)
		}
		if err := request.QualitySettings.validate(); err != nil {
			return nil, err
		}
	}
	t, u, _, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	j := &Job{ID: uuid.NewString(), TenantID: t, OwnerID: u, ResourceID: id, Kind: kind, Action: action, Status: "queued", Stage: "等待执行", Request: asJSON(request), CreatedAt: time.Now(), UpdatedAt: time.Now()}
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		current, err := lockResource(ctx, tx, kind, id)
		if err != nil {
			return err
		}
		// Generation is billed to its owner; administrator inspection and manual
		// edits must not silently enqueue billable work for somebody else.
		if current.OwnerID != u {
			return ErrForbidden
		}
		if current.Version != request.Version {
			return ErrConflict
		}
		if err = requireActionExecutable(current, action); err != nil {
			return err
		}
		if kind == "content-tasks" {
			if err = s.validateTaskAction(ctx, current, action); err != nil {
				return err
			}
		} else if err = s.validate(ctx, current); err != nil {
			return err
		}
		if err = requireIdle(tx, t, id); err != nil {
			return err
		}
		return tx.Create(j).Error
	})
	if e != nil {
		return nil, e
	}
	_ = s.event(ctx, id, action, "提交任务："+actionName(action), "info")
	return j, nil
}

// Start uses database leases so multiple service instances cannot claim the same
// job. Expired jobs require an explicit retry: a provider request may have succeeded.
func (s *Service) Start(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(2 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				s.tick(ctx)
			}
		}
	}()
}
func (s *Service) tick(ctx context.Context) {
	if err := s.releaseExpiredEvidence(ctx); err != nil {
		logJobFailure(ctx, nil, "evidence_recovery", err)
	}
	if err := s.releaseExpiredBalances(ctx); err != nil {
		logJobFailure(ctx, nil, "balance_recovery", err)
	}
	now := time.Now()
	if err := s.recoverExpiredJobs(ctx, now); err != nil {
		logJobFailure(ctx, nil, "lease_recovery", err)
		return
	}
	j, err := s.claim(ctx)
	if err != nil {
		logJobFailure(ctx, nil, "claim", err)
		return
	}
	if j == nil {
		return
	}
	s.run(ctx, j)
}

func (s *Service) claim(ctx context.Context) (*Job, error) {
	var j Job
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		q := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).Where("status = ?", "queued").Order("created_at ASC, id ASC").Limit(1).Find(&j)
		if q.Error != nil || q.RowsAffected == 0 {
			return q.Error
		}
		now := time.Now()
		until := now.Add(10 * time.Minute)
		j.Status = "running"
		j.Progress = 5
		j.Stage = "检查任务配置"
		j.LeaseToken = uuid.NewString()
		j.LeaseUntil = &until
		q = tx.Model(&Job{}).Where("id = ? AND status = ?", j.ID, "queued").Updates(map[string]any{"status": j.Status, "progress": j.Progress, "stage": j.Stage, "lease_token": j.LeaseToken, "lease_until": until, "updated_at": now})
		if q.Error != nil {
			return q.Error
		}
		if q.RowsAffected != 1 {
			return ErrConflict
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	if j.ID == "" {
		return nil, nil
	}
	return &j, nil
}

type jobLease struct{ ID, Token string }
type leaseContextKey struct{}

func (s *Service) progress(ctx context.Context, id string, percent int, stage string) error {
	lease, ok := ctx.Value(leaseContextKey{}).(jobLease)
	if !ok || lease.ID != id {
		return errors.New("missing job lease")
	}
	q := s.db.WithContext(ctx).Model(&Job{}).Where("id = ? AND lease_token = ? AND status = ? AND lease_until > ?", id, lease.Token, "running", time.Now()).Updates(map[string]any{"progress": percent, "stage": stage, "updated_at": time.Now()})
	if q.Error != nil {
		return q.Error
	}
	if q.RowsAffected != 1 {
		return ErrConflict
	}
	return nil
}
func (s *Service) run(parent context.Context, j *Job) {
	ctx, cancel := context.WithTimeout(parent, 8*time.Minute)
	defer cancel()
	ctx = context.WithValue(ctx, types.TenantIDContextKey, j.TenantID)
	ctx = context.WithValue(ctx, types.UserIDContextKey, j.OwnerID)
	ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{j.ID, j.LeaseToken})
	ctx = types.WithBackgroundTask(ctx)
	go s.watchJobLease(ctx, cancel, j)
	var runErr error
	defer func() {
		if panicValue := recover(); panicValue != nil {
			runErr = errors.New("后台任务发生异常，请重试")
			logger.GetLogger(ctx).WithFields(logger.Fields{"job_id": j.ID, "panic_type": fmt.Sprintf("%T", panicValue), "stack": string(debug.Stack())}).Error("工作台任务发生异常")
		}
		finishCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
		defer stop()
		status, stage, message := "succeeded", "已完成", ""
		updates := map[string]any{"status": status, "stage": stage, "progress": 100, "error": message, "updated_at": time.Now()}
		if runErr != nil {
			if err := s.releaseJobEvidence(finishCtx, j); err != nil {
				logJobFailure(finishCtx, j, "evidence_release", err)
			}
			status = "failed"
			stage = "执行失败"
			message = publicJobError(runErr)
			logJobFailure(ctx, j, "execute", runErr)
			updates["status"], updates["stage"], updates["error"] = status, stage, message
			// Preserve the latest persisted stage progress instead of resetting to 5%.
			delete(updates, "progress")
		}
		q := s.db.WithContext(finishCtx).Model(&Job{}).Where("id = ? AND lease_token = ? AND status = ? AND lease_until > ?", j.ID, j.LeaseToken, "running", time.Now()).Updates(updates)
		if q.Error != nil {
			logJobFailure(finishCtx, j, "finalize", q.Error)
			s.recordJobFailure(finishCtx, j, "finalize", q.Error)
			return
		}
		if q.RowsAffected != 1 {
			var current Job
			if err := s.db.WithContext(finishCtx).Where("id=?", j.ID).Take(&current).Error; err == nil && (current.Status == "paused" || current.Status == "cancelled") {
				return
			}
			logJobFailure(finishCtx, j, "finalize_lease", ErrConflict)
			return
		}
		level := "info"
		if runErr != nil {
			level = "error"
			s.recordJobFailure(finishCtx, j, "execute", runErr)
		}
		if err := s.event(finishCtx, j.ResourceID, j.Action, stage+"："+message, level); err != nil {
			logJobFailure(finishCtx, j, "event", err)
		}
	}()
	// Re-read membership on execution; queued work does not preserve revoked roles.
	var member types.TenantMember
	runErr = s.db.WithContext(ctx).Where("user_id = ? AND tenant_id = ? AND status = ?", j.OwnerID, j.TenantID, types.TenantMemberStatusActive).First(&member).Error
	if errors.Is(runErr, gorm.ErrRecordNotFound) || (runErr == nil && !member.Role.HasPermission(types.TenantRoleContributor)) {
		runErr = ErrForbidden
		return
	}
	if runErr != nil {
		return
	}
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, member.Role)
	ctx, runErr = s.restoreJobTenant(ctx, j.TenantID)
	if runErr != nil {
		return
	}
	if s.accessSnapshot != nil {
		var snapshot *accesspolicy.Snapshot
		snapshot, runErr = s.accessSnapshot(ctx)
		if runErr != nil {
			return
		}
		ctx = accesspolicy.WithSnapshot(ctx, snapshot)
		if !accesspolicy.Allowed(ctx, accesspolicy.WorkbenchActionPermission(j.Kind, j.Action)) {
			runErr = ErrForbidden
			return
		}
	}
	r, e := s.get(ctx, j.Kind, j.ResourceID)
	if e != nil {
		runErr = e
		return
	}
	request, e := decode[ActionRequest](j.Request)
	if e != nil {
		runErr = e
		return
	}
	if r.Version != request.Version {
		runErr = ErrConflict
		return
	}
	if runErr = requireActionExecutable(r, j.Action); runErr != nil {
		return
	}
	if r.Kind == "content-tasks" {
		if runErr = s.validateTaskAction(ctx, r, j.Action); runErr != nil {
			return
		}
	} else if runErr = s.validate(ctx, r); runErr != nil {
		return
	}
	switch r.Kind {
	case "training-maps":
		runErr = s.executeMap(ctx, j, r, request)
	case "content-tasks":
		runErr = s.executeTask(ctx, j, r, request)
	default:
		runErr = ErrInvalid
	}
	if runErr == nil {
		runErr = s.save(ctx, r, request.Version)
	}
}
func publicJobError(e error) string {
	_, message := jobErrorDetails(e)
	return message
}