stock-analysis/internal/job/runner.go

package job

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	"stock-analysis/internal/cache"
	"stock-analysis/internal/datasource"
	"stock-analysis/internal/model"
)

// fetchResult 单只转债拉取结果，通过 channel 发往 collector。
type fetchResult struct {
	code   string
	klines []model.Kline
	err    error
}

// Runner 单任务执行器：N worker 并发拉取 + 单 collector 串行批量写入。
//
// 并发模型：
//   - 主 goroutine 遍历候选列表，为每只转债启动一个 worker goroutine
//   - worker 从 gate 获取全局限速令牌后拉取数据，结果发往 resultCh
//   - collector 单 goroutine 从 resultCh 收集，每 batchSize 只批量写 SQLite
//   - SQLite SetMaxOpenConns(1) 下，只有 collector 写入，避免锁竞争
type Runner struct {
	store     Store
	cache     cache.Cache
	ds        datasource.MarketDataSource
	gate      <-chan struct{} // 全局限速令牌
	workers   int
	batchSize int
}

// NewRunner 创建执行器。
func NewRunner(store Store, c cache.Cache, ds datasource.MarketDataSource, gate <-chan struct{}, workers, batchSize int) *Runner {
	return &Runner{
		store:     store,
		cache:     c,
		ds:        ds,
		gate:      gate,
		workers:   workers,
		batchSize: batchSize,
	}
}

// run 执行任务（异步调用，不阻塞调用方）。
func (r *Runner) run(ctx context.Context, task *cache.Task, lister BondLister, since string, limit int) {
	// 1. 解析候选列表
	candidates, err := resolveCandidates(ctx, r.store, lister, task.Type, limit)
	if err != nil {
		r.failTask(ctx, task, err)
		return
	}
	task.Total = len(candidates)

	// 2. 预加载：全量模式加载已入库集合（跳过断点续跑）；增量模式加载最后日期映射
	var skipSet map[string]bool
	var lastDateMap map[string]string
	if task.Type == cache.TaskTypeFull {
		codes, err := r.store.ListCodes(ctx)
		if err == nil {
			skipSet = make(map[string]bool, len(codes))
			for _, c := range codes {
				skipSet[c] = true
			}
		}
	} else {
		lastDateMap, _ = r.store.LastKlineDates(ctx)
	}

	// 3. 更新状态为 running
	task.Status = cache.TaskStatusRunning
	r.reportProgress(ctx, task, true)

	log.Printf("[job] task %s started: type=%s, total=%d", task.ID, task.Type, task.Total)

	// 4. 启动 collector（唯一 SQLite 写者）
	resultCh := make(chan fetchResult, r.workers*2)
	collectorDone := make(chan struct{})
	go r.collect(ctx, task, resultCh, collectorDone)

	// 5. 并发拉取：主 goroutine 遍历候选，每只启动一个 worker
	sem := make(chan struct{}, r.workers)
	var wg sync.WaitGroup
	for _, b := range candidates {
		if ctx.Err() != nil {
			break
		}
		// 全量模式跳过已入库转债
		if task.Type == cache.TaskTypeFull && skipSet[b.Code] {
			task.Done++
			task.Success++
			r.reportProgress(ctx, task, false)
			continue
		}
		sem <- struct{}{}
		wg.Add(1)
		go func(bond model.ConvertibleBond) {
			defer wg.Done()
			defer func() { <-sem }()

			// 全局限速：获取令牌后才允许发请求
			select {
			case <-r.gate:
			case <-ctx.Done():
				return
			}

			klines, err := r.fetchOne(ctx, task.Type, bond, since, lastDateMap)
			select {
			case resultCh <- fetchResult{code: bond.Code, klines: klines, err: err}:
			case <-ctx.Done():
			}
		}(b)
	}

	// 6. 等待所有 worker 完成
	wg.Wait()
	close(resultCh)
	<-collectorDone

	// 7. 确定终态
	r.finishTask(ctx, task)
}

