Merge pull request #30 from krau/dev

feat: cancel task

bot/handle_add_task.go

@@ -153,7 +153,6 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
 		Status:         types.Pending,
 		File:           file,
 		StorageName:    storageName,
-		StoragePath:    path.Join(),
 		FileChatID:     record.ChatID,
 		ReplyMessageID: record.ReplyMessageID,
 		FileMessageID:  record.MessageID,
@@ -164,7 +163,7 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
 		task.StoragePath = path.Join(dir.Path, file.FileName)
 	}
 
-	queue.AddTask(task)
+	queue.AddTask(&task)
 
 	entityBuilder := entity.Builder{}
 	var entities []tg.MessageEntityClass

bot/handle_cancel_task.go

@@ -0,0 +1,27 @@
+package bot
+
+import (
+	"strings"
+
+	"github.com/celestix/gotgproto/dispatcher"
+	"github.com/celestix/gotgproto/ext"
+	"github.com/gotd/td/tg"
+	"github.com/krau/SaveAny-Bot/queue"
+)
+
+func cancelTask(ctx *ext.Context, update *ext.Update) error {
+	key := strings.Split(string(update.CallbackQuery.Data), " ")[1]
+	ok := queue.CancelTask(key)
+	if ok {
+		ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
+			QueryID: update.CallbackQuery.QueryID,
+			Message: "任务已取消",
+		})
+		return dispatcher.EndGroups
+	}
+	ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
+		QueryID: update.CallbackQuery.QueryID,
+		Message: "任务取消失败",
+	})
+	return dispatcher.EndGroups
+}

bot/handlers.go

@@ -22,5 +22,6 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
 	dispatcher.AddHandler(handlers.NewMessage(linkRegexFilter, handleLinkMessage))
 	dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("add"), AddToQueue))
 	dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("set_default"), setDefaultStorage))
+	dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("cancel"), cancelTask))
 	dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, handleFileMessage))
 }

bot/utils.go

@@ -264,7 +264,7 @@ func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *dao.User, t
 		})
 		return dispatcher.EndGroups
 	}
-	queue.AddTask(*task)
+	queue.AddTask(task)
 	ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
 		Message: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", task.FileName(), queue.Len()),
 		ID:      task.ReplyMessageID,

core/core.go

@@ -22,13 +22,12 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
 		switch task.Status {
 		case types.Pending:
 			logger.L.Infof("Processing task: %s", task.String())
-			if err := processPendingTask(&task); err != nil {
-				logger.L.Errorf("Failed to do task: %s", err)
+			if err := processPendingTask(task); err != nil {
 				task.Error = err
 				if errors.Is(err, context.Canceled) {
-					logger.L.Debugf("Task canceled: %s", task.String())
 					task.Status = types.Canceled
 				} else {
+					logger.L.Errorf("Failed to do task: %s", err)
 					task.Status = types.Failed
 				}
 			} else {
@@ -37,23 +36,43 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
 			queue.AddTask(task)
 		case types.Succeeded:
 			logger.L.Infof("Task succeeded: %s", task.String())
-			task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
+			extCtx, ok := task.Ctx.(*ext.Context)
+			if !ok {
+				logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
+			} else {
+				extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
 					Message: fmt.Sprintf("文件保存成功\n [%s]: %s", task.StorageName, task.StoragePath),
 					ID:      task.ReplyMessageID,
 				})
+			}
 		case types.Failed:
 			logger.L.Errorf("Task failed: %s", task.String())
-			task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
+			extCtx, ok := task.Ctx.(*ext.Context)
+			if !ok {
+				logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
+			} else {
+				extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
 					Message: "文件保存失败\n" + task.Error.Error(),
 					ID:      task.ReplyMessageID,
 				})
+			}
 		case types.Canceled:
 			logger.L.Infof("Task canceled: %s", task.String())
