mirror of https://github.com/cloudreve/Cloudreve
61 lines
1.0 KiB
Go
61 lines
1.0 KiB
Go
package queue
|
|
|
|
import "sync"
|
|
|
|
type (
|
|
// TaskRegistry is used in slave node to track in-memory stateful tasks.
|
|
TaskRegistry interface {
|
|
// NextID returns the next available Task ID.
|
|
NextID() int
|
|
// Get returns the Task by ID.
|
|
Get(id int) (Task, bool)
|
|
// Set sets the Task by ID.
|
|
Set(id int, t Task)
|
|
// Delete deletes the Task by ID.
|
|
Delete(id int)
|
|
}
|
|
|
|
taskRegistry struct {
|
|
tasks map[int]Task
|
|
current int
|
|
mu sync.Mutex
|
|
}
|
|
)
|
|
|
|
// NewTaskRegistry creates a new TaskRegistry.
|
|
func NewTaskRegistry() TaskRegistry {
|
|
return &taskRegistry{
|
|
tasks: make(map[int]Task),
|
|
}
|
|
}
|
|
|
|
func (r *taskRegistry) NextID() int {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
r.current++
|
|
return r.current
|
|
}
|
|
|
|
func (r *taskRegistry) Get(id int) (Task, bool) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
t, ok := r.tasks[id]
|
|
return t, ok
|
|
}
|
|
|
|
func (r *taskRegistry) Set(id int, t Task) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
r.tasks[id] = t
|
|
}
|
|
|
|
func (r *taskRegistry) Delete(id int) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
delete(r.tasks, id)
|
|
}
|