12 changed files with 397 additions and 73 deletions
@ -0,0 +1,169 @@
|
||||
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 |
||||
} |
||||
@ -0,0 +1,45 @@
|
||||
package progress |
||||
|
||||
import ( |
||||
"context" |
||||
"fmt" |
||||
"testing" |
||||
"time" |
||||
) |
||||
|
||||
func TestProgressManager(t *testing.T) { |
||||
c := make(chan struct{}) |
||||
|
||||
m := NewProgressManager() |
||||
taskId := m.StartTask(context.Background(), "testing", func(ctx context.Context, progress IProgressReporter) error { |
||||
defer close(c) |
||||
|
||||
progress.SetTotal(100) |
||||
for range 100 { |
||||
progress.AddProgress(1) |
||||
time.Sleep(110 * time.Millisecond) |
||||
} |
||||
return nil |
||||
}) |
||||
|
||||
go func() { |
||||
task, _ := m.GetTask(taskId) |
||||
for { |
||||
select { |
||||
case <-c: |
||||
return |
||||
default: |
||||
} |
||||
progress, total := task.Progress.Get() |
||||
if total != 0 { |
||||
fmt.Printf("task %s progress: %.2f%%\n", task.Name, float64(progress)/float64(total)*100) |
||||
if progress == total { |
||||
return |
||||
} |
||||
} |
||||
time.Sleep(300 * time.Millisecond) |
||||
} |
||||
}() |
||||
|
||||
<-c |
||||
} |
||||
Loading…
Reference in new issue