package download import ( "errors" "github.com/monkeyWie/gopeed/internal/controller" "github.com/monkeyWie/gopeed/internal/fetcher" "github.com/monkeyWie/gopeed/pkg/base" "github.com/monkeyWie/gopeed/pkg/util" "strings" "sync" "sync/atomic" "time" ) const ( bucketTask = "task" bucketSave = "save" 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 { fetchBuilders map[string]fetcher.FetcherBuilder storage Storage tasks []*Task listener Listener refreshInterval int lock *sync.Mutex closed atomic.Bool } func NewDownloader(cfg *DownloaderConfig) *Downloader { if cfg == nil { cfg = &DownloaderConfig{} } cfg.Init() d := &Downloader{ fetchBuilders: make(map[string]fetcher.FetcherBuilder), refreshInterval: cfg.RefreshInterval, storage: cfg.Storage, lock: &sync.Mutex{}, } for _, f := range cfg.FetchBuilders { for _, p := range f.Schemes() { d.fetchBuilders[strings.ToUpper(p)] = f } } return d } func (d *Downloader) Setup() error { // setup storage if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketConfig}); err != nil { return err } // 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 _, task := range tasks { initTask(task) if task.Status != base.DownloadStatusDone && task.Status != base.DownloadStatusError { task.Status = base.DownloadStatusPause } } } d.tasks = tasks // 每秒统计一次下载速度 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.DownloadStatusDone || task.Status == base.DownloadStatusError || task.Status == base.DownloadStatusPause { 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 = current - task.Progress.Downloaded 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.refreshInterval)) } }() return nil } func (d *Downloader) buildFetcher(url string) (fetcher.FetcherBuilder, fetcher.Fetcher, error) { schema := util.ParseSchema(url) fetchBuilder, ok := d.fetchBuilders[schema] if !ok { fetchBuilder, ok = d.fetchBuilders[util.FileSchema] } if ok { fetcher := fetchBuilder.Build() d.setupFetcher(fetcher) return fetchBuilder, fetcher, nil } return nil, 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) (*base.Resource, error) { _, fetcher, err := d.buildFetcher(req.URL) if err != nil { return nil, err } res, err := fetcher.Resolve(req) if err != nil { return nil, err } return res, nil } func (d *Downloader) Create(res *base.Resource, opts *base.Options) (taskId string, err error) { if opts == nil { opts = &base.Options{} } if len(opts.SelectFiles) == 0 { opts.SelectFiles = make([]int, 0) for i := range res.Files { opts.SelectFiles = append(opts.SelectFiles, i) } } if opts.Path == "" { exist, storeConfig, err := d.GetConfig() if err != nil { return "", err } if exist { opts.Path = storeConfig.DownloadDir } } fetcherBuilder, fetcher, err := d.buildFetcher(res.Req.URL) if err != nil { return } err = fetcher.Create(res, opts) if err != nil { return } task := NewTask() task.fetcherBuilder = fetcherBuilder task.fetcher = fetcher task.Res = res task.Opts = opts task.Progress = &Progress{} task.Status = base.DownloadStatusRunning // calculate select files size for _, selectIndex := range opts.SelectFiles { task.Size += res.Files[selectIndex].Size } initTask(task) task.timer.Start() err = func() error { d.lock.Lock() defer d.lock.Unlock() if err := d.storage.Put(bucketTask, task.ID, task.clone()); err != nil { return err } d.tasks = append(d.tasks, task) return nil }() if err != nil { return } d.emit(EventKeyStart, task) taskId = task.ID go func() { err := fetcher.Start() if err != nil { d.emit(EventKeyError, task, err) return } d.watch(task) }() return } func (d *Downloader) Pause(id string) (err error) { task := d.GetTask(id) if task == nil { return } task.lock.Lock() defer task.lock.Unlock() if task.Status == base.DownloadStatusRunning { task.Status = base.DownloadStatusPause task.timer.Pause() task.fetcher.Pause() d.storage.Put(bucketTask, task.ID, task.clone()) d.emit(EventKeyPause, task) } return } func (d *Downloader) Continue(id string) (err error) { task := d.GetTask(id) if task == nil { return } task.lock.Lock() defer task.lock.Unlock() if task.Status == base.DownloadStatusPause || task.Status == base.DownloadStatusError { task.Status = base.DownloadStatusRunning err = d.restoreFetcher(task) if err != nil { return } task.timer.Start() err = task.fetcher.Continue() if err != nil { return } d.storage.Put(bucketTask, task.ID, task.clone()) d.emit(EventKeyContinue, task) } 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 { names := make([]string, 0) for _, file := range task.Res.Files { names = append(names, util.Filepath(file.Path, file.Name, task.Opts.Name)) } util.SafeRemoveAll(task.Opts.Path, names) } d.emit(EventKeyDelete, task) 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.Pause(task.ID) } 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() (bool, *DownloaderStoreConfig, error) { var cfg DownloaderStoreConfig exist, err := d.storage.Get(bucketConfig, "config", &cfg) if err != nil { return false, nil, err } return exist, &cfg, nil } func (d *Downloader) PutConfig(v *DownloaderStoreConfig) error { return d.storage.Put(bucketConfig, "config", v) } func (d *Downloader) getProtocolConfig(name string, v any) (bool, error) { exist, cfg, err := d.GetConfig() if err != nil { return false, err } if !exist { return false, nil } 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 } 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.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) } func (d *Downloader) restoreFetcher(task *Task) error { if task.fetcher == nil { fetcherBuilder, _, err := d.buildFetcher(task.Res.Req.URL) if err != nil { return err } task.fetcherBuilder = fetcherBuilder err = func() error { v, f := fetcherBuilder.Restore() if v != nil { err := d.storage.Pop(bucketSave, task.ID, v) if err != nil { return err } task.fetcher = f(task.Res, task.Opts, v) } return nil }() if err != nil { // TODO log } if task.fetcher == nil { task.fetcher = fetcherBuilder.Build() } d.setupFetcher(task.fetcher) task.fetcher.Create(task.Res, task.Opts) go d.watch(task) } return nil } func initTask(task *Task) { task.timer = &util.Timer{} task.lock = &sync.Mutex{} } 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() (*base.Resource, 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) { res, err := b.Resolve() if err != nil { return "", err } defaultDownloader.Listener(b.listener) return defaultDownloader.Create(res, opts) } func Boot() *boot { err := defaultDownloader.Setup() if err != nil { panic(err) } return &boot{} }