package download import ( "errors" "github.com/GopeedLab/gopeed/internal/controller" "github.com/GopeedLab/gopeed/internal/fetcher" "github.com/GopeedLab/gopeed/pkg/base" "github.com/GopeedLab/gopeed/pkg/util" gonanoid "github.com/matoous/go-nanoid/v2" "github.com/virtuald/go-paniclog" "math" "os" "path" "path/filepath" "sort" "strings" "sync" "sync/atomic" "time" ) const ( // task info bucket bucketTask = "task" // task download data bucket bucketSave = "save" // downloader config bucket bucketConfig = "config" ) type Listener func(event *Event) type Progress struct { // 下载耗时(纳秒) Used int64 `json:"used"` // 每秒下载字节数 Speed int64 `json:"speed"` // 已下载的字节数 Downloaded int64 `json:"downloaded"` } type Downloader struct { cfg *DownloaderConfig fetchBuilders map[string]fetcher.FetcherBuilder fetcherMap map[string]fetcher.Fetcher storage Storage tasks []*Task listener Listener lock *sync.Mutex closed atomic.Bool } func NewDownloader(cfg *DownloaderConfig) *Downloader { if cfg == nil { cfg = &DownloaderConfig{} } cfg.Init() d := &Downloader{ cfg: cfg, fetchBuilders: make(map[string]fetcher.FetcherBuilder), fetcherMap: make(map[string]fetcher.Fetcher), storage: cfg.Storage, lock: &sync.Mutex{}, } for _, f := range cfg.FetchBuilders { for _, p := range f.Schemes() { d.fetchBuilders[strings.ToUpper(p)] = f } } logPanic(filepath.Join(cfg.StorageDir, "logs")) return d } func (d *Downloader) Setup() error { // setup storage if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketConfig}); err != nil { return err } // load config from storage var cfg DownloaderStoreConfig exist, err := d.storage.Get(bucketConfig, "config", &cfg) if err != nil { return err } if exist { d.cfg.DownloaderStoreConfig = &cfg } else { d.cfg.DownloaderStoreConfig = &DownloaderStoreConfig{ FirstLoad: true, } } // init default config d.cfg.DownloaderStoreConfig.Init() // load tasks from storage var tasks []*Task if err := d.storage.List(bucketTask, &tasks); err != nil { return err } if tasks == nil { tasks = make([]*Task, 0) } else { for i := len(tasks) - 1; i >= 0; i-- { task := tasks[i] // Remove broken tasks if task.Meta == nil { tasks = append(tasks[:i], tasks[i+1:]...) continue } initTask(task) if task.Status != base.DownloadStatusDone && task.Status != base.DownloadStatusError { task.Status = base.DownloadStatusPause } } // sort by create time sort.Slice(tasks, func(i, j int) bool { return tasks[i].CreatedAt.Before(tasks[j].CreatedAt) }) } d.tasks = tasks // 每个tick统计一次下载速度 go func() { for !d.closed.Load() { if len(d.tasks) > 0 { for _, task := range d.tasks { func() { task.lock.Lock() defer task.lock.Unlock() if task.Status != base.DownloadStatusRunning { return } // check if task is deleted if d.GetTask(task.ID) == nil { return } current := task.fetcher.Progress().TotalDownloaded() task.Progress.Used = task.timer.Used() task.Progress.Speed = task.calcSpeed(current-task.Progress.Downloaded, float64(d.cfg.RefreshInterval)/1000) task.Progress.Downloaded = current d.emit(EventKeyProgress, task) // store fetcher progress data, err := task.fetcherBuilder.Store(task.fetcher) if err != nil { return // TODO log } if err := d.storage.Put(bucketSave, task.ID, data); err != nil { return // TODO log } if err := d.storage.Put(bucketTask, task.ID, task); err != nil { return // TODO log } }() } } time.Sleep(time.Millisecond * time.Duration(d.cfg.RefreshInterval)) } }() return nil } func (d *Downloader) parseFb(url string) (fetcher.FetcherBuilder, error) { schema := util.ParseSchema(url) fetchBuilder, ok := d.fetchBuilders[schema] if ok { return fetchBuilder, nil } return nil, errors.New("unsupported protocol") } func (d *Downloader) setupFetcher(fetcher fetcher.Fetcher) { ctl := controller.NewController() ctl.GetConfig = func(v any) (bool, error) { return d.getProtocolConfig(fetcher.Name(), v) } fetcher.Setup(ctl) } func (d *Downloader) Resolve(req *base.Request) (rr *ResolveResult, err error) { fb, err := d.parseFb(req.URL) if err != nil { return } fetcher := fb.Build() d.setupFetcher(fetcher) err = fetcher.Resolve(req) if err != nil { return } rrId, err := gonanoid.New() if err != nil { return } d.fetcherMap[rrId] = fetcher rr = &ResolveResult{ ID: rrId, Res: fetcher.Meta().Res, } return } func (d *Downloader) CreateDirect(req *base.Request, opts *base.Options) (taskId string, err error) { rr, err := d.Resolve(req) if err != nil { return } return d.Create(rr.ID, opts) } func (d *Downloader) Create(rrId string, opts *base.Options) (taskId string, err error) { if opts == nil { opts = &base.Options{} } fetcher, ok := d.fetcherMap[rrId] if !ok { return "", errors.New("invalid resource id") } defer delete(d.fetcherMap, rrId) meta := fetcher.Meta() meta.Opts = opts res := meta.Res if len(opts.SelectFiles) == 0 { opts.SelectFiles = make([]int, 0) for i := range res.Files { opts.SelectFiles = append(opts.SelectFiles, i) } } if opts.Path == "" { storeConfig, err := d.GetConfig() if err != nil { return "", err } opts.Path = storeConfig.DownloadDir } fb, err := d.parseFb(fetcher.Meta().Req.URL) if err != nil { return } // check if the download file is duplicated and rename it automatically. files := res.Files if res.RootDir != "" { fullDirPath := path.Join(opts.Path, res.RootDir) newName, err := util.CheckDuplicateAndRename(fullDirPath) if err != nil { return "", err } res.RootDir = newName } else if len(files) == 1 { fullPath := meta.Filepath(files[0]) newName, err := util.CheckDuplicateAndRename(fullPath) if err != nil { return "", err } opts.Name = newName } err = fetcher.Create(opts) if err != nil { return } task := NewTask() task.fetcherBuilder = fb task.fetcher = fetcher task.Meta = fetcher.Meta() task.Progress = &Progress{} task.Status = base.DownloadStatusReady // calculate select files size for _, selectIndex := range opts.SelectFiles { task.Size += res.Files[selectIndex].Size } initTask(task) err = func() error { d.lock.Lock() defer d.lock.Unlock() d.tasks = append(d.tasks, task) return d.storage.Put(bucketTask, task.ID, task.clone()) }() if err != nil { return } taskId = task.ID d.tryRunning(func(remain int) { if remain > 0 { err = d.start(task) } }) return } func (d *Downloader) Pause(id string) (err error) { task := d.GetTask(id) if task == nil { return } if err = d.doPause(task, base.DownloadStatusPause); err != nil { return err } return d.notifyRunning() } func (d *Downloader) Continue(id string) (err error) { task := d.GetTask(id) if task == nil { return } d.tryRunning(func(remain int) { if remain == 0 { // no more task can be running, need to doPause a running task for _, t := range d.tasks { if t.Status == base.DownloadStatusRunning { err = d.doPause(t, base.DownloadStatusWait) break } } } if err != nil { return } }) if err != nil { return } return d.doContinue(task) } func (d *Downloader) PauseAll() (err error) { d.lock.Lock() defer d.lock.Unlock() for _, task := range d.tasks { if task.Status == base.DownloadStatusRunning { if err = d.doPause(task, base.DownloadStatusPause); err != nil { return } } } return } func (d *Downloader) ContinueAll() (err error) { d.lock.Lock() defer d.lock.Unlock() runningCount := 0 continueTasks := make([]*Task, 0) for _, task := range d.tasks { if task.Status == base.DownloadStatusRunning { runningCount++ } else if task.Status != base.DownloadStatusDone { continueTasks = append(continueTasks, task) } } // calculate how many tasks can be continued, can't exceed maxRunning continueCount := int(math.Min(float64(d.cfg.MaxRunning-runningCount), float64(len(continueTasks)))) for i := 0; i < continueCount; i++ { if err = d.doContinue(continueTasks[i]); err != nil { return } } return } func (d *Downloader) Delete(id string, force bool) (err error) { d.lock.Lock() defer d.lock.Unlock() task := d.GetTask(id) if task == nil { return } task.lock.Lock() defer task.lock.Unlock() if task.fetcher != nil { err = task.fetcher.Close() if err != nil { return } } for i, t := range d.tasks { if t.ID == id { d.tasks = append(d.tasks[:i], d.tasks[i+1:]...) break } } err = d.storage.Delete(bucketTask, id) if err != nil { return } err = d.storage.Delete(bucketSave, id) if err != nil { return } if force { if task.Meta.Res.RootDir != "" { if err = os.RemoveAll(path.Join(task.Meta.Opts.Path, task.Meta.Res.RootDir)); err != nil { return } } else { for _, file := range task.Meta.Res.Files { if err = util.SafeRemove(task.Meta.Filepath(file)); err != nil { return err } } } } d.emit(EventKeyDelete, task) if task.Status == base.DownloadStatusRunning { if err = d.doNotifyRunning(); err != nil { return err } } task = nil return } func (d *Downloader) Close() error { d.lock.Lock() defer d.lock.Unlock() d.closed.Store(true) for _, task := range d.tasks { d.doPause(task, base.DownloadStatusPause) } return d.storage.Close() } func (d *Downloader) Clear() error { if err := d.Close(); err != nil { return err } if err := d.storage.Clear(); err != nil { return err } return nil } func (d *Downloader) Listener(fn Listener) { d.listener = fn } func (d *Downloader) emit(eventKey EventKey, task *Task, errs ...error) { if d.listener != nil { var err error if len(errs) > 0 { err = errs[0] } d.listener(&Event{ Key: eventKey, Task: task, Err: err, }) } } func (d *Downloader) GetTask(id string) *Task { for _, task := range d.tasks { if task.ID == id { return task } } return nil } func (d *Downloader) GetTasks() []*Task { return d.tasks } func (d *Downloader) GetConfig() (*DownloaderStoreConfig, error) { return d.cfg.DownloaderStoreConfig, nil } func (d *Downloader) PutConfig(v *DownloaderStoreConfig) error { d.cfg.DownloaderStoreConfig = v return d.storage.Put(bucketConfig, "config", v) } func (d *Downloader) getProtocolConfig(name string, v any) (bool, error) { cfg, err := d.GetConfig() if err != nil { return false, err } if cfg.ProtocolConfig == nil || cfg.ProtocolConfig[name] == nil { return false, nil } if err := util.MapToStruct(cfg.ProtocolConfig[name], v); err != nil { return false, err } return true, nil } // wait task done func (d *Downloader) watch(task *Task) { err := task.fetcher.Wait() if err != nil { task.Status = base.DownloadStatusError d.emit(EventKeyError, task, err) } else { task.Progress.Used = task.timer.Used() if task.Size == 0 { task.Size = task.fetcher.Progress().TotalDownloaded() task.Meta.Res.Size = task.Size } used := task.Progress.Used / int64(time.Second) if used == 0 { used = 1 } task.Progress.Speed = task.Size / used task.Progress.Downloaded = task.Size task.Status = base.DownloadStatusDone d.storage.Put(bucketTask, task.ID, task.clone()) d.emit(EventKeyDone, task) } d.emit(EventKeyFinally, task, err) d.notifyRunning() } func (d *Downloader) restoreFetcher(task *Task) error { if task.fetcher == nil { fb, err := d.parseFb(task.Meta.Req.URL) if err != nil { return err } task.fetcherBuilder = fb err = func() error { v, f := fb.Restore() if v != nil { err := d.storage.Pop(bucketSave, task.ID, v) if err != nil { return err } } task.fetcher = f(task.Meta, v) return nil }() if err != nil { // TODO log } if task.fetcher == nil { task.fetcher = fb.Build() } d.setupFetcher(task.fetcher) task.fetcher.Create(task.Meta.Opts) } return nil } func (d *Downloader) tryRunning(cb func(remain int)) { d.lock.Lock() defer d.lock.Unlock() runningCount := 0 for _, task := range d.tasks { if task.Status == base.DownloadStatusRunning { runningCount++ } } cb(d.cfg.MaxRunning - runningCount) } func (d *Downloader) notifyRunning() error { d.lock.Lock() defer d.lock.Unlock() return d.doNotifyRunning() } func (d *Downloader) doNotifyRunning() error { for _, task := range d.tasks { if task.Status == base.DownloadStatusReady || task.Status == base.DownloadStatusWait { return d.doContinue(task) } } return nil } func (d *Downloader) start(task *Task) error { var event EventKey if task.Status == base.DownloadStatusReady { event = EventKeyStart } else { event = EventKeyContinue } task.Status = base.DownloadStatusRunning task.Progress.Speed = 0 task.timer.Start() if err := d.storage.Put(bucketTask, task.ID, task.clone()); err != nil { return err } var err error if event == EventKeyStart { err = task.fetcher.Start() } if event == EventKeyContinue { err = task.fetcher.Continue() } if err != nil { return err } go d.watch(task) d.emit(event, task) return nil } func (d *Downloader) doPause(task *Task, status base.Status) (err error) { task.lock.Lock() defer task.lock.Unlock() if task.Status == base.DownloadStatusRunning { task.Status = status task.timer.Pause() task.fetcher.Pause() d.storage.Put(bucketTask, task.ID, task.clone()) d.emit(EventKeyPause, task) } return } func (d *Downloader) doContinue(task *Task) (err error) { task.lock.Lock() defer task.lock.Unlock() if task.Status != base.DownloadStatusRunning && task.Status != base.DownloadStatusDone { err = d.restoreFetcher(task) if err != nil { return } err = d.start(task) } return } // redirect stderr to log file, when panic happened log it func logPanic(logDir string) { if err := util.CreateDirIfNotExist(logDir); err != nil { return } f, err := os.Create(filepath.Join(logDir, "crash.log")) if err != nil { return } paniclog.RedirectStderr(f) } func initTask(task *Task) { task.timer = util.NewTimer(task.Progress.Used) task.lock = &sync.Mutex{} task.speedArr = make([]int64, 0) } var defaultDownloader = NewDownloader(nil) type boot struct { url string extra interface{} listener Listener } func (b *boot) URL(url string) *boot { b.url = url return b } func (b *boot) Extra(extra interface{}) *boot { b.extra = extra return b } func (b *boot) Resolve() (*ResolveResult, error) { return defaultDownloader.Resolve(&base.Request{ URL: b.url, Extra: b.extra, }) } func (b *boot) Listener(listener Listener) *boot { b.listener = listener return b } func (b *boot) Create(opts *base.Options) (string, error) { rr, err := b.Resolve() if err != nil { return "", err } defaultDownloader.Listener(b.listener) return defaultDownloader.Create(rr.ID, opts) } func Boot() *boot { err := defaultDownloader.Setup() if err != nil { panic(err) } return &boot{} }