mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-19 19:33:58 +08:00
832eb27d4f
Use the resource fingerprint as the dedup key on insert and delete. Check and set processing entries atomically instead of TOCTOU. Count only successful downloads.
132 lines
2.8 KiB
Go
132 lines
2.8 KiB
Go
package directlinks
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/krau/SaveAny-Bot/config"
|
|
"github.com/krau/SaveAny-Bot/core"
|
|
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
|
|
"github.com/krau/SaveAny-Bot/storage"
|
|
)
|
|
|
|
type File struct {
|
|
Name string
|
|
URL string
|
|
Size int64
|
|
}
|
|
|
|
func (f *File) FileName() string {
|
|
return f.Name
|
|
}
|
|
|
|
func (f *File) FileSize() int64 {
|
|
return f.Size
|
|
}
|
|
|
|
var _ core.Executable = (*Task)(nil)
|
|
|
|
type Task struct {
|
|
ID string
|
|
ctx context.Context
|
|
files []*File
|
|
Storage storage.Storage
|
|
StorPath string
|
|
Progress ProgressTracker
|
|
|
|
client *http.Client // [TODO] parallel download
|
|
stream bool
|
|
totalBytes int64 // total bytes to download
|
|
downloadedBytes atomic.Int64 // downloaded bytes
|
|
totalFiles int64 // total files to download
|
|
downloaded atomic.Int64 // downloaded files count
|
|
processing map[string]*File // {"url": File}
|
|
processingMu sync.RWMutex
|
|
}
|
|
|
|
// Title implements core.Exectable.
|
|
func (t *Task) Title() string {
|
|
return fmt.Sprintf("[%s](%s...->%s:%s)", t.Type(), t.files[0].Name, t.Storage.Name(), t.StorPath)
|
|
}
|
|
|
|
// DownloadedBytes implements TaskInfo.
|
|
func (t *Task) DownloadedBytes() int64 {
|
|
return t.downloadedBytes.Load()
|
|
}
|
|
|
|
// Processing implements TaskInfo.
|
|
func (t *Task) Processing() []FileInfo {
|
|
t.processingMu.RLock()
|
|
defer t.processingMu.RUnlock()
|
|
infos := make([]FileInfo, 0, len(t.processing))
|
|
for _, f := range t.processing {
|
|
infos = append(infos, f)
|
|
}
|
|
return infos
|
|
}
|
|
|
|
// StorageName implements TaskInfo.
|
|
func (t *Task) StorageName() string {
|
|
return t.Storage.Name()
|
|
}
|
|
|
|
// StoragePath implements TaskInfo.
|
|
func (t *Task) StoragePath() string {
|
|
if len(t.files) == 1 {
|
|
return t.StorPath + "/" + t.files[0].Name
|
|
}
|
|
return t.StorPath
|
|
}
|
|
|
|
// TotalBytes implements TaskInfo.
|
|
func (t *Task) TotalBytes() int64 {
|
|
return t.totalBytes
|
|
}
|
|
|
|
// TotalFiles implements TaskInfo.
|
|
func (t *Task) TotalFiles() int {
|
|
return int(t.totalFiles)
|
|
}
|
|
|
|
func (t *Task) Type() tasktype.TaskType {
|
|
return tasktype.TaskTypeDirectlinks
|
|
}
|
|
|
|
func (t *Task) TaskID() string {
|
|
return t.ID
|
|
}
|
|
|
|
func NewTask(
|
|
id string,
|
|
ctx context.Context,
|
|
links []string,
|
|
stor storage.Storage,
|
|
storPath string,
|
|
progressTracker ProgressTracker,
|
|
) *Task {
|
|
_, ok := stor.(storage.StorageCannotStream)
|
|
stream := config.C().Stream && !ok
|
|
files := make([]*File, 0, len(links))
|
|
for _, link := range links {
|
|
files = append(files, &File{
|
|
URL: link,
|
|
})
|
|
}
|
|
return &Task{
|
|
ID: id,
|
|
ctx: ctx,
|
|
files: files,
|
|
Storage: stor,
|
|
StorPath: storPath,
|
|
Progress: progressTracker,
|
|
stream: stream,
|
|
client: http.DefaultClient,
|
|
processing: make(map[string]*File),
|
|
processingMu: sync.RWMutex{},
|
|
totalFiles: int64(len(files)),
|
|
}
|
|
}
|