ProjectHentai/alist
6
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(task): generic task manager

Browse Source
This commit is contained in:
Noah Hsu
2022-06-21 16:14:37 +08:00
parent 1b3387ca1a
commit 55d6434daa
6 changed files with 135 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
pkg/task/manager.go
View File

@ -1,27 +1,28 @@
package task
import (
log "github.com/sirupsen/logrus"
"sync/atomic"
"github.com/alist-org/alist/v3/pkg/generic_sync"
log "github.com/sirupsen/logrus"
)
type Manager struct {
workerC chan struct{}
curID uint64
tasks generic_sync.MapOf[uint64, *Task]
type Manager[K comparable, V any] struct {
workerC chan struct{}
curID K
updateID func(*K)
tasks generic_sync.MapOf[K, *Task[K, V]]
}
func (tm *Manager) Submit(name string, f Func, callbacks ...Callback) uint64 {
task := newTask(name, f, callbacks...)
tm.addTask(task)
tm.do(task.ID)
func (tm *Manager[K, V]) Submit(task *Task[K, V]) K {
if tm.updateID != nil {
task.ID = tm.curID
tm.updateID(&task.ID)
}
tm.tasks.Store(task.ID, task)
tm.do(task)
return task.ID
}
func (tm *Manager) do(tid uint64) {
task := tm.MustGet(tid)
func (tm *Manager[K, V]) do(task *Task[K, V]) {
go func() {
log.Debugf("task [%s] waiting for worker", task.Name)
select {
@ -30,39 +31,34 @@ func (tm *Manager) do(tid uint64) {
task.run()
log.Debugf("task [%s] ended", task.Name)
}
// return worker
tm.workerC <- struct{}{}
}()
}
func (tm *Manager) addTask(task *Task) {
task.ID = tm.curID
atomic.AddUint64(&tm.curID, 1)
tm.tasks.Store(task.ID, task)
}
func (tm *Manager) GetAll() []*Task {
func (tm *Manager[K, V]) GetAll() []*Task[K, V] {
return tm.tasks.Values()
}
func (tm *Manager) Get(tid uint64) (*Task, bool) {
func (tm *Manager[K, V]) Get(tid K) (*Task[K, V], bool) {
return tm.tasks.Load(tid)
}
func (tm *Manager) MustGet(tid uint64) *Task {
func (tm *Manager[K, V]) MustGet(tid K) *Task[K, V] {
task, _ := tm.Get(tid)
return task
}
func (tm *Manager) Retry(tid uint64) error {
func (tm *Manager[K, V]) Retry(tid K) error {
t, ok := tm.Get(tid)
if !ok {
return ErrTaskNotFound
}
tm.do(t.ID)
tm.do(t)
return nil
}
func (tm *Manager) Cancel(tid uint64) error {
func (tm *Manager[K, V]) Cancel(tid K) error {
t, ok := tm.Get(tid)
if !ok {
return ErrTaskNotFound
@ -71,17 +67,17 @@ func (tm *Manager) Cancel(tid uint64) error {
return nil
}
func (tm *Manager) Remove(tid uint64) {
func (tm *Manager[K, V]) Remove(tid K) {
tm.tasks.Delete(tid)
}
// RemoveAll removes all tasks from the manager, this maybe shouldn't be used
// because the task maybe still running.
func (tm *Manager) RemoveAll() {
func (tm *Manager[K, V]) RemoveAll() {
tm.tasks.Clear()
}
func (tm *Manager) RemoveFinished() {
func (tm *Manager[K, V]) RemoveFinished() {
tasks := tm.GetAll()
for _, task := range tasks {
if task.Status == FINISHED {
@ -90,7 +86,7 @@ func (tm *Manager) RemoveFinished() {
}
}
func (tm *Manager) RemoveError() {
func (tm *Manager[K, V]) RemoveError() {
tasks := tm.GetAll()
for _, task := range tasks {
if task.Error != nil {
@ -99,9 +95,16 @@ func (tm *Manager) RemoveError() {
}
}
func NewTaskManager() *Manager {
return &Manager{
tasks: generic_sync.MapOf[uint64, *Task]{},
curID: 0,
func NewTaskManager[K comparable, V any](maxWorker int, updateID ...func(*K)) *Manager[K, V] {
tm := &Manager[K, V]{
tasks: generic_sync.MapOf[K, *Task[K, V]]{},
workerC: make(chan struct{}, maxWorker),
}
for i := 0; i < maxWorker; i++ {
tm.workerC <- struct{}{}
}
if len(updateID) > 0 {
tm.updateID = updateID[0]
}
return tm
}