feat: init commit

2026-05-11 22:59:49 +08:00 · 2024-10-10 23:31:01 +08:00
commit 1cbcebc148
31 changed files with 1673 additions and 0 deletions
--- a/core/core.go
+++ b/core/core.go
@@ -0,0 +1,135 @@
+package core
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/imroc/req/v3"
+	"github.com/krau/SaveAny-Bot/bot"
+	"github.com/krau/SaveAny-Bot/common"
+	"github.com/krau/SaveAny-Bot/config"
+	"github.com/krau/SaveAny-Bot/dao"
+	"github.com/krau/SaveAny-Bot/logger"
+	"github.com/krau/SaveAny-Bot/queue"
+	"github.com/krau/SaveAny-Bot/storage"
+	"github.com/krau/SaveAny-Bot/types"
+	"github.com/mymmrac/telego"
+	"github.com/mymmrac/telego/telegoutil"
+)
+
+func processPendingTask(task types.Task) error {
+	fileRecord, err := dao.GetReceivedFileByFileID(task.FileID)
+	if err != nil {
+		return err
+	}
+	fileRecord.Processing = true
+	if err := dao.UpdateReceivedFile(fileRecord); err != nil {
+		return err
+	}
+
+	_, err = common.ReqClient.R().SetContext(task.Ctx).SetOutputFile(task.FileName).SetDownloadCallbackWithInterval(func(info req.DownloadInfo) {
+		if info.Response == nil || info.Response.Response == nil || info.Response.Response.StatusCode != 200 {
+			return
+		}
+		process := fmt.Sprintf("%.2f%%", float64(info.DownloadedSize)/float64(info.Response.ContentLength)*100.0)
+		bot.Bot.EditMessageText(&telego.EditMessageTextParams{
+			ChatID:    telegoutil.ID(task.ChatID),
+			MessageID: task.ReplyMessageID,
+			Text:      "正在下载文件: " + process,
+		})
+	}, time.Second*time.Duration(3*func() int {
+		if queue.Len() > 0 {
+			return queue.Len()
+		}
+		return 1
+	}())).Get(bot.FileDownloadURL(task.FilePath))
+	if err != nil {
+		return err
+	}
+
+	bot.Bot.EditMessageText(&telego.EditMessageTextParams{
+		ChatID:    telegoutil.ID(task.ChatID),
+		MessageID: task.ReplyMessageID,
+		Text:      "下载完成, 正在转存...",
+	})
+
+	defer func() {
+		if config.Cfg.Temp.CacheTTL > 0 {
+			common.PurgeFileAfter(common.GetDownloadedFilePath(task.FileName),
+				time.Duration(config.Cfg.Temp.CacheTTL)*time.Second)
+		} else {
+			if err := common.PurgeFile(common.GetDownloadedFilePath(task.FileName)); err != nil {
+				logger.L.Errorf("Failed to purge file: %s", err)
+			}
+		}
+	}()
+	if task.StoragePath == "" {
+		task.StoragePath = task.FileName
+	}
+
+	if err := storage.Save(task.Storage, task.Ctx, common.GetDownloadedFilePath(task.FileName), task.StoragePath); err != nil {
+		return err
+	}
+	return nil
+}
+
+func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
+	for {
+		semaphore <- struct{}{}
+		task := queue.GetTask()
+		logger.L.Debugf("Got task: %s", task.FileName)
+
+		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)
+				task.Error = err
+				if errors.Is(err, context.Canceled) {
+					task.Status = types.Canceled
+				} else {
+					task.Status = types.Failed
+				}
+			} else {
+				task.Status = types.Succeeded
+			}
+			queue.AddTask(task)
+		case types.Succeeded:
+			logger.L.Infof("Task succeeded: %s", task.String())
+			bot.Bot.EditMessageText(&telego.EditMessageTextParams{
+				ChatID:    telegoutil.ID(task.ChatID),
+				MessageID: task.ReplyMessageID,
+				Text:      "文件转存完成",
+			})
+			if err := dao.DeleteReceivedFileByFileID(task.FileID); err != nil {
+				logger.L.Errorf("Failed to delete received file: %s", err)
+			}
+		case types.Failed:
+			logger.L.Errorf("Task failed: %s", task.String())
+			bot.Bot.EditMessageText(&telego.EditMessageTextParams{
+				ChatID:    telegoutil.ID(task.ChatID),
+				MessageID: task.ReplyMessageID,
+				Text:      "文件转存失败:" + "\n" + task.Error.Error(),
+			})
+			if err := dao.DeleteReceivedFileByFileID(task.FileID); err != nil {
+				logger.L.Errorf("Failed to delete received file: %s", err)
+			}
+		case types.Canceled:
+			logger.L.Infof("Task canceled: %s", task.String())
+		default:
+			logger.L.Errorf("Unknown task status: %s", task.Status)
+		}
+		<-semaphore
+		logger.L.Debugf("Task done: %s", task.FileName)
+	}
+}
+
+func Run() {
+	logger.L.Info("Start processing tasks...")
+	semaphore := make(chan struct{}, 3)
+	for i := 0; i < 3; i++ {
+		go worker(queue.Queue, semaphore)
+	}
+}