sig-pub/pkg/utils/progress/progress.go

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
}