stock-analysis/internal/job/job.go

package job

import (
	"context"
	"encoding/json"
	"log"
	"sync"
	"time"

	"github.com/google/uuid"

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

// Config Manager 配置。
type Config struct {
	Workers      int           // 并发 worker 数，默认 8
	BatchSize    int           // 批量写入阈值，默认 20
	RateInterval time.Duration // 全局限速间隔，默认 400ms（≈2.5 req/s）
	RateBurst    int           // 限速突发额度，默认 4
}

// runningTask 记录运行中任务的取消句柄和状态快照。
type runningTask struct {
	cancel context.CancelFunc
	task   *cache.Task
}

// Manager 管理任务注册表和全局限速器，允许多任务并发运行。
type Manager struct {
	store     Store
	cache     cache.Cache
	ds        datasource.MarketDataSource
	lister    BondLister
	gate      <-chan struct{} // 全局限速令牌（跨任务共享）
	workers   int
	batchSize int

	mu      sync.RWMutex
	running map[string]*runningTask
}

// NewManager 创建任务管理器，内部启动全局限速器。
func NewManager(store Store, c cache.Cache, ds datasource.MarketDataSource, lister BondLister, cfg Config) *Manager {
	if cfg.Workers <= 0 {
		cfg.Workers = 8
	}
	if cfg.BatchSize <= 0 {
		cfg.BatchSize = 20
	}
	if cfg.RateInterval <= 0 {
		cfg.RateInterval = 400 * time.Millisecond
	}
	if cfg.RateBurst <= 0 {
		cfg.RateBurst = 4
	}

	return &Manager{
		store:     store,
		cache:     c,
		ds:        ds,
		lister:    lister,
		gate:      newRateLimiter(cfg.RateInterval, cfg.RateBurst),
		workers:   cfg.Workers,
		batchSize: cfg.BatchSize,
		running:   make(map[string]*runningTask),
	}
}

// Start 创建并异步启动一个任务，立即返回任务快照（pending 状态）。
func (m *Manager) Start(ctx context.Context, typ cache.TaskType, since string, limit int) (*cache.Task, error) {
	params, _ := json.Marshal(map[string]any{"since": since, "limit": limit})

	task := &cache.Task{
		ID:        uuid.New().String(),
		Type:      typ,
		Status:    cache.TaskStatusPending,
		StartedAt: time.Now().Format(time.RFC3339),
		Params:    string(params),
	}

	// 持久化到 SQLite + Redis
	if err := m.store.SaveTask(ctx, task); err != nil {
		return nil, err
	}
	_ = m.cache.SetTask(ctx, task, 24*time.Hour)
	_ = m.cache.AddActiveTask(ctx, task.ID)

	// 创建可取消 context
	taskCtx, cancel := context.WithCancel(context.Background())

	m.mu.Lock()
	m.running[task.ID] = &runningTask{cancel: cancel, task: task}
	m.mu.Unlock()

	// 异步执行
	runner := NewRunner(m.store, m.cache, m.ds, m.gate, m.workers, m.batchSize)
	go func() {
		defer cancel()
		runner.run(taskCtx, task, m.lister, since, limit)
		// 任务结束后从 running map 移除
		m.mu.Lock()
		delete(m.running, task.ID)
		m.mu.Unlock()
	}()

	log.Printf("[job] task %s created: type=%s, since=%s, limit=%d", task.ID, typ, since, limit)
	return task, nil
}

// Cancel 取消运行中任务，返回是否成功取消。
func (m *Manager) Cancel(id string) bool {
	m.mu.RLock()
	rt, ok := m.running[id]
	m.mu.RUnlock()
	if !ok {
		return false
	}
	rt.cancel()
	log.Printf("[job] task %s cancelled", id)
	return true
}

// Get 获取任务状态：Redis 优先（实时），回退 SQLite。
func (m *Manager) Get(ctx context.Context, id string) (*cache.Task, error) {
	// 1. 先查 Redis（实时状态）
	t, err := m.cache.GetTask(ctx, id)
	if err == nil && t != nil {
		return t, nil
	}
	// 2. 回退 SQLite
	return m.store.GetTask(ctx, id)
}

// List 返回任务列表：SQLite 历史记录 + Redis 实时覆盖运行中任务。
func (m *Manager) List(ctx context.Context, limit int) ([]*cache.Task, error) {
	tasks, err := m.store.ListTasks(ctx, limit)
	if err != nil {
		return nil, err
	}
	// 对运行中任务用 Redis 实时数据覆盖
	for _, t := range tasks {
		if t.Status == cache.TaskStatusRunning || t.Status == cache.TaskStatusPending {
			if rt, _ := m.cache.GetTask(ctx, t.ID); rt != nil {
				t.Done = rt.Done
				t.Success = rt.Success
				t.Failed = rt.Failed
				t.Status = rt.Status
				t.FailedCodes = rt.FailedCodes
			}
		}
	}
	return tasks, nil
}

// StopAll 优雅停止所有运行中任务（服务关闭时调用）。
func (m *Manager) StopAll() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for id, rt := range m.running {
		rt.cancel()
		log.Printf("[job] stopping task %s on shutdown", id)
	}
}

// newRateLimiter 创建全局限速器（令牌桶模式）。
// 返回一个只读 channel，每次接收表示获得一个请求令牌。
// 稳态速率 = 1 token / interval，初始突发 = burst。
func newRateLimiter(interval time.Duration, burst int) <-chan struct{} {
	ch := make(chan struct{}, burst)
	// 预填突发额度
	for i := 0; i < burst; i++ {
		ch <- struct{}{}
	}
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for range ticker.C {
			select {
			case ch <- struct{}{}:
			default: // 满了丢弃，保持稳态速率
			}
		}
	}()
	return ch
}