diff --git a/common/tdler/resume.go b/common/tdler/resume.go new file mode 100644 index 0000000..dfc2cdf --- /dev/null +++ b/common/tdler/resume.go @@ -0,0 +1,238 @@ +package tdler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "os" + "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<= 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 + } + 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" +} diff --git a/common/tdler/resume_test.go b/common/tdler/resume_test.go new file mode 100644 index 0000000..5eeea5e --- /dev/null +++ b/common/tdler/resume_test.go @@ -0,0 +1,160 @@ +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") + } +} + +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 +} diff --git a/core/tasks/batchtfile/download.go b/core/tasks/batchtfile/download.go new file mode 100644 index 0000000..c15dbf9 --- /dev/null +++ b/core/tasks/batchtfile/download.go @@ -0,0 +1,88 @@ +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 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" + localFile, err := fsutil.CreateFile(partPath) + 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 { + return fmt.Errorf("failed to remove resume state: %w", 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, + }) + } +} diff --git a/core/tasks/tfile/download.go b/core/tasks/tfile/download.go new file mode 100644 index 0000000..66360b6 --- /dev/null +++ b/core/tasks/tfile/download.go @@ -0,0 +1,71 @@ +package tfile + +import ( + "context" + "fmt" + "github.com/charmbracelet/log" + "os" + + "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 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" + localFile, err := fsutil.CreateFile(partPath) + 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 { + return fmt.Errorf("failed to remove resume state: %w", err) + } + logger.Info("File downloaded successfully") + return nil +}