refactor: refactor task logic for better scalability (#76)
* refactor: a big refactor. wip * refactor: port handle file * refactor: place all handlers * fix: task info nil pointer * feat: enhance task progress tracking and context management * feat: cancel task * feat: stream mode * feat: silent mode * feat: dir cmd * refactor: remove unused old file * feat: rule cmd * feat: handle silent mode * feat: batch task * fix: batch task progress and temp file cleanup * refactor: update file creation and cleanup methods for better resource management * feat: add save command with silent mode handling * feat: message link * feat: update message prompts to include file count in storage selection * feat: slient save links * refactor: reduce dup code * feat: rule type * feat: chose dir * feat: refactor file handling and storage rules, improve error handling and logging * feat: rule mode * feat: telegraph pics * fix: tphpics nil pointer and inaccurate dirpath * feat: silent save telegraph * feat: add suffix to avoid file overwrite * feat: new storage telegram * chore: tidy go mod
This commit is contained in:
Krau
241
pkg/queue/queue.go
Normal file
241
pkg/queue/queue.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type TaskQueue[T any] struct {
|
||||
tasks *list.List
|
||||
taskMap map[string]*Task[T]
|
||||
runningTaskMap map[string]*Task[T]
|
||||
mu sync.RWMutex
|
||||
cond *sync.Cond
|
||||
closed bool
|
||||
}
|
||||
|
||||
func NewTaskQueue[T any]() *TaskQueue[T] {
|
||||
tq := &TaskQueue[T]{
|
||||
tasks: list.New(),
|
||||
taskMap: make(map[string]*Task[T]),
|
||||
runningTaskMap: make(map[string]*Task[T]),
|
||||
}
|
||||
tq.cond = sync.NewCond(&tq.mu)
|
||||
return tq
|
||||
}
|
||||
|
||||
func (tq *TaskQueue[T]) Add(task *Task[T]) error {
|
||||
tq.mu.Lock()
|
||||
defer tq.mu.Unlock()
|
||||
|
||||
if tq.closed {
|
||||
return errors.New("queue is closed")
|
||||
}
|
||||
|
||||
if _, exists := tq.taskMap[task.ID]; exists {
|
||||
return fmt.Errorf("task with ID %s already exists", task.ID)
|
||||
}
|
||||
|
||||
if task.IsCancelled() {
|
||||
return fmt.Errorf("task %s has been cancelled", task.ID)
|
||||
}
|
||||
|
||||
element := tq.tasks.PushBack(task)
|
||||
task.element = element
|
||||
tq.taskMap[task.ID] = task
|
||||
|
||||
tq.cond.Signal()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (tq *TaskQueue[T]) Get() (*Task[T], error) {
|
||||
tq.mu.Lock()
|
||||
defer tq.mu.Unlock()
|
||||
|
||||
for tq.tasks.Len() == 0 && !tq.closed {
|
||||
tq.cond.Wait()
|
||||
}
|
||||
|
||||
if tq.closed && tq.tasks.Len() == 0 {
|
||||
return nil, fmt.Errorf("queue is closed and empty")
|
||||
}
|
||||
|
||||
for tq.tasks.Len() > 0 {
|
||||
element := tq.tasks.Front()
|
||||
task := element.Value.(*Task[T])
|
||||
|
||||
tq.tasks.Remove(element)
|
||||
task.element = nil
|
||||
|
||||
if !task.IsCancelled() {
|
||||
tq.runningTaskMap[task.ID] = task
|
||||
return task, nil
|
||||
}
|
||||
}
|
||||
|
||||
if !tq.closed {
|
||||
return tq.Get()
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("queue is closed and empty")
|
||||
}
|
||||
|
||||
func (tq *TaskQueue[T]) Done(taskID string) {
|
||||
tq.mu.Lock()
|
||||
defer tq.mu.Unlock()
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
func (tq *TaskQueue[T]) ActiveLength() int {
|
||||
tq.mu.RLock()
|
||||
defer tq.mu.RUnlock()
|
||||
|
||||
count := 0
|
||||
for element := tq.tasks.Front(); element != nil; element = element.Next() {
|
||||
task := element.Value.(*Task[T])
|
||||
if !task.IsCancelled() {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func (tq *TaskQueue[T]) CancelTask(taskID string) error {
|
||||
tq.mu.RLock()
|
||||
task, exists := tq.taskMap[taskID]
|
||||
if !exists {
|
||||
task, exists = tq.runningTaskMap[taskID]
|
||||
}
|
||||
tq.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return fmt.Errorf("task %s does not exist", taskID)
|
||||
}
|
||||
|
||||
task.Cancel()
|
||||
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()
|
||||
|
||||
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
|
||||
}
|
||||
172
pkg/queue/queue_test.go
Normal file
172
pkg/queue/queue_test.go
Normal file
@@ -0,0 +1,172 @@
|
||||
package queue_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/pkg/queue"
|
||||
)
|
||||
|
||||
// helper to create a simple Task with integer payload
|
||||
func newTask(id string) *queue.Task[int] {
|
||||
return queue.NewTask(context.Background(), id, 0)
|
||||
}
|
||||
|
||||
func TestAddAndLength(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
if q.Length() != 0 {
|
||||
t.Fatalf("expected length 0, got %d", q.Length())
|
||||
}
|
||||
t1 := newTask("t1")
|
||||
if err := q.Add(t1); err != nil {
|
||||
t.Fatalf("unexpected error on Add: %v", err)
|
||||
}
|
||||
if q.Length() != 1 {
|
||||
t.Fatalf("expected length 1, got %d", q.Length())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateAdd(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
t1 := newTask("dup")
|
||||
if err := q.Add(t1); err != nil {
|
||||
t.Fatalf("unexpected error on first Add: %v", err)
|
||||
}
|
||||
if err := q.Add(t1); err == nil {
|
||||
t.Fatal("expected error on duplicate Add, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
t2 := newTask("2")
|
||||
q.Add(t1)
|
||||
q.Add(t2)
|
||||
// Cancel t1
|
||||
if err := q.CancelTask("1"); err != nil {
|
||||
t.Fatalf("unexpected error on CancelTask: %v", err)
|
||||
}
|
||||
// Length counts all entries
|
||||
if q.Length() != 2 {
|
||||
t.Fatalf("expected total length 2, got %d", q.Length())
|
||||
}
|
||||
// ActiveLength skips cancelled
|
||||
if got := q.ActiveLength(); got != 1 {
|
||||
t.Fatalf("expected active length 1, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
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{})
|
||||
// consumer
|
||||
go func() {
|
||||
_, err := q.Get()
|
||||
if err == nil {
|
||||
t.Errorf("expected error when getting from closed empty queue, got nil")
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
// allow goroutine to block
|
||||
|
||||
// close queue
|
||||
q.Close()
|
||||
<-done
|
||||
}
|
||||
|
||||
func TestConcurrencySafety(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
var wg sync.WaitGroup
|
||||
n := 1000
|
||||
// producers
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < n; i++ {
|
||||
q.Add(newTask(fmt.Sprintf("p%d", i)))
|
||||
}
|
||||
}()
|
||||
// consumers
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
count := 0
|
||||
for count < n {
|
||||
_, err := q.Get()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
count++
|
||||
}
|
||||
}()
|
||||
wg.Wait()
|
||||
}
|
||||
44
pkg/queue/task.go
Normal file
44
pkg/queue/task.go
Normal file
@@ -0,0 +1,44 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Task[T any] struct {
|
||||
ID string
|
||||
Data T
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
created time.Time
|
||||
element *list.Element
|
||||
}
|
||||
|
||||
func NewTask[T any](ctx context.Context, id string, data T) *Task[T] {
|
||||
cancelCtx, cancel := context.WithCancel(ctx)
|
||||
return &Task[T]{
|
||||
ID: id,
|
||||
Data: data,
|
||||
ctx: cancelCtx,
|
||||
cancel: cancel,
|
||||
created: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task[T]) IsCancelled() bool {
|
||||
select {
|
||||
case <-t.ctx.Done():
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task[T]) Cancel() {
|
||||
t.cancel()
|
||||
}
|
||||
|
||||
func (t *Task[T]) Context() context.Context {
|
||||
return t.ctx
|
||||
}
|
||||
Reference in New Issue
Block a user