+			extCtx, ok := task.Ctx.(*ext.Context)
+			if !ok {
+				logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
+			} else {
+				extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
+					Message: "任务已取消",
+					ID:      task.ReplyMessageID,
+				})
+			}
 		default:
 			logger.L.Errorf("Unknown task status: %s", task.Status)
 		}
 		<-semaphore
-		logger.L.Debugf("Task done: %s", task.String())
+		logger.L.Debugf("Task done: %s; status: %s", task.String(), task.Status)
+		queue.DoneTask(task)
 	}
 }
 

core/download.go

@@ -1,6 +1,7 @@
 package core
 
 import (
+	"context"
 	"fmt"
 	"path/filepath"
 	"time"
@@ -48,11 +49,15 @@ func processPendingTask(task *types.Task) error {
 		return fmt.Errorf("context is not *ext.Context: %T", task.Ctx)
 	}
 
+	cancelCtx, cancel := context.WithCancel(ctx)
+	task.Cancel = cancel
+
 	text, entities := buildProgressMessageEntity(task, 0, task.StartTime, 0)
 	ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
 		Message:     text,
 		Entities:    entities,
 		ID:          task.ReplyMessageID,
+		ReplyMarkup: getCancelTaskMarkup(task),
 	})
 	progressCallback := buildProgressCallback(ctx, task, getProgressUpdateCount(task.File.FileSize))
 
@@ -63,7 +68,7 @@ func processPendingTask(task *types.Task) error {
 	defer dest.Close()
 	task.StartTime = time.Now()
 	downloadBuider := Downloader.Download(bot.Client.API(), task.File.Location).WithThreads(getTaskThreads(task.File.FileSize))
-	_, err = downloadBuider.Parallel(ctx, dest)
+	_, err = downloadBuider.Parallel(cancelCtx, dest)
 	if err != nil {
 		return fmt.Errorf("下载文件失败: %w", err)
 	}
@@ -78,5 +83,5 @@ func processPendingTask(task *types.Task) error {
 		ID:      task.ReplyMessageID,
 	})
 
-	return saveFileWithRetry(task, taskStorage, cacheDestPath)
+	return saveFileWithRetry(cancelCtx, task, taskStorage, cacheDestPath)
 }

core/utils.go

@@ -1,6 +1,7 @@
 package core
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path"
@@ -19,13 +20,21 @@ import (
 	"github.com/krau/SaveAny-Bot/types"
 )
 
-func saveFileWithRetry(task *types.Task, taskStorage storage.Storage, localFilePath string) error {
+func saveFileWithRetry(ctx context.Context, task *types.Task, taskStorage storage.Storage, localFilePath string) error {
 	for i := 0; i <= config.Cfg.Retry; i++ {
-		if err := taskStorage.Save(task.Ctx, localFilePath, task.StoragePath); err != nil {
+		if err := ctx.Err(); err != nil {
+			return fmt.Errorf("context canceled while saving file: %w", err)
+		}
+		if err := taskStorage.Save(ctx, localFilePath, task.StoragePath); err != nil {
 			if i == config.Cfg.Retry {
 				return fmt.Errorf("failed to save file: %w", err)
 			}
 			logger.L.Errorf("Failed to save file: %s, retrying...", err)
+			select {
+			case <-ctx.Done():
+				return fmt.Errorf("context canceled during retry delay: %w", ctx.Err())
+			case <-time.After(time.Duration(i*500) * time.Millisecond):
+			}
 			continue
 		}
 		return nil
@@ -56,7 +65,7 @@ func processPhoto(task *types.Task, taskStorage storage.Storage, cachePath strin
 
 	logger.L.Infof("Downloaded file: %s", cachePath)
 
-	return saveFileWithRetry(task, taskStorage, cachePath)
+	return saveFileWithRetry(task.Ctx, task, taskStorage, cachePath)
 }
 
 // func getProgressBar(progress float64, updateCount int) string {
@@ -142,10 +151,17 @@ func buildProgressCallback(ctx *ext.Context, task *types.Task, updateCount int)
 			Message:     text,
 			Entities:    entities,
 			ID:          task.ReplyMessageID,
+			ReplyMarkup: getCancelTaskMarkup(task),
 		})
 	}
 }
 
