mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-26 02:30:46 +08:00
Compare commits
58 Commits
v0.60.4
...
fix/qualit
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
942022969c | ||
|
|
ab89310325 | ||
|
|
b27bb734be | ||
|
|
fc0c8975a0 | ||
|
|
524a259ae4 | ||
|
|
e8257c46d2 | ||
|
|
6a11b8ac8d | ||
|
|
47d941d851 | ||
|
|
7e502442c5 | ||
|
|
a9425cb5ea | ||
|
|
6ad2a2c884 | ||
|
|
0d42c9f23d | ||
|
|
4c8f35ae80 | ||
|
|
0a67b8bb57 | ||
|
|
196bb9941f | ||
|
|
973f3499a9 | ||
|
|
d9cb5f9cfc | ||
|
|
ccb702148e | ||
|
|
85b1f35dc5 | ||
|
|
9f93a95258 | ||
|
|
189bf9c736 | ||
|
|
534ed7a7c2 | ||
|
|
1d4997ba2d | ||
|
|
f7e532ca7e | ||
|
|
0957d93da3 | ||
|
|
324d8c100f | ||
|
|
4cb23f04b6 | ||
|
|
a92ca75c4f | ||
|
|
4e99b4bdc9 | ||
|
|
389be59371 | ||
|
|
0d49ae94af | ||
|
|
48e739f06d | ||
|
|
b4fbbae068 | ||
|
|
29be3ee90c | ||
|
|
56660a7705 | ||
|
|
0dfb6af153 | ||
|
|
63f53fac7f | ||
|
|
607ec8aced | ||
|
|
cd9886d710 | ||
|
|
3f7f50133e | ||
|
|
c7911cc208 | ||
|
|
b31d628c19 | ||
|
|
60fd2e04d9 | ||
|
|
b72dd67be9 | ||
|
|
aa25eb1510 | ||
|
|
ba0deababc | ||
|
|
d1b2dbfe5f | ||
|
|
bd7da7c31e | ||
|
|
1981d8662b | ||
|
|
51fcd17922 | ||
|
|
610f586d45 | ||
|
|
6bc12d6feb | ||
|
|
ad41aec22c | ||
|
|
4389bfe046 | ||
|
|
832eb27d4f | ||
|
|
b6e981282d | ||
|
|
9ed5b5f962 | ||
|
|
4527022853 |
4
.github/workflows/build-release.yml
vendored
4
.github/workflows/build-release.yml
vendored
@@ -43,7 +43,7 @@ jobs:
|
||||
goarch: arm64
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Extract version from Git Ref
|
||||
id: extract_version
|
||||
@@ -64,7 +64,7 @@ jobs:
|
||||
ldflags: >-
|
||||
-s -w
|
||||
-X "github.com/krau/SaveAny-Bot/config.Version=${{ env.VERSION }}"
|
||||
-X "github.com/krau/SaveAny-Bot/config.BuildTime=${{ format(github.event.repository.updated_at, 'yyyy-MM-dd HH:mm:ss') }}"
|
||||
-X "github.com/krau/SaveAny-Bot/config.BuildTime=${{ github.event.repository.updated_at }}"
|
||||
-X "github.com/krau/SaveAny-Bot/config.GitCommit=${{ github.sha }}"
|
||||
binary_name: saveany-bot
|
||||
env:
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -1,5 +1,6 @@
|
||||
config.toml
|
||||
logs/
|
||||
/cache/
|
||||
tmp/
|
||||
data/
|
||||
downloads/
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"net/http"
|
||||
"strings"
|
||||
@@ -9,9 +8,6 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
)
|
||||
|
||||
// tokenContextKey 用于在 context 中存储 token
|
||||
type tokenContextKey struct{}
|
||||
|
||||
// AuthMiddleware 返回认证中间件
|
||||
func AuthMiddleware() func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
@@ -40,9 +36,7 @@ func AuthMiddleware() func(http.Handler) http.Handler {
|
||||
return
|
||||
}
|
||||
|
||||
// 将 token 添加到 context
|
||||
ctx := context.WithValue(r.Context(), tokenContextKey{}, token)
|
||||
next.ServeHTTP(w, r.WithContext(ctx))
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ func NewTaskFactory(ctx context.Context) *TaskFactory {
|
||||
// CreateTask 创建任务
|
||||
func (f *TaskFactory) CreateTask(req *CreateTaskRequest) (*CreateTaskResponse, error) {
|
||||
// 验证存储
|
||||
stor, ok := storage.Storages[req.Storage]
|
||||
stor, ok := storage.GetStorage(req.Storage)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("storage not found: %s", req.Storage)
|
||||
}
|
||||
@@ -327,12 +327,12 @@ func (f *TaskFactory) createTransferTask(taskID string, createdAt time.Time, req
|
||||
}
|
||||
|
||||
// 验证源存储和目标存储
|
||||
sourceStor, ok := storage.Storages[params.SourceStorage]
|
||||
sourceStor, ok := storage.GetStorage(params.SourceStorage)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("source storage not found: %s", params.SourceStorage)
|
||||
}
|
||||
|
||||
targetStor, ok := storage.Storages[params.TargetStorage]
|
||||
targetStor, ok := storage.GetStorage(params.TargetStorage)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("target storage not found: %s", params.TargetStorage)
|
||||
}
|
||||
|
||||
@@ -135,8 +135,9 @@ func (h *Handlers) ListStoragesHandler(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
storages := make([]StorageInfo, 0, len(storage.Storages))
|
||||
for name, stor := range storage.Storages {
|
||||
all := storage.AllStorages()
|
||||
storages := make([]StorageInfo, 0, len(all))
|
||||
for name, stor := range all {
|
||||
storages = append(storages, StorageInfo{
|
||||
Name: name,
|
||||
Type: string(stor.Type()),
|
||||
|
||||
@@ -198,19 +198,3 @@ func (t *TaskProgressInfo) Emit(e taskevent.Event) {
|
||||
SendWebhook(nil, payload)
|
||||
}
|
||||
}
|
||||
|
||||
// ProgressTracker is retained for compatibility but is no longer the primary
|
||||
// progress path; taskevent drives updates now. These methods are safe no-ops
|
||||
// when called on a nil receiver.
|
||||
type ProgressTracker struct{}
|
||||
|
||||
func NewProgressTracker(taskID, taskType, storage, path, title, webhook string) *ProgressTracker {
|
||||
return &ProgressTracker{}
|
||||
}
|
||||
|
||||
func (p *ProgressTracker) OnStart(totalBytes int64, totalFiles int) {}
|
||||
func (p *ProgressTracker) OnProgress(downloadedBytes int64, downloadedFiles int) {}
|
||||
func (p *ProgressTracker) OnDone(err error) {}
|
||||
func (p *ProgressTracker) GetInfo() *TaskProgressInfo { return nil }
|
||||
func (p *ProgressTracker) UpdateProgressBytes(bytes int64) {}
|
||||
func (p *ProgressTracker) UpdateProgressFiles(files int) {}
|
||||
|
||||
@@ -3,6 +3,7 @@ package api
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
@@ -90,9 +91,15 @@ func (s *Server) Start(ctx context.Context) error {
|
||||
|
||||
logger.Infof("Starting API server on %s", s.httpServer.Addr)
|
||||
|
||||
// Bind synchronously so listen failures are returned to the caller.
|
||||
ln, err := net.Listen("tcp", s.httpServer.Addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to listen on %s: %w", s.httpServer.Addr, err)
|
||||
}
|
||||
|
||||
// 在 goroutine 中启动服务器
|
||||
go func() {
|
||||
if err := s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
if err := s.httpServer.Serve(ln); err != nil && err != http.ErrServerClosed {
|
||||
logger.Errorf("API server error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -37,6 +37,9 @@ func SendWebhook(ctx context.Context, payload *WebhookPayload) {
|
||||
} else {
|
||||
logger = log.Default().With("task_id", payload.TaskID)
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
payloadBytes, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
@@ -44,10 +47,15 @@ func SendWebhook(ctx context.Context, payload *WebhookPayload) {
|
||||
return
|
||||
}
|
||||
|
||||
// 重试 3 次
|
||||
for i := range 3 {
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, webhookURL, bytes.NewBuffer(payloadBytes))
|
||||
// 重试 3 次, 指数退避 (100ms/400ms/1.6s)
|
||||
const maxAttempts = 3
|
||||
const requestTimeout = 30 * time.Second
|
||||
backoff := 100 * time.Millisecond
|
||||
for i := range maxAttempts {
|
||||
reqCtx, cancel := context.WithTimeout(ctx, requestTimeout)
|
||||
req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, webhookURL, bytes.NewBuffer(payloadBytes))
|
||||
if err != nil {
|
||||
cancel()
|
||||
logger.Errorf("Failed to create webhook request: %v", err)
|
||||
return
|
||||
}
|
||||
@@ -56,9 +64,13 @@ func SendWebhook(ctx context.Context, payload *WebhookPayload) {
|
||||
req.Header.Set("User-Agent", "SaveAny-Bot/1.0")
|
||||
|
||||
resp, err := webhookClient.Do(req)
|
||||
cancel()
|
||||
if err != nil {
|
||||
logger.Warnf("Webhook request failed (attempt %d/3): %v", i+1, err)
|
||||
time.Sleep(time.Second * time.Duration(i+1))
|
||||
logger.Warnf("Webhook request failed (attempt %d/%d): %v", i+1, maxAttempts, err)
|
||||
if i < maxAttempts-1 {
|
||||
time.Sleep(backoff)
|
||||
}
|
||||
backoff *= 4
|
||||
continue
|
||||
}
|
||||
resp.Body.Close()
|
||||
@@ -68,11 +80,14 @@ func SendWebhook(ctx context.Context, payload *WebhookPayload) {
|
||||
return
|
||||
}
|
||||
|
||||
logger.Warnf("Webhook returned non-2xx status (attempt %d/3): %d", i+1, resp.StatusCode)
|
||||
time.Sleep(time.Second * time.Duration(i+1))
|
||||
logger.Warnf("Webhook returned non-2xx status (attempt %d/%d): %d", i+1, maxAttempts, resp.StatusCode)
|
||||
if i < maxAttempts-1 {
|
||||
time.Sleep(backoff)
|
||||
}
|
||||
backoff *= 4
|
||||
}
|
||||
|
||||
logger.Errorf("Failed to send webhook after 3 attempts")
|
||||
logger.Errorf("Failed to send webhook after %d attempts", maxAttempts)
|
||||
}()
|
||||
}
|
||||
|
||||
|
||||
@@ -70,9 +70,6 @@ func Init(ctx context.Context) <-chan struct{} {
|
||||
}{nil, err}
|
||||
return
|
||||
}
|
||||
client.API().BotsSetBotCommands(ctx, &tg.BotsSetBotCommandsRequest{
|
||||
Scope: &tg.BotCommandScopeDefault{},
|
||||
})
|
||||
commands := make([]tg.BotCommand, 0, len(handlers.CommandHandlers))
|
||||
for _, info := range handlers.CommandHandlers {
|
||||
commands = append(commands, tg.BotCommand{Command: info.Cmd, Description: i18n.T(info.Desc)})
|
||||
|
||||
@@ -23,7 +23,11 @@ import (
|
||||
)
|
||||
|
||||
func handleAddCallback(ctx *ext.Context, update *ext.Update) error {
|
||||
dataid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
|
||||
dataParts := strings.Split(string(update.CallbackQuery.Data), " ")
|
||||
if len(dataParts) < 2 {
|
||||
return fmt.Errorf("invalid callback data: %q", update.CallbackQuery.Data)
|
||||
}
|
||||
dataid := dataParts[1]
|
||||
data, err := shortcut.GetCallbackDataWithAnswer[tcbdata.Add](ctx, update, dataid)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
@@ -14,7 +15,11 @@ import (
|
||||
)
|
||||
|
||||
func handleCancelCallback(ctx *ext.Context, update *ext.Update) error {
|
||||
taskid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
|
||||
dataParts := strings.Split(string(update.CallbackQuery.Data), " ")
|
||||
if len(dataParts) < 2 {
|
||||
return fmt.Errorf("invalid callback data: %q", update.CallbackQuery.Data)
|
||||
}
|
||||
taskid := dataParts[1]
|
||||
if err := core.CancelTask(ctx, taskid); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to cancel task %s: %v", taskid, err)
|
||||
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(update.CallbackQuery.GetQueryID(), i18n.T(i18nk.BotMsgCancelErrorCancelFailed, map[string]any{
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/database"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
@@ -42,7 +43,9 @@ func handleDirCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
if _, err := storage.GetStorageByUserIDAndName(ctx, user.ChatID, args[2]); err != nil {
|
||||
ctx.Reply(update, ext.ReplyTextString(err.Error()), nil)
|
||||
logger.Errorf("Failed to get storage %q: %s", args[2], err)
|
||||
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgCommonErrorGetStorageFailed,
|
||||
tgutil.EscapeHTMLTemplateData(map[string]any{"Error": err.Error()}))), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
|
||||
@@ -20,9 +20,17 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
|
||||
// mediaGroupKey uniquely identifies a media group by chat, sender, and group
|
||||
// ID so files from different users in the same chat can never be mixed.
|
||||
type mediaGroupKey struct {
|
||||
chatID int64
|
||||
userID int64
|
||||
groupID int64
|
||||
}
|
||||
|
||||
type MediaGroupHandler struct {
|
||||
groups map[int64][]tfile.TGFileMessage
|
||||
timers map[int64]*time.Timer
|
||||
groups map[mediaGroupKey][]tfile.TGFileMessage
|
||||
timers map[mediaGroupKey]*time.Timer
|
||||
mu sync.Mutex
|
||||
timeout time.Duration
|
||||
setupOnce sync.Once
|
||||
@@ -39,8 +47,8 @@ func (m *MediaGroupHandler) SetupTimeout(timeoutSec int) {
|
||||
|
||||
var (
|
||||
mediaGroupHandler = &MediaGroupHandler{
|
||||
groups: make(map[int64][]tfile.TGFileMessage),
|
||||
timers: make(map[int64]*time.Timer),
|
||||
groups: make(map[mediaGroupKey][]tfile.TGFileMessage),
|
||||
timers: make(map[mediaGroupKey]*time.Timer),
|
||||
mu: sync.Mutex{},
|
||||
}
|
||||
)
|
||||
@@ -66,32 +74,37 @@ func handleGroupMediaMessage(ctx *ext.Context, update *ext.Update, message *tg.M
|
||||
}
|
||||
mediaGroupHandler.mu.Lock()
|
||||
defer mediaGroupHandler.mu.Unlock()
|
||||
if mediaGroupHandler.groups[groupID] == nil {
|
||||
mediaGroupHandler.groups[groupID] = make([]tfile.TGFileMessage, 0)
|
||||
key := mediaGroupKey{
|
||||
chatID: update.EffectiveChat().GetID(),
|
||||
userID: userId,
|
||||
groupID: groupID,
|
||||
}
|
||||
mediaGroupHandler.groups[groupID] = append(mediaGroupHandler.groups[groupID], file)
|
||||
if mediaGroupHandler.groups[key] == nil {
|
||||
mediaGroupHandler.groups[key] = make([]tfile.TGFileMessage, 0)
|
||||
}
|
||||
mediaGroupHandler.groups[key] = append(mediaGroupHandler.groups[key], file)
|
||||
|
||||
if timer, exists := mediaGroupHandler.timers[groupID]; exists {
|
||||
if timer, exists := mediaGroupHandler.timers[key]; exists {
|
||||
timer.Stop()
|
||||
}
|
||||
mediaGroupHandler.timers[groupID] = time.AfterFunc(mediaGroupHandler.timeout, func() {
|
||||
processMediaGroup(ctx, update, groupID)
|
||||
mediaGroupHandler.timers[key] = time.AfterFunc(mediaGroupHandler.timeout, func() {
|
||||
processMediaGroup(ctx, update, key)
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
func processMediaGroup(ctx *ext.Context, update *ext.Update, groupID int64) {
|
||||
func processMediaGroup(ctx *ext.Context, update *ext.Update, key mediaGroupKey) {
|
||||
logger := log.FromContext(ctx)
|
||||
mediaGroupHandler.mu.Lock()
|
||||
items := mediaGroupHandler.groups[groupID]
|
||||
delete(mediaGroupHandler.groups, groupID)
|
||||
delete(mediaGroupHandler.timers, groupID)
|
||||
items := mediaGroupHandler.groups[key]
|
||||
delete(mediaGroupHandler.groups, key)
|
||||
delete(mediaGroupHandler.timers, key)
|
||||
mediaGroupHandler.mu.Unlock()
|
||||
if len(items) == 0 {
|
||||
logger.Warn("No media items to process for group", "groupID", groupID)
|
||||
logger.Warn("No media items to process for group", "groupID", key.groupID)
|
||||
return
|
||||
}
|
||||
logger.Debugf("Processing media group %d with %d items", groupID, len(items))
|
||||
logger.Debugf("Processing media group %d with %d items", key.groupID, len(items))
|
||||
|
||||
userId := update.GetUserChat().GetID()
|
||||
msg, err := ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgMediaGroupInfoSavingFiles, nil)), nil)
|
||||
|
||||
@@ -22,6 +22,17 @@ func checkPermission(ctx *ext.Context, update *ext.Update) error {
|
||||
return dispatcher.ContinueGroups
|
||||
}
|
||||
|
||||
// withPermission wraps a callback handler with the same whitelist check used
|
||||
// for message handlers (checkPermission).
|
||||
func withPermission(handler func(*ext.Context, *ext.Update) error) func(*ext.Context, *ext.Update) error {
|
||||
return func(ctx *ext.Context, update *ext.Update) error {
|
||||
if err := checkPermission(ctx, update); err != nil {
|
||||
return err
|
||||
}
|
||||
return handler(ctx, update)
|
||||
}
|
||||
}
|
||||
|
||||
func handleSilentMode(next func(*ext.Context, *ext.Update) error, handler func(*ext.Context, *ext.Update) error) func(*ext.Context, *ext.Update) error {
|
||||
return func(ctx *ext.Context, update *ext.Update) error {
|
||||
userID := update.GetUserChat().GetID()
|
||||
|
||||
@@ -56,11 +56,11 @@ func Register(disp dispatcher.Dispatcher) {
|
||||
for _, info := range CommandHandlers {
|
||||
disp.AddHandler(handlers.NewCommand(info.Cmd, info.handler))
|
||||
}
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("update"), handleUpdateCallback))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeAdd), handleAddCallback))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeSetDefault), handleSetDefaultCallback))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeCancel), handleCancelCallback))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeConfig), handleConfigCallback))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("update"), withPermission(handleUpdateCallback)))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeAdd), withPermission(handleAddCallback)))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeSetDefault), withPermission(handleSetDefaultCallback)))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeCancel), withPermission(handleCancelCallback)))
|
||||
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeConfig), withPermission(handleConfigCallback)))
|
||||
disp.AddHandler(handlers.NewMessage(sabotfilters.RegexUrl(regexp.MustCompile(re.TgMessageLinkRegexString)), handleSilentMode(handleMessageLink, handleSilentSaveLink)))
|
||||
disp.AddHandler(handlers.NewMessage(sabotfilters.RegexUrl(regexp.MustCompile(re.TelegraphUrlRegexString)), handleSilentMode(handleTelegraphUrlMessage, handleSilentSaveTelegraph)))
|
||||
disp.AddHandler(handlers.NewMessage(filters.Message.Media, handleSilentMode(handleMediaMessage, handleSilentSaveMedia)))
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
@@ -43,7 +44,11 @@ func handleSilentCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
|
||||
func handleSetDefaultCallback(ctx *ext.Context, update *ext.Update) error {
|
||||
dataid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
|
||||
dataParts := strings.Split(string(update.CallbackQuery.Data), " ")
|
||||
if len(dataParts) < 2 {
|
||||
return fmt.Errorf("invalid callback data: %q", update.CallbackQuery.Data)
|
||||
}
|
||||
dataid := dataParts[1]
|
||||
data, ok := cache.Get[tcbdata.SetDefaultStorage](dataid)
|
||||
|
||||
failedAnswer := func(message string) error {
|
||||
|
||||
@@ -89,7 +89,11 @@ func showQueuedTasks(ctx *ext.Context, update *ext.Update) {
|
||||
styling.Bold(i18n.T(i18nk.BotMsgTasksQueuedTitle)),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgTasksTotalPrefix, map[string]any{"Count": len(tasks)})),
|
||||
)
|
||||
for _, t := range tasks {
|
||||
const maxShown = 10
|
||||
for i, t := range tasks {
|
||||
if i >= maxShown {
|
||||
break
|
||||
}
|
||||
created := t.Created.In(time.Local).Format("2006-01-02 15:04:05")
|
||||
status := i18n.T(i18nk.BotMsgTasksStatusQueued)
|
||||
if t.Cancelled {
|
||||
@@ -105,10 +109,9 @@ func showQueuedTasks(ctx *ext.Context, update *ext.Update) {
|
||||
styling.Plain("\n"+i18n.T(i18nk.BotMsgTasksFieldStatus)),
|
||||
styling.Code(status),
|
||||
)
|
||||
if len(tasks) > 10 {
|
||||
opts = append(opts, styling.Plain("\n"+i18n.T(i18nk.BotMsgTasksTruncatedNote, map[string]any{"Count": len(tasks)})))
|
||||
break
|
||||
}
|
||||
}
|
||||
if len(tasks) > maxShown {
|
||||
opts = append(opts, styling.Plain("\n"+i18n.T(i18nk.BotMsgTasksTruncatedNote, map[string]any{"Count": len(tasks)})))
|
||||
}
|
||||
ctx.Reply(update, ext.ReplyTextStyledTextArray(opts), nil)
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/gotd/td/telegram/message/html"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -100,7 +101,10 @@ func handleUpdateCmd(ctx *ext.Context, u *ext.Update) error {
|
||||
func handleUpdateCallback(ctx *ext.Context, u *ext.Update) error {
|
||||
currentV, err := semver.Parse(config.Version)
|
||||
if err != nil {
|
||||
return err
|
||||
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(u.CallbackQuery.GetQueryID(), i18n.T(i18nk.BotMsgUpdateErrorVersionVarInvalid, map[string]any{
|
||||
"Error": err.Error(),
|
||||
})))
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
ctx.EditMessage(u.GetUserChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: u.CallbackQuery.GetMsgID(),
|
||||
|
||||
@@ -2,6 +2,7 @@ package user
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/celestix/gotgproto"
|
||||
@@ -20,17 +21,18 @@ import (
|
||||
)
|
||||
|
||||
var uc *gotgproto.Client
|
||||
var ectx *ext.Context
|
||||
|
||||
// getEctx lazily creates the user-client ext.Context exactly once. Guarded by
|
||||
// sync.OnceValue so concurrent GetCtx calls cannot race on ectx creation.
|
||||
var getEctx = sync.OnceValue(func() *ext.Context {
|
||||
return uc.CreateContext()
|
||||
})
|
||||
|
||||
func GetCtx() *ext.Context {
|
||||
if ectx != nil {
|
||||
return ectx
|
||||
}
|
||||
if uc == nil {
|
||||
return nil
|
||||
}
|
||||
ectx = uc.CreateContext()
|
||||
return ectx
|
||||
return getEctx()
|
||||
}
|
||||
|
||||
func Login(ctx context.Context) (*gotgproto.Client, error) {
|
||||
|
||||
@@ -22,7 +22,13 @@ func main() {
|
||||
pkg := flag.String("pkg", "i18nk", "Package name for generated file")
|
||||
flag.Parse()
|
||||
|
||||
type localeFile struct {
|
||||
path string
|
||||
keys map[string]struct{}
|
||||
}
|
||||
|
||||
keys := make(map[string]struct{})
|
||||
var localeFiles []localeFile
|
||||
|
||||
err := filepath.WalkDir(*dir, func(path string, d fs.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
@@ -42,7 +48,12 @@ func main() {
|
||||
return fmt.Errorf("failed to parse yaml %s: %w", path, err)
|
||||
}
|
||||
|
||||
collectKeys(content, "", keys)
|
||||
fileKeys := make(map[string]struct{})
|
||||
collectKeys(content, "", fileKeys)
|
||||
localeFiles = append(localeFiles, localeFile{path: path, keys: fileKeys})
|
||||
for k := range fileKeys {
|
||||
keys[k] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
@@ -50,6 +61,25 @@ func main() {
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// 一致性校验: 每个语言文件必须包含全部 key
|
||||
invalid := false
|
||||
for _, f := range localeFiles {
|
||||
var missing []string
|
||||
for k := range keys {
|
||||
if _, ok := f.keys[k]; !ok {
|
||||
missing = append(missing, k)
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
invalid = true
|
||||
sort.Strings(missing)
|
||||
fmt.Fprintf(os.Stderr, "Error: locale file %s is missing %d key(s): %s\n", f.path, len(missing), strings.Join(missing, ", "))
|
||||
}
|
||||
}
|
||||
if invalid {
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var list []string
|
||||
for k := range keys {
|
||||
list = append(list, k)
|
||||
|
||||
@@ -61,6 +61,7 @@ func Run(cmd *cobra.Command, _ []string) {
|
||||
<-ctx.Done()
|
||||
logger.Info("Exiting...")
|
||||
defer logger.Info("Exit complete")
|
||||
core.Close()
|
||||
cleanCache()
|
||||
}
|
||||
|
||||
@@ -87,7 +88,7 @@ func initAll(ctx context.Context) (<-chan struct{}, error) {
|
||||
}
|
||||
}
|
||||
if err := api.Start(ctx); err != nil {
|
||||
logger.Error("Failed to start API server", "error", err)
|
||||
logger.Fatal("Failed to start API server", "error", err)
|
||||
}
|
||||
return bot.Init(ctx), nil
|
||||
}
|
||||
|
||||
33
common/cache/ristretto.go
vendored
33
common/cache/ristretto.go
vendored
@@ -2,6 +2,7 @@ package cache
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
@@ -9,24 +10,26 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
)
|
||||
|
||||
var cache *ristretto.Cache[string, any]
|
||||
var (
|
||||
cache *ristretto.Cache[string, any]
|
||||
initOnce sync.Once
|
||||
)
|
||||
|
||||
func Init() {
|
||||
if cache != nil {
|
||||
panic("cache already initialized")
|
||||
}
|
||||
c, err := ristretto.NewCache(&ristretto.Config[string, any]{
|
||||
NumCounters: config.C().Cache.NumCounters,
|
||||
MaxCost: config.C().Cache.MaxCost,
|
||||
BufferItems: 64,
|
||||
OnReject: func(item *ristretto.Item[any]) {
|
||||
log.Warnf("Cache item rejected: key=%d, value=%v", item.Key, item.Value)
|
||||
},
|
||||
initOnce.Do(func() {
|
||||
c, err := ristretto.NewCache(&ristretto.Config[string, any]{
|
||||
NumCounters: config.C().Cache.NumCounters,
|
||||
MaxCost: config.C().Cache.MaxCost,
|
||||
BufferItems: 64,
|
||||
OnReject: func(item *ristretto.Item[any]) {
|
||||
log.Warnf("Cache item rejected: key=%d, value=%v", item.Key, item.Value)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("failed to create ristretto cache: %v", err)
|
||||
}
|
||||
cache = c
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("failed to create ristretto cache: %v", err)
|
||||
}
|
||||
cache = c
|
||||
}
|
||||
|
||||
func Set(key string, value any) error {
|
||||
|
||||
@@ -192,8 +192,10 @@ const (
|
||||
BotMsgProgressSingleDownloadingUnknown Key = "bot.msg.progress.single_downloading_unknown"
|
||||
BotMsgProgressSingleFailed Key = "bot.msg.progress.single_failed"
|
||||
BotMsgProgressSingleStatusHeader Key = "bot.msg.progress.single_status_header"
|
||||
BotMsgProgressSingleUploading Key = "bot.msg.progress.single_uploading"
|
||||
BotMsgProgressSingleUploadRetrying Key = "bot.msg.progress.single_upload_retrying"
|
||||
BotMsgProgressSingleUploading Key = "bot.msg.progress.single_uploading"
|
||||
BotMsgProgressSizeWithFiles Key = "bot.msg.progress.size_with_files"
|
||||
BotMsgProgressSizeWithResources Key = "bot.msg.progress.size_with_resources"
|
||||
BotMsgProgressTaskCanceledWithId Key = "bot.msg.progress.task_canceled_with_id"
|
||||
BotMsgProgressTaskFailedWithError Key = "bot.msg.progress.task_failed_with_error"
|
||||
BotMsgProgressTelegraphDonePrefix Key = "bot.msg.progress.telegraph_done_prefix"
|
||||
@@ -244,7 +246,6 @@ const (
|
||||
BotMsgSaveHelpText Key = "bot.msg.save_help_text"
|
||||
BotMsgStorageInfoFilenamePrefix Key = "bot.msg.storage.info_filename_prefix"
|
||||
BotMsgStorageInfoPromptSelectStorage Key = "bot.msg.storage.info_prompt_select_storage"
|
||||
BotMsgSyncpeersDone Key = "bot.msg.syncpeers.done"
|
||||
BotMsgSyncpeersFailed Key = "bot.msg.syncpeers.failed"
|
||||
BotMsgSyncpeersStart Key = "bot.msg.syncpeers.start"
|
||||
BotMsgSyncpeersSuccess Key = "bot.msg.syncpeers.success"
|
||||
|
||||
@@ -384,6 +384,8 @@ bot:
|
||||
single_canceled: "<b>🚫 Task canceled</b>\n\nFilename: <code>{{.Name}}</code>"
|
||||
single_failed: "<b>❌ Processing failed</b>\n\nFilename: <code>{{.Name}}</code>\nReason: <code>{{.Reason}}</code>"
|
||||
downloading_prefix: "Downloading\nTotal size: "
|
||||
size_with_files: "{{.Size}} ({{.Count}} files)"
|
||||
size_with_resources: "{{.Size}} ({{.Count}} resources)"
|
||||
processing_list_prefix: "\nProcessing:\n"
|
||||
processing_none: " - None"
|
||||
avg_speed_prefix: "\nAverage speed: "
|
||||
@@ -424,7 +426,7 @@ bot:
|
||||
transfer_failed_files_prefix: "\nFailed files: "
|
||||
syncpeers:
|
||||
start: "Starting to sync peers..."
|
||||
done: "Peer sync completed, total {{.Count}} chats synced"
|
||||
success: "Peer sync completed, total {{.Count}} chats synced"
|
||||
failed: "Peer sync failed: {{.Error}}"
|
||||
aria2:
|
||||
error_aria2_not_enabled: "Aria2 feature is not enabled in the configuration"
|
||||
|
||||
@@ -238,9 +238,9 @@ bot:
|
||||
info_install_plugin_success: "插件安装成功: {{.Name}}"
|
||||
parse:
|
||||
info_parsing: "正在解析..."
|
||||
error_parse_text_failed: "Failed to parse text: {{.Error}}"
|
||||
error_build_storage_select_keyboard_failed: "Failed to build storage selection keyboard: {{.Error}}"
|
||||
error_build_parsed_text_entity_failed: "Failed to build parsed text entity: {{.Error}}"
|
||||
error_parse_text_failed: "解析文本失败: {{.Error}}"
|
||||
error_build_storage_select_keyboard_failed: "构建存储选择键盘失败: {{.Error}}"
|
||||
error_build_parsed_text_entity_failed: "构建解析文本实体失败: {{.Error}}"
|
||||
info_link_prefix: "\n链接: "
|
||||
info_author_prefix: "\n作者: "
|
||||
info_description_prefix: "\n描述: "
|
||||
@@ -385,6 +385,8 @@ bot:
|
||||
single_canceled: "<b>🚫 任务已取消</b>\n\n文件名:<code>{{.Name}}</code>"
|
||||
single_failed: "<b>❌ 处理失败</b>\n\n文件名:<code>{{.Name}}</code>\n原因:<code>{{.Reason}}</code>"
|
||||
downloading_prefix: "正在下载\n总大小: "
|
||||
size_with_files: "{{.Size}} ({{.Count}} 个文件)"
|
||||
size_with_resources: "{{.Size}} ({{.Count}} 个资源)"
|
||||
processing_list_prefix: "\n正在处理:\n"
|
||||
processing_none: " - 无"
|
||||
avg_speed_prefix: "\n平均速度: "
|
||||
|
||||
27
common/utils/fsutil/unique.go
Normal file
27
common/utils/fsutil/unique.go
Normal file
@@ -0,0 +1,27 @@
|
||||
package fsutil
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path"
|
||||
"strings"
|
||||
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
// UniquePath returns a non-taken path under basePath: the name itself, then
|
||||
// numbered variants, then a random suffix.
|
||||
func UniquePath(basePath, name string, exists func(candidate string) bool, maxAttempts int) string {
|
||||
candidate := path.Join(basePath, name)
|
||||
if !exists(candidate) {
|
||||
return candidate
|
||||
}
|
||||
ext := path.Ext(name)
|
||||
stem := strings.TrimSuffix(name, ext)
|
||||
for i := 1; i <= maxAttempts; i++ {
|
||||
candidate = path.Join(basePath, fmt.Sprintf("%s_%d%s", stem, i, ext))
|
||||
if !exists(candidate) {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return path.Join(basePath, fmt.Sprintf("%s_%s%s", stem, xid.New().String(), ext))
|
||||
}
|
||||
50
common/utils/progressutil/progress.go
Normal file
50
common/utils/progressutil/progress.go
Normal file
@@ -0,0 +1,50 @@
|
||||
// Package progressutil provides shared progress-update throttling for task
|
||||
// progress trackers.
|
||||
package progressutil
|
||||
|
||||
// updatesLevels picks the percent step by file size.
|
||||
var updatesLevels = []struct {
|
||||
size int64 // file size threshold
|
||||
stepPercent int // minimum percent step between updates
|
||||
}{
|
||||
{10 << 20, 100},
|
||||
{50 << 20, 20},
|
||||
{200 << 20, 10},
|
||||
{500 << 20, 5},
|
||||
}
|
||||
|
||||
// ShouldUpdate reports whether a byte-based progress update should be shown.
|
||||
func ShouldUpdate(total, downloaded int64, lastUpdatePercent int) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
percent := int((downloaded * 100) / total)
|
||||
if percent <= lastUpdatePercent {
|
||||
return false
|
||||
}
|
||||
|
||||
step := updatesLevels[len(updatesLevels)-1].stepPercent
|
||||
for _, lvl := range updatesLevels {
|
||||
if total < lvl.size {
|
||||
step = lvl.stepPercent
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return percent >= lastUpdatePercent+step
|
||||
}
|
||||
|
||||
// ShouldUpdateCount reports whether a count-based progress update (e.g. files
|
||||
// downloaded so far) should be shown: every 10 units, or when finished.
|
||||
func ShouldUpdateCount(downloaded, total int64) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
const step = int64(10)
|
||||
if downloaded < step {
|
||||
return downloaded == total
|
||||
}
|
||||
return downloaded%step == 0 || downloaded == total
|
||||
}
|
||||
@@ -194,97 +194,6 @@ func getMessagesRange(ctx *ext.Context, chatID int64, minId, maxId int) ([]*tg.M
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// [TODO]
|
||||
// type MessageItem struct {
|
||||
// Message *tg.Message
|
||||
// Error error
|
||||
// }
|
||||
|
||||
// func IterMessages(ctx *ext.Context, chatID int64, minId, maxId int) (<-chan MessageItem, error) {
|
||||
// total := maxId - minId + 1
|
||||
// ch := make(chan MessageItem, 100)
|
||||
|
||||
// go func() {
|
||||
// defer close(ch)
|
||||
// if !ctx.Self.Bot {
|
||||
// perr := ctx.PeerStorage.GetInputPeerById(chatID)
|
||||
// if perr == nil || perr.(*tg.InputPeerEmpty) != nil {
|
||||
// ch <- MessageItem{
|
||||
// Error: fmt.Errorf("peer not found: %d", chatID),
|
||||
// }
|
||||
// return
|
||||
// }
|
||||
|
||||
// for i := 0; i < total; i += 100 {
|
||||
// start := minId + i
|
||||
// end := min(start+100, maxId)
|
||||
// msgs, err := ctx.Raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
|
||||
// Peer: perr,
|
||||
// OffsetID: start,
|
||||
// AddOffset: start - end,
|
||||
// Limit: 100,
|
||||
// })
|
||||
// if err != nil {
|
||||
// ch <- MessageItem{
|
||||
// Error: fmt.Errorf("failed to get messages: %w", err),
|
||||
// }
|
||||
// return
|
||||
// }
|
||||
// var msgClass []tg.MessageClass
|
||||
// switch msgsv := msgs.(type) {
|
||||
// case *tg.MessagesMessages:
|
||||
// msgClass = msgsv.GetMessages()
|
||||
// case *tg.MessagesMessagesSlice:
|
||||
// msgClass = msgsv.GetMessages()
|
||||
// case *tg.MessagesChannelMessages:
|
||||
// msgClass = msgsv.GetMessages()
|
||||
// default:
|
||||
// ch <- MessageItem{
|
||||
// Error: fmt.Errorf("unsupported message type: %T", msgsv),
|
||||
// }
|
||||
// continue
|
||||
// }
|
||||
// for _, msg := range msgClass {
|
||||
// msg, ok := msg.AsNotEmpty()
|
||||
// if !ok {
|
||||
// continue
|
||||
// }
|
||||
// switch msg := msg.(type) {
|
||||
// case *tg.Message:
|
||||
// key := fmt.Sprintf("tgmsg:%d:%d:%d", ctx.Self.ID, chatID, msg.GetID())
|
||||
// cache.Set(key, msg)
|
||||
// ch <- MessageItem{
|
||||
// Message: msg,
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// } else {
|
||||
// for i := 0; i < total; i += 100 {
|
||||
// start := minId + i
|
||||
// end := min(start+100, maxId)
|
||||
// msgs, err := GetMessagesRange(ctx, chatID, start, end)
|
||||
// if err != nil {
|
||||
// ch <- MessageItem{
|
||||
// Error: fmt.Errorf("failed to get messages: %w", err),
|
||||
// }
|
||||
// return
|
||||
// }
|
||||
// for _, msg := range msgs {
|
||||
// if msg == nil {
|
||||
// continue
|
||||
// }
|
||||
// ch <- MessageItem{
|
||||
// Message: msg,
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }()
|
||||
|
||||
// return ch, nil
|
||||
// }
|
||||
|
||||
func getMessageByID(ctx *ext.Context, chatID int64, msgID int) (*tg.Message, error) {
|
||||
key := fmt.Sprintf("tgmsg:%d:%d:%d", ctx.Self.ID, chatID, msgID)
|
||||
if msg, ok := cache.Get[*tg.Message](key); ok {
|
||||
|
||||
@@ -44,6 +44,25 @@ format = ""
|
||||
# 下载后转封装的视频容器格式, 留空则不转封装. 默认 mp4
|
||||
recode = "mp4"
|
||||
|
||||
# 解析器配置
|
||||
[parser]
|
||||
# 启用 JS 解析器插件 (Go 内置解析器默认启用)
|
||||
plugin_enable = false
|
||||
# 插件目录, 可以是多个目录
|
||||
plugin_dirs = ["./plugins"]
|
||||
# 解析器默认代理
|
||||
proxy = ""
|
||||
|
||||
# Twitter/X 解析器配置
|
||||
[parser.twitter]
|
||||
# 自定义 API 域名
|
||||
api_domain = "api.fxtwitter.com"
|
||||
# 单独为此解析器指定代理 (留空则使用 [parser] 中的 proxy)
|
||||
# proxy = "http://127.0.0.1:7890"
|
||||
|
||||
# Kemono 解析器配置 (暂无可配置项, 留空即可)
|
||||
[parser.kemono]
|
||||
|
||||
# HTTP API 配置
|
||||
[api]
|
||||
# 启用 HTTP API
|
||||
|
||||
@@ -10,13 +10,4 @@ type hookExecConfig struct {
|
||||
TaskSuccess string `toml:"task_success" mapstructure:"task_success" json:"task_success"`
|
||||
TaskFail string `toml:"task_fail" mapstructure:"task_fail" json:"task_fail"`
|
||||
TaskCancel string `toml:"task_cancel" mapstructure:"task_cancel" json:"task_cancel"`
|
||||
|
||||
// TaskTypes map[string]hookExecOnTypeConfig `toml:"task_types" mapstructure:"task_types" json:"task_types"` // [TODO]
|
||||
}
|
||||
|
||||
// type hookExecOnTypeConfig struct {
|
||||
// TaskBeforeStart string `toml:"task_before_start" mapstructure:"task_before_start" json:"task_before_start"`
|
||||
// TaskSuccess string `toml:"task_success" mapstructure:"task_success" json:"task_success"`
|
||||
// TaskFail string `toml:"task_fail" mapstructure:"task_fail" json:"task_fail"`
|
||||
// TaskCancel string `toml:"task_cancel" mapstructure:"task_cancel" json:"task_cancel"`
|
||||
// }
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/duke-git/lancet/v2/slice"
|
||||
"github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/spf13/viper"
|
||||
@@ -68,6 +69,13 @@ func (c Config) GetStorageByName(name string) storage.StorageConfig {
|
||||
}
|
||||
|
||||
func Init(ctx context.Context, configFile ...string) error {
|
||||
logger := log.FromContext(ctx)
|
||||
|
||||
// Reset side tables for re-init.
|
||||
storages = nil
|
||||
userIDs = nil
|
||||
userStorages = make(map[int64][]string)
|
||||
|
||||
viper.SetConfigType("toml")
|
||||
viper.SetEnvPrefix("SAVEANY")
|
||||
viper.AutomaticEnv()
|
||||
@@ -76,11 +84,13 @@ func Init(ctx context.Context, configFile ...string) error {
|
||||
|
||||
// 如果指定了配置文件路径,则使用指定的配置文件
|
||||
// 配置文件支持传入一个 http(s) URL 地址
|
||||
loadedFromURL := false
|
||||
if len(configFile) > 0 && configFile[0] != "" {
|
||||
cfg := configFile[0]
|
||||
if strings.HasPrefix(cfg, "http://") || strings.HasPrefix(cfg, "https://") {
|
||||
// 使用远程配置文件
|
||||
resp, err := http.Get(cfg)
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
resp, err := client.Get(cfg)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to fetch remote config file: %w", err)
|
||||
}
|
||||
@@ -91,6 +101,7 @@ func Init(ctx context.Context, configFile ...string) error {
|
||||
if err := viper.ReadConfig(resp.Body); err != nil {
|
||||
return fmt.Errorf("failed to read remote config file: %w", err)
|
||||
}
|
||||
loadedFromURL = true
|
||||
} else {
|
||||
viper.SetConfigFile(cfg)
|
||||
}
|
||||
@@ -141,13 +152,15 @@ func Init(ctx context.Context, configFile ...string) error {
|
||||
viper.SetDefault(key, value)
|
||||
}
|
||||
|
||||
if err := viper.ReadInConfig(); err != nil {
|
||||
fmt.Println("Error reading config file, ", err)
|
||||
return err
|
||||
if !loadedFromURL {
|
||||
if err := viper.ReadInConfig(); err != nil {
|
||||
logger.Errorf("Error reading config file: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if err := viper.Unmarshal(cfg); err != nil {
|
||||
fmt.Println("Error unmarshalling config file, ", err)
|
||||
logger.Errorf("Error unmarshalling config file: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
|
||||
29
core/core.go
29
core/core.go
@@ -3,6 +3,7 @@ package core
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -11,7 +12,18 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
)
|
||||
|
||||
var queueInstance *queue.TaskQueue[Executable]
|
||||
var (
|
||||
queueOnce sync.Once
|
||||
queueInstance *queue.TaskQueue[Executable]
|
||||
)
|
||||
|
||||
// initQueue lazily creates the shared task queue.
|
||||
func initQueue() *queue.TaskQueue[Executable] {
|
||||
queueOnce.Do(func() {
|
||||
queueInstance = queue.NewTaskQueue[Executable]()
|
||||
})
|
||||
return queueInstance
|
||||
}
|
||||
|
||||
type Executable interface {
|
||||
Type() tasktype.TaskType
|
||||
@@ -65,17 +77,22 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan
|
||||
func Run(ctx context.Context) {
|
||||
log.FromContext(ctx).Info("Start processing tasks...")
|
||||
semaphore := make(chan struct{}, config.C().Workers)
|
||||
if queueInstance == nil {
|
||||
queueInstance = queue.NewTaskQueue[Executable]()
|
||||
}
|
||||
q := initQueue()
|
||||
for range config.C().Workers {
|
||||
go worker(ctx, queueInstance, semaphore)
|
||||
go worker(ctx, q, semaphore)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// Close stops the queue and unblocks workers in Get.
|
||||
func Close() {
|
||||
if q := initQueue(); q != nil {
|
||||
q.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func AddTask(ctx context.Context, task Executable) error {
|
||||
return queueInstance.Add(queue.NewTask(ctx, task.TaskID(), task.Title(), task))
|
||||
return initQueue().Add(queue.NewTask(ctx, task.TaskID(), task.Title(), task))
|
||||
}
|
||||
|
||||
func CancelTask(ctx context.Context, id string) error {
|
||||
|
||||
@@ -2,10 +2,12 @@ package batchtfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
@@ -33,7 +35,9 @@ func (g executionGroup) usesBatchSaver() bool {
|
||||
func (t *Task) Execute(ctx context.Context) error {
|
||||
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("batch_file[%s]", t.ID))
|
||||
logger.Info("Starting batch file task")
|
||||
t.Progress.OnStart(ctx, t)
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnStart(ctx, t)
|
||||
}
|
||||
groups := t.executionGroups()
|
||||
var err error
|
||||
for i := 0; i < len(groups); {
|
||||
@@ -53,7 +57,11 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
i = end
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
if !t.IgnoreErrors || errors.Is(err, context.Canceled) {
|
||||
break
|
||||
}
|
||||
logger.Warnf("Group processing failed (ignored): %v", err)
|
||||
err = nil
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
@@ -62,10 +70,19 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
logger.Info("Batch file task completed successfully")
|
||||
}
|
||||
t.finishItems(err)
|
||||
t.Progress.OnDone(ctx, t, err)
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnDone(ctx, t, err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// notifyProgress reports a progress update to the optional tracker.
|
||||
func (t *Task) notifyProgress(ctx context.Context) {
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) executionGroups() []executionGroup {
|
||||
groups := make([]executionGroup, 0, len(t.elems))
|
||||
for i := 0; i < len(t.elems); {
|
||||
@@ -104,7 +121,13 @@ func (t *Task) processElements(ctx context.Context, elems []*TaskElement) error
|
||||
return err
|
||||
}
|
||||
defer t.unmarkProcessing(elem.ID)
|
||||
return t.processElement(gctx, *elem)
|
||||
err := t.processElement(gctx, *elem)
|
||||
if err != nil && t.IgnoreErrors && !errors.Is(err, context.Canceled) {
|
||||
// Per-item failure: keep siblings running.
|
||||
log.FromContext(ctx).Warnf("Element %s failed (ignored): %v", elem.ID, err)
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
})
|
||||
}
|
||||
return eg.Wait()
|
||||
@@ -119,23 +142,51 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
}
|
||||
}()
|
||||
|
||||
type downloadResult struct {
|
||||
elem *TaskElement
|
||||
err error
|
||||
}
|
||||
results := make([]downloadResult, len(group.elems))
|
||||
var resultsMu sync.Mutex
|
||||
|
||||
eg, gctx := errgroup.WithContext(ctx)
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for _, elem := range group.elems {
|
||||
for i, elem := range group.elems {
|
||||
eg.Go(func() error {
|
||||
if err := t.markProcessing(ctx, elem); err != nil {
|
||||
return err
|
||||
}
|
||||
defer t.unmarkProcessing(elem.ID)
|
||||
return t.downloadElement(gctx, elem)
|
||||
err := t.downloadElement(gctx, elem)
|
||||
// Store by original index.
|
||||
resultsMu.Lock()
|
||||
results[i] = downloadResult{elem: elem, err: err}
|
||||
resultsMu.Unlock()
|
||||
if err != nil && t.IgnoreErrors && !errors.Is(err, context.Canceled) {
|
||||
// Per-item failure: keep siblings running.
|
||||
log.FromContext(ctx).Warnf("Element %s failed (ignored): %v", elem.ID, err)
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
})
|
||||
}
|
||||
if err := eg.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
items := make([]storagetypes.BatchItem, 0, len(group.elems))
|
||||
openFiles := make([]*os.File, 0, len(group.elems))
|
||||
// Upload only successfully downloaded elements.
|
||||
successElems := make([]*TaskElement, 0, len(group.elems))
|
||||
for _, r := range results {
|
||||
if r.err == nil {
|
||||
successElems = append(successElems, r.elem)
|
||||
}
|
||||
}
|
||||
if len(successElems) == 0 {
|
||||
return fmt.Errorf("all elements failed to download")
|
||||
}
|
||||
|
||||
items := make([]storagetypes.BatchItem, 0, len(successElems))
|
||||
openFiles := make([]*os.File, 0, len(successElems))
|
||||
defer func() {
|
||||
for _, file := range openFiles {
|
||||
if err := file.Close(); err != nil {
|
||||
@@ -143,7 +194,7 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
}
|
||||
}
|
||||
}()
|
||||
for _, elem := range group.elems {
|
||||
for _, elem := range successElems {
|
||||
file, err := os.Open(elem.localPath)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
@@ -166,28 +217,28 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
})
|
||||
}
|
||||
for index, item := range items {
|
||||
t.recordDownloadComplete(group.elems[index].ID, item.Size)
|
||||
t.recordDownloadComplete(successElems[index].ID, item.Size)
|
||||
}
|
||||
return t.saveBatchItems(ctx, group, items)
|
||||
return t.saveBatchItems(ctx, successElems, items)
|
||||
}
|
||||
|
||||
func (t *Task) saveBatchItems(ctx context.Context, group executionGroup, items []storagetypes.BatchItem) error {
|
||||
func (t *Task) saveBatchItems(ctx context.Context, successElems []*TaskElement, items []storagetypes.BatchItem) error {
|
||||
t.startUpload(ctx)
|
||||
if progressSaver, ok := group.batchSaver.(storage.StorageBatchProgressSaver); ok {
|
||||
if progressSaver, ok := successElems[0].Storage.(storage.StorageBatchProgressSaver); ok {
|
||||
err := progressSaver.SaveBatchWithProgress(ctx, items, func(index int, uploaded, total int64) {
|
||||
if index < 0 || index >= len(group.elems) {
|
||||
if index < 0 || index >= len(successElems) {
|
||||
return
|
||||
}
|
||||
t.uploadCallback(ctx, group.elems[index].ID)(uploaded, total)
|
||||
t.uploadCallback(ctx, successElems[index].ID)(uploaded, total)
|
||||
})
|
||||
if err != nil {
|
||||
for _, elem := range group.elems {
|
||||
for _, elem := range successElems {
|
||||
t.markItemFailed(elem.ID, FailureStageBatchUpload, err)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to save batch: %w", err)
|
||||
}
|
||||
for index, elem := range group.elems {
|
||||
for index, elem := range successElems {
|
||||
t.uploadCallback(ctx, elem.ID)(items[index].Size, items[index].Size)
|
||||
t.markItemCompleted(elem.ID)
|
||||
}
|
||||
@@ -198,17 +249,17 @@ func (t *Task) saveBatchItems(ctx context.Context, group executionGroup, items [
|
||||
items[i].Reader = ioutil.NewProgressReader(
|
||||
items[i].Reader,
|
||||
items[i].Size,
|
||||
t.uploadCallback(ctx, group.elems[i].ID),
|
||||
t.uploadCallback(ctx, successElems[i].ID),
|
||||
)
|
||||
}
|
||||
if err := group.batchSaver.SaveBatch(ctx, items); err != nil {
|
||||
for _, elem := range group.elems {
|
||||
if err := successElems[0].Storage.(storage.StorageBatchSaver).SaveBatch(ctx, items); err != nil {
|
||||
for _, elem := range successElems {
|
||||
t.markItemFailed(elem.ID, FailureStageBatchUpload, err)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to save batch: %w", err)
|
||||
}
|
||||
for index, elem := range group.elems {
|
||||
for index, elem := range successElems {
|
||||
t.uploadCallback(ctx, elem.ID)(items[index].Size, items[index].Size)
|
||||
t.markItemCompleted(elem.ID)
|
||||
}
|
||||
@@ -225,7 +276,7 @@ func (t *Task) markProcessing(ctx context.Context, elem *TaskElement) error {
|
||||
t.processing[elem.ID] = elem
|
||||
t.processingMu.Unlock()
|
||||
t.markItemActive(elem.ID, elem.stream, time.Now())
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
t.notifyProgress(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -247,7 +298,7 @@ func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
t.notifyProgress(ctx)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
@@ -274,7 +325,7 @@ func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||
}
|
||||
}
|
||||
t.markItemDownloaded(elem.ID)
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
t.notifyProgress(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -295,7 +346,7 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
wr := ioutil.NewProgressWriter(pw, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
t.notifyProgress(ctx)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
@@ -318,7 +369,12 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
if err := errg.Wait(); err != nil {
|
||||
return fmt.Errorf("failed to download file in stream mode: %w", err)
|
||||
}
|
||||
t.recordDownloadComplete(elem.ID, 0)
|
||||
// Streamed bytes are the uploaded bytes.
|
||||
var streamedBytes int64
|
||||
t.updateItem(elem.ID, func(item *itemProgressState) {
|
||||
streamedBytes = item.downloaded
|
||||
})
|
||||
t.recordDownloadComplete(elem.ID, streamedBytes)
|
||||
t.markItemCompleted(elem.ID)
|
||||
t.notifyStateChange(ctx)
|
||||
logger.Info("File downloaded successfully in stream mode")
|
||||
@@ -339,7 +395,7 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
t.notifyProgress(ctx)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
|
||||
35
core/tasks/batchtfile/notify_test.go
Normal file
35
core/tasks/batchtfile/notify_test.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package batchtfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type recordingTracker struct {
|
||||
calls atomic.Int64
|
||||
}
|
||||
|
||||
func (r *recordingTracker) OnStart(context.Context, TaskInfo) {}
|
||||
func (r *recordingTracker) OnDone(context.Context, TaskInfo, error) {
|
||||
}
|
||||
func (r *recordingTracker) OnProgress(context.Context, TaskInfo) {
|
||||
r.calls.Add(1)
|
||||
}
|
||||
|
||||
// Regression: notifyProgress must call the tracker, not itself. The previous
|
||||
// self-call recursed until stack overflow on any batch task with a tracker.
|
||||
func TestNotifyProgressCallsTracker(t *testing.T) {
|
||||
tracker := &recordingTracker{}
|
||||
task := &Task{Progress: tracker}
|
||||
task.notifyProgress(t.Context())
|
||||
if tracker.calls.Load() != 1 {
|
||||
t.Fatalf("expected 1 tracker call, got %d", tracker.calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// The nil tracker path must stay a no-op.
|
||||
func TestNotifyProgressNilTracker(t *testing.T) {
|
||||
task := &Task{}
|
||||
task.notifyProgress(t.Context()) // must not panic
|
||||
}
|
||||
@@ -194,10 +194,13 @@ func buildBatchDoneMarkup(info TaskInfo, skipped []string, err error) string {
|
||||
totalSize = info.TotalSize()
|
||||
}
|
||||
if err == nil {
|
||||
if len(skipped) > 0 {
|
||||
completed, _, _, failed := itemCounts(items)
|
||||
// Report per-element failures instead of full completion.
|
||||
totalSkipped := len(skipped) + failed
|
||||
if totalSkipped > 0 {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchDoneWithSkipped, map[string]any{
|
||||
"Success": len(items),
|
||||
"Skipped": len(skipped),
|
||||
"Success": completed,
|
||||
"Skipped": totalSkipped,
|
||||
"Size": dlutil.FormatSize(totalSize),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -47,7 +47,6 @@ type Task struct {
|
||||
uploadOnce sync.Once
|
||||
uploadMu sync.Mutex
|
||||
uploaded map[string]int64
|
||||
failed map[string]error // [TODO] errors for each element
|
||||
}
|
||||
|
||||
// Title implements core.Exectable.
|
||||
@@ -136,7 +135,6 @@ func NewBatchTGFileTask(
|
||||
uploaded: make(map[string]int64),
|
||||
IgnoreErrors: ignoreErrors,
|
||||
processingMu: sync.RWMutex{},
|
||||
failed: make(map[string]error),
|
||||
}
|
||||
return task
|
||||
}
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
package batchtfile
|
||||
|
||||
var progressUpdatesLevels = []struct {
|
||||
size int64 // 文件大小阈值
|
||||
stepPercent int // 每多少 % 更新一次
|
||||
}{
|
||||
{10 << 20, 100},
|
||||
{50 << 20, 20},
|
||||
{200 << 20, 10},
|
||||
{500 << 20, 5},
|
||||
}
|
||||
|
||||
func shouldUpdateProgress(total, downloaded int64, lastUpdatePercent int) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
percent := int((downloaded * 100) / total)
|
||||
if percent <= lastUpdatePercent {
|
||||
return false
|
||||
}
|
||||
|
||||
step := progressUpdatesLevels[len(progressUpdatesLevels)-1].stepPercent
|
||||
for _, lvl := range progressUpdatesLevels {
|
||||
if total < lvl.size {
|
||||
step = lvl.stepPercent
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return percent >= lastUpdatePercent+step
|
||||
}
|
||||
@@ -76,12 +76,11 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for _, file := range t.files {
|
||||
eg.Go(func() error {
|
||||
t.processingMu.RLock()
|
||||
t.processingMu.Lock()
|
||||
if _, ok := t.processing[file.URL]; ok {
|
||||
t.processingMu.Unlock()
|
||||
return fmt.Errorf("file %s is already being processed", file.URL)
|
||||
}
|
||||
t.processingMu.RUnlock()
|
||||
t.processingMu.Lock()
|
||||
t.processing[file.URL] = file
|
||||
t.processingMu.Unlock()
|
||||
defer func() {
|
||||
@@ -90,7 +89,6 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
t.processingMu.Unlock()
|
||||
}()
|
||||
err := t.processLink(gctx, file)
|
||||
t.downloaded.Add(1)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
logger.Debug("Link processing canceled")
|
||||
return err
|
||||
@@ -99,6 +97,7 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
logger.Errorf("Error processing link %s: %v", file.URL, err)
|
||||
return fmt.Errorf("failed to process link %s: %w", file.URL, err)
|
||||
}
|
||||
t.downloaded.Add(1)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/progressutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
)
|
||||
|
||||
@@ -102,7 +103,7 @@ func (p *Progress) OnDone(ctx context.Context, info TaskInfo, err error) {
|
||||
|
||||
// OnProgress implements ProgressTracker.
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
if !shouldUpdateProgress(info.TotalBytes(), info.DownloadedBytes(), int(p.lastUpdatePercent.Load())) {
|
||||
if !progressutil.ShouldUpdate(info.TotalBytes(), info.DownloadedBytes(), int(p.lastUpdatePercent.Load())) {
|
||||
return
|
||||
}
|
||||
percent := int((info.DownloadedBytes() * 100) / info.TotalBytes())
|
||||
@@ -115,7 +116,10 @@ func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressDownloadingPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB (%d个文件)", float64(info.TotalBytes())/(1024*1024), info.TotalFiles())),
|
||||
styling.Code(i18n.T(i18nk.BotMsgProgressSizeWithFiles, map[string]any{
|
||||
"Size": fmt.Sprintf("%.2f MB", float64(info.TotalBytes())/(1024*1024)),
|
||||
"Count": info.TotalFiles(),
|
||||
})),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressProcessingListPrefix, nil)),
|
||||
func() styling.StyledTextOption {
|
||||
var lines []string
|
||||
|
||||
@@ -45,7 +45,6 @@ type Task struct {
|
||||
downloaded atomic.Int64 // downloaded files count
|
||||
processing map[string]*File // {"url": File}
|
||||
processingMu sync.RWMutex
|
||||
failed map[string]error // [TODO] errors for each file
|
||||
}
|
||||
|
||||
// Title implements core.Exectable.
|
||||
@@ -127,7 +126,6 @@ func NewTask(
|
||||
client: http.DefaultClient,
|
||||
processing: make(map[string]*File),
|
||||
processingMu: sync.RWMutex{},
|
||||
failed: make(map[string]error),
|
||||
totalFiles: int64(len(files)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -207,34 +207,3 @@ func parseFilenameFallback(cd string) string {
|
||||
|
||||
return decodeFilenameParam(value)
|
||||
}
|
||||
|
||||
var progressUpdatesLevels = []struct {
|
||||
size int64 // 文件大小阈值
|
||||
stepPercent int // 每多少 % 更新一次
|
||||
}{
|
||||
{10 << 20, 100},
|
||||
{50 << 20, 50},
|
||||
{200 << 20, 20},
|
||||
{500 << 20, 10},
|
||||
}
|
||||
|
||||
func shouldUpdateProgress(total, downloaded int64, lastUpdatePercent int) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
percent := int((downloaded * 100) / total)
|
||||
if percent <= lastUpdatePercent {
|
||||
return false
|
||||
}
|
||||
|
||||
step := progressUpdatesLevels[len(progressUpdatesLevels)-1].stepPercent
|
||||
for _, lvl := range progressUpdatesLevels {
|
||||
if total < lvl.size {
|
||||
step = lvl.stepPercent
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return percent >= lastUpdatePercent+step
|
||||
}
|
||||
|
||||
@@ -30,21 +30,20 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for _, resource := range t.item.Resources {
|
||||
eg.Go(func() error {
|
||||
t.processingMu.RLock()
|
||||
if t.processing[resource.ID()] != nil {
|
||||
return fmt.Errorf("resource %s is already being processed", resource.ID())
|
||||
}
|
||||
t.processingMu.RUnlock()
|
||||
resourceID := resource.ID()
|
||||
t.processingMu.Lock()
|
||||
t.processing[resource.ID()] = &resource
|
||||
if t.processing[resourceID] != nil {
|
||||
t.processingMu.Unlock()
|
||||
return fmt.Errorf("resource %s is already being processed", resourceID)
|
||||
}
|
||||
t.processing[resourceID] = &resource
|
||||
t.processingMu.Unlock()
|
||||
defer func() {
|
||||
t.processingMu.Lock()
|
||||
delete(t.processing, resource.URL)
|
||||
delete(t.processing, resourceID)
|
||||
t.processingMu.Unlock()
|
||||
}()
|
||||
err := t.processResource(gctx, resource)
|
||||
t.downloaded.Add(1)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
logger.Debug("Resource processing canceled")
|
||||
return err
|
||||
@@ -53,6 +52,7 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
logger.Errorf("Error processing resource %s: %v", resource.URL, err)
|
||||
return fmt.Errorf("failed to process resource %s: %w", resource.URL, err)
|
||||
}
|
||||
t.downloaded.Add(1)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -15,40 +15,10 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/progressutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
)
|
||||
|
||||
var progressUpdatesLevels = []struct {
|
||||
size int64 // 文件大小阈值
|
||||
stepPercent int // 每多少 % 更新一次
|
||||
}{
|
||||
{10 << 20, 100},
|
||||
{50 << 20, 50},
|
||||
{200 << 20, 20},
|
||||
{500 << 20, 10},
|
||||
}
|
||||
|
||||
func shouldUpdateProgress(total, downloaded int64, lastUpdatePercent int) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
percent := int((downloaded * 100) / total)
|
||||
if percent <= lastUpdatePercent {
|
||||
return false
|
||||
}
|
||||
|
||||
step := progressUpdatesLevels[len(progressUpdatesLevels)-1].stepPercent
|
||||
for _, lvl := range progressUpdatesLevels {
|
||||
if total < lvl.size {
|
||||
step = lvl.stepPercent
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return percent >= lastUpdatePercent+step
|
||||
}
|
||||
|
||||
type ProgressTracker interface {
|
||||
OnStart(ctx context.Context, info TaskInfo)
|
||||
OnProgress(ctx context.Context, info TaskInfo)
|
||||
@@ -73,7 +43,10 @@ func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressParsedStartPrefix, map[string]any{
|
||||
"Site": info.Site(),
|
||||
})),
|
||||
styling.Code(fmt.Sprintf("%.2f MB (%d个资源)", float64(info.TotalBytes())/(1024*1024), info.TotalResources())),
|
||||
styling.Code(i18n.T(i18nk.BotMsgProgressSizeWithResources, map[string]any{
|
||||
"Size": fmt.Sprintf("%.2f MB", float64(info.TotalBytes())/(1024*1024)),
|
||||
"Count": info.TotalResources(),
|
||||
})),
|
||||
); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
|
||||
return
|
||||
@@ -101,7 +74,7 @@ func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
}
|
||||
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
if !shouldUpdateProgress(info.TotalBytes(), info.DownloadedBytes(), int(p.lastUpdatePercent.Load())) {
|
||||
if !progressutil.ShouldUpdate(info.TotalBytes(), info.DownloadedBytes(), int(p.lastUpdatePercent.Load())) {
|
||||
return
|
||||
}
|
||||
percent := int((info.DownloadedBytes() * 100) / info.TotalBytes())
|
||||
@@ -114,7 +87,10 @@ func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressDownloadingPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB (%d个文件)", float64(info.TotalBytes())/(1024*1024), info.TotalResources())),
|
||||
styling.Code(i18n.T(i18nk.BotMsgProgressSizeWithResources, map[string]any{
|
||||
"Size": fmt.Sprintf("%.2f MB", float64(info.TotalBytes())/(1024*1024)),
|
||||
"Count": info.TotalResources(),
|
||||
})),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressProcessingListPrefix, nil)),
|
||||
func() styling.StyledTextOption {
|
||||
var lines []string
|
||||
|
||||
@@ -33,7 +33,6 @@ type Task struct {
|
||||
downloadedBytes atomic.Int64 // downloaded bytes count
|
||||
processing map[string]ResourceInfo
|
||||
processingMu sync.RWMutex
|
||||
failed map[string]error // [TODO] errors for each resource
|
||||
}
|
||||
|
||||
// Title implements core.Exectable.
|
||||
@@ -84,6 +83,5 @@ func NewTask(
|
||||
progress: progressTracker,
|
||||
processing: make(map[string]ResourceInfo),
|
||||
processingMu: sync.RWMutex{},
|
||||
failed: make(map[string]error),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,7 +18,9 @@ import (
|
||||
func (t *Task) Execute(ctx context.Context) error {
|
||||
logger := log.FromContext(ctx)
|
||||
logger.Infof("Starting Telegraph task %s", t.PhPath)
|
||||
t.progress.OnStart(ctx, t)
|
||||
if t.progress != nil {
|
||||
t.progress.OnStart(ctx, t)
|
||||
}
|
||||
eg, gctx := errgroup.WithContext(ctx)
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for i, pic := range t.Pics {
|
||||
@@ -29,7 +31,9 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
return fmt.Errorf("failed to process picture %s: %w", pic, err)
|
||||
}
|
||||
downloaded := t.downloaded.Add(1)
|
||||
t.progress.OnProgress(gctx, t)
|
||||
if t.progress != nil {
|
||||
t.progress.OnProgress(gctx, t)
|
||||
}
|
||||
taskevent.Emit(gctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
@@ -45,7 +49,9 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
} else {
|
||||
logger.Infof("Telegraph task %s completed successfully", t.PhPath)
|
||||
}
|
||||
t.progress.OnDone(ctx, t, err)
|
||||
if t.progress != nil {
|
||||
t.progress.OnDone(ctx, t, err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
|
||||
31
core/tasks/telegraph/nil_tracker_test.go
Normal file
31
core/tasks/telegraph/nil_tracker_test.go
Normal file
@@ -0,0 +1,31 @@
|
||||
package telegraph_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
storconfig "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/core/tasks/telegraph"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
|
||||
type mockStorage struct{}
|
||||
|
||||
func (mockStorage) Init(context.Context, storconfig.StorageConfig) error { return nil }
|
||||
func (mockStorage) Type() storenum.StorageType { return storenum.Local }
|
||||
func (mockStorage) Name() string { return "mock" }
|
||||
func (mockStorage) Save(context.Context, io.Reader, string) error { return nil }
|
||||
func (mockStorage) Exists(context.Context, string) bool { return false }
|
||||
|
||||
var _ storage.Storage = mockStorage{}
|
||||
|
||||
// Regression: API-created tasks run with a nil ProgressTracker (progress goes
|
||||
// through taskevent); Execute must not panic on the tracker callbacks.
|
||||
func TestExecuteWithNilTracker(t *testing.T) {
|
||||
task := telegraph.NewTask("id", t.Context(), "/page", nil, mockStorage{}, "/out", nil, nil)
|
||||
if err := task.Execute(t.Context()); err != nil {
|
||||
t.Fatalf("Execute returned error: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/progressutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
)
|
||||
|
||||
@@ -60,7 +61,7 @@ func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
}
|
||||
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
if !shouldUpdateProgress(info.Downloaded(), int64(info.TotalPics())) {
|
||||
if !progressutil.ShouldUpdateCount(info.Downloaded(), int64(info.TotalPics())) {
|
||||
return
|
||||
}
|
||||
log.FromContext(ctx).Debugf("Progress update: %s, %d/%d", info.TaskID(), info.Downloaded(), info.TotalPics())
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
package telegraph
|
||||
|
||||
func shouldUpdateProgress(downloaded int64, total int64) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
step := int64(10)
|
||||
if downloaded < step {
|
||||
return downloaded == total
|
||||
}
|
||||
return downloaded%step == 0 || downloaded == total
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/progressutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
)
|
||||
|
||||
@@ -111,7 +112,7 @@ func (p *Progress) OnProgress(ctx context.Context, info TaskInfo, downloaded, to
|
||||
|
||||
func shouldUpdateSingleDownloadProgress(total, downloaded int64, lastPercent int, elapsed time.Duration) bool {
|
||||
if total > 0 {
|
||||
return shouldUpdateProgress(total, downloaded, lastPercent)
|
||||
return progressutil.ShouldUpdate(total, downloaded, lastPercent)
|
||||
}
|
||||
return downloaded > 0 && elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
@@ -183,7 +184,7 @@ func shouldUpdateUploadProgress(total, uploaded int64, lastPercent int, elapsed
|
||||
if percent == lastPercent {
|
||||
return elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
return shouldUpdateProgress(total, uploaded, lastPercent) || elapsed >= uploadProgressMaxInterval
|
||||
return progressutil.ShouldUpdate(total, uploaded, lastPercent) || elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
|
||||
func singleUploadPhase(attempt int) singleProgressPhase {
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
package tfile
|
||||
|
||||
var progressUpdatesLevels = []struct {
|
||||
size int64 // 文件大小阈值
|
||||
stepPercent int // 每多少 % 更新一次
|
||||
}{
|
||||
{10 << 20, 100},
|
||||
{50 << 20, 20},
|
||||
{200 << 20, 10},
|
||||
{500 << 20, 5},
|
||||
}
|
||||
|
||||
func shouldUpdateProgress(total, downloaded int64, lastUpdatePercent int) bool {
|
||||
if total <= 0 || downloaded <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
percent := int((downloaded * 100) / total)
|
||||
if percent <= lastUpdatePercent {
|
||||
return false
|
||||
}
|
||||
|
||||
step := progressUpdatesLevels[len(progressUpdatesLevels)-1].stepPercent
|
||||
for _, lvl := range progressUpdatesLevels {
|
||||
if total < lvl.size {
|
||||
step = lvl.stepPercent
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return percent >= lastUpdatePercent+step
|
||||
}
|
||||
75
core/tasks/transfer/cancel_test.go
Normal file
75
core/tasks/transfer/cancel_test.go
Normal file
@@ -0,0 +1,75 @@
|
||||
package transfer_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
storconfig "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/core/tasks/transfer"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
|
||||
// initConfig seeds the global config so task execution reads a sane Workers
|
||||
// value (the zero default would deadlock errgroup.SetLimit(0)).
|
||||
func initConfig(t *testing.T) {
|
||||
t.Helper()
|
||||
path := filepath.Join(t.TempDir(), "config.toml")
|
||||
if err := os.WriteFile(path, []byte("workers = 2\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := config.Init(t.Context(), path); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
type cancelSource struct{}
|
||||
|
||||
func (cancelSource) Init(context.Context, storconfig.StorageConfig) error { return nil }
|
||||
func (cancelSource) Type() storenum.StorageType { return storenum.Local }
|
||||
func (cancelSource) Name() string { return "cancel-source" }
|
||||
func (cancelSource) Save(context.Context, io.Reader, string) error { return nil }
|
||||
func (cancelSource) Exists(context.Context, string) bool { return false }
|
||||
func (cancelSource) ListFiles(context.Context, string) ([]storagetypes.FileInfo, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// OpenFile reports cancellation, as the task context would be cancelled.
|
||||
func (cancelSource) OpenFile(ctx context.Context, path string) (io.ReadCloser, int64, error) {
|
||||
return nil, 0, ctx.Err()
|
||||
}
|
||||
|
||||
type voidTarget struct{}
|
||||
|
||||
func (voidTarget) Init(context.Context, storconfig.StorageConfig) error { return nil }
|
||||
func (voidTarget) Type() storenum.StorageType { return storenum.Local }
|
||||
func (voidTarget) Name() string { return "void-target" }
|
||||
func (voidTarget) Save(context.Context, io.Reader, string) error { return nil }
|
||||
func (voidTarget) Exists(context.Context, string) bool { return false }
|
||||
|
||||
var (
|
||||
_ storage.StorageReadable = cancelSource{}
|
||||
_ storage.Storage = voidTarget{}
|
||||
)
|
||||
|
||||
// Regression: IgnoreErrors must swallow element failures but never a task
|
||||
// cancellation; Execute must surface context.Canceled.
|
||||
func TestIgnoreErrorsPropagatesCancel(t *testing.T) {
|
||||
initConfig(t)
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
cancel()
|
||||
|
||||
elem := transfer.NewTaskElement(cancelSource{}, storagetypes.FileInfo{Path: "/a", Name: "a", Size: 1}, voidTarget{}, "/out")
|
||||
task := transfer.NewTransferTask("id", ctx, []transfer.TaskElement{*elem}, nil, true)
|
||||
|
||||
err := task.Execute(ctx)
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("expected context.Canceled, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package transfer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
@@ -20,7 +21,9 @@ import (
|
||||
func (t *Task) Execute(ctx context.Context) error {
|
||||
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("transfer[%s]", t.ID))
|
||||
logger.Info("Starting transfer task")
|
||||
t.Progress.OnStart(ctx, t)
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnStart(ctx, t)
|
||||
}
|
||||
|
||||
workers := config.C().Workers
|
||||
eg, gctx := errgroup.WithContext(ctx)
|
||||
@@ -28,14 +31,11 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
|
||||
for _, elem := range t.elems {
|
||||
eg.Go(func() error {
|
||||
t.processingMu.RLock()
|
||||
t.processingMu.Lock()
|
||||
if t.processing[elem.ID] != nil {
|
||||
t.processingMu.RUnlock()
|
||||
t.processingMu.Unlock()
|
||||
return fmt.Errorf("element with ID %s is already being processed", elem.ID)
|
||||
}
|
||||
t.processingMu.RUnlock()
|
||||
|
||||
t.processingMu.Lock()
|
||||
t.processing[elem.ID] = &elem
|
||||
t.processingMu.Unlock()
|
||||
|
||||
@@ -46,7 +46,7 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
}()
|
||||
|
||||
err := t.processElement(gctx, elem)
|
||||
if err != nil && !t.IgnoreErrors {
|
||||
if err != nil && (!t.IgnoreErrors || errors.Is(err, context.Canceled)) {
|
||||
return err
|
||||
}
|
||||
if err != nil {
|
||||
@@ -66,7 +66,9 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
logger.Info("Transfer task completed successfully")
|
||||
}
|
||||
|
||||
t.Progress.OnDone(ctx, t, err)
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnDone(ctx, t, err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -116,7 +118,9 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
}
|
||||
|
||||
t.uploaded.Add(size)
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
}
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/progressutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
)
|
||||
|
||||
@@ -83,7 +84,7 @@ func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
}
|
||||
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
if !shouldUpdateProgress(info.TotalSize(), info.Uploaded(), int(p.lastUpdatePercent.Load())) {
|
||||
if !progressutil.ShouldUpdate(info.TotalSize(), info.Uploaded(), int(p.lastUpdatePercent.Load())) {
|
||||
return
|
||||
}
|
||||
percent := int((info.Uploaded() * 100) / info.TotalSize())
|
||||
@@ -221,14 +222,6 @@ func (p *Progress) OnDone(ctx context.Context, info TaskInfo, err error) {
|
||||
}
|
||||
}
|
||||
|
||||
func shouldUpdateProgress(total, current int64, lastPercent int) bool {
|
||||
if total == 0 {
|
||||
return false
|
||||
}
|
||||
currentPercent := int((current * 100) / total)
|
||||
return currentPercent > lastPercent && currentPercent%5 == 0
|
||||
}
|
||||
|
||||
func formatDuration(d time.Duration) string {
|
||||
d = d.Round(time.Second)
|
||||
h := d / time.Hour
|
||||
|
||||
@@ -44,7 +44,10 @@ func jsRegisterParser(vm *goja.Runtime) func(call goja.FunctionCall) goja.Value
|
||||
return vm.NewGoError(errors.New("metadata cannot be null or undefined"))
|
||||
}
|
||||
|
||||
pluginV := semver.MustParse(metadata.Version)
|
||||
pluginV, err := semver.Parse(metadata.Version)
|
||||
if err != nil {
|
||||
return vm.NewGoError(fmt.Errorf("invalid parser version %q: %w", metadata.Version, err))
|
||||
}
|
||||
if pluginV.LT(MinimumParserVersion) {
|
||||
return vm.NewGoError(fmt.Errorf("parser version %s is not supported, must be at least %s", metadata.Version, MinimumParserVersion))
|
||||
}
|
||||
@@ -57,6 +60,9 @@ func jsRegisterParser(vm *goja.Runtime) func(call goja.FunctionCall) goja.Value
|
||||
if parseFn == nil || goja.IsUndefined(parseFn) {
|
||||
return vm.NewGoError(errors.New("parser must provide a parse function"))
|
||||
}
|
||||
if handleFn == nil || goja.IsUndefined(handleFn) {
|
||||
return vm.NewGoError(errors.New("parser must provide a canHandle function"))
|
||||
}
|
||||
parsers.Add(newJSParser(vm, handleFn, parseFn, metadata))
|
||||
return goja.Undefined()
|
||||
}
|
||||
|
||||
107
parsers/js/js.go
107
parsers/js/js.go
@@ -9,6 +9,7 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/dop251/goja"
|
||||
@@ -34,11 +35,23 @@ type jsParserResp struct {
|
||||
err error
|
||||
}
|
||||
|
||||
const canHandleTimeout = 10 * time.Second
|
||||
|
||||
func (p *jsParser) CanHandle(url string) bool {
|
||||
respCh := make(chan jsParserResp, 1)
|
||||
p.reqCh <- jsParserReq{method: ParserMethodCanHandle, url: url, respCh: respCh}
|
||||
resp := <-respCh
|
||||
return resp.ok && resp.err == nil
|
||||
timer := time.NewTimer(canHandleTimeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case p.reqCh <- jsParserReq{method: ParserMethodCanHandle, url: url, respCh: respCh}:
|
||||
case <-timer.C:
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case resp := <-respCh:
|
||||
return resp.ok && resp.err == nil
|
||||
case <-timer.C:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (p *jsParser) Parse(ctx context.Context, url string) (*parser.Item, error) {
|
||||
@@ -61,41 +74,57 @@ func newJSParser(vm *goja.Runtime, canHandleFunc, parseFunc goja.Value, metadata
|
||||
|
||||
go func() {
|
||||
for req := range p.reqCh {
|
||||
switch req.method {
|
||||
case ParserMethodCanHandle:
|
||||
fn, _ := goja.AssertFunction(canHandleFunc)
|
||||
res, err := fn(goja.Undefined(), p.vm.ToValue(req.url))
|
||||
if err != nil {
|
||||
req.respCh <- jsParserResp{ok: false, err: err}
|
||||
continue
|
||||
}
|
||||
req.respCh <- jsParserResp{ok: res.ToBoolean()}
|
||||
case ParserMethodParse:
|
||||
fn, _ := goja.AssertFunction(parseFunc)
|
||||
result, err := fn(goja.Undefined(), p.vm.ToValue(req.url))
|
||||
if err != nil {
|
||||
req.respCh <- jsParserResp{err: err}
|
||||
continue
|
||||
}
|
||||
|
||||
var item parser.Item
|
||||
if exported := result.Export(); exported != nil {
|
||||
data, err := json.Marshal(exported)
|
||||
func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Errorf("JS parser %q panicked while handling method %d: %v", p.meta.Name, req.method, r)
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("JS parser %q panicked: %v", p.meta.Name, r)}
|
||||
}
|
||||
}()
|
||||
switch req.method {
|
||||
case ParserMethodCanHandle:
|
||||
fn, ok := goja.AssertFunction(canHandleFunc)
|
||||
if !ok {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("canHandle is not a function")}
|
||||
return
|
||||
}
|
||||
res, err := fn(goja.Undefined(), p.vm.ToValue(req.url))
|
||||
if err != nil {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("failed to marshal result to JSON: %w", err)}
|
||||
continue
|
||||
req.respCh <- jsParserResp{ok: false, err: err}
|
||||
return
|
||||
}
|
||||
req.respCh <- jsParserResp{ok: res.ToBoolean()}
|
||||
case ParserMethodParse:
|
||||
fn, ok := goja.AssertFunction(parseFunc)
|
||||
if !ok {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("parse is not a function")}
|
||||
return
|
||||
}
|
||||
result, err := fn(goja.Undefined(), p.vm.ToValue(req.url))
|
||||
if err != nil {
|
||||
req.respCh <- jsParserResp{err: err}
|
||||
return
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(data, &item); err != nil {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("failed to unmarshal JSON to Item: %w", err)}
|
||||
continue
|
||||
var item parser.Item
|
||||
if exported := result.Export(); exported != nil {
|
||||
data, err := json.Marshal(exported)
|
||||
if err != nil {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("failed to marshal result to JSON: %w", err)}
|
||||
return
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(data, &item); err != nil {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("failed to unmarshal JSON to Item: %w", err)}
|
||||
return
|
||||
}
|
||||
} else {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("JS function returned null or undefined")}
|
||||
return
|
||||
}
|
||||
} else {
|
||||
req.respCh <- jsParserResp{err: fmt.Errorf("JS function returned null or undefined")}
|
||||
continue
|
||||
req.respCh <- jsParserResp{item: &item}
|
||||
}
|
||||
req.respCh <- jsParserResp{item: &item}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -116,7 +145,8 @@ func LoadPlugins(ctx context.Context, dir string) error {
|
||||
scriptPath := filepath.Join(dir, e.Name())
|
||||
code, err := os.ReadFile(scriptPath)
|
||||
if err != nil {
|
||||
return err
|
||||
log.Warnf("Failed to read plugin file %s: %v", e.Name(), err)
|
||||
continue
|
||||
}
|
||||
|
||||
vm := goja.New()
|
||||
@@ -130,7 +160,8 @@ func LoadPlugins(ctx context.Context, dir string) error {
|
||||
vm.Set("playwright", jsPlaywright(vm, logger))
|
||||
|
||||
if _, err := vm.RunString(string(code)); err != nil {
|
||||
return fmt.Errorf("error loading plugin %s: %w", e.Name(), err)
|
||||
logger.Warnf("Failed to load plugin %s: %v", e.Name(), err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
@@ -164,8 +195,12 @@ func addPlugin(ctx context.Context, code string, name string) error {
|
||||
if len(configuredDirs) > 0 {
|
||||
dir = configuredDirs[0]
|
||||
}
|
||||
fileName := filepath.Base(name)
|
||||
if fileName == "" || fileName == "." || fileName == ".." {
|
||||
return fmt.Errorf("invalid plugin name %q", name)
|
||||
}
|
||||
if err := os.MkdirAll(dir, 0755); err == nil {
|
||||
pluginPath := filepath.Join(dir, name)
|
||||
pluginPath := filepath.Join(dir, fileName)
|
||||
if err := os.WriteFile(pluginPath, []byte(code), 0644); err != nil {
|
||||
logger.Warn("Failed to save plugin file: " + err.Error())
|
||||
}
|
||||
|
||||
35
parsers/js/plugin_version_test.go
Normal file
35
parsers/js/plugin_version_test.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package js_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/parsers/js"
|
||||
)
|
||||
|
||||
// Regression: a plugin with an invalid semver version must be rejected
|
||||
// without panicking the process (previously semver.MustParse crashed the bot).
|
||||
func TestLoadPluginsRejectsInvalidVersion(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
bad := `registerParser({
|
||||
metadata: { name: "probe", version: "not-a-version", description: "", author: "" },
|
||||
canHandle: function(url) { return true; },
|
||||
parse: async function(url) { return { resources: [] }; }
|
||||
});`
|
||||
if err := os.WriteFile(filepath.Join(dir, "bad.js"), []byte(bad), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
good := `registerParser({
|
||||
metadata: { name: "good", version: "1.0.0", description: "", author: "" },
|
||||
canHandle: function(url) { return false; },
|
||||
parse: async function(url) { return { resources: [] }; }
|
||||
});`
|
||||
if err := os.WriteFile(filepath.Join(dir, "good.js"), []byte(good), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Must not panic; the bad plugin is skipped and remaining plugins load.
|
||||
if err := js.LoadPlugins(t.Context(), dir); err != nil {
|
||||
t.Fatalf("LoadPlugins returned error: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"path"
|
||||
"strings"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/duke-git/lancet/v2/strutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/netutil"
|
||||
"github.com/krau/SaveAny-Bot/pkg/parser"
|
||||
@@ -135,9 +136,17 @@ func (k *KemonoParser) parseOne(ctx context.Context, info *DownloadInfo) (*parse
|
||||
if preview.Type == nil || *preview.Type != "thumbnail" {
|
||||
continue
|
||||
}
|
||||
if preview.Path == nil || preview.Server == nil {
|
||||
log.FromContext(ctx).Warnf("Skipping kemono preview with missing path or server: post %s", info.PostID)
|
||||
continue
|
||||
}
|
||||
picCdnMap[*preview.Path] = *preview.Server
|
||||
}
|
||||
for _, attachment := range postInfo.Post.Attachments {
|
||||
if attachment.Path == nil || attachment.Name == nil {
|
||||
log.FromContext(ctx).Warnf("Skipping kemono post attachment with missing path or name: post %s", info.PostID)
|
||||
continue
|
||||
}
|
||||
if !isImageExt(*attachment.Path) {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
package kemono
|
||||
|
||||
type PostLegacy struct {
|
||||
Props Props `json:"props"`
|
||||
Results []Result `json:"results"`
|
||||
}
|
||||
|
||||
type Props struct {
|
||||
Count uint `json:"count"`
|
||||
Limit uint `json:"limit"`
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
package kemono
|
||||
|
||||
type UserProfile struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Service string `json:"service"`
|
||||
PublicID *string `json:"public_id,omitempty"`
|
||||
}
|
||||
@@ -1,101 +1,5 @@
|
||||
package twitter
|
||||
|
||||
// type AutoGenerated struct {
|
||||
// Code int `json:"code"`
|
||||
// Message string `json:"message"`
|
||||
// Tweet struct {
|
||||
// URL string `json:"url"`
|
||||
// ID string `json:"id"`
|
||||
// Text string `json:"text"`
|
||||
// RawText struct {
|
||||
// Text string `json:"text"`
|
||||
// Facets []struct {
|
||||
// Type string `json:"type"`
|
||||
// Indices []int `json:"indices"`
|
||||
// Original string `json:"original"`
|
||||
// ID string `json:"id,omitempty"`
|
||||
// Display string `json:"display,omitempty"`
|
||||
// Replacement string `json:"replacement,omitempty"`
|
||||
// } `json:"facets"`
|
||||
// } `json:"raw_text"`
|
||||
// Author struct {
|
||||
// ID string `json:"id"`
|
||||
// Name string `json:"name"`
|
||||
// ScreenName string `json:"screen_name"`
|
||||
// AvatarURL string `json:"avatar_url"`
|
||||
// BannerURL interface{} `json:"banner_url"`
|
||||
// Description string `json:"description"`
|
||||
// Location string `json:"location"`
|
||||
// URL string `json:"url"`
|
||||
// Followers int `json:"followers"`
|
||||
// Following int `json:"following"`
|
||||
// Joined string `json:"joined"`
|
||||
// Likes int `json:"likes"`
|
||||
// MediaCount int `json:"media_count"`
|
||||
// Protected bool `json:"protected"`
|
||||
// Website struct {
|
||||
// URL string `json:"url"`
|
||||
// DisplayURL string `json:"display_url"`
|
||||
// } `json:"website"`
|
||||
// Tweets int `json:"tweets"`
|
||||
// AvatarColor interface{} `json:"avatar_color"`
|
||||
// } `json:"author"`
|
||||
// Replies int `json:"replies"`
|
||||
// Retweets int `json:"retweets"`
|
||||
// Likes int `json:"likes"`
|
||||
// Bookmarks int `json:"bookmarks"`
|
||||
// CreatedAt string `json:"created_at"`
|
||||
// CreatedTimestamp int `json:"created_timestamp"`
|
||||
// PossiblySensitive bool `json:"possibly_sensitive"`
|
||||
// Views int `json:"views"`
|
||||
// IsNoteTweet bool `json:"is_note_tweet"`
|
||||
// CommunityNote interface{} `json:"community_note"`
|
||||
// Lang string `json:"lang"`
|
||||
// ReplyingTo interface{} `json:"replying_to"`
|
||||
// ReplyingToStatus interface{} `json:"replying_to_status"`
|
||||
// Media struct {
|
||||
// All []struct {
|
||||
// URL string `json:"url"`
|
||||
// ThumbnailURL string `json:"thumbnail_url"`
|
||||
// Duration int `json:"duration"`
|
||||
// Width int `json:"width"`
|
||||
// Height int `json:"height"`
|
||||
// Format string `json:"format"`
|
||||
// Type string `json:"type"`
|
||||
// Variants []struct {
|
||||
// Bitrate int `json:"bitrate"`
|
||||
// ContentType string `json:"content_type"`
|
||||
// URL string `json:"url"`
|
||||
// } `json:"variants"`
|
||||
// } `json:"all"`
|
||||
// Photos []struct {
|
||||
// Type string `json:"type"`
|
||||
// URL string `json:"url"`
|
||||
// Width int `json:"width"`
|
||||
// Height int `json:"height"`
|
||||
// } `json:"photos"`
|
||||
// Videos []struct {
|
||||
// URL string `json:"url"`
|
||||
// ThumbnailURL string `json:"thumbnail_url"`
|
||||
// Duration int `json:"duration"`
|
||||
// Width int `json:"width"`
|
||||
// Height int `json:"height"`
|
||||
// Format string `json:"format"`
|
||||
// Type string `json:"type"`
|
||||
// Variants []struct {
|
||||
// Bitrate int `json:"bitrate"`
|
||||
// ContentType string `json:"content_type"`
|
||||
// URL string `json:"url"`
|
||||
// } `json:"variants"`
|
||||
// } `json:"videos"`
|
||||
// } `json:"media"`
|
||||
// Source string `json:"source"`
|
||||
// TwitterCard string `json:"twitter_card"`
|
||||
// Color interface{} `json:"color"`
|
||||
// Provider string `json:"provider"`
|
||||
// } `json:"tweet"`
|
||||
// }
|
||||
|
||||
type FxTwitterApiResp struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
|
||||
@@ -2,6 +2,7 @@ package parsers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -39,5 +40,5 @@ func Get() []parser.Parser {
|
||||
configOnce.Do(configParsers)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
return parsers
|
||||
return slices.Clone(parsers)
|
||||
}
|
||||
|
||||
@@ -8,4 +8,7 @@ const (
|
||||
MaxPartSize = 1024 * 1024
|
||||
MaxUploadPartSize = uploader.MaximumPartSize
|
||||
MaxPhotoSize = 10 * 1024 * 1024
|
||||
// MaxAlbumItems is the Telegram media-album item cap used for batching
|
||||
// uploads and lossless video splitting.
|
||||
MaxAlbumItems = 10
|
||||
)
|
||||
|
||||
@@ -4,6 +4,8 @@ import (
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"fmt"
|
||||
"maps"
|
||||
"slices"
|
||||
)
|
||||
|
||||
type Parser interface {
|
||||
@@ -57,14 +59,15 @@ func (r *Resource) ID() string {
|
||||
h.Write([]byte(r.Extension))
|
||||
fmt.Fprintf(h, "%d", r.Size)
|
||||
|
||||
for k, v := range r.Hash {
|
||||
// Sort keys for a stable fingerprint.
|
||||
for _, k := range slices.Sorted(maps.Keys(r.Hash)) {
|
||||
h.Write([]byte(k))
|
||||
h.Write([]byte(v))
|
||||
h.Write([]byte(r.Hash[k]))
|
||||
}
|
||||
|
||||
for k, v := range r.Headers {
|
||||
for _, k := range slices.Sorted(maps.Keys(r.Headers)) {
|
||||
h.Write([]byte(k))
|
||||
h.Write([]byte(v))
|
||||
h.Write([]byte(r.Headers[k]))
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%x", h.Sum(nil))
|
||||
|
||||
@@ -50,38 +50,41 @@ func (tq *TaskQueue[T]) Add(task *Task[T]) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ErrQueueClosed is returned by Get when the queue is closed and no tasks remain.
|
||||
var ErrQueueClosed = errors.New("queue is closed and empty")
|
||||
|
||||
// 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()
|
||||
|
||||
for tq.tasks.Len() == 0 && !tq.closed {
|
||||
tq.cond.Wait()
|
||||
}
|
||||
for {
|
||||
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])
|
||||
|
||||
for tq.tasks.Len() > 0 {
|
||||
element := tq.tasks.Front()
|
||||
task := element.Value.(*Task[T])
|
||||
tq.tasks.Remove(element)
|
||||
task.element = nil
|
||||
|
||||
tq.tasks.Remove(element)
|
||||
task.element = nil
|
||||
if task.Cancelled() {
|
||||
// Skip cancelled tasks and release their IDs.
|
||||
delete(tq.taskMap, task.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
if !task.Cancelled() {
|
||||
tq.runningTaskMap[task.ID] = task
|
||||
return task, nil
|
||||
}
|
||||
}
|
||||
|
||||
if !tq.closed {
|
||||
return tq.Get()
|
||||
if tq.closed {
|
||||
return nil, ErrQueueClosed
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("queue is closed and empty")
|
||||
}
|
||||
|
||||
// Done stops(cancels) and removes the task from the running tasks.
|
||||
|
||||
@@ -2,9 +2,11 @@ package queue_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/pkg/queue"
|
||||
)
|
||||
@@ -65,8 +67,8 @@ func TestCloseBehavior(t *testing.T) {
|
||||
// consumer
|
||||
go func() {
|
||||
_, err := q.Get()
|
||||
if err == nil {
|
||||
t.Errorf("expected error when getting from closed empty queue, got nil")
|
||||
if !errors.Is(err, queue.ErrQueueClosed) {
|
||||
t.Errorf("expected ErrQueueClosed from closed empty queue, got %v", err)
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
@@ -77,6 +79,86 @@ func TestCloseBehavior(t *testing.T) {
|
||||
<-done
|
||||
}
|
||||
|
||||
// Regression: Get() must not deadlock when every queued task was cancelled
|
||||
// before a worker picked it up (previously recursed while holding the mutex).
|
||||
func TestGetAfterAllCancelled(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
for i := range 3 {
|
||||
if err := q.Add(newTask(fmt.Sprintf("c%d", i))); err != nil {
|
||||
t.Fatalf("unexpected error on Add: %v", err)
|
||||
}
|
||||
}
|
||||
for i := range 3 {
|
||||
if err := q.CancelTask(fmt.Sprintf("c%d", i)); err != nil {
|
||||
t.Fatalf("unexpected error on CancelTask: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A task added after the cancelled ones must still be delivered.
|
||||
if err := q.Add(newTask("late")); err != nil {
|
||||
t.Fatalf("unexpected error on Add after cancel: %v", err)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
var got *queue.Task[int]
|
||||
go func() {
|
||||
var err error
|
||||
got, err = q.Get()
|
||||
if err != nil {
|
||||
t.Errorf("unexpected error on Get: %v", err)
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
if got == nil || got.ID != "late" {
|
||||
t.Fatalf("expected task 'late', got %v", got)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Get() deadlocked after all queued tasks were cancelled")
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelled queued tasks must be dropped from the task map so their TaskID
|
||||
// can be reused (previously they leaked until process exit).
|
||||
func TestCancelledQueuedTaskIDReusable(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
if err := q.Add(newTask("dup")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := q.CancelTask("dup"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
done := make(chan struct{})
|
||||
var gotID string
|
||||
go func() {
|
||||
task, err := q.Get()
|
||||
if err != nil {
|
||||
t.Errorf("unexpected error on Get: %v", err)
|
||||
close(done)
|
||||
return
|
||||
}
|
||||
gotID = task.ID
|
||||
close(done)
|
||||
}()
|
||||
// The first Get skips the cancelled task and blocks; a live task unblocks it.
|
||||
if err := q.Add(newTask("late")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
if gotID != "late" {
|
||||
t.Fatalf("expected 'late', got %q", gotID)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Get did not return the live task")
|
||||
}
|
||||
// The cancelled task was dropped from the map: its ID is reusable.
|
||||
if err := q.Add(newTask("dup")); err != nil {
|
||||
t.Fatalf("expected cancelled task ID to be reusable, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrencySafety(t *testing.T) {
|
||||
q := queue.NewTaskQueue[int]()
|
||||
var wg sync.WaitGroup
|
||||
|
||||
@@ -9,23 +9,35 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
type Alist struct {
|
||||
client *http.Client
|
||||
token string
|
||||
baseURL string
|
||||
loginInfo *loginRequest
|
||||
config config.AlistStorageConfig
|
||||
logger *log.Logger
|
||||
client *http.Client
|
||||
tokenMu sync.RWMutex
|
||||
token string
|
||||
lastLoginAt time.Time
|
||||
tokenFlight singleflight.Group
|
||||
baseURL string
|
||||
loginInfo *loginRequest
|
||||
config config.AlistStorageConfig
|
||||
logger *log.Logger
|
||||
}
|
||||
|
||||
// authHeader returns the current token for use in API requests.
|
||||
func (a *Alist) authHeader() string {
|
||||
a.tokenMu.RLock()
|
||||
defer a.tokenMu.RUnlock()
|
||||
return a.token
|
||||
}
|
||||
|
||||
func (a *Alist) Init(ctx context.Context, cfg config.StorageConfig) error {
|
||||
@@ -42,39 +54,35 @@ func (a *Alist) Init(ctx context.Context, cfg config.StorageConfig) error {
|
||||
a.logger = log.FromContext(ctx).WithPrefix(fmt.Sprintf("alist[%s]", alistConfig.Name))
|
||||
|
||||
if alistConfig.Token != "" {
|
||||
a.tokenMu.Lock()
|
||||
a.token = alistConfig.Token
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
|
||||
a.tokenMu.Unlock()
|
||||
tokenCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, a.baseURL+"/api/me", nil)
|
||||
req, err := http.NewRequestWithContext(tokenCtx, http.MethodGet, a.baseURL+"/api/me", nil)
|
||||
if err != nil {
|
||||
a.logger.Fatalf("Failed to create request: %v", err)
|
||||
return err
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", a.token)
|
||||
req.Header.Set("Authorization", a.authHeader())
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
a.logger.Fatalf("Failed to send request: %v", err)
|
||||
return err
|
||||
return fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
a.logger.Fatalf("Failed to get alist user info: %s", resp.Status)
|
||||
return err
|
||||
return fmt.Errorf("failed to get alist user info: %s", resp.Status)
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
a.logger.Fatalf("Failed to read response body: %v", err)
|
||||
return err
|
||||
return fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
var meResp meResponse
|
||||
if err := json.Unmarshal(body, &meResp); err != nil {
|
||||
a.logger.Fatalf("Failed to unmarshal me response: %v", err)
|
||||
return err
|
||||
return fmt.Errorf("failed to unmarshal me response: %w", err)
|
||||
}
|
||||
if meResp.Code != http.StatusOK {
|
||||
a.logger.Fatalf("Failed to get alist user info: %s", meResp.Message)
|
||||
return err
|
||||
return fmt.Errorf("failed to get alist user info: %s", meResp.Message)
|
||||
}
|
||||
a.logger.Debugf("Logged in Alist as %s", meResp.Data.Username)
|
||||
return nil
|
||||
@@ -85,12 +93,15 @@ func (a *Alist) Init(ctx context.Context, cfg config.StorageConfig) error {
|
||||
}
|
||||
|
||||
if err := a.getToken(ctx); err != nil {
|
||||
a.logger.Fatalf("Failed to login to Alist: %v", err)
|
||||
return err
|
||||
return fmt.Errorf("failed to login to Alist: %w", err)
|
||||
}
|
||||
// The init login must not satisfy the refresh dedup window.
|
||||
a.tokenMu.Lock()
|
||||
a.lastLoginAt = time.Time{}
|
||||
a.tokenMu.Unlock()
|
||||
a.logger.Debug("Logged in to Alist")
|
||||
|
||||
go a.refreshToken(*alistConfig)
|
||||
go a.refreshToken(ctx, *alistConfig)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -104,33 +115,40 @@ func (a *Alist) Name() string {
|
||||
|
||||
func (a *Alist) Save(ctx context.Context, reader io.Reader, storagePath string) error {
|
||||
a.logger.Infof("Saving file to %s", storagePath)
|
||||
storagePath = a.JoinStoragePath(storagePath)
|
||||
ext := path.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
candidate := a.JoinStoragePath(storagePath)
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
for i := 1; a.existsPath(ctx, candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
}
|
||||
candidate = fsutil.UniquePath(a.config.BasePath, storagePath, func(c string) bool {
|
||||
return a.existsPath(ctx, c)
|
||||
}, 1000)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, a.baseURL+"/api/fs/put", reader)
|
||||
resp, err := a.putFile(ctx, reader, candidate)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", a.token)
|
||||
req.Header.Set("File-Path", url.PathEscape(candidate))
|
||||
req.Header.Set("Content-Type", "application/octet-stream")
|
||||
if length := ctx.Value(ctxkey.ContentLength); length != nil {
|
||||
length, ok := length.(int64)
|
||||
if ok {
|
||||
req.ContentLength = length
|
||||
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||
status := resp.Status
|
||||
resp.Body.Close()
|
||||
// Token-only storage cannot refresh: surface the auth error.
|
||||
if a.loginInfo == nil {
|
||||
return fmt.Errorf("failed to save file to Alist: %s", status)
|
||||
}
|
||||
if err := a.getToken(ctx); err != nil {
|
||||
return fmt.Errorf("failed to refresh alist token: %w", err)
|
||||
}
|
||||
rs, seekable := reader.(io.ReadSeeker)
|
||||
if !seekable {
|
||||
a.logger.Warnf("Upload rejected with %s; reader is not seekable, cannot retry", status)
|
||||
return fmt.Errorf("failed to save file to Alist: %s (streaming reader cannot be replayed for retry)", status)
|
||||
}
|
||||
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to rewind reader before retry: %w", err)
|
||||
}
|
||||
a.logger.Info("Retrying upload with refreshed token")
|
||||
resp, err = a.putFile(ctx, reader, candidate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
@@ -155,6 +173,29 @@ func (a *Alist) Save(ctx context.Context, reader io.Reader, storagePath string)
|
||||
return nil
|
||||
}
|
||||
|
||||
// putFile performs a single PUT upload with the given token and returns the response.
|
||||
func (a *Alist) putFile(ctx context.Context, reader io.Reader, storagePath string) (*http.Response, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, a.baseURL+"/api/fs/put", reader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", a.authHeader())
|
||||
req.Header.Set("File-Path", url.PathEscape(storagePath))
|
||||
req.Header.Set("Content-Type", "application/octet-stream")
|
||||
if length := ctx.Value(ctxkey.ContentLength); length != nil {
|
||||
length, ok := length.(int64)
|
||||
if ok {
|
||||
req.ContentLength = length
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (a *Alist) JoinStoragePath(p string) string {
|
||||
return path.Join(a.config.BasePath, p)
|
||||
}
|
||||
@@ -189,7 +230,7 @@ func (a *Alist) existsPath(ctx context.Context, storagePath string) bool {
|
||||
a.logger.Errorf("Failed to create request: %v", err)
|
||||
return false
|
||||
}
|
||||
req.Header.Set("Authorization", a.token)
|
||||
req.Header.Set("Authorization", a.authHeader())
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
@@ -244,7 +285,7 @@ func (a *Alist) ListFiles(ctx context.Context, dirPath string) ([]storagetypes.F
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", a.token)
|
||||
req.Header.Set("Authorization", a.authHeader())
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
@@ -319,7 +360,7 @@ func (a *Alist) OpenFile(ctx context.Context, filePath string) (io.ReadCloser, i
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", a.token)
|
||||
req.Header.Set("Authorization", a.authHeader())
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
|
||||
152
storage/alist/alist_test.go
Normal file
152
storage/alist/alist_test.go
Normal file
@@ -0,0 +1,152 @@
|
||||
package alist_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
storconfig "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/storage/alist"
|
||||
)
|
||||
|
||||
// newAlistServer starts a fake alist whose login endpoint issues sequential
|
||||
// tokens and whose PUT endpoint rejects the given token (simulating an expired
|
||||
// credential) while accepting refreshed ones.
|
||||
func newAlistServer(t *testing.T, rejectedToken string) (*httptest.Server, *sync.Mutex, *int, *[]putRecord) {
|
||||
t.Helper()
|
||||
var mu sync.Mutex
|
||||
loginCount := 0
|
||||
tokenSeq := 0
|
||||
var puts []putRecord
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/api/auth/login", func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
loginCount++
|
||||
tokenSeq++
|
||||
token := fmt.Sprintf("token-%d", tokenSeq)
|
||||
mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]any{"code": 200, "message": "ok", "data": map[string]any{"token": token}})
|
||||
})
|
||||
mux.HandleFunc("/api/me", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]any{"code": 200, "message": "ok", "data": map[string]any{"username": "probe"}})
|
||||
})
|
||||
mux.HandleFunc("/api/fs/put", func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
rejected := r.Header.Get("Authorization") == rejectedToken
|
||||
puts = append(puts, putRecord{auth: r.Header.Get("Authorization"), rejected: rejected})
|
||||
mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if rejected {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
json.NewEncoder(w).Encode(map[string]any{"code": 401, "message": "unauthorized"})
|
||||
return
|
||||
}
|
||||
json.NewEncoder(w).Encode(map[string]any{"code": 200, "message": "ok"})
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv, &mu, &loginCount, &puts
|
||||
}
|
||||
|
||||
type putRecord struct {
|
||||
auth string
|
||||
rejected bool
|
||||
}
|
||||
|
||||
// Regression: concurrent uploads hitting 401 must share a single re-login
|
||||
// (singleflight) and retry with the refreshed token. Init performs login #1
|
||||
// (token-1); the server rejects it, so the concurrent uploads must trigger a
|
||||
// second, merged login (token-2).
|
||||
func TestConcurrent401RetrySingleLogin(t *testing.T) {
|
||||
srv, mu, loginCount, putAuths := newAlistServer(t, "token-1")
|
||||
|
||||
cfg := &storconfig.AlistStorageConfig{}
|
||||
cfg.Name = "probe"
|
||||
cfg.URL = srv.URL
|
||||
cfg.Username = "user"
|
||||
cfg.Password = "pass"
|
||||
cfg.BasePath = "/probe"
|
||||
|
||||
stor := &alist.Alist{}
|
||||
if err := stor.Init(t.Context(), cfg); err != nil {
|
||||
t.Fatalf("Init failed: %v", err)
|
||||
}
|
||||
|
||||
const workers = 10
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, workers)
|
||||
for range workers {
|
||||
wg.Go(func() {
|
||||
errs <- stor.Save(t.Context(), bytes.NewReader([]byte("data")), "dir/file.txt")
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
t.Fatalf("Save failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
// One login for init, one merged login for the 401 storm.
|
||||
if *loginCount != 2 {
|
||||
t.Fatalf("expected 2 logins (init + merged retry), got %d", *loginCount)
|
||||
}
|
||||
accepted := 0
|
||||
for _, put := range *putAuths {
|
||||
if put.rejected {
|
||||
if put.auth != "token-1" {
|
||||
t.Fatalf("expected rejected uploads to use the expired token-1, got %q", put.auth)
|
||||
}
|
||||
continue
|
||||
}
|
||||
accepted++
|
||||
if put.auth != "token-2" {
|
||||
t.Fatalf("expected accepted uploads to use the refreshed token-2, got %q", put.auth)
|
||||
}
|
||||
}
|
||||
if accepted != workers {
|
||||
t.Fatalf("expected %d accepted uploads, got %d", workers, accepted)
|
||||
}
|
||||
}
|
||||
|
||||
// A token-only storage receives 401 and must return the auth error without
|
||||
// attempting a login (it has no credentials to refresh with).
|
||||
func TestTokenOnlyNoLoginOn401(t *testing.T) {
|
||||
srv, mu, loginCount, _ := newAlistServer(t, "token-0")
|
||||
|
||||
cfg := &storconfig.AlistStorageConfig{}
|
||||
cfg.Name = "probe"
|
||||
cfg.URL = srv.URL
|
||||
cfg.Token = "token-0"
|
||||
cfg.BasePath = "/probe"
|
||||
|
||||
stor := &alist.Alist{}
|
||||
if err := stor.Init(t.Context(), cfg); err != nil {
|
||||
t.Fatalf("Init failed: %v", err)
|
||||
}
|
||||
|
||||
err := stor.Save(t.Context(), bytes.NewReader([]byte("data")), "dir/file.txt")
|
||||
if err == nil {
|
||||
t.Fatal("expected auth error from token-only storage, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "401") {
|
||||
t.Fatalf("expected 401 auth error, got: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if *loginCount != 0 {
|
||||
t.Fatalf("expected no login attempts for token-only storage, got %d", *loginCount)
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,42 @@ import (
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
)
|
||||
|
||||
// minTokenRefreshInterval deduplicates login storms; it is capped by the
|
||||
// configured token expiry.
|
||||
const minTokenRefreshInterval = 30 * time.Second
|
||||
|
||||
func (a *Alist) tokenRefreshWindow() time.Duration {
|
||||
window := minTokenRefreshInterval
|
||||
if exp := time.Duration(a.config.TokenExp) * time.Second; exp > 0 && exp < window {
|
||||
window = exp
|
||||
}
|
||||
return window
|
||||
}
|
||||
|
||||
// getToken refreshes the JWT, merging concurrent calls.
|
||||
func (a *Alist) getToken(ctx context.Context) error {
|
||||
a.tokenMu.RLock()
|
||||
fresh := !a.lastLoginAt.IsZero() && time.Since(a.lastLoginAt) < a.tokenRefreshWindow()
|
||||
a.tokenMu.RUnlock()
|
||||
if fresh {
|
||||
return nil
|
||||
}
|
||||
_, err, _ := a.tokenFlight.Do("token", func() (any, error) {
|
||||
a.tokenMu.RLock()
|
||||
fresh := !a.lastLoginAt.IsZero() && time.Since(a.lastLoginAt) < a.tokenRefreshWindow()
|
||||
a.tokenMu.RUnlock()
|
||||
if fresh {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, a.fetchToken(ctx)
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (a *Alist) fetchToken(ctx context.Context) error {
|
||||
if a.loginInfo == nil {
|
||||
return fmt.Errorf("token-only alist storage cannot refresh credentials")
|
||||
}
|
||||
loginBody, err := json.Marshal(a.loginInfo)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal login request: %w", err)
|
||||
@@ -44,22 +79,31 @@ func (a *Alist) getToken(ctx context.Context) error {
|
||||
return fmt.Errorf("%w: %s", ErrAlistLoginFailed, loginResp.Message)
|
||||
}
|
||||
|
||||
a.tokenMu.Lock()
|
||||
a.token = loginResp.Data.Token
|
||||
a.lastLoginAt = time.Now()
|
||||
a.tokenMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Alist) refreshToken(cfg config.AlistStorageConfig) {
|
||||
func (a *Alist) refreshToken(ctx context.Context, cfg config.AlistStorageConfig) {
|
||||
tokenExp := cfg.TokenExp
|
||||
if tokenExp <= 0 {
|
||||
a.logger.Warn("Invalid token expiration time, using default value")
|
||||
tokenExp = 3600
|
||||
}
|
||||
ticker := time.NewTicker(time.Duration(tokenExp) * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
time.Sleep(time.Duration(tokenExp) * time.Second)
|
||||
if err := a.getToken(context.Background()); err != nil {
|
||||
a.logger.Errorf("Failed to refresh jwt token: %v", err)
|
||||
continue
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := a.getToken(ctx); err != nil {
|
||||
a.logger.Errorf("Failed to refresh jwt token: %v", err)
|
||||
continue
|
||||
}
|
||||
a.logger.Info("Refreshed Alist jwt token")
|
||||
}
|
||||
a.logger.Info("Refreshed Alist jwt token")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,13 +3,44 @@ package storage
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"maps"
|
||||
"sync"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
var UserStorages = make(map[int64][]Storage)
|
||||
var (
|
||||
storageMu sync.RWMutex
|
||||
// Storages maps storage names to initialized storage instances.
|
||||
Storages = make(map[string]Storage)
|
||||
|
||||
userStoragesMu sync.RWMutex
|
||||
// UserStorages maps user IDs to their available storage instances.
|
||||
UserStorages = make(map[int64][]Storage)
|
||||
|
||||
initFlight singleflight.Group
|
||||
)
|
||||
|
||||
// GetStorage returns the initialized storage instance for name, without
|
||||
// creating one on demand.
|
||||
func GetStorage(name string) (Storage, bool) {
|
||||
storageMu.RLock()
|
||||
defer storageMu.RUnlock()
|
||||
s, ok := Storages[name]
|
||||
return s, ok
|
||||
}
|
||||
|
||||
// AllStorages returns a snapshot copy of all initialized storages.
|
||||
func AllStorages() map[string]Storage {
|
||||
storageMu.RLock()
|
||||
defer storageMu.RUnlock()
|
||||
out := make(map[string]Storage, len(Storages))
|
||||
maps.Copy(out, Storages)
|
||||
return out
|
||||
}
|
||||
|
||||
// GetStorageByName returns storage by name from cache or creates new one
|
||||
// It should NOT be used to get storage for user, use GetStorageByUserIDAndName instead
|
||||
@@ -18,21 +49,41 @@ func GetStorageByName(ctx context.Context, name string) (Storage, error) {
|
||||
return nil, ErrStorageNameEmpty
|
||||
}
|
||||
|
||||
storageMu.RLock()
|
||||
storage, ok := Storages[name]
|
||||
storageMu.RUnlock()
|
||||
if ok {
|
||||
return storage, nil
|
||||
}
|
||||
cfg := config.C().GetStorageByName(name)
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("未找到存储 %s", name)
|
||||
return nil, fmt.Errorf("storage %s not found", name)
|
||||
}
|
||||
|
||||
storage, err := NewStorage(ctx, cfg)
|
||||
// Merge concurrent first-time initializations.
|
||||
v, err, _ := initFlight.Do("storage:"+name, func() (any, error) {
|
||||
storageMu.RLock()
|
||||
if existing, ok := Storages[name]; ok {
|
||||
storageMu.RUnlock()
|
||||
return existing, nil
|
||||
}
|
||||
storageMu.RUnlock()
|
||||
storage, err := NewStorage(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
storageMu.Lock()
|
||||
defer storageMu.Unlock()
|
||||
if existing, ok := Storages[name]; ok {
|
||||
return existing, nil
|
||||
}
|
||||
Storages[name] = storage
|
||||
return storage, nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
Storages[name] = storage
|
||||
return storage, nil
|
||||
return v.(Storage), nil
|
||||
}
|
||||
|
||||
// 检查 user 是否可用指定的 storage, 若不可用则返回未找到错误
|
||||
@@ -52,8 +103,11 @@ func GetUserStorages(ctx context.Context, chatID int64) []Storage {
|
||||
if chatID <= 0 {
|
||||
return nil
|
||||
}
|
||||
if storages, ok := UserStorages[chatID]; ok {
|
||||
return storages
|
||||
userStoragesMu.RLock()
|
||||
cached, ok := UserStorages[chatID]
|
||||
userStoragesMu.RUnlock()
|
||||
if ok {
|
||||
return cached
|
||||
}
|
||||
var storages []Storage
|
||||
for _, name := range config.C().GetStorageNamesByUserID(chatID) {
|
||||
@@ -75,9 +129,16 @@ func LoadStorages(ctx context.Context) {
|
||||
logger.Errorf("failed to load storage %s: %v", storage.GetName(), err)
|
||||
}
|
||||
}
|
||||
logger.Infof("successfully loaded %d storages", len(Storages))
|
||||
storageMu.RLock()
|
||||
loaded := len(Storages)
|
||||
storageMu.RUnlock()
|
||||
logger.Infof("successfully loaded %d storages", loaded)
|
||||
for user := range config.C().GetUsersID() {
|
||||
UserStorages[int64(user)] = GetUserStorages(ctx, int64(user))
|
||||
uid := int64(user)
|
||||
storages := GetUserStorages(ctx, uid)
|
||||
userStoragesMu.Lock()
|
||||
UserStorages[uid] = storages
|
||||
userStoragesMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/duke-git/lancet/v2/fileutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
@@ -52,15 +53,17 @@ func (l *Local) JoinStoragePath(path string) string {
|
||||
|
||||
func (l *Local) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||
l.logger.Infof("Saving file to %s", storagePath)
|
||||
storagePath = l.JoinStoragePath(storagePath)
|
||||
storagePath = filepath.Clean(storagePath)
|
||||
if filepath.IsAbs(storagePath) {
|
||||
return fmt.Errorf("local: storage path must be relative: %s", storagePath)
|
||||
}
|
||||
if storagePath == ".." || strings.HasPrefix(storagePath, ".."+string(filepath.Separator)) {
|
||||
return fmt.Errorf("local: storage path escapes base directory: %s", storagePath)
|
||||
}
|
||||
|
||||
ext := filepath.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
candidate := l.JoinStoragePath(storagePath)
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
for i := 1; l.existsPath(candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
}
|
||||
candidate = fsutil.UniquePath(l.config.BasePath, storagePath, l.existsPath, 1000)
|
||||
}
|
||||
|
||||
absPath, err := filepath.Abs(candidate)
|
||||
@@ -68,13 +71,17 @@ func (l *Local) Save(ctx context.Context, r io.Reader, storagePath string) error
|
||||
return err
|
||||
}
|
||||
if err := fileutil.CreateDir(filepath.Dir(absPath)); err != nil {
|
||||
return err
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
file, err := os.Create(absPath)
|
||||
if err != nil {
|
||||
return err
|
||||
return fmt.Errorf("failed to create file: %w", err)
|
||||
}
|
||||
defer file.Close()
|
||||
defer func() {
|
||||
if err := file.Close(); err != nil {
|
||||
l.logger.Errorf("Failed to close file %s: %v", absPath, err)
|
||||
}
|
||||
}()
|
||||
_, err = io.Copy(file, r)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -11,12 +11,12 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/minio/minio-go/v7"
|
||||
"github.com/minio/minio-go/v7/pkg/credentials"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -77,19 +77,11 @@ func (m *Minio) JoinStoragePath(p string) string {
|
||||
|
||||
func (m *Minio) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||
m.logger.Infof("Saving file from reader to %s", storagePath)
|
||||
storagePath = m.JoinStoragePath(storagePath)
|
||||
ext := path.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
candidate := m.JoinStoragePath(storagePath)
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
for i := 1; m.existsObject(ctx, candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
if i > 10 {
|
||||
m.logger.Errorf("Too many attempts to find a unique filename for %s", storagePath)
|
||||
candidate = fmt.Sprintf("%s_%s%s", base, xid.New().String(), ext)
|
||||
break
|
||||
}
|
||||
}
|
||||
candidate = fsutil.UniquePath(strings.TrimPrefix(m.config.BasePath, "/"), storagePath, func(c string) bool {
|
||||
return m.existsObject(ctx, c)
|
||||
}, 10)
|
||||
}
|
||||
size := int64(-1)
|
||||
if length := ctx.Value(ctxkey.ContentLength); length != nil {
|
||||
|
||||
@@ -8,7 +8,4 @@ var (
|
||||
ErrFailedToSaveFile = errors.New("rclone: failed to save file")
|
||||
ErrFailedToListFiles = errors.New("rclone: failed to list files")
|
||||
ErrFailedToOpenFile = errors.New("rclone: failed to open file")
|
||||
ErrFailedToCheckFile = errors.New("rclone: failed to check file exists")
|
||||
ErrFailedToCreateDir = errors.New("rclone: failed to create directory")
|
||||
ErrCommandFailed = errors.New("rclone: command execution failed")
|
||||
)
|
||||
|
||||
@@ -13,11 +13,11 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
type Rclone struct {
|
||||
@@ -51,9 +51,6 @@ func (r *Rclone) Init(ctx context.Context, cfg config.StorageConfig) error {
|
||||
}
|
||||
|
||||
remoteName := strings.TrimSuffix(r.config.Remote, ":")
|
||||
if !strings.HasSuffix(r.config.Remote, ":") {
|
||||
remoteName = r.config.Remote
|
||||
}
|
||||
|
||||
found := false
|
||||
scanner := bufio.NewScanner(bytes.NewReader(output))
|
||||
@@ -105,18 +102,11 @@ func (r *Rclone) getRemotePath(storagePath string) string {
|
||||
func (r *Rclone) Save(ctx context.Context, reader io.Reader, storagePath string) error {
|
||||
r.logger.Infof("Saving file to %s", storagePath)
|
||||
|
||||
ext := path.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
for i := 1; r.Exists(ctx, candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
if i > 100 {
|
||||
r.logger.Errorf("Too many attempts to find a unique filename for %s", storagePath)
|
||||
candidate = fmt.Sprintf("%s_%s%s", base, xid.New().String(), ext)
|
||||
break
|
||||
}
|
||||
}
|
||||
candidate = fsutil.UniquePath("", storagePath, func(c string) bool {
|
||||
return r.Exists(ctx, c)
|
||||
}, 100)
|
||||
}
|
||||
|
||||
remotePath := r.getRemotePath(candidate)
|
||||
@@ -285,6 +275,12 @@ func (r *rcloneCatReader) Close() error {
|
||||
if err := r.reader.Close(); err != nil {
|
||||
r.logger.Warnf("Failed to close reader: %v", err)
|
||||
}
|
||||
// Kill the cat process so Wait cannot block on a hung pipe.
|
||||
if r.cmd.Process != nil {
|
||||
if err := r.cmd.Process.Kill(); err != nil {
|
||||
r.logger.Warnf("Failed to kill rclone cat process: %v", err)
|
||||
}
|
||||
}
|
||||
if err := r.cmd.Wait(); err != nil {
|
||||
r.logger.Warnf("rclone cat process exited with error: %v", err)
|
||||
}
|
||||
|
||||
@@ -8,11 +8,11 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
storconfig "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/s3"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
type S3 struct {
|
||||
@@ -65,21 +65,13 @@ func (m *S3) JoinStoragePath(p string) string {
|
||||
|
||||
func (m *S3) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||
m.logger.Infof("Saving file from reader to %s", storagePath)
|
||||
storagePath = m.JoinStoragePath(storagePath)
|
||||
ext := path.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
candidate := m.JoinStoragePath(storagePath)
|
||||
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
// Unique filename
|
||||
for i := 1; m.existsKey(ctx, candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
if i > 10 {
|
||||
m.logger.Errorf("Too many attempts for unique filename: %s", storagePath)
|
||||
candidate = fmt.Sprintf("%s_%s%s", base, xid.New().String(), ext)
|
||||
break
|
||||
}
|
||||
}
|
||||
candidate = fsutil.UniquePath(strings.TrimPrefix(m.config.BasePath, "/"), storagePath, func(c string) bool {
|
||||
return m.existsKey(ctx, c)
|
||||
}, 10)
|
||||
}
|
||||
|
||||
// Determine content length
|
||||
|
||||
@@ -75,11 +75,18 @@ type StorageReadable interface {
|
||||
OpenFile(ctx context.Context, filePath string) (io.ReadCloser, int64, error)
|
||||
}
|
||||
|
||||
var Storages = make(map[string]Storage)
|
||||
|
||||
var _ StorageProgressSaver = (*telegram.Telegram)(nil)
|
||||
var _ StorageBatchProgressSaver = (*telegram.Telegram)(nil)
|
||||
|
||||
var _ StorageListable = (*alist.Alist)(nil)
|
||||
var _ StorageReadable = (*alist.Alist)(nil)
|
||||
var _ StorageListable = (*local.Local)(nil)
|
||||
var _ StorageReadable = (*local.Local)(nil)
|
||||
var _ StorageListable = (*rclone.Rclone)(nil)
|
||||
var _ StorageReadable = (*rclone.Rclone)(nil)
|
||||
var _ StorageListable = (*webdav.Webdav)(nil)
|
||||
var _ StorageReadable = (*webdav.Webdav)(nil)
|
||||
|
||||
type StorageConstructor func() Storage
|
||||
|
||||
var storageConstructors = map[storenum.StorageType]StorageConstructor{
|
||||
|
||||
@@ -83,6 +83,9 @@ func (t *Telegram) Name() string {
|
||||
return t.config.Name
|
||||
}
|
||||
|
||||
// Exists always reports false: Telegram offers no reliable way to query
|
||||
// whether a file already exists in a chat, so conflict policies do not apply
|
||||
// to this backend.
|
||||
func (t *Telegram) Exists(ctx context.Context, storagePath string) bool {
|
||||
return false
|
||||
}
|
||||
@@ -447,7 +450,7 @@ func planMediaGroups(items []batchMediaItem) [][]batchMediaItem {
|
||||
continue
|
||||
}
|
||||
end := i + 1
|
||||
for end < len(items) && end-i < 10 {
|
||||
for end < len(items) && end-i < tglimit.MaxAlbumItems {
|
||||
next := items[end]
|
||||
if next.useSingleSave || !next.albumEligible || next.chatID != item.chatID || next.item.SourceGroupKey != item.item.SourceGroupKey {
|
||||
break
|
||||
@@ -623,16 +626,16 @@ func (t *Telegram) splitUpload(
|
||||
|
||||
sender := ctx.Sender
|
||||
|
||||
if len(multiMedia) <= 10 {
|
||||
if len(multiMedia) <= tglimit.MaxAlbumItems {
|
||||
_, err = sender.WithUploader(upler).
|
||||
To(peer).
|
||||
Album(ctx, multiMedia[0], multiMedia[1:]...)
|
||||
return err
|
||||
}
|
||||
|
||||
// more than 10 parts, send in batches, each batch up to 10 parts
|
||||
for i := 0; i < len(multiMedia); i += 10 {
|
||||
end := min(i+10, len(multiMedia))
|
||||
// more than MaxAlbumItems parts, send in batches, each batch up to MaxAlbumItems parts
|
||||
for i := 0; i < len(multiMedia); i += tglimit.MaxAlbumItems {
|
||||
end := min(i+tglimit.MaxAlbumItems, len(multiMedia))
|
||||
batch := multiMedia[i:end]
|
||||
_, err = sender.WithUploader(upler).
|
||||
To(peer).
|
||||
|
||||
@@ -19,13 +19,14 @@ import (
|
||||
"github.com/rs/xid"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
|
||||
)
|
||||
|
||||
const (
|
||||
videoPartTargetRatio = 0.95
|
||||
videoSplitAttempts = 4
|
||||
minSegmentDuration = 1.0
|
||||
maxLosslessVideoParts = 10
|
||||
maxLosslessVideoParts = tglimit.MaxAlbumItems
|
||||
)
|
||||
|
||||
type losslessVideoPart struct {
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
package webdav
|
||||
|
||||
import "errors"
|
||||
|
||||
var (
|
||||
ErrFailedToCreateDirectory = errors.New("webdav: failed to create directory")
|
||||
ErrFailedToWriteFile = errors.New("webdav: failed to write file")
|
||||
ErrFailedToCheckFileExists = errors.New("webdav: failed to check if file exists")
|
||||
)
|
||||
@@ -11,11 +11,11 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
type Webdav struct {
|
||||
@@ -54,28 +54,18 @@ func (w *Webdav) JoinStoragePath(p string) string {
|
||||
|
||||
func (w *Webdav) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||
w.logger.Infof("Saving file to %s", storagePath)
|
||||
storagePath = w.JoinStoragePath(storagePath)
|
||||
ext := path.Ext(storagePath)
|
||||
base := strings.TrimSuffix(storagePath, ext)
|
||||
candidate := storagePath
|
||||
candidate := w.JoinStoragePath(storagePath)
|
||||
if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
|
||||
for i := 1; w.existsPath(ctx, candidate); i++ {
|
||||
candidate = fmt.Sprintf("%s_%d%s", base, i, ext)
|
||||
if i > 1000 {
|
||||
w.logger.Errorf("Too many attempts to find a unique filename for %s", storagePath)
|
||||
candidate = fmt.Sprintf("%s_%s%s", base, xid.New().String(), ext)
|
||||
break
|
||||
}
|
||||
}
|
||||
candidate = fsutil.UniquePath(w.config.BasePath, storagePath, func(c string) bool {
|
||||
return w.existsPath(ctx, c)
|
||||
}, 1000)
|
||||
}
|
||||
|
||||
if err := w.client.MkDir(ctx, path.Dir(candidate)); err != nil {
|
||||
w.logger.Errorf("Failed to create directory %s: %v", path.Dir(candidate), err)
|
||||
return ErrFailedToCreateDirectory
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
if err := w.client.WriteFile(ctx, candidate, r); err != nil {
|
||||
w.logger.Errorf("Failed to write file %s: %v", candidate, err)
|
||||
return ErrFailedToWriteFile
|
||||
return fmt.Errorf("failed to write file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -136,9 +126,6 @@ func (w *Webdav) ListFiles(ctx context.Context, dirPath string) ([]storagetypes.
|
||||
|
||||
isDir := resp.Propstat.Prop.ResourceType.IsCollection()
|
||||
|
||||
filePath := strings.TrimPrefix(decodedHref, path.Join("/", strings.Trim(path.Dir(fullPath), "/")))
|
||||
filePath = strings.TrimPrefix(filePath, "/")
|
||||
|
||||
fileInfo := storagetypes.FileInfo{
|
||||
Name: name,
|
||||
Path: path.Join(dirPath, name),
|
||||
|
||||
Reference in New Issue
Block a user