feat(tdler): resumable chunked download with block bitmap

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<<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 {
+		return nil, fmt.Errorf("parse resume bitmap: %w", err)
+	}
+	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
+	}
+	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"
+}

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
+}

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,
+		})
+	}
+}

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
+}