cmsp/internal/task/task.go

// Package task 提供批量取视频任务的队列、进度和停止状态机。
package task

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"time"
)

const (
	StateRunning       = "running"
	StateStopped       = "stopped"
	StateCompleted     = "completed"
	StateLoginRequired = "login_required"
)

// ErrLoginRequired 是全局停止门。Prepare 返回此错误后，Runner 不再启动后续商品。
var ErrLoginRequired = errors.New("淘宝登录已失效")

// Progress 是批量任务对界面公开的只读进度。
type Progress struct {
	Total       int    `json:"total"`
	Done        int    `json:"done"`
	Current     string `json:"current"`
	CurrentName string `json:"currentName"`
	Downloaded  int    `json:"downloaded"`
	Failed      int    `json:"failed"`
	Skipped     int    `json:"skipped"`
	ElapsedSec  int    `json:"elapsedSec"`
	State       string `json:"state"`
}

// Product 是 Runner 判断断点与显示当前商品所需的最小信息。
type Product struct {
	ID             string
	ItemID         string
	Name           string
	DownloadStatus string
	VideoStatus    string
}

// DownloadFunc 是一段完全不接触 CDP 的纯下载工作。
type DownloadFunc func(context.Context) error

// Work 是一个商品在串行准备完成后交给并行下载阶段的工作。
// Finalize 在全部下载结束后调用，用来集中写回商品和视频状态。
type Work struct {
	Skipped   bool
	Downloads []DownloadFunc
	Finalize  func([]error) error
}

type LoadFunc func(context.Context, string) (Product, error)
type PrepareFunc func(context.Context, Product) (Work, error)
type ProgressFunc func(Progress)

type Options struct {
	Load        LoadFunc
	Prepare     PrepareFunc
	Concurrency int
	WaitMin     time.Duration
	WaitMax     time.Duration
	OnProgress  ProgressFunc
	OnFinished  ProgressFunc
}

// Runner 管理一次批量任务。同一时间只允许一个任务在跑。
type Runner struct {
	mu       sync.Mutex
	options  Options
	progress Progress
	started  time.Time
	active   bool
	stop     chan struct{}
	stopOnce *sync.Once
}

func NewRunner(options Options) *Runner {
	if options.Concurrency < 1 {
		options.Concurrency = 1
	}
	if options.WaitMin < 0 {
		options.WaitMin = 0
	}
	if options.WaitMax < options.WaitMin {
		options.WaitMax = options.WaitMin
	}
	return &Runner{options: options}
}

type queuedProduct struct {
	product Product
	err     error
}

