GithubBackup/SaveAny-Bot
1
0
Fork 0
mirror of https://github.com/krau/SaveAny-Bot.git synced 2026-05-11 22:59:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: refactor queue to remove unused methods and add comments

Browse Source
This commit is contained in:
krau
2025-12-15 10:49:40 +08:00
parent 45c978980c
commit 651835c467
3 changed files with 60 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
pkg/queue/queue.go
View File

@@ -38,7 +38,7 @@ func (tq *TaskQueue[T]) Add(task *Task[T]) error {
return fmt.Errorf("task with ID %s already exists", task.ID)
}
if task.IsCancelled() {
if task.Cancelled() {
return fmt.Errorf("task %s has been cancelled", task.ID)
}
@@ -50,6 +50,8 @@ func (tq *TaskQueue[T]) Add(task *Task[T]) error {
return nil
}
// Get retrieves and removes the next non-cancelled task from the queue, adding it to the running tasks.
// Blocks until a task is available or the queue is closed.
func (tq *TaskQueue[T]) Get() (*Task[T], error) {
tq.mu.Lock()
defer tq.mu.Unlock()
@@ -69,7 +71,7 @@ func (tq *TaskQueue[T]) Get() (*Task[T], error) {
tq.tasks.Remove(element)
task.element = nil
if !task.IsCancelled() {
if !task.Cancelled() {
tq.runningTaskMap[task.ID] = task
return task, nil
}
@@ -82,38 +84,23 @@ func (tq *TaskQueue[T]) Get() (*Task[T], error) {
return nil, fmt.Errorf("queue is closed and empty")
}
// Done stops(cancels) and removes the task from the running tasks.
func (tq *TaskQueue[T]) Done(taskID string) {
tq.mu.Lock()
defer tq.mu.Unlock()
tq.CancelTask(taskID) // ensure it's cancelled
delete(tq.taskMap, taskID)
delete(tq.runningTaskMap, taskID)
}
func (tq *TaskQueue[T]) Peek() (*Task[T], error) {
tq.mu.RLock()
defer tq.mu.RUnlock()
if tq.tasks.Len() == 0 {
return nil, fmt.Errorf("queue is empty")
}
for element := tq.tasks.Front(); element != nil; element = element.Next() {
task := element.Value.(*Task[T])
if !task.IsCancelled() {
return task, nil
}
}
return nil, fmt.Errorf("queue has no valid tasks")
}
func (tq *TaskQueue[T]) Length() int {
tq.mu.RLock()
defer tq.mu.RUnlock()
return tq.tasks.Len()
}
// ActiveLength returns the number of non-cancelled tasks in the queue.
func (tq *TaskQueue[T]) ActiveLength() int {
tq.mu.RLock()
defer tq.mu.RUnlock()
@@ -121,13 +108,56 @@ func (tq *TaskQueue[T]) ActiveLength() int {
count := 0
for element := tq.tasks.Front(); element != nil; element = element.Next() {
task := element.Value.(*Task[T])
if !task.IsCancelled() {
if !task.Cancelled() {
count++
}
}
return count
}
// RunningTasks returns the currently running tasks' info.
func (tq *TaskQueue[T]) RunningTasks() []TaskInfo {
tq.mu.RLock()
defer tq.mu.RUnlock()
tasks := make([]TaskInfo, 0, len(tq.runningTaskMap))
for _, task := range tq.runningTaskMap {
if task.Cancelled() {
continue
}
tasks = append(tasks, TaskInfo{
ID: task.ID,
Created: task.created,
Cancelled: task.Cancelled(),
})
}
return tasks
}
// QueuedTasks returns the queued (not yet running) tasks' info.
// The sorting is in the order of addition.
func (tq *TaskQueue[T]) QueuedTasks() []TaskInfo {
tq.mu.RLock()
defer tq.mu.RUnlock()
tasks := make([]TaskInfo, 0, tq.tasks.Len())
for element := tq.tasks.Front(); element != nil; element = element.Next() {
task := element.Value.(*Task[T])
if !task.Cancelled() {
tasks = append(tasks, TaskInfo{
ID: task.ID,
Created: task.created,
Cancelled: task.Cancelled(),
})
}
}
return tasks
}
// CancelTask cancels a task by its ID.
// It looks for the task in both queued and running tasks.
// [NOTE] Cancelled tasks will not be removed from the queue, but marked as cancelled. Use Done to remove them.
// [WARN] Cancelling a running task relies on the task's implementation to respect the cancellation. If the task does not check for cancellation, it may continue running.
func (tq *TaskQueue[T]) CancelTask(taskID string) error {
tq.mu.RLock()
task, exists := tq.taskMap[taskID]
@@ -144,52 +174,6 @@ func (tq *TaskQueue[T]) CancelTask(taskID string) error {
return nil
}
func (tq *TaskQueue[T]) RemoveTask(taskID string) error {
tq.mu.Lock()
defer tq.mu.Unlock()
task, exists := tq.taskMap[taskID]
if !exists {
_, exists = tq.runningTaskMap[taskID]
if exists {
delete(tq.runningTaskMap, taskID)
}
return fmt.Errorf("task %s is already running, cannot remove from queue", taskID)
}
if task.element != nil {
tq.tasks.Remove(task.element)
}
delete(tq.taskMap, taskID)
task.Cancel()
return nil
}
func (tq *TaskQueue[T]) CancelAll() {
tq.mu.RLock()
tasks := make([]*Task[T], 0, tq.tasks.Len())
for element := tq.tasks.Front(); element != nil; element = element.Next() {
tasks = append(tasks, element.Value.(*Task[T]))
}
tq.mu.RUnlock()
for _, task := range tasks {
task.Cancel()
}
}
func (tq *TaskQueue[T]) GetTask(taskID string) (*Task[T], error) {
tq.mu.RLock()
defer tq.mu.RUnlock()
task, exists := tq.taskMap[taskID]
if !exists {
return nil, fmt.Errorf("task %s does not exist", taskID)
}
return task, nil
}
func (tq *TaskQueue[T]) Close() {
tq.mu.Lock()
defer tq.mu.Unlock()
@@ -197,45 +181,3 @@ func (tq *TaskQueue[T]) Close() {
tq.closed = true
tq.cond.Broadcast()
}
func (tq *TaskQueue[T]) IsClosed() bool {
tq.mu.RLock()
defer tq.mu.RUnlock()
return tq.closed
}
func (tq *TaskQueue[T]) Clear() {
tq.mu.Lock()
defer tq.mu.Unlock()
for element := tq.tasks.Front(); element != nil; element = element.Next() {
task := element.Value.(*Task[T])
task.Cancel()
}
tq.tasks.Init()
tq.taskMap = make(map[string]*Task[T])
}
func (tq *TaskQueue[T]) CleanupCancelled() int {
tq.mu.Lock()
defer tq.mu.Unlock()
removed := 0
element := tq.tasks.Front()
for element != nil {
next := element.Next()
task := element.Value.(*Task[T])
if task.IsCancelled() {
tq.tasks.Remove(element)
delete(tq.taskMap, task.ID)
removed++
}
element = next
}
return removed
}

68
pkg/queue/queue_test.go
View File

@@ -39,37 +39,6 @@ func TestDuplicateAdd(t *testing.T) {
}
}
func TestGetAndPeek(t *testing.T) {
q := queue.NewTaskQueue[int]()
t1 := newTask("a")
t2 := newTask("b")
q.Add(t1)
q.Add(t2)
// Peek should return t1
peeked, err := q.Peek()
if err != nil {
t.Fatalf("unexpected error on Peek: %v", err)
}
if peeked.ID != "a" {
t.Fatalf("expected Peek ID 'a', got '%s'", peeked.ID)
}
// Get should return t1 then t2
first, err := q.Get()
if err != nil {
t.Fatalf("unexpected error on Get: %v", err)
}
if first.ID != "a" {
t.Fatalf("expected first Get ID 'a', got '%s'", first.ID)
}
second, err := q.Get()
if err != nil {
t.Fatalf("unexpected error on second Get: %v", err)
}
if second.ID != "b" {
t.Fatalf("expected second Get ID 'b', got '%s'", second.ID)
}
}
func TestCancelAndActiveLength(t *testing.T) {
q := queue.NewTaskQueue[int]()
t1 := newTask("1")
@@ -90,41 +59,6 @@ func TestCancelAndActiveLength(t *testing.T) {
}
}
func TestRemoveTask(t *testing.T) {
q := queue.NewTaskQueue[int]()
t1 := newTask("r1")
q.Add(t1)
if err := q.RemoveTask("r1"); err != nil {
t.Fatalf("unexpected error on RemoveTask: %v", err)
}
if q.Length() != 0 {
t.Fatalf("expected length 0 after remove, got %d", q.Length())
}
}
func TestClearAndCleanupCancelled(t *testing.T) {
q := queue.NewTaskQueue[int]()
tasks := []*queue.Task[int]{newTask("c1"), newTask("c2"), newTask("c3")}
for _, tsk := range tasks {
q.Add(tsk)
}
// Cancel one
q.CancelTask("c2")
// Cleanup cancelled
removed := q.CleanupCancelled()
if removed != 1 {
t.Fatalf("expected removed 1, got %d", removed)
}
if q.ActiveLength() != 2 {
t.Fatalf("expected active length 2 after cleanup, got %d", q.ActiveLength())
}
// Clear all
q.Clear()
if q.Length() != 0 {
t.Fatalf("expected length 0 after clear, got %d", q.Length())
}
}
func TestCloseBehavior(t *testing.T) {
q := queue.NewTaskQueue[int]()
done := make(chan struct{})
@@ -169,4 +103,4 @@ func TestConcurrencySafety(t *testing.T) {
}
}()
wg.Wait()
}
}

9
pkg/queue/task.go
View File

@@ -15,6 +15,13 @@ type Task[T any] struct {
element *list.Element
}
// Read-only info about a task
type TaskInfo struct {
ID string
Created time.Time
Cancelled bool
}
func NewTask[T any](ctx context.Context, id string, data T) *Task[T] {
cancelCtx, cancel := context.WithCancel(ctx)
return &Task[T]{
@@ -26,7 +33,7 @@ func NewTask[T any](ctx context.Context, id string, data T) *Task[T] {
}
}
func (t *Task[T]) IsCancelled() bool {
func (t *Task[T]) Cancelled() bool {
select {
case <-t.ctx.Done():
return true