+func getCancelTaskMarkup(task *types.Task) *tg.ReplyInlineMarkup {
+	return &tg.ReplyInlineMarkup{
+		Rows: []tg.KeyboardButtonRow{{Buttons: []tg.KeyboardButtonClass{&tg.KeyboardButtonCallback{Text: "取消任务", Data: fmt.Appendf(nil, "cancel %s", task.Key())}}}},
+	}
+}
+
 func fixTaskFileExt(task *types.Task, localFilePath string) {
 	if path.Ext(task.FileName()) == "" {
 		mimeType, err := mimetype.DetectFile(localFilePath)

queue/queue.go

@@ -11,27 +11,62 @@ type TaskQueue struct {
 	list      *list.List
 	cond      *sync.Cond
 	mutex     *sync.Mutex
+	activeMap map[string]*types.Task
 }
 
-func (q *TaskQueue) AddTask(task types.Task) {
+func (q *TaskQueue) AddTask(task *types.Task) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 	q.list.PushBack(task)
 	q.cond.Signal()
+	if task.Status != types.Pending {
+		delete(q.activeMap, task.Key())
+	}
 }
 
-func (q *TaskQueue) GetTask() types.Task {
+func (q *TaskQueue) GetTask() *types.Task {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 	for q.list.Len() == 0 {
 		q.cond.Wait()
 	}
 	e := q.list.Front()
-	task := e.Value.(types.Task)
+	task := e.Value.(*types.Task)
 	q.list.Remove(e)
+	if task.Status == types.Pending {
+		q.activeMap[task.Key()] = task
+	}
 	return task
 }
 
+func (q *TaskQueue) DoneTask(task *types.Task) {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+	delete(q.activeMap, task.Key())
+}
+
+func (q *TaskQueue) CancelTask(key string) bool {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+	if task, ok := q.activeMap[key]; ok {
+		if task.Cancel != nil {
+			task.Cancel()
+			return true
+		}
+	}
+	for e := q.list.Front(); e != nil; e = e.Next() {
+		task := e.Value.(*types.Task)
+		if task.Key() == key {
+			if task.Cancel != nil {
+				task.Cancel()
+			}
+			q.list.Remove(e)
+			return true
+		}
+	}
+	return false
+}
+
 func (q *TaskQueue) Len() int {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
@@ -50,17 +85,26 @@ func NewQueue() *TaskQueue {
 		list:      list.New(),
 		cond:      sync.NewCond(m),
 		mutex:     m,
+		activeMap: make(map[string]*types.Task),
 	}
 }
 
-func AddTask(task types.Task) {
+func AddTask(task *types.Task) {
 	Queue.AddTask(task)
 }
 
-func GetTask() types.Task {
+func GetTask() *types.Task {
 	return Queue.GetTask()
 }
 
 func Len() int {
 	return Queue.Len()
 }
+
+func CancelTask(key string) bool {
+	return Queue.CancelTask(key)
+}
+
+func DoneTask(task *types.Task) {
+	Queue.DoneTask(task)
+}

types/types.go

@@ -36,6 +36,7 @@ var StorageTypeDisplay = map[StorageType]string{
 
 type Task struct {
 	Ctx         context.Context
+	Cancel      context.CancelFunc
 	Error       error
 	Status      TaskStatus
 	File        *File
@@ -52,6 +53,10 @@ type Task struct {
 	UserID int64
 }
 
+func (t Task) Key() string {
+	return fmt.Sprintf("%d:%d", t.FileChatID, t.FileMessageID)
+}
+
 func (t Task) String() string {
 	return fmt.Sprintf("[%d:%d]:%s", t.FileChatID, t.FileMessageID, t.File.FileName)
 }