mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-09-06 08:06:56 +08:00
feat(tdler): resumable chunked download with block bitmap
This commit is contained in:
@@ -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"
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user