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