GithubBackup/SaveAny-Bot
1
0
Fork 0
mirror of https://github.com/krau/SaveAny-Bot.git synced 2026-09-02 14:16:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): persist tasks and recover them after restart

Browse Source
This commit is contained in:
krau
2026-08-25 15:05:55 +08:00
parent 54dc4caafe
commit 991f454096
14 changed files with 938 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/tasks/batchtfile/codec.go
+106
View File
@@ -0,0 +1,106 @@
package batchtfile
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/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 {
ID string `json:"id"`
Elements []elementPayload `json:"elements"`
ChatID int64 `json:"chat_id"`
MessageID int `json:"message_id"`
IgnoreErrors bool `json:"ignore_errors"`
}
type taskCodec struct{}
func init() {
core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, 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,
}
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 (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")
}
elems := make([]TaskElement, 0, len(p.Elements))
for _, ep := range p.Elements {
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)
}
return NewBatchTGFileTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors), nil
}
core/tasks/batchtfile/execute.go
+20 -43
View File
@@ -134,7 +134,12 @@ 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)
@@ -219,7 +224,11 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
for index, item := range items {
t.recordDownloadComplete(successElems[index].ID, item.Size)
}
return t.saveBatchItems(ctx, successElems, items)
err := t.saveBatchItems(ctx, successElems, items)
if err == nil {
uploaded = true
}
return err
}
func (t *Task) saveBatchItems(ctx context.Context, successElems []*TaskElement, items []storagetypes.BatchItem) error {
@@ -289,34 +298,10 @@ 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")
localFile, err := fsutil.CreateFile(elem.localPath)
if err != nil {
t.markItemFailed(elem.ID, FailureStageCache, err)
if err := t.downloadToCache(ctx, elem); err != nil {
t.markItemFailed(elem.ID, FailureStageDownload, err)
t.notifyStateChange(ctx)
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)
return fmt.Errorf("failed to download file: %w", err)
}
logger.Info("File downloaded successfully")
if path.Ext(elem.FileName()) == "" {
@@ -387,24 +372,15 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
t.notifyStateChange(ctx)
return fmt.Errorf("failed to create local file: %w", err)
}
success := false
defer func() {
if err := localFile.CloseAndRemove(); err != nil {
logger.Errorf("Failed to close local file: %v", err)
if success {
if err := localFile.CloseAndRemove(); err != nil {
logger.Errorf("Failed to close local file: %v", 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,
})
})
_, err = tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
if err != nil {
if err := t.downloadToCache(ctx, &elem); err != nil {
t.markItemFailed(elem.ID, FailureStageDownload, err)
t.notifyStateChange(ctx)
return fmt.Errorf("failed to download file: %w", err)
@@ -460,6 +436,7 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
onProgress(fileStat.Size(), fileStat.Size())
t.markItemCompleted(elem.ID)
t.notifyStateChange(vctx)
success = true
} else {
t.markItemFailed(elem.ID, lastFailureStage, err)
t.notifyStateChange(vctx)
core/tasks/tfile/codec.go
+85
View File
@@ -0,0 +1,85 @@
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/tasktype"
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
)
type taskPayload struct {
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"`
}
type taskCodec struct{}
func init() {
core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, 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{
ID: t.ID,
Storage: t.Storage.Name(),
Path: t.Path,
File: filePayload,
}
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,
}, nil
}
core/tasks/tfile/execute.go
+12 -23
View File
@@ -9,7 +9,6 @@ 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"
@@ -19,8 +18,13 @@ import (
"github.com/krau/SaveAny-Bot/storage"
)
func (t *Task) Execute(ctx context.Context) error {
func (t *Task) Execute(ctx context.Context) (err 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.Progress != nil {
t.Progress.OnStart(ctx, t)
}
@@ -29,35 +33,16 @@ func (t *Task) Execute(ctx context.Context) error {
}
logger.Info("Starting file download")
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 {
if err := t.download(ctx); 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
}
}
var fileStat os.FileInfo
fileStat, err = os.Stat(t.localPath)
fileStat, err := os.Stat(t.localPath)
if err != nil {
return fmt.Errorf("failed to get file stat: %w", err)
}
@@ -97,6 +82,10 @@ func (t *Task) Execute(ctx context.Context) 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
}
core/tasks/transfer/codec.go
+92
View File
@@ -0,0 +1,92 @@
package transfer
import (
"context"
"encoding/json"
"fmt"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
"github.com/krau/SaveAny-Bot/storage"
)
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"`
}
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,
}
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)
}
return NewTransferTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors), nil
}