diff --git a/cmd/run.go b/cmd/run.go index 7499296..33020ba 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -10,6 +10,7 @@ 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" @@ -56,6 +57,14 @@ func Run(cmd *cobra.Command, _ []string) { cancel() }() + core.SetDownloaderProvider(func() downloader.Client { + if ectx := bot.ExtContext(); ectx != nil { + return ectx.Raw + } + return nil + }) + core.RecoverTasks(ctx) + core.Run(ctx) <-ctx.Done() @@ -102,6 +111,15 @@ 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) diff --git a/core/core.go b/core/core.go index 95561a5..975377b 100644 --- a/core/core.go +++ b/core/core.go @@ -7,6 +7,7 @@ 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" @@ -45,6 +46,9 @@ 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) @@ -70,6 +74,9 @@ 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) + if err := database.DeleteTask(ctx, exe.TaskID()); err != nil { + logger.Errorf("Failed to delete persisted task %s: %v", exe.TaskID(), err) + } <-semaphore } } @@ -92,12 +99,21 @@ 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) - return err + 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 } func GetLength(ctx context.Context) int { diff --git a/core/persist.go b/core/persist.go new file mode 100644 index 0000000..9379039 --- /dev/null +++ b/core/persist.go @@ -0,0 +1,114 @@ +package core + +import ( + "context" + "sync" + + "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), + }) +} + +// RecoverTasks re-enqueues tasks that were unfinished when the process last +// exited. Must be called after storages are loaded and before Run. Tasks of +// types without a registered codec are dropped with a warning. +func RecoverTasks(ctx context.Context) { + logger := log.FromContext(ctx) + 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("Dropping unrecoverable task %s (type %s): no codec registered", t.ID, t.Type) + if err := database.DeleteTask(ctx, t.ID); err != nil { + logger.Errorf("Failed to delete task %s: %v", t.ID, err) + } + continue + } + task, err := codec.Unmarshal(t.Payload) + if err != nil { + logger.Errorf("Dropping task %s: failed to rebuild: %v", t.ID, err) + if err := database.DeleteTask(ctx, t.ID); err != nil { + logger.Errorf("Failed to delete task %s: %v", t.ID, err) + } + continue + } + if err := AddTask(ctx, task); err != nil { + logger.Errorf("Dropping task %s: failed to re-enqueue: %v", t.ID, err) + if err := database.DeleteTask(ctx, t.ID); err != nil { + logger.Errorf("Failed to delete task %s: %v", t.ID, err) + } + continue + } + logger.Infof("Recovered task %s (%s)", t.ID, t.Type) + } +} diff --git a/core/persist_test.go b/core/persist_test.go new file mode 100644 index 0000000..3db8784 --- /dev/null +++ b/core/persist_test.go @@ -0,0 +1,124 @@ +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 TestRecoverTasksReenqueuesAndDropsUnknown(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) + } + } +} + +func TestRecoverTasksDropsInvalidPayload(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 nor remain in the database. + 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) + } +} diff --git a/core/tasks/batchtfile/codec.go b/core/tasks/batchtfile/codec.go new file mode 100644 index 0000000..489ada7 --- /dev/null +++ b/core/tasks/batchtfile/codec.go @@ -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 +} diff --git a/core/tasks/batchtfile/execute.go b/core/tasks/batchtfile/execute.go index a00c602..7b298a0 100644 --- a/core/tasks/batchtfile/execute.go +++ b/core/tasks/batchtfile/execute.go @@ -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) diff --git a/core/tasks/tfile/codec.go b/core/tasks/tfile/codec.go new file mode 100644 index 0000000..f3cc206 --- /dev/null +++ b/core/tasks/tfile/codec.go @@ -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 +} diff --git a/core/tasks/tfile/execute.go b/core/tasks/tfile/execute.go index 8eea959..3f8db41 100644 --- a/core/tasks/tfile/execute.go +++ b/core/tasks/tfile/execute.go @@ -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 } diff --git a/core/tasks/transfer/codec.go b/core/tasks/transfer/codec.go new file mode 100644 index 0000000..06740b8 --- /dev/null +++ b/core/tasks/transfer/codec.go @@ -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 +} diff --git a/database/db.go b/database/db.go index 3c4300e..ceb9c8a 100644 --- a/database/db.go +++ b/database/db.go @@ -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{}); err != nil { + if err := db.AutoMigrate(&User{}, &Dir{}, &Rule{}, &WatchChat{}, &Task{}); 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 { diff --git a/database/task.go b/database/task.go new file mode 100644 index 0000000..c03c62d --- /dev/null +++ b/database/task.go @@ -0,0 +1,92 @@ +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 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 +} diff --git a/database/task_test.go b/database/task_test.go new file mode 100644 index 0000000..7855696 --- /dev/null +++ b/database/task_test.go @@ -0,0 +1,110 @@ +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) + } +} diff --git a/pkg/tfile/payload.go b/pkg/tfile/payload.go new file mode 100644 index 0000000..f65a192 --- /dev/null +++ b/pkg/tfile/payload.go @@ -0,0 +1,86 @@ +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, + } +} diff --git a/pkg/tfile/payload_test.go b/pkg/tfile/payload_test.go new file mode 100644 index 0000000..21d5372 --- /dev/null +++ b/pkg/tfile/payload_test.go @@ -0,0 +1,61 @@ +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()) + } +}