hpz937/zurg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Reapply download map and unrestrict map

Browse Source
This commit is contained in:
Ben Sarmiento
2024-05-25 21:09:12 +02:00
parent 7886ed69a2
commit e6a47be656
7 changed files with 82 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
internal/torrent/manager.go
View File

@@ -1,6 +1,7 @@
package torrent
import (
"context"
"io"
"os"
"path/filepath"
@@ -15,6 +16,7 @@ import (
mapset "github.com/deckarep/golang-set/v2"
cmap "github.com/orcaman/concurrent-map/v2"
"github.com/panjf2000/ants/v2"
"gopkg.in/vansante/go-ffprobe.v2"
)
const (
@@ -30,8 +32,9 @@ type TorrentManager struct {
log *logutil.Logger
repairLog *logutil.Logger
DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
DownloadMap cmap.ConcurrentMap[string, *realdebrid.Download]
DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
DownloadMap cmap.ConcurrentMap[string, *realdebrid.Download]
UnrestrictMap cmap.ConcurrentMap[string, *realdebrid.Download]
RootNode *fs.FileNode
@@ -63,8 +66,9 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
log: log,
repairLog: repairLog,
DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
DownloadMap: cmap.New[*realdebrid.Download](),
DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
DownloadMap: cmap.New[*realdebrid.Download](),
UnrestrictMap: cmap.New[*realdebrid.Download](),
RootNode: fs.NewFileNode("root", true),
@@ -77,34 +81,55 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
t.initializeBins()
t.initializeDirectoryMaps()
var wg sync.WaitGroup
wg.Add(2)
t.workerPool.Submit(func() {
t.refreshTorrents()
wg.Done()
t.setNewLatestState(t.getCurrentState())
t.StartRefreshJob()
t.StartRepairJob()
t.TriggerRepair(nil)
})
t.workerPool.Submit(func() {
t.mountDownloads()
t.mountNewDownloads()
wg.Done()
t.StartDownloadsJob()
})
t.workerPool.Submit(func() {
wg.Wait()
t.log.Info("Applying media info details to all torrents")
allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
allTorrents.IterCb(func(_ string, torrent *Torrent) {
t.applyMediaInfoDetails(torrent)
t.writeTorrentToFile(torrent)
t.log.Debugf("Applied media info details to torrent %s", torrent.Hash)
})
})
return t
}
// proxy function
func (t *TorrentManager) UnrestrictLinkUntilOk(link string, checkFirstByte bool) *realdebrid.Download {
if strings.HasPrefix(link, "https://real-debrid.com/d/") && t.DownloadMap.Has(link[0:39]) {
ret, _ := t.DownloadMap.Get(link[0:39])
isRealDebrid := strings.HasPrefix(link, "https://real-debrid.com/d/")
if isRealDebrid && t.UnrestrictMap.Has(link[0:39]) {
ret, _ := t.UnrestrictMap.Get(link[0:39])
return ret
} else if !isRealDebrid && t.UnrestrictMap.Has(link) {
ret, _ := t.UnrestrictMap.Get(link)
return ret
}
ret, err := t.api.UnrestrictLink(link, checkFirstByte)
t.DownloadMap.Set(ret.Link[0:39], ret)
if err != nil {
t.log.Warnf("Cannot unrestrict link %s: %v", link, err)
return nil
}
t.DownloadMap.Set(ret.Filename, ret)
if isRealDebrid {
t.UnrestrictMap.Set(ret.Link[0:39], ret)
} else {
t.UnrestrictMap.Set(ret.Link, ret)
}
return ret
}
@@ -152,7 +177,7 @@ func (t *TorrentManager) getTorrentFiles(parentDir string) mapset.Set[string] {
}
func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) {
filePath := "data/" + torrent.Hash + ".zurgtorrent"
filePath := "data/" + t.GetKey(torrent) + ".zurgtorrent"
file, err := os.Create(filePath)
if err != nil {
t.log.Warnf("Cannot create file %s: %v", filePath, err)
@@ -176,6 +201,28 @@ func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) {
t.log.Debugf("Saved torrent %s (hash=%s) to file", t.GetKey(torrent), torrent.Hash)
}
func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) {
torrent.SelectedFiles.IterCb(func(_ string, file *File) {
if file.MediaInfo != nil {
return
}
unrestrict := t.UnrestrictFileUntilOk(file, false)
if unrestrict == nil {
file.State.Event(context.Background(), "break_file")
return
}
ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second)
defer cancelFn()
data, err := ffprobe.ProbeURL(ctx, unrestrict.Download)
if err != nil {
t.log.Warnf("Cannot probe file %s: %v", file.Path, err)
file.State.Event(context.Background(), "break_file")
return
}
file.MediaInfo = data
})
}
func (t *TorrentManager) readTorrentFromFile(filePath string) *Torrent {
file, err := os.Open(filePath)
if err != nil {
@@ -268,15 +315,20 @@ func (t *TorrentManager) deleteInfoFile(torrentID string) {
/// end info functions
func (t *TorrentManager) mountDownloads() {
t.DownloadMap.Clear()
_ = t.workerPool.Submit(func() {
downloads := t.api.GetDownloads()
for i := range downloads {
idx := i
t.DownloadMap.Set(downloads[idx].Filename, &downloads[idx])
func (t *TorrentManager) mountNewDownloads() {
downloads := t.api.GetDownloads()
for _, download := range downloads {
isRealDebrid := strings.HasPrefix(download.Link, "https://real-debrid.com/d/")
if isRealDebrid && !t.UnrestrictMap.Has(download.Link[0:39]) {
t.UnrestrictMap.Set(download.Link[0:39], &download)
} else if !isRealDebrid {
if !t.UnrestrictMap.Has(download.Link) {
t.UnrestrictMap.Set(download.Link, &download)
}
filename := filepath.Base(download.Filename)
t.DownloadMap.Set(filename, &download)
}
})
}
}
func (t *TorrentManager) StartDownloadsJob() {
@@ -287,9 +339,10 @@ func (t *TorrentManager) StartDownloadsJob() {
for {
select {
case <-remountTicker.C:
t.mountDownloads()
t.mountNewDownloads()
case <-t.RemountTrigger:
t.mountDownloads()
t.DownloadMap.Clear()
t.mountNewDownloads()
}
}
})