Compare commits

..
Author SHA1 Message Date
krau 942022969c style: format codes 2026-08-17 19:01:51 +08:00
krau ab89310325 fix: drop unbounded saved-path cache in telegram storage
Telegram cannot reliably query remote file existence, so the
cache answered a question it could not answer and grew without
bound. Exists returns false again, as before.
2026-08-16 22:05:08 +08:00
krau b27bb734be style: gofmt test files 2026-08-16 21:56:01 +08:00
krau fc0c8975a0 style: trim verbose comments
Drop process-style explanations; keep one-line behavior notes.
2026-08-16 21:56:01 +08:00
krau 524a259ae4 test: exercise the alist 401 refresh path
Concurrent uploads now reject the init token, share one
refresh and retry with the new token; token-only stays inert.
2026-08-16 21:51:11 +08:00
krau e8257c46d2 fix: refresh alist token after startup 401s
The init login no longer satisfies the dedup window, so an
early 401 triggers a real refresh; later 401s reuse it.
Clarify that streaming uploads cannot replay their body.
2026-08-16 21:51:11 +08:00
krau 6a11b8ac8d test: cover cancellation with IgnoreErrors 2026-08-16 21:44:40 +08:00
krau 47d941d851 test: cover notifyProgress tracker call 2026-08-16 21:44:40 +08:00
krau 7e502442c5 fix: propagate cancellation past IgnoreErrors
Cancelled tasks must not be reported as successful: only
ordinary element failures are ignored.
2026-08-16 21:44:35 +08:00
krau a9425cb5ea fix: call tracker from notifyProgress instead of recursing
The helper called itself, overflowing the stack on any batch
task with a progress tracker.
2026-08-16 21:44:35 +08:00
krau 6ad2a2c884 test: cover alist refresh semantics
Concurrent refresh uses username/password; token-only storage
never attempts a login on 401.
2026-08-16 21:35:21 +08:00
krau 0d42c9f23d fix: skip login refresh for token-only alist storage
401 responses surface the auth error instead of sending a
credential-less login; the refresh window never exceeds TokenExp.
2026-08-16 21:35:21 +08:00
krau 4c8f35ae80 fix: report partial failure in batch done message
IgnoreErrors runs with failed elements show success and failed
counts instead of claiming every file completed.
2026-08-16 21:35:14 +08:00
krau 0a67b8bb57 test: cover nil-tracker task execution 2026-08-16 21:35:14 +08:00
krau 196bb9941f fix: keep album order when filtering failed batch items
Download results are stored by original index so surviving
elements keep their source order.
2026-08-16 21:35:14 +08:00
krau 973f3499a9 fix: guard nil progress trackers in api-created tasks
Batch, telegraph and transfer tasks run without a Telegram tracker
when created through the API; their callbacks must not panic.
2026-08-16 21:34:52 +08:00
krau d9cb5f9cfc fix: deduplicate concurrent storage initialization
singleflight merges first-time inits so side effects are not duplicated.
2026-08-16 21:20:12 +08:00
krau ccb702148e fix: localize storage lookup errors in /dir
Use the shared i18n key and escape the dynamic error.
2026-08-16 21:20:12 +08:00
krau 85b1f35dc5 fix: count parsed resources in progress text 2026-08-16 21:20:12 +08:00
krau 9f93a95258 test: cover concurrent alist 401 retry
Ten parallel uploads share a single re-login under -race.
2026-08-16 21:20:12 +08:00
krau 189bf9c736 fix: deduplicate alist token refreshes
Guard token access with a mutex and merge concurrent logins.
Reuse a recent refresh to avoid login storms.
2026-08-16 21:20:11 +08:00
krau 534ed7a7c2 fix: record telegram saved paths only after upload
Skip-large returns a sentinel so skipped files are not marked as saved.
2026-08-16 21:20:11 +08:00
krau 1d4997ba2d fix: upload only downloaded batch elements
Failed elements keep partial cache files and never reach the backend.
Successfully downloaded siblings still upload when one element fails.
2026-08-16 21:20:11 +08:00
krau f7e532ca7e fix: make cache init idempotent 2026-08-16 20:30:06 +08:00
krau 0957d93da3 chore: ignore cache directory 2026-08-16 20:30:06 +08:00
krau 324d8c100f ci: fix BuildTime formatting and align checkout
Actions format does not format dates; pass the raw timestamp.
2026-08-16 20:28:45 +08:00
krau 4cb23f04b6 docs: document parser plugin config 2026-08-16 20:28:45 +08:00
krau a92ca75c4f refactor: drop unused hook config 2026-08-16 20:28:45 +08:00
krau 4e99b4bdc9 fix: load remote config without local lookup
Skip the local file search after reading a config URL and add a timeout.
2026-08-16 20:28:45 +08:00
krau 389be59371 refactor: remove dead code from api and bot
Drop the empty ProgressTracker shim, unused token context key and a redundant SetBotCommands call.
2026-08-16 20:28:39 +08:00
krau 0d49ae94af fix: add timeouts and backoff to webhook delivery 2026-08-16 20:28:39 +08:00
krau 48e739f06d fix: fail fast when API listen fails
Bind synchronously and surface errors instead of logging them.
2026-08-16 20:28:39 +08:00
krau b4fbbae068 fix: notify users on invalid update version 2026-08-16 20:28:33 +08:00
krau 29be3ee90c fix: avoid leaking raw errors in /dir reply 2026-08-16 20:28:33 +08:00
krau 56660a7705 fix: initialize userbot context once
Replace the racy lazy init with sync.OnceValue.
2026-08-16 20:28:33 +08:00
krau 0dfb6af153 fix: require permission for callback handlers 2026-08-16 20:28:33 +08:00
krau 63f53fac7f fix: isolate media groups per user
Key pending groups by chat, user and group id.
2026-08-16 20:28:33 +08:00
krau 607ec8aced fix: guard callback data parsing
Reject malformed callback payloads before indexing split parts.
2026-08-16 20:28:33 +08:00
krau cd9886d710 fix: show all queued tasks in /task list
Render up to ten tasks and append the truncation note once.
2026-08-16 20:28:33 +08:00
krau 3f7f50133e test: cover invalid plugin version rejection 2026-08-16 20:28:27 +08:00
krau c7911cc208 refactor: remove commented-out dead code
Drop unused kemono legacy types and commented response structs.
2026-08-16 20:28:27 +08:00
krau b31d628c19 fix: guard kemono parser against nil fields
Skip sparse preview and attachment entries instead of panicking.
2026-08-16 20:28:27 +08:00
krau 60fd2e04d9 fix: harden JS parser plugin runtime
Validate semver instead of panicking and require canHandle.
Recover plugin worker panics and time out CanHandle calls.
Return a copy from the registry and sanitize install filenames.
2026-08-16 20:28:27 +08:00
krau b72dd67be9 fix: guard storage registry maps
Protect Storages and UserStorages with mutexes and expose read accessors.
2026-08-16 20:28:21 +08:00
krau aa25eb1510 fix: enforce telegram album limits and track saved paths
Use tglimit.MaxAlbumItems for album batching and video splitting.
Exists now reports previously saved paths instead of always false.
2026-08-16 20:28:21 +08:00
krau ba0deababc fix: fail alist init instead of exiting
Replace log.Fatalf with wrapped errors so a bad alist cannot kill the bot.
Cancel token refresh with the init context and re-login on 401.
2026-08-16 20:28:21 +08:00
krau d1b2dbfe5f fix: kill rclone subprocess on reader close
Prevent cat processes from hanging after the pipe is closed.
2026-08-16 20:28:15 +08:00
krau bd7da7c31e fix: preserve webdav error causes
Wrap mkdir and write failures with %w and drop dead error values.
2026-08-16 20:28:15 +08:00
krau 1981d8662b fix: sanitize local storage paths
Reject absolute paths and dot-dot escapes in Save.
Check close errors and wrap creation failures.
2026-08-16 20:28:15 +08:00
krau 51fcd17922 refactor: share unique filename logic
Storage backends use fsutil.UniquePath instead of local loops.
2026-08-16 20:28:15 +08:00
krau 610f586d45 refactor: share progress throttling helpers
Move size-tiered and count-based throttling into progressutil.
Localize hardcoded Chinese progress strings.
Drop five duplicated implementations and dead local copies.
2026-08-16 20:28:07 +08:00
krau 6bc12d6feb fix: validate i18n key parity across locales
geni18n fails when a language file misses any key.
Align the syncpeers completion key between en and zh-Hans.
Translate three untranslated parse keys in zh-Hans.
2026-08-16 20:28:03 +08:00
krau ad41aec22c fix: report streamed upload bytes in batch tasks
Stream downloads report their byte count as the upload total.
2026-08-16 20:27:53 +08:00
krau 4389bfe046 fix: honor IgnoreErrors in batch tasks
Element failures no longer cancel sibling elements or stop later groups.
2026-08-16 20:27:42 +08:00
krau 832eb27d4f fix: harden per-item processing tracking in tasks
Use the resource fingerprint as the dedup key on insert and delete.
Check and set processing entries atomically instead of TOCTOU.
Count only successful downloads.
2026-08-16 20:27:42 +08:00
krau b6e981282d fix: make resource fingerprints deterministic
Sort map keys before hashing so Resource.ID is stable.
2026-08-16 20:27:42 +08:00
krau 9ed5b5f962 fix: init task queue lazily and close on shutdown
AddTask is safe before Run by initializing the queue once.
Close unblocks workers waiting in Get.
2026-08-16 20:27:26 +08:00
krau 4527022853 fix: prevent queue deadlock after cancelling tasks
Get no longer recurses while holding the mutex.
Cancelled queued tasks leave the map, making their IDs reusable.
Closed empty queues return ErrQueueClosed.
2026-08-16 20:27:26 +08:00
34 changed files with 86 additions and 2232 deletions
+1 -2
View File
@@ -1,7 +1,6 @@
package api
import (
"context"
"sync"
"time"
@@ -196,6 +195,6 @@ func (t *TaskProgressInfo) Emit(e taskevent.Event) {
if notify {
payload := CreateWebhookPayload(t.TaskID, t.Type, t.Status, t.Storage, t.Path, e.Err)
SendWebhook(context.Background(), payload)
SendWebhook(nil, payload)
}
}
+4 -22
View File
@@ -1,13 +1,10 @@
package handlers
import (
"errors"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/dirutil"
"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"
@@ -15,23 +12,10 @@ import (
"github.com/krau/SaveAny-Bot/storage"
)
// responsibleUserID returns the sender's ID. Callback queries carry it
// natively; message updates resolve it through the entity map.
func responsibleUserID(u *ext.Update) int64 {
if u.CallbackQuery != nil {
return u.CallbackQuery.GetUserID()
}
return u.GetUserChat().GetID()
}
func checkPermission(ctx *ext.Context, update *ext.Update) error {
userID := responsibleUserID(update)
userID := update.GetUserChat().GetID()
if !slice.Contain(config.C().GetUsersID(), userID) {
if cbq := update.CallbackQuery; cbq != nil {
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(cbq.GetQueryID(), i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)))
} else {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)), nil)
}
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)), nil)
return dispatcher.EndGroups
}
@@ -39,12 +23,10 @@ func checkPermission(ctx *ext.Context, update *ext.Update) error {
}
// withPermission wraps a callback handler with the same whitelist check used
// for message handlers (checkPermission). ContinueGroups is the dispatcher's
// success sentinel, not an error: only real failures and EndGroups stop the
// chain before the wrapped handler runs.
// 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 && !errors.Is(err, dispatcher.ContinueGroups) {
if err := checkPermission(ctx, update); err != nil {
return err
}
return handler(ctx, update)
-85
View File
@@ -1,85 +0,0 @@
package handlers
import (
"os"
"path/filepath"
"testing"
"github.com/celestix/gotgproto/ext"
"github.com/celestix/gotgproto/types"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/config"
)
// Regression: callback queries usually arrive as updateShort without entity
// maps, so resolving the sender through the entity map yields ID 0 and every
// click was denied by the whitelist check. Callback updates must use the
// native UserID field.
func TestResponsibleUserID(t *testing.T) {
tests := []struct {
name string
update *ext.Update
want int64
}{
{
name: "callback query uses native user id",
update: &ext.Update{CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 42}},
want: 42,
},
{
name: "message resolves through entity map",
update: &ext.Update{
EffectiveMessage: &types.Message{Message: &tg.Message{PeerID: &tg.PeerUser{UserID: 7}}},
Entities: &tg.Entities{Users: map[int64]*tg.User{7: {ID: 7}}},
},
want: 7,
},
{
name: "callback query ignores entity map",
update: &ext.Update{
CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 9},
Entities: &tg.Entities{Users: map[int64]*tg.User{8: {ID: 8}}},
},
want: 9,
},
{
name: "unresolvable update yields zero",
update: &ext.Update{},
want: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := responsibleUserID(tt.update); got != tt.want {
t.Fatalf("responsibleUserID() = %d, want %d", got, tt.want)
}
})
}
}
// Regression: withPermission must treat ContinueGroups (the dispatcher's
// success sentinel) as a pass and invoke the wrapped handler. v0.60.1 treated
// it as an error, so every permitted callback was swallowed before the real
// handler ran.
func TestWithPermissionInvokesHandler(t *testing.T) {
path := filepath.Join(t.TempDir(), "config.toml")
if err := os.WriteFile(path, []byte("workers = 2\n\n[[users]]\nid = 42\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := config.Init(t.Context(), path); err != nil {
t.Fatal(err)
}
update := &ext.Update{CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 42}}
called := false
handler := withPermission(func(ctx *ext.Context, u *ext.Update) error {
called = true
return nil
})
if err := handler(&ext.Context{}, update); err != nil {
t.Fatalf("withPermission returned error: %v", err)
}
if !called {
t.Fatal("withPermission did not invoke the wrapped handler")
}
}
+1 -1
View File
@@ -64,7 +64,7 @@ func handleWatchCmd(ctx *ext.Context, update *ext.Update) error {
filter := ""
if len(args) > 2 {
filterArg := strings.Join(args[2:], " ")
filterType, _, _ := strings.Cut(filterArg, ":")
filterType := strings.Split(filterArg, ":")[0]
filterData := strings.Split(filterArg, ":")[1]
if filterType == "" || filterData == "" {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgWatchErrorFilterFormatInvalid)), nil)
-24
View File
@@ -10,14 +10,12 @@ import (
"slices"
"github.com/charmbracelet/log"
"github.com/gotd/td/telegram/downloader"
"github.com/krau/SaveAny-Bot/api"
"github.com/krau/SaveAny-Bot/client/bot"
userclient "github.com/krau/SaveAny-Bot/client/user"
"github.com/krau/SaveAny-Bot/common/cache"
"github.com/krau/SaveAny-Bot/common/i18n"
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/database"
@@ -58,19 +56,6 @@ func Run(cmd *cobra.Command, _ []string) {
cancel()
}()
core.SetDownloaderProvider(func() downloader.Client {
if ectx := bot.ExtContext(); ectx != nil {
return ectx.Raw
}
return nil
})
// 恢复任务携带 ext 上下文, 让进度编辑/取消按钮在恢复后继续工作。
recoverCtx := context.Background()
if ectx := bot.ExtContext(); ectx != nil {
recoverCtx = tgutil.ExtWithContext(recoverCtx, ectx)
}
core.RecoverTasks(recoverCtx)
core.Run(ctx)
<-ctx.Done()
@@ -117,15 +102,6 @@ func cleanCache() {
log.Error("Invalid cache directory", "path", config.C().Temp.BasePath)
return
}
unfinished, err := database.CountUnfinishedTasks(context.Background())
if err != nil {
log.Error("Failed to count unfinished tasks, skipping cache cleanup", "error", err)
return
}
if unfinished > 0 {
log.Info("Skipping cache cleanup: unfinished tasks need their cache files for recovery", "tasks", unfinished)
return
}
currentDir, err := os.Getwd()
if err != nil {
log.Error("Failed to get working directory", "error", err)
+1 -1
View File
@@ -351,7 +351,7 @@ bot:
info_filename_prefix: "Filename: "
info_prompt_select_storage: "\nPlease select storage"
progress:
batch_status_header: "<b>📦 Processing</b>\n\nFiles: <code>{{.Total}}</code> | Total size: <code>{{.TotalSize}}</code>\nStatus: ✅ <code>{{.Completed}}</code> | 📥 <code>{{.Downloaded}}</code> | ⏳ <code>{{.Waiting}}</code>\nTotal speed: ⬇️ <code>{{.DownloadSpeed}}</code> | ⬆️ <code>{{.UploadSpeed}}</code>"
batch_status_header: "<b>📦 Processing</b>\n\nFiles: <code>{{.Total}}</code>\nStatus: ✅ <code>{{.Completed}}</code> | 📥 <code>{{.Downloaded}}</code> | ⏳ <code>{{.Waiting}}</code>\nTotal speed: ⬇️ <code>{{.DownloadSpeed}}</code> | ⬆️ <code>{{.UploadSpeed}}</code>"
batch_item_downloading: "<blockquote><b>⬇️ {{.Index}}/{{.Total}} Downloading</b>\n<code>{{.Name}}</code>\n{{.Bar}} <code>{{.Progress}}%</code>\nSpeed: <code>{{.Speed}}</code>\nSize: <code>{{.Current}}</code> / <code>{{.Size}}</code></blockquote>"
batch_item_downloading_unknown: "<blockquote><b>⬇️ {{.Index}}/{{.Total}} Downloading</b>\n<code>{{.Name}}</code>\nSpeed: <code>{{.Speed}}</code>\nSize: <code>{{.Current}}</code> / unknown</blockquote>"
batch_item_transferring: "<blockquote><b>↕️ {{.Index}}/{{.Total}} Transferring</b>\n<code>{{.Name}}</code>\n{{.Bar}} <code>{{.Progress}}%</code>\nSpeed: <code>{{.Speed}}</code>\nSize: <code>{{.Current}}</code> / <code>{{.Size}}</code></blockquote>"
+1 -1
View File
@@ -352,7 +352,7 @@ bot:
info_filename_prefix: "文件名: "
info_prompt_select_storage: "\n请选择存储位置"
progress:
batch_status_header: "<b>📦 正在处理</b>\n\n文件:<code>{{.Total}}</code> 总大小:<code>{{.TotalSize}}</code>\n状态:✅ <code>{{.Completed}}</code> 📥 <code>{{.Downloaded}}</code> ⏳ <code>{{.Waiting}}</code>\n总速度:⬇️ <code>{{.DownloadSpeed}}</code> ⬆️ <code>{{.UploadSpeed}}</code>"
batch_status_header: "<b>📦 正在处理</b>\n\n文件:<code>{{.Total}}</code>\n状态:✅ <code>{{.Completed}}</code> 📥 <code>{{.Downloaded}}</code> ⏳ <code>{{.Waiting}}</code>\n总速度:⬇️ <code>{{.DownloadSpeed}}</code> ⬆️ <code>{{.UploadSpeed}}</code>"
batch_item_downloading: "<blockquote><b>⬇️ {{.Index}}/{{.Total}} 下载中</b>\n<code>{{.Name}}</code>\n{{.Bar}} <code>{{.Progress}}%</code>\n速度:<code>{{.Speed}}</code>\n大小:<code>{{.Current}}</code> / <code>{{.Size}}</code></blockquote>"
batch_item_downloading_unknown: "<blockquote><b>⬇️ {{.Index}}/{{.Total}} 下载中</b>\n<code>{{.Name}}</code>\n速度:<code>{{.Speed}}</code>\n大小:<code>{{.Current}}</code> / 未知</blockquote>"
batch_item_transferring: "<blockquote><b>↕️ {{.Index}}/{{.Total}} 传输中</b>\n<code>{{.Name}}</code>\n{{.Bar}} <code>{{.Progress}}%</code>\n速度:<code>{{.Speed}}</code>\n大小:<code>{{.Current}}</code> / <code>{{.Size}}</code></blockquote>"
+1 -22
View File
@@ -1,10 +1,7 @@
package tdler
import (
"context"
"github.com/gotd/td/telegram/downloader"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
@@ -13,23 +10,5 @@ import (
func NewDownloader(file tfile.TGFile) *downloader.Builder {
return downloader.NewDownloader().WithPartSize(tglimit.MaxPartSize).
Download(eofAwareClient{Client: file.Dler(), size: file.Size()}, file.Location()).
WithThreads(dlutil.BestThreads(file.Size(), config.C().Threads))
}
// eofAwareClient answers upload.getFile requests at or past the end of the
// file with an empty chunk. gotd's downloader is size-unaware: for files
// whose size is an exact multiple of the part size it issues one final
// request at offset == size and expects an empty chunk, but Telegram rejects
// it with 400 OFFSET_INVALID and the whole download fails.
type eofAwareClient struct {
downloader.Client
size int64
}
func (c eofAwareClient) UploadGetFile(ctx context.Context, req *tg.UploadGetFileRequest) (tg.UploadFileClass, error) {
if req.Offset >= c.size {
return &tg.UploadFile{}, nil
}
return c.Client.UploadGetFile(ctx, req)
Download(file.Dler(), file.Location()).WithThreads(dlutil.BestThreads(file.Size(), config.C().Threads))
}
-112
View File
@@ -1,112 +0,0 @@
package tdler
import (
"bytes"
"context"
"sync"
"testing"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
// serverLikeClient mimics real Telegram upload.getFile behavior: it returns
// up to limit bytes per chunk, and answers any offset at or past the end of
// the file with 400 OFFSET_INVALID.
type serverLikeClient struct {
data []byte
mu sync.Mutex
maxOffset int64
}
func (c *serverLikeClient) UploadGetFile(_ context.Context, req *tg.UploadGetFileRequest) (tg.UploadFileClass, error) {
c.mu.Lock()
if req.Offset > c.maxOffset {
c.maxOffset = req.Offset
}
c.mu.Unlock()
if req.Offset >= int64(len(c.data)) {
return nil, tgerr.New(400, "OFFSET_INVALID")
}
end := min(len(c.data), int(req.Offset)+req.Limit)
return &tg.UploadFile{Bytes: c.data[req.Offset:end]}, nil
}
func (c *serverLikeClient) UploadGetFileHashes(context.Context, *tg.UploadGetFileHashesRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadReuploadCDNFile(context.Context, *tg.UploadReuploadCDNFileRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadGetCDNFileHashes(context.Context, *tg.UploadGetCDNFileHashesRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadGetWebFile(context.Context, *tg.UploadGetWebFileRequest) (*tg.UploadWebFile, error) {
return nil, nil
}
type memWriterAt struct {
b []byte
}
func (w *memWriterAt) WriteAt(p []byte, off int64) (int, error) {
copy(w.b[off:], p)
return len(p), nil
}
func TestDownloadServerLikeEOF(t *testing.T) {
const partSize = 1024 * 1024
tests := []struct {
name string
size int
parallel bool
}{
{"stream exact multiple of part size", 2 * partSize, false},
{"stream non-multiple", 2*partSize + 12345, false},
{"stream smaller than part size", 1234, false},
{"parallel exact multiple of part size", 2 * partSize, true},
{"parallel non-multiple", 2*partSize + 12345, true},
{"parallel smaller than part size", 1234, true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
data := make([]byte, tt.size)
for i := range data {
data[i] = byte(i % 251)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(
&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2},
client, int64(tt.size), "test.bin",
)
dl := NewDownloader(file)
var got []byte
var err error
if tt.parallel {
buf := make([]byte, tt.size)
_, err = dl.WithThreads(4).Parallel(context.Background(), &memWriterAt{b: buf})
got = buf
} else {
var buf bytes.Buffer
_, err = dl.Stream(context.Background(), &buf)
got = buf.Bytes()
}
if err != nil {
t.Fatalf("download failed: %v", err)
}
if !bytes.Equal(got, data) {
t.Fatalf("downloaded %d bytes, want %d matching bytes", len(got), len(data))
}
if client.maxOffset >= int64(tt.size) {
t.Fatalf("requested offset %d at or past EOF (size %d)", client.maxOffset, tt.size)
}
})
}
}
-257
View File
@@ -1,257 +0,0 @@
package tdler
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"os"
"strings"
"sync"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"golang.org/x/sync/errgroup"
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
const maxChunkRetries = 20
// resumeBitmap records which partSize-aligned blocks of a download have been
// durably written, so an interrupted download can continue after a restart.
// The bitmap file is rewritten atomically after every completed block.
type resumeBitmap struct {
PartSize int `json:"part_size"`
Size int64 `json:"size"`
Blocks []uint64 `json:"blocks"`
mu sync.Mutex
}
func newResumeBitmap(size int64) *resumeBitmap {
b := &resumeBitmap{PartSize: tglimit.MaxPartSize, Size: size}
b.ensureBlocks()
return b
}
func (b *resumeBitmap) ensureBlocks() {
if need := (b.blockCount() + 63) / 64; len(b.Blocks) < need {
b.Blocks = make([]uint64, need)
}
}
func (b *resumeBitmap) blockCount() int {
return int((b.Size + int64(b.PartSize) - 1) / int64(b.PartSize))
}
func (b *resumeBitmap) isDone(block int) bool {
return b.Blocks[block/64]&(1<<uint(block%64)) != 0
}
func (b *resumeBitmap) markDone(block int) {
b.Blocks[block/64] |= 1 << uint(block%64)
}
func (b *resumeBitmap) complete() bool {
for block := 0; block < b.blockCount(); block++ {
if !b.isDone(block) {
return false
}
}
return true
}
func (b *resumeBitmap) missingBlocks() []int {
missing := make([]int, 0, b.blockCount())
for block := 0; block < b.blockCount(); block++ {
if !b.isDone(block) {
missing = append(missing, block)
}
}
return missing
}
func loadResumeBitmap(path string) (*resumeBitmap, error) {
data, err := os.ReadFile(path)
if errors.Is(err, os.ErrNotExist) {
return nil, nil
}
if err != nil {
return nil, fmt.Errorf("read resume bitmap: %w", err)
}
var b resumeBitmap
if err := json.Unmarshal(data, &b); err != nil {
// 无法解析的位图 (外部损坏): 删除并视为不存在, 全量重下自愈。
_ = os.Remove(path)
return nil, nil
}
if b.Size <= 0 || b.PartSize <= 0 {
// 无效位图 (损坏或旧格式), 视为不存在, 全量重下。
_ = os.Remove(path)
return nil, nil
}
b.ensureBlocks()
return &b, nil
}
func (b *resumeBitmap) save(path string) error {
b.mu.Lock()
defer b.mu.Unlock()
return b.saveLocked(path)
}
func (b *resumeBitmap) saveLocked(path string) error {
data, err := json.Marshal(b)
if err != nil {
return fmt.Errorf("marshal resume bitmap: %w", err)
}
tmp := path + ".tmp"
if err := os.WriteFile(tmp, data, 0o644); err != nil {
return fmt.Errorf("write resume bitmap: %w", err)
}
return os.Rename(tmp, path)
}
func (b *resumeBitmap) markAndSave(block int, path string) error {
b.mu.Lock()
defer b.mu.Unlock()
b.markDone(block)
return b.saveLocked(path)
}
func isRetryableTimeout(ctx context.Context, err error) bool {
if err == nil || ctx.Err() != nil {
return false
}
if tgerr.Is(err, tg.ErrTimeout) || errors.Is(err, context.DeadlineExceeded) {
return true
}
var netErr net.Error
return errors.As(err, &netErr) && netErr.Timeout()
}
// fetchChunk downloads one partSize-aligned chunk, retrying flood waits and
// transient timeouts like gotd's downloader does.
func fetchChunk(ctx context.Context, file tfile.TGFile, offset int64, limit int) ([]byte, error) {
req := &tg.UploadGetFileRequest{
Location: file.Location(),
Offset: offset,
Limit: limit,
}
timeoutRetries := 0
for {
res, err := file.Dler().UploadGetFile(ctx, req)
if err == nil {
switch r := res.(type) {
case *tg.UploadFile:
return r.Bytes, nil
case *tg.UploadFileCDNRedirect:
return nil, fmt.Errorf("CDN redirect is not supported (dc %d)", r.DCID)
default:
return nil, fmt.Errorf("unexpected upload.getFile response %T", res)
}
}
if flood, ferr := tgerr.FloodWait(ctx, err); ferr != nil {
if flood {
// FloodWait already slept; retry.
continue
}
if isRetryableTimeout(ctx, ferr) {
timeoutRetries++
if timeoutRetries >= maxChunkRetries {
return nil, fmt.Errorf("get chunk at %d: retry limit reached: %w", offset, ferr)
}
continue
}
return nil, fmt.Errorf("get chunk at %d: %w", offset, ferr)
}
}
}
// DownloadResumable downloads file to w in partSize chunks, skipping blocks
// already recorded as complete in bitmapPath and persisting every completed
// block so an interrupted download can resume. A missing or incompatible
// bitmap starts a full download. Requires a known, non-zero file size.
func DownloadResumable(
ctx context.Context,
file tfile.TGFile,
w io.WriterAt,
threads int,
bitmapPath string,
) error {
if file.Size() <= 0 {
return fmt.Errorf("resumable download requires a known size")
}
bm, err := loadResumeBitmap(bitmapPath)
if err != nil {
return err
}
// 位图描述的数据文件 (bitmapPath 去掉 .bitmap 后缀) 必须存在且非空:
// 若缺失或为空, 已标记完成的块字节已丢失, 必须重置位图全量重下。
if bm != nil {
partPath := strings.TrimSuffix(bitmapPath, ".bitmap")
if stat, err := os.Stat(partPath); err != nil || stat.Size() == 0 {
if err := os.Remove(bitmapPath); err != nil && !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("reset stale resume bitmap: %w", err)
}
bm = nil
}
}
if bm == nil || bm.PartSize != tglimit.MaxPartSize || bm.Size != file.Size() {
bm = newResumeBitmap(file.Size())
if err := bm.save(bitmapPath); err != nil {
return err
}
}
missing := bm.missingBlocks()
if len(missing) == 0 {
return nil
}
eg, gctx := errgroup.WithContext(ctx)
eg.SetLimit(threads)
for _, block := range missing {
block := block
eg.Go(func() error {
offset := int64(block) * int64(bm.PartSize)
data, err := fetchChunk(gctx, file, offset, bm.PartSize)
if err != nil {
return err
}
if len(data) == 0 {
return fmt.Errorf("file ended early at offset %d (expected size %d)", offset, bm.Size)
}
if _, err := w.WriteAt(data, offset); err != nil {
return fmt.Errorf("write chunk at offset %d: %w", offset, err)
}
return bm.markAndSave(block, bitmapPath)
})
}
if err := eg.Wait(); err != nil {
return err
}
if !bm.complete() {
return fmt.Errorf("download finished with missing blocks")
}
return nil
}
// RemoveResumeState deletes the bitmap file of a completed download.
func RemoveResumeState(bitmapPath string) error {
if err := os.Remove(bitmapPath); err != nil && !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("remove resume bitmap: %w", err)
}
if err := os.Remove(bitmapPath + ".tmp"); err != nil && !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("remove resume bitmap temp: %w", err)
}
return nil
}
// ResumeStatePath returns the bitmap path for a download cache file.
func ResumeStatePath(cachePath string) string {
return cachePath + ".bitmap"
}
-271
View File
@@ -1,271 +0,0 @@
package tdler
import (
"context"
"errors"
"os"
"path/filepath"
"testing"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
// failAfterClient serves the first failAfter chunks, then returns err.
type failAfterClient struct {
*serverLikeClient
failAfter int
calls int
err error
}
func (c *failAfterClient) UploadGetFile(ctx context.Context, req *tg.UploadGetFileRequest) (tg.UploadFileClass, error) {
c.calls++
if c.calls > c.failAfter {
return nil, c.err
}
return c.serverLikeClient.UploadGetFile(ctx, req)
}
func TestDownloadResumableFull(t *testing.T) {
data := make([]byte, 3*1024*1024+123)
for i := range data {
data[i] = byte(i % 251)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
dir := t.TempDir()
bitmapPath := filepath.Join(dir, "test.bin.bitmap")
w := &memWriterAt{b: make([]byte, len(data))}
if err := DownloadResumable(context.Background(), file, w, 4, bitmapPath); err != nil {
t.Fatalf("download failed: %v", err)
}
if !bytesEqual(w.b, data) {
t.Fatalf("downloaded data mismatch")
}
bm, err := loadResumeBitmap(bitmapPath)
if err != nil {
t.Fatalf("load bitmap: %v", err)
}
if bm == nil || !bm.complete() {
t.Fatalf("bitmap not complete after full download")
}
}
func TestDownloadResumableInterrupted(t *testing.T) {
data := make([]byte, 5*1024*1024) // exactly 5 blocks
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
bitmapPath := filepath.Join(dir, "test.bin.bitmap")
w := &memWriterAt{b: make([]byte, len(data))}
// First run: 3 blocks complete, 4th request fails.
flaky := &failAfterClient{
serverLikeClient: &serverLikeClient{data: data},
failAfter: 3,
err: tgerr.New(500, "INTERNAL_SERVER_ERROR"),
}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, flaky, int64(len(data)), "test.bin")
err := DownloadResumable(context.Background(), file, w, 1, bitmapPath)
if err == nil {
t.Fatalf("expected first run to fail")
}
bm, err := loadResumeBitmap(bitmapPath)
if err != nil {
t.Fatalf("load bitmap after interruption: %v", err)
}
if bm == nil {
t.Fatalf("bitmap missing after interruption")
}
if got := bm.blockCount() - len(bm.missingBlocks()); got != 3 {
t.Fatalf("expected 3 completed blocks, got %d", got)
}
// Second run: only the missing blocks are requested.
healthy := &serverLikeClient{data: data}
file = tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, healthy, int64(len(data)), "test.bin")
if err := DownloadResumable(context.Background(), file, w, 1, bitmapPath); err != nil {
t.Fatalf("resume failed: %v", err)
}
if !bytesEqual(w.b, data) {
t.Fatalf("resumed data mismatch")
}
if healthy.maxOffset >= int64(len(data)) {
t.Fatalf("resume requested offset %d at or past EOF", healthy.maxOffset)
}
if bm, err = loadResumeBitmap(bitmapPath); err != nil || bm == nil || !bm.complete() {
t.Fatalf("bitmap not complete after resume: %v", err)
}
}
func TestDownloadResumableBitmapResetOnSizeChange(t *testing.T) {
data := make([]byte, 2*1024*1024)
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
bitmapPath := filepath.Join(dir, "test.bin.bitmap")
w := &memWriterAt{b: make([]byte, len(data))}
// Record a bitmap claiming the old, larger file is fully downloaded.
stale := newResumeBitmap(int64(4 * 1024 * 1024))
if err := stale.save(bitmapPath); err != nil {
t.Fatalf("save stale bitmap: %v", err)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
if err := DownloadResumable(context.Background(), file, w, 1, bitmapPath); err != nil {
t.Fatalf("download with stale bitmap failed: %v", err)
}
if !bytesEqual(w.b, data) {
t.Fatalf("data mismatch with stale bitmap")
}
}
// TestDownloadResumablePartMissingOrTruncated resets the bitmap: skipped
// blocks would otherwise be zero-filled (caller recreates the part file
// without its bytes), or the download would wedge forever on a stale
// complete bitmap.
func TestDownloadResumablePartMissingOrTruncated(t *testing.T) {
data := make([]byte, 5*1024*1024)
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
partPath := filepath.Join(dir, "test.bin.part")
bitmapPath := ResumeStatePath(partPath)
tests := []struct {
name string
doneBlocks []int
createPart bool
truncate bool
}{
{"part missing, partial bitmap", []int{0, 1, 2}, false, false},
{"part empty, partial bitmap", []int{0, 1, 2}, true, true},
{"part missing, complete bitmap", []int{0, 1, 2, 3, 4}, false, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
os.Remove(partPath)
os.Remove(bitmapPath)
bm := newResumeBitmap(int64(len(data)))
for _, block := range tt.doneBlocks {
bm.markDone(block)
}
if err := bm.save(bitmapPath); err != nil {
t.Fatal(err)
}
if tt.createPart {
// Simulate the caller re-creating the part file (truncating).
if err := os.WriteFile(partPath, nil, 0o644); err != nil {
t.Fatal(err)
}
if tt.truncate {
if err := os.WriteFile(partPath, make([]byte, 0), 0o644); err != nil {
t.Fatal(err)
}
}
}
partFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
t.Fatal(err)
}
defer partFile.Close()
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
if err := DownloadResumable(context.Background(), file, partFile, 1, bitmapPath); err != nil {
t.Fatalf("download failed: %v", err)
}
got := make([]byte, len(data))
if _, err := partFile.ReadAt(got, 0); err != nil {
t.Fatal(err)
}
if !bytesEqual(got, data) {
t.Fatalf("downloaded data mismatch (blocks not reset)")
}
})
}
}
// TestDownloadResumableInvalidBitmap treats a corrupt bitmap as absent.
func TestDownloadResumableInvalidBitmap(t *testing.T) {
data := make([]byte, 1024*1024+7)
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
partPath := filepath.Join(dir, "test.bin.part")
bitmapPath := ResumeStatePath(partPath)
for _, content := range []string{
`{"part_size":1048576,"size":-1,"blocks":[]}`,
`{"part_size":1048576,"size":9223372036854775807,"blocks":[]}`,
`not json`,
} {
os.Remove(partPath)
if err := os.WriteFile(bitmapPath, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
partFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
t.Fatal(err)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
err = DownloadResumable(context.Background(), file, partFile, 1, bitmapPath)
partFile.Close()
if err != nil {
t.Fatalf("download with corrupt bitmap %q failed: %v", content, err)
}
got := make([]byte, len(data))
f, err := os.Open(partPath)
if err != nil {
t.Fatal(err)
}
if _, err := f.ReadAt(got, 0); err != nil {
t.Fatal(err)
}
f.Close()
if !bytesEqual(got, data) {
t.Fatalf("downloaded data mismatch with corrupt bitmap %q", content)
}
}
}
func TestRemoveResumeState(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "x.bitmap")
if err := os.WriteFile(path, []byte("{}"), 0o644); err != nil {
t.Fatal(err)
}
if err := RemoveResumeState(path); err != nil {
t.Fatalf("RemoveResumeState: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("bitmap still exists: %v", err)
}
// Removing again must be a no-op.
if err := RemoveResumeState(path); err != nil {
t.Fatalf("RemoveResumeState second call: %v", err)
}
}
func bytesEqual(a, b []byte) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
+1 -19
View File
@@ -7,7 +7,6 @@ import (
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/queue"
"github.com/krau/SaveAny-Bot/pkg/taskevent"
@@ -46,9 +45,6 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan
exe := qtask.Data
taskCtx := qtask.Context()
logger.Infof("Processing task: %s", exe.TaskID())
if err := database.UpdateTaskStatus(taskCtx, exe.TaskID(), database.TaskStatusRunning, ""); err != nil {
logger.Errorf("Failed to mark task %s as running: %v", exe.TaskID(), err)
}
taskevent.Emit(taskCtx, taskevent.Event{TaskID: exe.TaskID(), Phase: taskevent.PhaseStart})
if err := ExecCommandString(taskCtx, execHooks.TaskBeforeStart); err != nil {
logger.Errorf("Failed to execute before start hook for task %s: %v", exe.TaskID(), err)
@@ -74,11 +70,6 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan
}
taskevent.Emit(taskCtx, taskevent.Event{TaskID: exe.TaskID(), Phase: taskevent.PhaseDone, Err: err})
qe.Done(qtask.ID)
// 用独立 ctx 删除: 优雅关停时 run ctx 已被取消, 会留下已完成任务的行,
// 导致重启后重复执行 (重复上传)。
if err := database.DeleteTask(context.Background(), exe.TaskID()); err != nil {
logger.Errorf("Failed to delete persisted task %s: %v", exe.TaskID(), err)
}
<-semaphore
}
}
@@ -101,21 +92,12 @@ func Close() {
}
func AddTask(ctx context.Context, task Executable) error {
if err := persistTask(ctx, task); err != nil {
log.FromContext(ctx).Errorf("Failed to persist task %s: %v", task.TaskID(), err)
}
return initQueue().Add(queue.NewTask(ctx, task.TaskID(), task.Title(), task))
}
func CancelTask(ctx context.Context, id string) error {
err := queueInstance.CancelTask(id)
if err != nil {
return err
}
if err := database.DeleteTask(ctx, id); err != nil {
log.FromContext(ctx).Errorf("Failed to delete persisted task %s: %v", id, err)
}
return nil
return err
}
func GetLength(ctx context.Context) int {
-145
View File
@@ -1,145 +0,0 @@
package core
import (
"context"
"sync"
"time"
"fmt"
"github.com/charmbracelet/log"
"github.com/gotd/td/telegram/downloader"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
)
// TaskCodec serializes and rebuilds a task from its persisted payload.
// Task types without a registered codec are dropped with a warning on
// recovery instead of being silently re-enqueued.
type TaskCodec interface {
Marshal(task Executable) ([]byte, error)
Unmarshal(payload []byte) (Executable, error)
}
var (
taskCodecsMu sync.RWMutex
taskCodecs = make(map[tasktype.TaskType]TaskCodec)
dlerMu sync.RWMutex
dlerProvider func() downloader.Client
)
func RegisterTaskCodec(t tasktype.TaskType, codec TaskCodec) {
taskCodecsMu.Lock()
defer taskCodecsMu.Unlock()
taskCodecs[t] = codec
}
func TaskCodecFor(t tasktype.TaskType) (TaskCodec, bool) {
taskCodecsMu.RLock()
defer taskCodecsMu.RUnlock()
codec, ok := taskCodecs[t]
return codec, ok
}
// SetDownloaderProvider registers the download client factory used to
// rebuild tfile.TGFile values when recovering tasks.
func SetDownloaderProvider(f func() downloader.Client) {
dlerMu.Lock()
defer dlerMu.Unlock()
dlerProvider = f
}
// DownloaderClient returns the registered download client, or nil.
func DownloaderClient() downloader.Client {
dlerMu.RLock()
defer dlerMu.RUnlock()
if dlerProvider == nil {
return nil
}
return dlerProvider()
}
func persistTask(ctx context.Context, task Executable) error {
codec, ok := TaskCodecFor(task.Type())
if !ok {
return nil
}
payload, err := codec.Marshal(task)
if err != nil {
return err
}
return database.UpsertTask(ctx, &database.Task{
ID: task.TaskID(),
Type: string(task.Type()),
Payload: payload,
Status: string(database.TaskStatusQueued),
})
}
// UpdateTaskPayload atomically mutates the persisted payload of a running
// task (e.g. recording per-element upload progress for recovery).
func UpdateTaskPayload(ctx context.Context, id string, mutate func(payload []byte) ([]byte, error)) error {
row, err := database.GetTask(ctx, id)
if err != nil {
return err
}
updated, err := mutate(row.Payload)
if err != nil {
return fmt.Errorf("mutate payload: %w", err)
}
return database.UpdateTaskPayload(ctx, id, updated)
}
// RecoverTasks re-enqueues tasks that were unfinished when the process last
// exited. Must be called after storages are loaded and before Run. Tasks
// that cannot be recovered are marked failed and kept for visibility.
func RecoverTasks(ctx context.Context) {
logger := log.FromContext(ctx)
if err := database.DeleteStaleFailedTasks(ctx, 24*time.Hour); err != nil {
logger.Warnf("Failed to clean stale failed tasks: %v", err)
}
tasks, err := database.GetUnfinishedTasks(ctx)
if err != nil {
logger.Errorf("Failed to load unfinished tasks: %v", err)
return
}
for _, t := range tasks {
codec, ok := TaskCodecFor(tasktype.TaskType(t.Type))
if !ok {
logger.Warnf("Task %s (type %s) cannot be recovered: no codec registered", t.ID, t.Type)
markRecoverFailed(ctx, t, "no codec registered")
continue
}
task, err := codec.Unmarshal(t.Payload)
if err != nil {
logger.Errorf("Task %s cannot be recovered: failed to rebuild: %v", t.ID, err)
markRecoverFailed(ctx, t, err.Error())
continue
}
if initQueue().Contains(task.TaskID()) {
// Already live in the queue (e.g. submitted via API during
// startup); keep the row as-is.
logger.Infof("Task %s already queued, keeping row", t.ID)
continue
}
if err := AddTask(ctx, task); err != nil {
logger.Errorf("Task %s cannot be recovered: failed to re-enqueue: %v", t.ID, err)
markRecoverFailed(ctx, t, err.Error())
continue
}
// Upsert cleared the original creation time; restore it so
// GetUnfinishedTasks ordering stays stable across restarts.
if err := database.RestoreTaskCreatedAt(ctx, t.ID, t.CreatedAt); err != nil {
logger.Warnf("Failed to restore created_at for task %s: %v", t.ID, err)
}
logger.Infof("Recovered task %s (%s)", t.ID, t.Type)
}
}
func markRecoverFailed(ctx context.Context, t database.Task, reason string) {
if err := database.UpdateTaskStatus(ctx, t.ID, database.TaskStatusFailed, reason); err != nil {
log.FromContext(ctx).Errorf("Failed to mark task %s as failed: %v", t.ID, err)
}
}
-162
View File
@@ -1,162 +0,0 @@
package core
import (
"context"
"fmt"
"os"
"path/filepath"
"testing"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
)
const testRecoverType = tasktype.TaskType("test-recover")
type stubTask struct {
id string
}
func (s *stubTask) Type() tasktype.TaskType { return testRecoverType }
func (s *stubTask) Title() string { return s.id }
func (s *stubTask) TaskID() string { return s.id }
func (s *stubTask) Execute(context.Context) error { return nil }
type stubCodec struct{}
func (stubCodec) Marshal(task Executable) ([]byte, error) {
return []byte(task.TaskID()), nil
}
func (stubCodec) Unmarshal(payload []byte) (Executable, error) {
if len(payload) == 0 {
return nil, fmt.Errorf("empty payload")
}
return &stubTask{id: string(payload)}, nil
}
func initRecoveryEnv(t *testing.T) context.Context {
t.Helper()
dir := t.TempDir()
cfgPath := filepath.Join(dir, "config.toml")
content := fmt.Sprintf("[db]\npath = %q\n", filepath.Join(dir, "test.db"))
if err := os.WriteFile(cfgPath, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
if err := config.Init(context.Background(), cfgPath); err != nil {
t.Fatalf("config init: %v", err)
}
database.Init(context.Background())
RegisterTaskCodec(testRecoverType, stubCodec{})
return context.Background()
}
func TestRecoverTasksReenqueuesAndMarksUnknownFailed(t *testing.T) {
ctx := initRecoveryEnv(t)
if err := database.CreateTask(ctx, &database.Task{
ID: "rec-1", Type: string(testRecoverType), Payload: []byte("rec-1"), Status: string(database.TaskStatusQueued),
}); err != nil {
t.Fatal(err)
}
if err := database.CreateTask(ctx, &database.Task{
ID: "rec-2", Type: string(testRecoverType), Payload: []byte("rec-2"), Status: string(database.TaskStatusRunning),
}); err != nil {
t.Fatal(err)
}
if err := database.CreateTask(ctx, &database.Task{
ID: "drop-1", Type: "unregistered", Payload: nil, Status: string(database.TaskStatusQueued),
}); err != nil {
t.Fatal(err)
}
RecoverTasks(ctx)
ids := map[string]bool{}
for _, info := range GetQueuedTasks(ctx) {
ids[info.ID] = true
}
if !ids["rec-1"] || !ids["rec-2"] {
t.Fatalf("recovered task ids = %v, want rec-1 and rec-2", ids)
}
unfinished, err := database.GetUnfinishedTasks(ctx)
if err != nil {
t.Fatal(err)
}
if len(unfinished) != 2 {
t.Fatalf("unfinished rows = %d, want 2", len(unfinished))
}
for _, task := range unfinished {
if task.ID == "drop-1" {
t.Fatalf("unregistered task record was not dropped")
}
if task.Status != string(database.TaskStatusQueued) {
t.Fatalf("recovered task status = %s, want queued", task.Status)
}
}
// The unrecoverable task must be kept and marked failed, not silently deleted.
drop, err := database.GetTask(ctx, "drop-1")
if err != nil {
t.Fatalf("dropped task row missing: %v", err)
}
if drop.Status != string(database.TaskStatusFailed) {
t.Fatalf("dropped task status = %s, want failed", drop.Status)
}
if drop.Error == "" {
t.Fatalf("dropped task has no failure reason")
}
}
func TestRecoverTasksMarksInvalidPayloadFailed(t *testing.T) {
ctx := initRecoveryEnv(t)
if err := database.CreateTask(ctx, &database.Task{
ID: "bad-1", Type: string(testRecoverType), Payload: nil, Status: string(database.TaskStatusQueued),
}); err != nil {
t.Fatal(err)
}
RecoverTasks(ctx)
// bad-1 must not be enqueued; its row is kept as failed.
for _, info := range GetQueuedTasks(ctx) {
if info.ID == "bad-1" {
t.Fatalf("task with invalid payload was enqueued")
}
}
count, err := database.CountUnfinishedTasks(ctx)
if err != nil {
t.Fatal(err)
}
if count != 0 {
t.Fatalf("unfinished rows = %d, want 0", count)
}
bad, err := database.GetTask(ctx, "bad-1")
if err != nil {
t.Fatalf("failed task row missing: %v", err)
}
if bad.Status != string(database.TaskStatusFailed) {
t.Fatalf("bad task status = %s, want failed", bad.Status)
}
}
func TestRecoverTasksSkipsAlreadyQueued(t *testing.T) {
ctx := initRecoveryEnv(t)
// A task submitted during startup is both persisted and in the queue.
task := &stubTask{id: "live-1"}
if err := AddTask(ctx, task); err != nil {
t.Fatal(err)
}
RecoverTasks(ctx)
// The row must survive with its original status.
row, err := database.GetTask(ctx, "live-1")
if err != nil {
t.Fatalf("row missing for queued task: %v", err)
}
if row.Status != string(database.TaskStatusQueued) {
t.Fatalf("row status = %s, want queued", row.Status)
}
}
-176
View File
@@ -1,176 +0,0 @@
package batchtfile
import (
"context"
"encoding/json"
"fmt"
"path/filepath"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/core"
tftask "github.com/krau/SaveAny-Bot/core/tasks/tfile"
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
)
type elementPayload struct {
ID string `json:"id"`
Storage string `json:"storage"`
Path string `json:"path"`
File tfilepkg.FilePayload `json:"file"`
SourceGroupKey string `json:"source_group_key"`
SourceCaption string `json:"source_caption"`
PreserveCaption bool `json:"preserve_caption"`
}
type taskPayload struct {
Kind string `json:"kind"` // "batch"
ID string `json:"id"`
Elements []elementPayload `json:"elements"`
ChatID int64 `json:"chat_id"`
MessageID int `json:"message_id"`
IgnoreErrors bool `json:"ignore_errors"`
Overwrite bool `json:"overwrite"`
// Done lists element IDs whose upload completed; they are skipped on recovery.
Done []string `json:"done"`
}
// tgfilesCodec is the single codec registered for TaskTypeTgfiles: it
// dispatches between single-file and batch tasks by concrete type on marshal
// and by payload shape on unmarshal. Registering one codec per task class
// under the shared TaskTypeTgfiles key would let the last init() win and
// silently disable persistence for the other class.
type tgfilesCodec struct{}
func init() {
core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, tgfilesCodec{})
}
func (tgfilesCodec) Marshal(task core.Executable) ([]byte, error) {
switch t := task.(type) {
case *tftask.Task:
return tftask.TaskCodec.Marshal(t)
case *Task:
return batchCodec{}.Marshal(t)
default:
return nil, fmt.Errorf("unexpected task type %T", task)
}
}
// detectTaskKind returns "batch" or "file" for a persisted tgfiles payload.
// New payloads carry an explicit kind; legacy payloads are detected by shape.
func detectTaskKind(data []byte) (string, error) {
var shape struct {
Kind string `json:"kind"`
Elements []json.RawMessage `json:"elements"`
File json.RawMessage `json:"file"`
}
if err := json.Unmarshal(data, &shape); err != nil {
return "", fmt.Errorf("invalid task payload: %w", err)
}
switch {
case shape.Kind == "batch", shape.Kind == "" && shape.Elements != nil:
return "batch", nil
case shape.Kind == "file", shape.Kind == "" && shape.File != nil:
return "file", nil
default:
return "", fmt.Errorf("unrecognized task payload")
}
}
func (tgfilesCodec) Unmarshal(data []byte) (core.Executable, error) {
kind, err := detectTaskKind(data)
if err != nil {
return nil, err
}
if kind == "batch" {
return batchCodec{}.Unmarshal(data)
}
return tftask.TaskCodec.Unmarshal(data)
}
type batchCodec struct{}
func (batchCodec) Marshal(task core.Executable) ([]byte, error) {
t, ok := task.(*Task)
if !ok {
return nil, fmt.Errorf("unexpected task type %T", task)
}
p := taskPayload{
Kind: "batch",
ID: t.ID,
IgnoreErrors: t.IgnoreErrors,
Done: t.completedElementIDs(),
}
if overwrite, ok := t.ctx.Value(ctxkey.OverwriteExisting).(bool); ok {
p.Overwrite = overwrite
}
for _, elem := range t.elems {
filePayload, ok := tfilepkg.FilePayloadOf(elem.File)
if !ok {
return nil, fmt.Errorf("file %T is not serializable", elem.File)
}
p.Elements = append(p.Elements, elementPayload{
ID: elem.ID,
Storage: elem.Storage.Name(),
Path: elem.Path,
File: filePayload,
SourceGroupKey: elem.sourceGroupKey,
SourceCaption: elem.sourceCaption,
PreserveCaption: elem.preserveCaption,
})
}
if progress, ok := t.Progress.(*Progress); ok {
p.ChatID = progress.ChatID
p.MessageID = progress.MessageID
}
return json.Marshal(p)
}
func (batchCodec) Unmarshal(data []byte) (core.Executable, error) {
var p taskPayload
if err := json.Unmarshal(data, &p); err != nil {
return nil, fmt.Errorf("invalid task payload: %w", err)
}
dler := core.DownloaderClient()
if dler == nil {
return nil, fmt.Errorf("no downloader client available")
}
done := make(map[string]struct{}, len(p.Done))
for _, id := range p.Done {
done[id] = struct{}{}
}
elems := make([]TaskElement, 0, len(p.Elements))
for _, ep := range p.Elements {
if _, ok := done[ep.ID]; ok {
continue // upload already completed; do not re-run
}
stor, err := storage.GetStorageByName(context.Background(), ep.Storage)
if err != nil {
return nil, fmt.Errorf("storage %q: %w", ep.Storage, err)
}
localPath, err := filepath.Abs(filepath.Join(config.C().Temp.BasePath, fmt.Sprintf("%s_%s", ep.ID, ep.File.Name)))
if err != nil {
return nil, fmt.Errorf("failed to build cache path: %w", err)
}
elems = append(elems, TaskElement{
ID: ep.ID,
Storage: stor,
Path: ep.Path,
File: tfilepkg.FileFromPayload(ep.File, dler),
localPath: localPath,
sourceGroupKey: ep.SourceGroupKey,
sourceCaption: ep.SourceCaption,
preserveCaption: ep.PreserveCaption,
})
}
var progress ProgressTracker
if p.ChatID != 0 {
progress = NewProgressTracker(p.MessageID, p.ChatID)
}
task := NewBatchTGFileTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors)
task.overwrite = p.Overwrite
return task, nil
}
-39
View File
@@ -1,39 +0,0 @@
package batchtfile
import (
"testing"
)
func TestDetectTaskKind(t *testing.T) {
tests := []struct {
name string
payload string
want string
wantErr bool
}{
{"batch with kind", `{"kind":"batch","id":"1","elements":[]}`, "batch", false},
{"file with kind", `{"kind":"file","id":"1","file":{}}`, "file", false},
{"legacy batch by shape", `{"id":"1","elements":[]}`, "batch", false},
{"legacy file by shape", `{"id":"1","file":{}}`, "file", false},
{"legacy batch with element", `{"id":"1","elements":[{"id":"e"}]}`, "batch", false},
{"no discriminator", `{"id":"1"}`, "", true},
{"invalid json", `not json`, "", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := detectTaskKind([]byte(tt.payload))
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got kind %q", got)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != tt.want {
t.Fatalf("kind = %q, want %q", got, tt.want)
}
})
}
}
-92
View File
@@ -1,92 +0,0 @@
package batchtfile
import (
"context"
"fmt"
"os"
"time"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/common/tdler"
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
"github.com/krau/SaveAny-Bot/common/utils/ioutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/pkg/taskevent"
)
// downloadToCache fetches elem.File into the element cache path. It resumes
// from a partial .part download tracked by a resume bitmap, and reuses a
// complete cache file (e.g. when the previous run was interrupted during
// upload).
func (t *Task) downloadToCache(ctx context.Context, elem *TaskElement) error {
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name()))
if elem.File.Size() > 0 {
if stat, err := os.Stat(elem.localPath); err == nil && stat.Size() == elem.File.Size() {
logger.Info("Cache file already complete, skipping download")
return nil
}
}
onProgress := t.downloadCallback(ctx, elem)
if elem.File.Size() <= 0 {
// Unknown size (e.g. photos) cannot be resumed; use the plain downloader.
localFile, err := fsutil.CreateFile(elem.localPath)
if err != nil {
return fmt.Errorf("failed to create local file: %w", err)
}
defer localFile.Close()
wrAt := ioutil.NewProgressWriterAt(localFile, onProgress)
if _, err := tdler.NewDownloader(elem.File).Parallel(ctx, wrAt); err != nil {
return err
}
return nil
}
partPath := elem.localPath + ".part"
// 不截断已存在的 .part: 位图标记的已完成块依赖既有字节。
localFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
return fmt.Errorf("failed to create local file: %w", err)
}
wrAt := ioutil.NewProgressWriterAt(localFile, onProgress)
err = tdler.DownloadResumable(
ctx, elem.File, wrAt,
dlutil.BestThreads(elem.File.Size(), config.C().Threads),
tdler.ResumeStatePath(partPath),
)
closeErr := localFile.Close()
if err != nil {
return err
}
if closeErr != nil {
return fmt.Errorf("failed to close cache file: %w", closeErr)
}
stat, err := os.Stat(partPath)
if err != nil {
return fmt.Errorf("failed to stat downloaded file: %w", err)
}
if stat.Size() != elem.File.Size() {
return fmt.Errorf("downloaded size %d does not match expected %d", stat.Size(), elem.File.Size())
}
if err := os.Rename(partPath, elem.localPath); err != nil {
return fmt.Errorf("failed to finalize download: %w", err)
}
// 清理位图是尽力而为: 下载已完成, 清理失败不应使任务失败。
if err := tdler.RemoveResumeState(tdler.ResumeStatePath(partPath)); err != nil {
logger.Warnf("Failed to remove resume state: %v", err)
}
return nil
}
func (t *Task) downloadCallback(ctx context.Context, elem *TaskElement) func(int) {
return func(n int) {
t.recordItemDownload(elem.ID, int64(n), time.Now())
downloaded := t.downloaded.Add(int64(n))
t.notifyProgress(ctx)
taskevent.Emit(ctx, taskevent.Event{
TaskID: t.ID,
Phase: taskevent.PhaseProgress,
TotalBytes: t.totalSize,
DownloadedBytes: downloaded,
})
}
}
+51 -29
View File
@@ -35,9 +35,6 @@ 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")
if t.overwrite {
ctx = storage.WithOverwrite(ctx)
}
if t.Progress != nil {
t.Progress.OnStart(ctx, t)
}
@@ -137,12 +134,7 @@ func (t *Task) processElements(ctx context.Context, elems []*TaskElement) error
}
func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
// Cache files are kept on failure so a later restart can resume upload.
uploaded := false
defer func() {
if !uploaded {
return
}
for _, elem := range group.elems {
if err := os.Remove(elem.localPath); err != nil && !os.IsNotExist(err) {
log.FromContext(ctx).Warnf("Failed to cleanup batch cache file %s: %v", elem.localPath, err)
@@ -227,14 +219,7 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
for index, item := range items {
t.recordDownloadComplete(successElems[index].ID, item.Size)
}
err := t.saveBatchItems(ctx, successElems, items)
if err == nil {
uploaded = true
for _, elem := range successElems {
t.persistElementDone(ctx, elem.ID)
}
}
return err
return t.saveBatchItems(ctx, successElems, items)
}
func (t *Task) saveBatchItems(ctx context.Context, successElems []*TaskElement, items []storagetypes.BatchItem) error {
@@ -304,10 +289,34 @@ func (t *Task) unmarkProcessing(id string) {
func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name()))
logger.Info("Starting file download")
if err := t.downloadToCache(ctx, elem); err != nil {
t.markItemFailed(elem.ID, FailureStageDownload, err)
localFile, err := fsutil.CreateFile(elem.localPath)
if err != nil {
t.markItemFailed(elem.ID, FailureStageCache, err)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to download file: %w", err)
return fmt.Errorf("failed to create local file: %w", err)
}
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
t.recordItemDownload(elem.ID, int64(n), time.Now())
downloaded := t.downloaded.Add(int64(n))
t.notifyProgress(ctx)
taskevent.Emit(ctx, taskevent.Event{
TaskID: t.ID,
Phase: taskevent.PhaseProgress,
TotalBytes: t.totalSize,
DownloadedBytes: downloaded,
})
})
_, downloadErr := tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
closeErr := localFile.Close()
if downloadErr != nil {
t.markItemFailed(elem.ID, FailureStageDownload, downloadErr)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to download file: %w", downloadErr)
}
if closeErr != nil {
t.markItemFailed(elem.ID, FailureStageCache, closeErr)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to close cache file: %w", closeErr)
}
logger.Info("File downloaded successfully")
if path.Ext(elem.FileName()) == "" {
@@ -372,16 +381,30 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
return nil
}
logger.Info("Starting file download")
// 不预创建缓存文件: 预创建会截断上次运行保留的完整缓存, 使复用失效。
success := false
localFile, err := fsutil.CreateFile(elem.localPath)
if err != nil {
t.markItemFailed(elem.ID, FailureStageCache, err)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to create local file: %w", err)
}
defer func() {
if success {
if err := os.Remove(elem.localPath); err != nil {
logger.Errorf("Failed to remove cache file: %v", err)
}
if err := localFile.CloseAndRemove(); err != nil {
logger.Errorf("Failed to close local file: %v", err)
}
}()
if err := t.downloadToCache(ctx, &elem); err != nil {
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
t.recordItemDownload(elem.ID, int64(n), time.Now())
downloaded := t.downloaded.Add(int64(n))
t.notifyProgress(ctx)
taskevent.Emit(ctx, taskevent.Event{
TaskID: t.ID,
Phase: taskevent.PhaseProgress,
TotalBytes: t.totalSize,
DownloadedBytes: downloaded,
})
})
_, err = tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
if err != nil {
t.markItemFailed(elem.ID, FailureStageDownload, err)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to download file: %w", err)
@@ -393,7 +416,8 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
elem.Path = elem.Path + ext
}
}
fileStat, err := os.Stat(elem.localPath)
var fileStat os.FileInfo
fileStat, err = os.Stat(elem.localPath)
if err != nil {
t.markItemFailed(elem.ID, FailureStageCache, err)
t.notifyStateChange(ctx)
@@ -436,8 +460,6 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
onProgress(fileStat.Size(), fileStat.Size())
t.markItemCompleted(elem.ID)
t.notifyStateChange(vctx)
t.persistElementDone(ctx, elem.ID)
success = true
} else {
t.markItemFailed(elem.ID, lastFailureStage, err)
t.notifyStateChange(vctx)
-1
View File
@@ -152,7 +152,6 @@ func buildBatchProgressMessage(info TaskInfo, skipped []string, activeLimit int)
uploadSpeedText := formatSpeed(uploadSpeed)
header := localizedProgressMarkup(i18nk.BotMsgProgressBatchStatusHeader, map[string]any{
"Total": total,
"TotalSize": dlutil.FormatSize(info.ActualTotalSize()),
"Completed": completed,
"Downloaded": downloaded,
"Waiting": waiting,
@@ -95,30 +95,6 @@ func TestBatchProgressShowsTransferSpeedAndSize(t *testing.T) {
}
}
func TestBatchProgressHeaderShowsTotalSize(t *testing.T) {
useProgressRegressionLocale(t)
task := newProgressRegressionTask(nil,
progressRegressionFile{"first", 1024},
progressRegressionFile{"second", 1024},
)
message := buildBatchProgressMessage(task, nil, 2)
if message.Err != nil {
t.Fatalf("buildBatchProgressMessage() failed: %v", message.Err)
}
assertProgressRegressionContains(t, message.Text,
"文件:2 总大小:2.00 KB",
)
i18n.Init("en")
english := buildBatchProgressMessage(task, nil, 2)
if english.Err != nil {
t.Fatalf("English batch template failed: %v", english.Err)
}
assertProgressRegressionContains(t, english.Text,
"Files: 2 | Total size: 2.00 KB",
)
}
func TestBatchProgressLimitsRowsWithoutHidingActiveUpload(t *testing.T) {
useProgressRegressionLocale(t)
task := newProgressRegressionTask(nil,
-38
View File
@@ -2,13 +2,11 @@ package batchtfile
import (
"context"
"encoding/json"
"fmt"
"path/filepath"
"sync"
"sync/atomic"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/core"
@@ -49,7 +47,6 @@ type Task struct {
uploadOnce sync.Once
uploadMu sync.Mutex
uploaded map[string]int64
overwrite bool // recovered: overwrite storage targets instead of uniquifying
}
// Title implements core.Exectable.
@@ -61,41 +58,6 @@ func (t *Task) Type() tasktype.TaskType {
return tasktype.TaskTypeTgfiles
}
// completedElementIDs returns the element IDs whose upload finished, for
// persisting upload progress so recovery can skip them.
func (t *Task) completedElementIDs() []string {
t.itemMu.RLock()
defer t.itemMu.RUnlock()
var ids []string
for _, item := range t.itemStates {
if item.phase == ItemPhaseCompleted {
ids = append(ids, item.id)
}
}
return ids
}
// persistElementDone records an element's completed upload in the persisted
// payload so a restart does not re-upload it.
func (t *Task) persistElementDone(ctx context.Context, elemID string) {
err := core.UpdateTaskPayload(ctx, t.ID, func(payload []byte) ([]byte, error) {
var p taskPayload
if err := json.Unmarshal(payload, &p); err != nil {
return nil, err
}
for _, id := range p.Done {
if id == elemID {
return payload, nil
}
}
p.Done = append(p.Done, elemID)
return json.Marshal(p)
})
if err != nil {
log.FromContext(ctx).Warnf("Failed to persist element completion %s: %v", elemID, err)
}
}
func NewTaskElement(
stor storage.Storage,
path string,
-97
View File
@@ -1,97 +0,0 @@
package tfile
import (
"context"
"encoding/json"
"fmt"
"path/filepath"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
)
type taskPayload struct {
Kind string `json:"kind"` // "file"
ID string `json:"id"`
Storage string `json:"storage"`
Path string `json:"path"`
File tfilepkg.FilePayload `json:"file"`
ChatID int64 `json:"chat_id"`
MessageID int `json:"message_id"`
Overwrite bool `json:"overwrite"`
Caption string `json:"caption"`
}
type taskCodec struct{}
// TaskCodec serializes single-file tasks. It is registered together with the
// batch codec under TaskTypeTgfiles (see core/tasks/batchtfile/codec.go).
var TaskCodec core.TaskCodec = taskCodec{}
func (taskCodec) Marshal(task core.Executable) ([]byte, error) {
t, ok := task.(*Task)
if !ok {
return nil, fmt.Errorf("unexpected task type %T", task)
}
filePayload, ok := tfilepkg.FilePayloadOf(t.File)
if !ok {
return nil, fmt.Errorf("file %T is not serializable", t.File)
}
p := taskPayload{
Kind: "file",
ID: t.ID,
Storage: t.Storage.Name(),
Path: t.Path,
File: filePayload,
}
if overwrite, ok := t.Ctx.Value(ctxkey.OverwriteExisting).(bool); ok {
p.Overwrite = overwrite
}
if caption, ok := sourceCaption(t.File); ok {
p.Caption = caption
}
if progress, ok := t.Progress.(*Progress); ok {
p.ChatID = progress.ChatID
p.MessageID = progress.MessageID
}
return json.Marshal(p)
}
func (taskCodec) Unmarshal(data []byte) (core.Executable, error) {
var p taskPayload
if err := json.Unmarshal(data, &p); err != nil {
return nil, fmt.Errorf("invalid task payload: %w", err)
}
dler := core.DownloaderClient()
if dler == nil {
return nil, fmt.Errorf("no downloader client available")
}
file := tfilepkg.FileFromPayload(p.File, dler)
stor, err := storage.GetStorageByName(context.Background(), p.Storage)
if err != nil {
return nil, fmt.Errorf("storage %q: %w", p.Storage, err)
}
var progress ProgressTracker
if p.ChatID != 0 {
progress = NewProgressTrack(p.MessageID, p.ChatID)
}
localPath, err := filepath.Abs(filepath.Join(config.C().Temp.BasePath, fmt.Sprintf("%s_%s", p.ID, file.Name())))
if err != nil {
return nil, fmt.Errorf("failed to build cache path: %w", err)
}
return &Task{
ID: p.ID,
Ctx: context.Background(),
File: file,
Storage: stor,
Path: p.Path,
Progress: progress,
stream: false, // recovered tasks always download to cache first
localPath: localPath,
overwrite: p.Overwrite,
caption: p.Caption,
}, nil
}
-75
View File
@@ -1,75 +0,0 @@
package tfile
import (
"context"
"fmt"
"os"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/common/tdler"
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
"github.com/krau/SaveAny-Bot/config"
)
// download fetches the file into the cache path. It resumes from a partial
// .part download tracked by a resume bitmap, and reuses a complete cache
// file (e.g. when the previous run was interrupted during upload).
func (t *Task) download(ctx context.Context) error {
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", t.File.Name()))
if t.File.Size() > 0 {
if stat, err := os.Stat(t.localPath); err == nil && stat.Size() == t.File.Size() {
logger.Info("Cache file already complete, skipping download")
return nil
}
}
if t.File.Size() <= 0 {
// Unknown size (e.g. photos) cannot be resumed; use the plain downloader.
localFile, err := fsutil.CreateFile(t.localPath)
if err != nil {
return fmt.Errorf("failed to create local file: %w", err)
}
defer localFile.Close()
wrAt := newWriterAt(ctx, localFile, t.Progress, t)
if _, err := tdler.NewDownloader(t.File).Parallel(ctx, wrAt); err != nil {
return err
}
logger.Info("File downloaded successfully")
return nil
}
partPath := t.localPath + ".part"
// 不截断已存在的 .part: 位图标记的已完成块依赖既有字节。
localFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
return fmt.Errorf("failed to create local file: %w", err)
}
wrAt := newWriterAt(ctx, localFile, t.Progress, t)
err = tdler.DownloadResumable(
ctx, t.File, wrAt,
dlutil.BestThreads(t.File.Size(), config.C().Threads),
tdler.ResumeStatePath(partPath),
)
closeErr := localFile.Close()
if err != nil {
return err
}
if closeErr != nil {
return fmt.Errorf("failed to close cache file: %w", closeErr)
}
stat, err := os.Stat(partPath)
if err != nil {
return fmt.Errorf("failed to stat downloaded file: %w", err)
}
if stat.Size() != t.File.Size() {
return fmt.Errorf("downloaded size %d does not match expected %d", stat.Size(), t.File.Size())
}
if err := os.Rename(partPath, t.localPath); err != nil {
return fmt.Errorf("failed to finalize download: %w", err)
}
// 清理位图是尽力而为: 下载已完成, 清理失败不应使任务失败。
if err := tdler.RemoveResumeState(tdler.ResumeStatePath(partPath)); err != nil {
logger.Warnf("Failed to remove resume state: %v", err)
}
logger.Info("File downloaded successfully")
return nil
}
+24 -18
View File
@@ -9,6 +9,7 @@ import (
"github.com/charmbracelet/log"
"github.com/duke-git/lancet/v2/retry"
"github.com/krau/SaveAny-Bot/common/tdler"
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
"github.com/krau/SaveAny-Bot/common/utils/ioutil"
"github.com/krau/SaveAny-Bot/config"
@@ -18,16 +19,8 @@ import (
"github.com/krau/SaveAny-Bot/storage"
)
func (t *Task) Execute(ctx context.Context) (err error) {
func (t *Task) Execute(ctx context.Context) error {
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", t.File.Name()))
defer func() {
if t.Progress != nil {
t.Progress.OnDone(ctx, t, err)
}
}()
if t.overwrite {
ctx = storage.WithOverwrite(ctx)
}
if t.Progress != nil {
t.Progress.OnStart(ctx, t)
}
@@ -36,23 +29,40 @@ func (t *Task) Execute(ctx context.Context) (err error) {
}
logger.Info("Starting file download")
if err := t.download(ctx); err != nil {
localFile, err := fsutil.CreateFile(t.localPath)
if err != nil {
return fmt.Errorf("failed to create local file: %w", err)
}
defer func() {
if err := localFile.CloseAndRemove(); err != nil {
logger.Errorf("Failed to close local file: %v", err)
}
}()
wrAt := newWriterAt(ctx, localFile, t.Progress, t)
defer func() {
if t.Progress != nil {
t.Progress.OnDone(ctx, t, err)
}
}()
_, err = tdler.NewDownloader(t.File).Parallel(ctx, wrAt)
if err != nil {
return fmt.Errorf("failed to download file: %w", err)
}
logger.Infof("File downloaded successfully")
if path.Ext(t.File.Name()) == "" {
ext := fsutil.DetectFileExt(t.localPath)
if ext != "" {
t.Path = t.Path + ext
}
}
fileStat, err := os.Stat(t.localPath)
var fileStat os.FileInfo
fileStat, err = os.Stat(t.localPath)
if err != nil {
return fmt.Errorf("failed to get file stat: %w", err)
}
vctx := context.WithValue(ctx, ctxkey.ContentLength, fileStat.Size())
if t.caption != "" {
vctx = storagetypes.WithSourceCaption(vctx, t.caption)
} else if caption, ok := sourceCaption(t.File); ok {
if caption, ok := sourceCaption(t.File); ok {
vctx = storagetypes.WithSourceCaption(vctx, caption)
}
err = retry.Retry(func() error {
@@ -87,10 +97,6 @@ func (t *Task) Execute(ctx context.Context) (err error) {
if err != nil {
return fmt.Errorf("failed to save file after retries: %w", err)
}
// Cache file is kept on failure so a later restart can resume upload.
if err := os.Remove(t.localPath); err != nil {
logger.Errorf("Failed to remove cache file: %v", err)
}
return nil
}
-2
View File
@@ -23,8 +23,6 @@ type Task struct {
Progress ProgressTracker
stream bool // true if the file should be downloaded in stream mode
localPath string
overwrite bool // recovered: overwrite the storage target instead of uniquifying
caption string // recovered: source caption for the telegram backend
}
// Title implements core.Exectable.
-102
View File
@@ -1,102 +0,0 @@
package transfer
import (
"context"
"encoding/json"
"fmt"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
"github.com/krau/SaveAny-Bot/storage"
)
func ctxOverwrite(ctx context.Context) bool {
overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool)
return overwrite
}
type elementPayload struct {
ID string `json:"id"`
SourceStorage string `json:"source_storage"`
SourcePath string `json:"source_path"`
FileInfo storagetypes.FileInfo `json:"file_info"`
TargetStorage string `json:"target_storage"`
TargetPath string `json:"target_path"`
}
type taskPayload struct {
ID string `json:"id"`
Elements []elementPayload `json:"elements"`
ChatID int64 `json:"chat_id"`
MessageID int `json:"message_id"`
IgnoreErrors bool `json:"ignore_errors"`
Overwrite bool `json:"overwrite"`
}
type taskCodec struct{}
func init() {
core.RegisterTaskCodec(tasktype.TaskTypeTransfer, taskCodec{})
}
func (taskCodec) Marshal(task core.Executable) ([]byte, error) {
t, ok := task.(*Task)
if !ok {
return nil, fmt.Errorf("unexpected task type %T", task)
}
p := taskPayload{
ID: t.ID,
IgnoreErrors: t.IgnoreErrors,
Overwrite: ctxOverwrite(t.ctx),
}
for _, elem := range t.elems {
p.Elements = append(p.Elements, elementPayload{
ID: elem.ID,
SourceStorage: elem.SourceStorage.Name(),
SourcePath: elem.SourcePath,
FileInfo: elem.FileInfo,
TargetStorage: elem.TargetStorage.Name(),
TargetPath: elem.TargetPath,
})
}
if progress, ok := t.Progress.(*Progress); ok {
p.ChatID = progress.ChatID
p.MessageID = progress.MessageID
}
return json.Marshal(p)
}
func (taskCodec) Unmarshal(data []byte) (core.Executable, error) {
var p taskPayload
if err := json.Unmarshal(data, &p); err != nil {
return nil, fmt.Errorf("invalid task payload: %w", err)
}
elems := make([]TaskElement, 0, len(p.Elements))
for _, ep := range p.Elements {
source, err := storage.GetStorageByName(context.Background(), ep.SourceStorage)
if err != nil {
return nil, fmt.Errorf("source storage %q: %w", ep.SourceStorage, err)
}
target, err := storage.GetStorageByName(context.Background(), ep.TargetStorage)
if err != nil {
return nil, fmt.Errorf("target storage %q: %w", ep.TargetStorage, err)
}
elems = append(elems, TaskElement{
ID: ep.ID,
SourceStorage: source,
SourcePath: ep.SourcePath,
FileInfo: ep.FileInfo,
TargetStorage: target,
TargetPath: ep.TargetPath,
})
}
var progress ProgressTracker
if p.ChatID != 0 {
progress = NewProgressTracker(p.MessageID, p.ChatID)
}
task := NewTransferTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors)
task.overwrite = p.Overwrite
return task, nil
}
-3
View File
@@ -21,9 +21,6 @@ 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")
if t.overwrite {
ctx = storage.WithOverwrite(ctx)
}
if t.Progress != nil {
t.Progress.OnStart(ctx, t)
}
-1
View File
@@ -35,7 +35,6 @@ type Task struct {
processing map[string]TaskElementInfo
processingMu sync.RWMutex
failed map[string]error
overwrite bool // recovered: overwrite storage targets instead of uniquifying
}
// Title implements core.Executable.
+1 -1
View File
@@ -35,7 +35,7 @@ func Init(ctx context.Context) {
logger.Fatal("Failed to open database: ", err)
}
logger.Debug("Database connected")
if err := db.AutoMigrate(&User{}, &Dir{}, &Rule{}, &WatchChat{}, &Task{}); err != nil {
if err := db.AutoMigrate(&User{}, &Dir{}, &Rule{}, &WatchChat{}); err != nil {
logger.Fatal("Database migration failed; if upgrading from an old version, try deleting the database file and retrying", "error", err)
}
if err := syncUsers(ctx); err != nil {
-137
View File
@@ -1,137 +0,0 @@
package database
import (
"context"
"errors"
"time"
)
var errNotInitialized = errors.New("database not initialized")
type TaskStatus string
const (
TaskStatusQueued TaskStatus = "queued"
TaskStatusRunning TaskStatus = "running"
TaskStatusFailed TaskStatus = "failed"
TaskStatusCancelled TaskStatus = "cancelled"
)
// Task is the persisted record of a queued or running task, used to recover
// unfinished work after a process restart. Completed tasks are deleted on
// finish, so the table only ever holds queued/running rows.
type Task struct {
ID string `gorm:"primaryKey;size:64"`
Type string `gorm:"size:32;index"`
Payload []byte
Status string `gorm:"size:16;index"`
Error string
CreatedAt time.Time
UpdatedAt time.Time
}
func CreateTask(ctx context.Context, task *Task) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Create(task).Error
}
// UpsertTask inserts the task or replaces the existing row with the same ID.
func UpsertTask(ctx context.Context, task *Task) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Save(task).Error
}
func UpdateTaskStatus(ctx context.Context, id string, status TaskStatus, errMsg string) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Model(&Task{}).
Where("id = ?", id).
Updates(map[string]any{
"status": status,
"error": errMsg,
"updated_at": time.Now(),
}).Error
}
func GetTask(ctx context.Context, id string) (*Task, error) {
if db == nil {
return nil, errNotInitialized
}
var task Task
if err := db.WithContext(ctx).First(&task, "id = ?", id).Error; err != nil {
return nil, err
}
return &task, nil
}
// UpdateTaskPayload replaces the payload of an existing task row.
func UpdateTaskPayload(ctx context.Context, id string, payload []byte) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Model(&Task{}).
Where("id = ?", id).
Updates(map[string]any{
"payload": payload,
"updated_at": time.Now(),
}).Error
}
// RestoreTaskCreatedAt restores the original creation time after a
// re-enqueue overwrote it.
func RestoreTaskCreatedAt(ctx context.Context, id string, createdAt time.Time) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Model(&Task{}).
Where("id = ?", id).
Update("created_at", createdAt).Error
}
// DeleteStaleFailedTasks removes failed rows older than the given age.
func DeleteStaleFailedTasks(ctx context.Context, maxAge time.Duration) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).
Where("status = ? AND updated_at < ?", string(TaskStatusFailed), time.Now().Add(-maxAge)).
Delete(&Task{}).Error
}
func DeleteTask(ctx context.Context, id string) error {
if db == nil {
return errNotInitialized
}
return db.WithContext(ctx).Delete(&Task{}, "id = ?", id).Error
}
// GetUnfinishedTasks returns all tasks that were not finished when the
// process stopped, i.e. tasks that must be re-enqueued on startup.
func GetUnfinishedTasks(ctx context.Context) ([]Task, error) {
if db == nil {
return nil, errNotInitialized
}
var tasks []Task
err := db.WithContext(ctx).
Where("status IN ?", []string{string(TaskStatusQueued), string(TaskStatusRunning)}).
Order("created_at").
Find(&tasks).Error
return tasks, err
}
func CountUnfinishedTasks(ctx context.Context) (int64, error) {
if db == nil {
return 0, errNotInitialized
}
var count int64
err := db.WithContext(ctx).
Model(&Task{}).
Where("status IN ?", []string{string(TaskStatusQueued), string(TaskStatusRunning)}).
Count(&count).Error
return count, err
}
-110
View File
@@ -1,110 +0,0 @@
package database
import (
"context"
"path/filepath"
"testing"
"github.com/ncruces/go-sqlite3/gormlite"
"gorm.io/gorm"
)
func newTestDB(t *testing.T) {
t.Helper()
d, err := gorm.Open(gormlite.Open(filepath.Join(t.TempDir(), "test.db")), &gorm.Config{})
if err != nil {
t.Fatalf("open test db: %v", err)
}
if err := d.AutoMigrate(&Task{}); err != nil {
t.Fatalf("migrate: %v", err)
}
old := db
db = d
t.Cleanup(func() { db = old })
}
func TestTaskCRUD(t *testing.T) {
newTestDB(t)
ctx := context.Background()
task := &Task{
ID: "task-1",
Type: "tfile",
Payload: []byte(`{"file":"x"}`),
Status: string(TaskStatusQueued),
}
if err := CreateTask(ctx, task); err != nil {
t.Fatalf("create: %v", err)
}
unfinished, err := GetUnfinishedTasks(ctx)
if err != nil {
t.Fatalf("get unfinished: %v", err)
}
if len(unfinished) != 1 || unfinished[0].ID != "task-1" {
t.Fatalf("got %+v, want 1 task task-1", unfinished)
}
if err := UpdateTaskStatus(ctx, "task-1", TaskStatusRunning, ""); err != nil {
t.Fatalf("update: %v", err)
}
unfinished, err = GetUnfinishedTasks(ctx)
if err != nil {
t.Fatalf("get unfinished after update: %v", err)
}
if len(unfinished) != 1 || unfinished[0].Status != string(TaskStatusRunning) {
t.Fatalf("running status not persisted: %+v", unfinished)
}
if err := DeleteTask(ctx, "task-1"); err != nil {
t.Fatalf("delete: %v", err)
}
count, err := CountUnfinishedTasks(ctx)
if err != nil {
t.Fatalf("count: %v", err)
}
if count != 0 {
t.Fatalf("count = %d, want 0", count)
}
}
func TestTaskUpsert(t *testing.T) {
newTestDB(t)
ctx := context.Background()
task := &Task{ID: "task-2", Type: "tfile", Status: string(TaskStatusQueued)}
if err := UpsertTask(ctx, task); err != nil {
t.Fatalf("upsert create: %v", err)
}
task.Status = string(TaskStatusRunning)
task.Payload = []byte("new")
if err := UpsertTask(ctx, task); err != nil {
t.Fatalf("upsert update: %v", err)
}
unfinished, err := GetUnfinishedTasks(ctx)
if err != nil {
t.Fatalf("get unfinished: %v", err)
}
if len(unfinished) != 1 || unfinished[0].Status != string(TaskStatusRunning) || string(unfinished[0].Payload) != "new" {
t.Fatalf("upsert did not replace: %+v", unfinished)
}
}
func TestGetUnfinishedTasksExcludesFinished(t *testing.T) {
newTestDB(t)
ctx := context.Background()
if err := CreateTask(ctx, &Task{ID: "done", Type: "tfile", Status: string(TaskStatusFailed)}); err != nil {
t.Fatal(err)
}
if err := CreateTask(ctx, &Task{ID: "pending", Type: "tfile", Status: string(TaskStatusQueued)}); err != nil {
t.Fatal(err)
}
unfinished, err := GetUnfinishedTasks(ctx)
if err != nil {
t.Fatal(err)
}
if len(unfinished) != 1 || unfinished[0].ID != "pending" {
t.Fatalf("got %+v, want only pending", unfinished)
}
}
-16
View File
@@ -116,22 +116,6 @@ func (tq *TaskQueue[T]) ActiveLength() int {
return count
}
// Contains reports whether a task with the given ID is queued or running.
func (tq *TaskQueue[T]) Contains(taskID string) bool {
tq.mu.RLock()
defer tq.mu.RUnlock()
if _, ok := tq.runningTaskMap[taskID]; ok {
return true
}
for element := tq.tasks.Front(); element != nil; element = element.Next() {
task := element.Value.(*Task[T])
if task.ID == taskID && !task.Cancelled() {
return true
}
}
return false
}
// RunningTasks returns the currently running tasks' info.
func (tq *TaskQueue[T]) RunningTasks() []TaskInfo {
tq.mu.RLock()
-86
View File
@@ -1,86 +0,0 @@
package tfile
import (
"github.com/gotd/td/telegram/downloader"
"github.com/gotd/td/tg"
)
// Payloadable is implemented by TGFile implementations that can serialize
// themselves for task recovery.
type Payloadable interface {
Payload() FilePayload
}
// FilePayloadOf returns the serializable form of f.
func FilePayloadOf(f TGFile) (FilePayload, bool) {
p, ok := f.(Payloadable)
if !ok {
return FilePayload{}, false
}
return p.Payload(), true
}
// FilePayload is the minimal serializable representation of a TGFile,
// used to rebuild tasks after a process restart.
type FilePayload struct {
Kind string `json:"kind"` // "document" | "photo"
ID int64 `json:"id"`
AccessHash int64 `json:"access_hash"`
FileReference []byte `json:"file_reference"`
ThumbSize string `json:"thumb_size"`
Size int64 `json:"size"`
Name string `json:"name"`
}
// Payload returns the serializable representation of the file.
func (f *tgFile) Payload() FilePayload {
p := FilePayload{
Size: f.size,
Name: f.name,
}
switch loc := f.location.(type) {
case *tg.InputDocumentFileLocation:
p.Kind = "document"
p.ID = loc.ID
p.AccessHash = loc.AccessHash
p.FileReference = loc.FileReference
p.ThumbSize = loc.ThumbSize
case *tg.InputPhotoFileLocation:
p.Kind = "photo"
p.ID = loc.ID
p.AccessHash = loc.AccessHash
p.FileReference = loc.FileReference
p.ThumbSize = loc.ThumbSize
}
return p
}
// Location rebuilds the Telegram file location from the payload.
func (p FilePayload) Location() tg.InputFileLocationClass {
switch p.Kind {
case "photo":
return &tg.InputPhotoFileLocation{
ID: p.ID,
AccessHash: p.AccessHash,
FileReference: p.FileReference,
ThumbSize: p.ThumbSize,
}
default:
return &tg.InputDocumentFileLocation{
ID: p.ID,
AccessHash: p.AccessHash,
FileReference: p.FileReference,
ThumbSize: p.ThumbSize,
}
}
}
// FileFromPayload rebuilds a TGFile from its serialized payload.
func FileFromPayload(p FilePayload, dler downloader.Client) TGFile {
return &tgFile{
location: p.Location(),
dler: dler,
size: p.Size,
name: p.Name,
}
}
-61
View File
@@ -1,61 +0,0 @@
package tfile
import (
"reflect"
"testing"
"github.com/gotd/td/tg"
)
func TestFilePayloadDocumentRoundTrip(t *testing.T) {
file := NewTGFile(
&tg.InputDocumentFileLocation{
ID: 6287403840090150101,
AccessHash: -8452541528324991878,
FileReference: []byte{0x02, 0x0e, 0x80, 0xd6},
ThumbSize: "",
},
nil,
4194304000,
"常轨脱离Creative凸.7z.001",
)
p, ok := FilePayloadOf(file)
if !ok {
t.Fatalf("FilePayloadOf failed")
}
rebuilt := FileFromPayload(p, nil)
if !reflect.DeepEqual(rebuilt.Location(), file.Location()) {
t.Fatalf("location mismatch:\n got %#v\nwant %#v", rebuilt.Location(), file.Location())
}
if rebuilt.Size() != file.Size() || rebuilt.Name() != file.Name() {
t.Fatalf("size/name mismatch: got %d %q, want %d %q", rebuilt.Size(), rebuilt.Name(), file.Size(), file.Name())
}
if p.Kind != "document" {
t.Fatalf("kind = %q, want document", p.Kind)
}
}
func TestFilePayloadPhotoRoundTrip(t *testing.T) {
file := NewTGFile(
&tg.InputPhotoFileLocation{
ID: 123,
AccessHash: 456,
FileReference: []byte{0xaa, 0xbb},
ThumbSize: "y",
},
nil,
0,
"photo_123.png",
)
p, ok := FilePayloadOf(file)
if !ok {
t.Fatalf("FilePayloadOf failed")
}
if p.Kind != "photo" {
t.Fatalf("kind = %q, want photo", p.Kind)
}
rebuilt := FileFromPayload(p, nil)
if !reflect.DeepEqual(rebuilt.Location(), file.Location()) {
t.Fatalf("location mismatch:\n got %#v\nwant %#v", rebuilt.Location(), file.Location())
}
}