// Start 启动后台任务。断点直接复用 SQLite 中的商品状态，不另建文件或表，
// 从而不会产生两份断点互相不一致的问题。
func (r *Runner) Start(ctx context.Context, productIDs []string) error {
	r.mu.Lock()
	if r.active {
		r.mu.Unlock()
		return fmt.Errorf("已有任务在运行")
	}
	if r.options.Load == nil || r.options.Prepare == nil {
		r.mu.Unlock()
		return fmt.Errorf("任务处理器未配置")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	r.active = true
	r.started = time.Now()
	r.stop = make(chan struct{})
	r.stopOnce = &sync.Once{}
	r.progress = Progress{State: StateRunning}
	r.mu.Unlock()

	queue := make([]queuedProduct, 0, len(productIDs))
	for _, rawID := range productIDs {
		id := strings.TrimSpace(rawID)
		if id == "" {
			queue = append(queue, queuedProduct{product: Product{ID: id}, err: fmt.Errorf("商品 ID 不能为空")})
			continue
		}
		product, err := r.options.Load(ctx, id)
		if product.ID == "" {
			product.ID = id
		}
		if err == nil && (product.DownloadStatus == "done" || product.VideoStatus == "none") {
			continue
		}
		queue = append(queue, queuedProduct{product: product, err: err})
	}
	r.mu.Lock()
	r.progress.Total = len(queue)
	r.mu.Unlock()

	go r.run(ctx, queue)
	return nil
}

// Stop 请求在当前商品的下载完成后停止，不会取消正在进行的下载。
func (r *Runner) Stop() {
	r.mu.Lock()
	if !r.active || r.stop == nil || r.stopOnce == nil {
		r.mu.Unlock()
		return
	}
	stop := r.stop
	once := r.stopOnce
	r.mu.Unlock()
	once.Do(func() { close(stop) })
}

func (r *Runner) Snapshot() Progress {
	r.mu.Lock()
	defer r.mu.Unlock()
	progress := r.progress
	if r.active {
		progress.ElapsedSec = int(time.Since(r.started).Seconds())
	}
	return progress
}

func (r *Runner) run(ctx context.Context, queue []queuedProduct) {
	for index, item := range queue {
		if r.stopRequested() || ctx.Err() != nil {
			r.finish(StateStopped)
			return
		}
		r.setCurrent(item.product)
		if item.err != nil {
			r.completeProduct(0, false, true)
		} else {
			work, err := r.options.Prepare(ctx, item.product)
			if errors.Is(err, ErrLoginRequired) {
				r.finish(StateLoginRequired)
				return
			}
			if err != nil {
				r.completeProduct(0, false, true)
			} else if work.Skipped {
				if work.Finalize != nil {
					err = work.Finalize(nil)
				}
				r.completeProduct(0, err == nil, err != nil)
			} else {
				downloaded, downloadErr := RunDownloads(ctx, r.options.Concurrency, work)
				r.completeProduct(downloaded, false, downloadErr != nil)
			}
		}

		if index+1 < len(queue) {
			if r.stopRequested() || !r.wait(ctx, randomDuration(r.options.WaitMin, r.options.WaitMax)) {
				r.finish(StateStopped)
				return
			}
		}
	}
	if r.stopRequested() || ctx.Err() != nil {
		r.finish(StateStopped)
		return
	}
	r.finish(StateCompleted)
}

func (r *Runner) setCurrent(product Product) {
	r.mu.Lock()
	r.progress.Current = product.ItemID
	if r.progress.Current == "" {
		r.progress.Current = product.ID
	}
	r.progress.CurrentName = product.Name
	r.mu.Unlock()
}

func (r *Runner) completeProduct(downloaded int, skipped, failed bool) {
	r.mu.Lock()
	r.progress.Done++
	r.progress.Downloaded += downloaded
	if skipped {
		r.progress.Skipped++
	}
	if failed {
		r.progress.Failed++
	}
	r.progress.ElapsedSec = int(time.Since(r.started).Seconds())
	progress := r.progress
	callback := r.options.OnProgress
	r.mu.Unlock()
	if callback != nil {
		callback(progress)
	}
}

func (r *Runner) finish(state string) {
	r.mu.Lock()
	r.progress.State = state
	r.progress.ElapsedSec = int(time.Since(r.started).Seconds())
	r.active = false
	progress := r.progress
	callback := r.options.OnFinished
	r.mu.Unlock()
	if callback != nil {
		callback(progress)
	}
}

func (r *Runner) stopRequested() bool {
	r.mu.Lock()
	stop := r.stop
	r.mu.Unlock()
	if stop == nil {
		return false
	}
	select {
	case <-stop:
		return true
	default:
		return false
	}
}

func (r *Runner) wait(ctx context.Context, duration time.Duration) bool {
	if duration <= 0 {
		return !r.stopRequested() && ctx.Err() == nil
	}
	r.mu.Lock()
	stop := r.stop
	r.mu.Unlock()
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-stop:
		return false
	case <-ctx.Done():
		return false
	}
}

func randomDuration(minimum, maximum time.Duration) time.Duration {
	if maximum <= minimum {
		return minimum
	}
	return minimum + time.Duration(rand.Int63n(int64(maximum-minimum)+1))
}

// RunDownloads 按给定上限执行一个商品的纯下载工作，并在全部结束后收口。
// 单商品入口与批量 Runner 共用它，避免出现两套下载行为。
func RunDownloads(ctx context.Context, concurrency int, work Work) (int, error) {
	if concurrency < 1 {
		concurrency = 1
	}
	errs := make([]error, len(work.Downloads))
	sem := make(chan struct{}, concurrency)
	var wg sync.WaitGroup
	for index, job := range work.Downloads {
		index, job := index, job
		wg.Add(1)
		go func() {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()
			if job == nil {
				errs[index] = fmt.Errorf("下载任务未配置")
				return
			}
			errs[index] = job(ctx)
		}()
	}
	wg.Wait()
	downloaded := 0
	var firstError error
	for _, err := range errs {
		if err == nil {
			downloaded++
		} else if firstError == nil {
			firstError = err
		}
	}
	if work.Finalize != nil {
		if err := work.Finalize(errs); err != nil {
			return downloaded, err
		}
	}
	return downloaded, firstError
}