package progress import ( "context" "fmt" "sig-pub/pkg/types" "sig-pub/pkg/zlog" "sync" "sync/atomic" "time" ) // ProgressReporter 进度上报接口 type IProgressReporter interface { SetTotal(total int64) SetProgress(progress int64) AddProgress(progress int64) Get() (progress, total int64) } type reporter struct { total atomic.Int64 progress atomic.Int64 } func (r *reporter) SetTotal(total int64) { r.total.Store(total) } func (r *reporter) SetProgress(progress int64) { r.progress.Store(progress) } func (r *reporter) AddProgress(progress int64) { r.progress.Add(progress) } func (r *reporter) Get() (progress, total int64) { return r.progress.Load(), r.total.Load() } // ==================== 任务状态 ==================== type Status int32 const ( // StatusPending Status = "pending" StatusRunning Status = 1 //"running" StatusSuccess Status = 2 //"success" StatusFailed Status = 3 //"failed" StatusCancelled Status = 4 //"cancelled" ) // Task 任务结构体(带锁) type Task struct { ID string `json:"id"` // 任务id Name string `json:"name"` Status atomic.Int64 `json:"status"` Error error `json:"error,omitempty"` Ctime int64 `json:"ctime"` Etime int64 `json:"etime,omitempty"` Progress IProgressReporter cancel context.CancelFunc } // ProgressManager 管理器核心 type ProgressManager struct { nextID atomic.Int64 tasks sync.Map // string -> *Task completeTasks *types.RingSeries[*Task] mu sync.RWMutex } // NewProgressManager 创建管理器实例 func NewProgressManager() *ProgressManager { return &ProgressManager{ completeTasks: types.NewRingSeries[*Task](100, 0), } } func (m *ProgressManager) generateID() string { id := m.nextID.Add(1) return fmt.Sprintf("task_%d_%d", id, time.Now().UnixNano()%100000) } func (m *ProgressManager) StartTask(ctx context.Context, name string, fn func(ctx context.Context, progress IProgressReporter) error) string { id := m.generateID() ctx, cancel := context.WithCancel(ctx) task := &Task{ ID: id, Name: name, Status: atomic.Int64{}, Error: nil, Ctime: time.Now().UnixMilli(), Etime: 0, Progress: &reporter{}, cancel: cancel, } m.tasks.Store(id, task) go m.execTask(ctx, task, fn) return id } func (m *ProgressManager) execTask(ctx context.Context, task *Task, fn func(ctx context.Context, progress IProgressReporter) error) { defer func() { if r := recover(); r != nil { if !task.Status.CompareAndSwap(int64(StatusRunning), int64(StatusFailed)) { zlog.Error("task status not running, can not set panic failed: %s, name=%s, status=%d", task.ID, task.Name, task.Status.Load()) } task.Error = fmt.Errorf("panic: %v", r) task.Etime = time.Now().UnixMilli() } }() defer func() { m.tasks.Delete(task.ID) m.mu.Lock() defer m.mu.Unlock() m.completeTasks.Push(task) }() task.Status.Store(int64(StatusRunning)) err := fn(ctx, task.Progress) if err != nil { if !task.Status.CompareAndSwap(int64(StatusRunning), int64(StatusFailed)) { zlog.Error("task status not running, can not set failed: %s, name=%s, status=%d", task.ID, task.Name, task.Status.Load()) } task.Error = err task.Etime = time.Now().UnixMilli() return } task.Etime = time.Now().UnixMilli() if !task.Status.CompareAndSwap(int64(StatusRunning), int64(StatusSuccess)) { zlog.Error("task status not running, can not set success: %s, name=%s, status=%d", task.ID, task.Name, task.Status.Load()) } } func (m *ProgressManager) GetTask(id string) (*Task, bool) { task, ok := m.tasks.Load(id) if !ok { return nil, false } return task.(*Task), true } func (m *ProgressManager) ListTask() ([]*Task, bool) { runningTasks := m.ListRunningTask() m.mu.RLock() defer m.mu.RUnlock() tasks, ok := m.completeTasks.Series(0, m.completeTasks.Length()) if !ok { return nil, false } return append(runningTasks, tasks...), true } func (m *ProgressManager) ListRunningTask() (tasks []*Task) { m.tasks.Range(func(key, value any) bool { task := value.(*Task) if task.Status.Load() == int64(StatusRunning) { tasks = append(tasks, task) } return true }) return }