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 }