// fetchOne 拉取单只转债K线。
func (r *Runner) fetchOne(ctx context.Context, typ cache.TaskType, bond model.ConvertibleBond, since string, lastDateMap map[string]string) ([]model.Kline, error) {
	switch typ {
	case cache.TaskTypeFull:
		beg := fmt.Sprintf("%04d0101", yearOf(bond.ListDate))
		return r.ds.FetchKline(ctx, bond.Code, "101", beg, "20500101")
	case cache.TaskTypeIncremental:
		beg, ok := resolveIncrementalBeg(bond.Code, since, lastDateMap)
		if !ok {
			return nil, nil // 无存档，跳过
		}
		return r.ds.FetchKline(ctx, bond.Code, "101", beg, "20500101")
	default:
		return nil, fmt.Errorf("未知任务类型: %s", typ)
	}
}

// collect 单 goroutine 收集结果并批量写入 SQLite（唯一写者）。
func (r *Runner) collect(ctx context.Context, task *cache.Task, resultCh <-chan fetchResult, done chan<- struct{}) {
	defer close(done)
	batch := make(map[string][]model.Kline)

	for res := range resultCh {
		task.Done++
		if res.err != nil || len(res.klines) == 0 {
			task.Failed++
			if res.err != nil {
				task.FailedCodes = append(task.FailedCodes, res.code)
			}
		} else {
			task.Success++
			batch[res.code] = res.klines
			if len(batch) >= r.batchSize {
				if err := r.store.SaveKlinesBatch(ctx, batch); err != nil {
					log.Printf("[job] batch save error: %v", err)
				}
				batch = make(map[string][]model.Kline)
			}
		}
		r.reportProgress(ctx, task, false)
	}

	// flush 残余 batch
	if len(batch) > 0 {
		if err := r.store.SaveKlinesBatch(ctx, batch); err != nil {
			log.Printf("[job] final batch save error: %v", err)
		}
	}
}

// reportProgress 上报进度：Redis 每次（廉价），SQLite 按策略写（减少写压力）。
// force=true（状态变更）时用 SaveTask 写入完整状态（含 status/total），
// 确保 Redis 不可用时 SQLite 也能反映运行状态。
func (r *Runner) reportProgress(ctx context.Context, task *cache.Task, force bool) {
	_ = r.cache.SetTask(ctx, task, 24*time.Hour)
	if force {
		_ = r.store.SaveTask(ctx, task)
	} else if task.Done%5 == 0 {
		_ = r.store.UpdateTaskProgress(ctx, task.ID, task.Total, task.Done, task.Success, task.Failed, task.FailedCodes)
	}
}

// finishTask 确定任务终态并持久化。
// 使用 SaveTask 写入完整任务状态（含 total/done/success/failed），
// 确保 SQLite 中最终记录与内存一致（Redis 不可用时尤为重要）。
func (r *Runner) finishTask(ctx context.Context, task *cache.Task) {
	var errMsg string
	switch {
	case ctx.Err() != nil:
		task.Status = cache.TaskStatusCancelled
	case task.Success == 0 && task.Failed > 0:
		task.Status = cache.TaskStatusFailed
		errMsg = fmt.Sprintf("全部 %d 只失败", task.Failed)
	default:
		task.Status = cache.TaskStatusSuccess
		if task.Failed > 0 {
			errMsg = fmt.Sprintf("成功 %d，失败 %d", task.Success, task.Failed)
		}
	}
	task.Error = errMsg
	task.FinishedAt = time.Now().Format(time.RFC3339)

	_ = r.store.SaveTask(ctx, task)
	_ = r.cache.SetTask(ctx, task, 24*time.Hour)
	_ = r.cache.RemoveActiveTask(ctx, task.ID)

	log.Printf("[job] task %s finished: status=%s, total=%d, success=%d, failed=%d",
		task.ID, task.Status, task.Total, task.Success, task.Failed)
}

// failTask 标记任务失败（启动阶段错误）。
func (r *Runner) failTask(ctx context.Context, task *cache.Task, err error) {
	task.Status = cache.TaskStatusFailed
	task.Error = err.Error()
	task.FinishedAt = time.Now().Format(time.RFC3339)
	_ = r.store.SaveTask(ctx, task)
	_ = r.cache.SetTask(ctx, task, 24*time.Hour)
	_ = r.cache.RemoveActiveTask(ctx, task.ID)
	log.Printf("[job] task %s failed: %v", task.ID, err)
}