smartmate/backend/memory/repo/job_repo.go

package repo

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	memorymodel "github.com/LoveLosita/smartflow/backend/memory/model"
	"github.com/LoveLosita/smartflow/backend/model"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// JobRepo 封装 memory_jobs 的数据访问。
type JobRepo struct {
	db *gorm.DB
}

func NewJobRepo(db *gorm.DB) *JobRepo {
	return &JobRepo{db: db}
}

func (r *JobRepo) WithTx(tx *gorm.DB) *JobRepo {
	return &JobRepo{db: tx}
}

// CreatePendingExtractJob 创建“待抽取”任务（幂等写入）。
//
// 失败语义：
// 1. 参数非法直接返回 error，由上游决定 dead 或重试；
// 2. 同幂等键重复写入采用 DoNothing，保证无副作用。
func (r *JobRepo) CreatePendingExtractJob(
	ctx context.Context,
	payload memorymodel.ExtractJobPayload,
	sourceEventID string,
) error {
	if r == nil || r.db == nil {
		return errors.New("memory job repo is nil")
	}
	if payload.UserID <= 0 {
		return errors.New("invalid user_id")
	}
	if payload.IdempotencyKey == "" {
		return errors.New("idempotency_key is empty")
	}

	rawPayload, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	now := time.Now()
	job := model.MemoryJob{
		UserID:          payload.UserID,
		ConversationID:  strPtrOrNil(payload.ConversationID),
		SourceMessageID: int64PtrOrNil(payload.SourceMessageID),
		SourceEventID:   strPtrOrNil(sourceEventID),
		JobType:         model.MemoryJobTypeExtract,
		IdempotencyKey:  payload.IdempotencyKey,
		PayloadJSON:     string(rawPayload),
		Status:          model.MemoryJobStatusPending,
		RetryCount:      0,
		MaxRetry:        6,
		NextRetryAt:     &now,
	}

	return r.db.WithContext(ctx).
		Clauses(clause.OnConflict{
			Columns:   []clause.Column{{Name: "idempotency_key"}},
			DoNothing: true,
		}).
		Create(&job).Error
}

// ClaimNextRunnableExtractJob 抢占一个可执行的 extract 任务。
//
// 抢占规则：
// 1. 只从 pending/failed 中挑 next_retry_at 已到期任务；
// 2. 用行锁避免多个 worker 抢到同一条任务；
// 3. 抢占成功后立即置为 processing，防止重复执行。
func (r *JobRepo) ClaimNextRunnableExtractJob(ctx context.Context, now time.Time) (*model.MemoryJob, error) {
	if r == nil || r.db == nil {
		return nil, errors.New("memory job repo is nil")
	}

	var claimed *model.MemoryJob
	err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var job model.MemoryJob
		query := tx.
			Clauses(clause.Locking{Strength: "UPDATE"}).
			Where("job_type = ?", model.MemoryJobTypeExtract).
			Where("status IN ?", []string{model.MemoryJobStatusPending, model.MemoryJobStatusFailed}).
			Where("(next_retry_at IS NULL OR next_retry_at <= ?)", now).
			Order("id ASC").
			Limit(1).
			Find(&job)
		if query.Error != nil {
			return query.Error
		}
		if query.RowsAffected == 0 {
			return nil
		}

		updates := map[string]any{
			"status":     model.MemoryJobStatusProcessing,
			"updated_at": now,
			"last_error": nil,
		}
		if updateErr := tx.Model(&model.MemoryJob{}).Where("id = ?", job.ID).Updates(updates).Error; updateErr != nil {
			return updateErr
		}

		job.Status = model.MemoryJobStatusProcessing
		job.UpdatedAt = &now
		claimed = &job
		return nil
	})
	if err != nil {
		return nil, err
	}
	return claimed, nil
}

// MarkSuccess 把任务推进为 success 最终态。
func (r *JobRepo) MarkSuccess(ctx context.Context, jobID int64) error {
	if r == nil || r.db == nil {
		return errors.New("memory job repo is nil")
	}
	now := time.Now()
	updates := map[string]any{
		"status":        model.MemoryJobStatusSuccess,
		"last_error":    nil,
		"next_retry_at": nil,
		"updated_at":    now,
	}
	return r.db.WithContext(ctx).Model(&model.MemoryJob{}).Where("id = ?", jobID).Updates(updates).Error
}

// MarkFailed 按重试策略推进任务到 failed/dead。
//
// 规则：
// 1. retry_count +1 后若超上限，直接 dead；
// 2. 未超上限则写 failed 并设置 next_retry_at。
func (r *JobRepo) MarkFailed(ctx context.Context, jobID int64, reason string) error {
	if r == nil || r.db == nil {
		return errors.New("memory job repo is nil")
	}

	return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var job model.MemoryJob
		queryErr := tx.
			Clauses(clause.Locking{Strength: "UPDATE"}).
			Where("id = ?", jobID).
			First(&job).Error
		if queryErr != nil {
			return queryErr
		}
		if job.Status == model.MemoryJobStatusSuccess || job.Status == model.MemoryJobStatusDead {
			return nil
		}

		maxRetry := job.MaxRetry
		if maxRetry <= 0 {
			maxRetry = 6
		}
		nextRetryCount := job.RetryCount + 1
		now := time.Now()
		status := model.MemoryJobStatusFailed
		var nextRetryAt *time.Time
		if nextRetryCount >= maxRetry {
			status = model.MemoryJobStatusDead
			nextRetryAt = nil
		} else {
			t := now.Add(calcRetryBackoff(nextRetryCount))
			nextRetryAt = &t
		}

		lastErr := truncateError(reason)
		updates := map[string]any{
			"status":        status,
			"retry_count":   nextRetryCount,
			"last_error":    &lastErr,
			"next_retry_at": nextRetryAt,
			"updated_at":    now,
		}
		return tx.Model(&model.MemoryJob{}).Where("id = ?", jobID).Updates(updates).Error
	})
}

func calcRetryBackoff(retryCount int) time.Duration {
	if retryCount <= 0 {
		return time.Second
	}
	if retryCount > 6 {
		retryCount = 6
	}
	return time.Second * time.Duration(1<<(retryCount-1))
}

func truncateError(reason string) string {
	if len(reason) <= 2000 {
		return reason
	}
	return reason[:2000]
}

func strPtrOrNil(v string) *string {
	if v == "" {
		return nil
	}
	value := v
	return &value
}

func int64PtrOrNil(v int64) *int64 {
	if v <= 0 {
		return nil
	}
	value := v
	return &value
}