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

1
go.mod
View File

@@ -23,6 +23,7 @@ require (
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
gopkg.in/vansante/go-ffprobe.v2 v2.1.1 // indirect
)
require (

2
go.sum
View File

@@ -52,6 +52,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
gopkg.in/vansante/go-ffprobe.v2 v2.1.1 h1:DIh5fMn+tlBvG7pXyUZdemVmLdERnf2xX6XOFF+0BBU=
gopkg.in/vansante/go-ffprobe.v2 v2.1.1/go.mod h1:qF0AlAjk7Nqzqf3y333Ly+KxN3cKF2JqA3JT5ZheUGE=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

2
internal/dav/infuse.go
View File

@@ -114,7 +114,7 @@ func ServeDownloadsListForInfuse(torMgr *torrent.TorrentManager) ([]byte, error)
for _, filename := range filenames {
download, ok := torMgr.DownloadMap.Get(filename)
if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") {
if !ok {
continue
}
buf.WriteString(dav.File(download.Filename, download.Filesize, download.Generated))

2
internal/dav/listing.go
View File

@@ -125,7 +125,7 @@ func ServeDownloadsList(torMgr *torrent.TorrentManager) ([]byte, error) {
sort.Strings(filenames)
for _, filename := range filenames {
download, ok := torMgr.DownloadMap.Get(filename)
if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") {
if !ok {
continue
}
buf.WriteString(dav.File(download.Filename, download.Filesize, download.Generated))

2
internal/http/listing.go
View File

@@ -98,7 +98,7 @@ func ServeDownloadsList(torMgr *torrent.TorrentManager) ([]byte, error) {
sort.Strings(filenames)
for _, filename := range filenames {
download, ok := torMgr.DownloadMap.Get(filename)
if !ok || strings.HasPrefix(download.Link, "https://real-debrid.com/d/") {
if !ok {
continue
}
filePath := filepath.Join(config.DOWNLOADS, url.PathEscape(filename))

2
internal/torrent/file_types.go
View File

@@ -5,6 +5,7 @@ import (
"github.com/debridmediamanager/zurg/pkg/realdebrid"
"github.com/looplab/fsm"
"gopkg.in/vansante/go-ffprobe.v2"
)
type File struct {
@@ -15,6 +16,7 @@ type File struct {
State *fsm.FSM `json:"-"`
Rename string `json:"Rename"`
MediaInfo *ffprobe.ProbeData `json:"MediaInfo"`
}
func (f *File) MarshalJSON() ([]byte, error) {

83
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 (
@@ -32,6 +34,7 @@ type TorrentManager struct {
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
@@ -65,6 +68,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
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() {
func (t *TorrentManager) mountNewDownloads() {
downloads := t.api.GetDownloads()
for i := range downloads {
idx := i
t.DownloadMap.Set(downloads[idx].Filename, &downloads[idx])
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()
}
}
})