This commit is contained in:
Ben Sarmiento
2023-11-27 21:50:00 +01:00
parent a7623de410
commit c8334ecb3b
15 changed files with 277 additions and 221 deletions

View File

@@ -12,7 +12,6 @@ import (
"time"
"github.com/debridmediamanager.com/zurg/internal/config"
"github.com/debridmediamanager.com/zurg/pkg/logutil"
"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
"github.com/debridmediamanager.com/zurg/pkg/utils"
cmap "github.com/orcaman/concurrent-map/v2"
@@ -21,8 +20,9 @@ import (
)
const (
ALL_TORRENTS = "__all__"
DATA_DIR = "data"
INT_ALL = "int__all__"
INT_INFO_CACHE = "int__info__"
DATA_DIR = "data"
)
type TorrentManager struct {
@@ -33,6 +33,7 @@ type TorrentManager struct {
cfg config.ConfigInterface
api *realdebrid.RealDebrid
antsPool *ants.Pool
unrestrictPool *ants.Pool
log *zap.SugaredLogger
mu *sync.Mutex
}
@@ -40,21 +41,30 @@ type TorrentManager struct {
// NewTorrentManager creates a new torrent manager
// it will fetch all torrents and their info in the background
// and store them in-memory and cached in files
func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p *ants.Pool) *TorrentManager {
func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p *ants.Pool, log *zap.SugaredLogger) *TorrentManager {
t := &TorrentManager{
cfg: cfg,
DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
requiredVersion: "18.11.2023",
api: api,
antsPool: p,
log: logutil.NewLogger().Named("manager"),
log: log,
mu: &sync.Mutex{},
}
unrestrictPool, err := ants.NewPool(10)
if err != nil {
t.unrestrictPool = t.antsPool
} else {
t.unrestrictPool = unrestrictPool
}
ensureDir(DATA_DIR)
// create special directory
t.DirectoryMap.Set("__all__", cmap.New[*Torrent]()) // key is AccessKey
// create internal directories
t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is AccessKey
t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID
// create directory maps
for _, directory := range cfg.GetDirectories() {
t.DirectoryMap.Set(directory, cmap.New[*Torrent]())
@@ -82,18 +92,18 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
noInfoCount := 0
allCt := 0
allTorrents, _ := t.DirectoryMap.Get(ALL_TORRENTS)
allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
for info := range torrentsChan {
allCt++
if info == nil {
noInfoCount++
continue
}
if torrent, exists := allTorrents.Get(info.AccessKey); exists {
if torrent, exists := allTorrents.Get(info.AccessKey); !exists {
allTorrents.Set(info.AccessKey, info)
} else {
mainTorrent := t.mergeToMain(torrent, info)
allTorrents.Set(info.AccessKey, mainTorrent)
} else {
allTorrents.Set(info.AccessKey, info)
}
}
@@ -134,6 +144,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
go t.startRefreshJob()
t.latestAdded = newTorrents[0].Added // set the latest added to the first torrent's added
t.log.Info("Finished initializing torrent manager")
return t
}
@@ -168,8 +179,14 @@ func (t *TorrentManager) mergeToMain(mainTorrent, torrentToMerge *Torrent) *Torr
// proxy
func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.UnrestrictResponse {
ret := t.api.UnrestrictUntilOk(link)
return ret
retChan := make(chan *realdebrid.UnrestrictResponse, 1)
t.unrestrictPool.Submit(func() {
retChan <- t.api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
time.Sleep(1 * time.Second)
})
defer close(retChan)
return <-retChan
// return t.api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
}
type torrentsResponse struct {
@@ -178,9 +195,9 @@ type torrentsResponse struct {
}
func (t *TorrentManager) SetChecksum(checksum string) {
t.mu.Lock()
// t.mu.Lock()
t.checksum = checksum
t.mu.Unlock()
// t.mu.Unlock()
}
// generates a checksum based on the number of torrents, the first torrent id and the number of active torrents
@@ -251,6 +268,23 @@ func (t *TorrentManager) startRefreshJob() {
}
t.log.Infof("Detected changes! Refreshing %d torrents", len(newTorrents))
// handle deleted torrents in info cache
keep := make(map[string]bool)
for _, torrent := range newTorrents {
keep[torrent.ID] = true
}
var toDelete []string
infoCache, _ := t.DirectoryMap.Get(INT_INFO_CACHE)
infoCache.IterCb(func(torrentID string, torrent *Torrent) {
if _, ok := keep[torrentID]; !ok {
toDelete = append(toDelete, torrentID)
}
})
for _, torrentID := range toDelete {
infoCache.Remove(torrentID)
}
// end info cache cleanup
torrentsChan := make(chan *Torrent, len(newTorrents))
var wg sync.WaitGroup
for i := range newTorrents {
@@ -267,20 +301,20 @@ func (t *TorrentManager) startRefreshJob() {
t.log.Infof("Fetched info for %d torrents", len(newTorrents))
noInfoCount := 0
oldTorrents, _ := t.DirectoryMap.Get(ALL_TORRENTS)
oldTorrents, _ := t.DirectoryMap.Get(INT_ALL)
newSet := cmap.New[*Torrent]()
for info := range torrentsChan {
if info == nil {
noInfoCount++
continue
}
if torrent, exists := oldTorrents.Get(info.AccessKey); exists {
if torrent, exists := oldTorrents.Get(info.AccessKey); !exists {
oldTorrents.Set(info.AccessKey, info)
newSet.Set(info.AccessKey, info)
} else {
mainTorrent := t.mergeToMain(torrent, info)
oldTorrents.Set(info.AccessKey, mainTorrent)
newSet.Set(info.AccessKey, mainTorrent)
} else {
oldTorrents.Set(info.AccessKey, info)
newSet.Set(info.AccessKey, info)
}
}
@@ -333,15 +367,24 @@ func (t *TorrentManager) startRefreshJob() {
t.log.Info("Checking for torrents to repair")
t.repairAll()
t.log.Info("Finished checking for torrents to repair")
} else {
t.log.Info("Repair is disabled, skipping repair check")
}
go OnLibraryUpdateHook(updatedPaths, t.cfg, t.log)
t.latestAdded = newTorrents[0].Added
t.log.Info("Finished refreshing torrents")
}
}
// getMoreInfo gets original name, size and files for a torrent
func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *Torrent {
infoCache, _ := t.DirectoryMap.Get(INT_INFO_CACHE)
if infoCache.Has(rdTorrent.ID) {
tor, _ := infoCache.Get(rdTorrent.ID)
return tor
}
var info *realdebrid.TorrentInfo
var err error
// file cache
@@ -399,6 +442,9 @@ func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *Torrent {
if len(selectedFiles) > 0 && torrentFromFile == nil {
t.writeToFile(info) // only when there are selected files, else it's useless
}
infoCache.Set(rdTorrent.ID, &torrent)
return &torrent
}
@@ -470,10 +516,11 @@ func (t *TorrentManager) organizeChaos(links []string, selectedFiles []*File) ([
wg.Add(1)
link := link // redeclare to avoid closure on loop variable
// Use the existing worker pool to submit tasks
_ = t.antsPool.Submit(func() {
_ = t.unrestrictPool.Submit(func() {
defer wg.Done()
resp := t.api.UnrestrictUntilOk(link)
resp := t.api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
resultsChan <- Result{Response: resp}
time.Sleep(1 * time.Second)
})
}
@@ -522,7 +569,8 @@ func (t *TorrentManager) repairAll() {
return
}
allTorrents, _ := t.DirectoryMap.Get(ALL_TORRENTS)
var toDelete []string
allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
allTorrents.IterCb(func(_ string, torrent *Torrent) {
if torrent.AnyInProgress() {
t.log.Debugf("Skipping %s for repairs because it is in progress", torrent.AccessKey)
@@ -531,7 +579,7 @@ func (t *TorrentManager) repairAll() {
forRepair := false
unselected := 0
torrent.SelectedFiles.IterCb(func(_ string, file *File) {
if file.Link == "repair" {
if file.Link == "repair" && !forRepair {
t.log.Debugf("Found a file to repair for torrent %s", torrent.AccessKey)
forRepair = true
}
@@ -545,20 +593,29 @@ func (t *TorrentManager) repairAll() {
}
if unselected == torrent.SelectedFiles.Count() && unselected > 0 {
t.log.Infof("Deleting %s", torrent.AccessKey)
t.Delete(torrent.AccessKey)
toDelete = append(toDelete, torrent.AccessKey)
}
})
for _, accessKey := range toDelete {
t.Delete(accessKey)
}
}
func (t *TorrentManager) Delete(accessKey string) {
infoCache, _ := t.DirectoryMap.Get(INT_INFO_CACHE)
t.log.Infof("Deleting torrent %s", accessKey)
allTorrents, _ := t.DirectoryMap.Get(ALL_TORRENTS)
allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
if torrent, ok := allTorrents.Get(accessKey); ok {
for _, instance := range torrent.Instances {
infoCache.Remove(instance.ID)
t.api.DeleteTorrent(instance.ID)
}
allTorrents.Remove(accessKey)
}
t.DirectoryMap.IterCb(func(_ string, torrents cmap.ConcurrentMap[string, *Torrent]) {
if _, ok := torrents.Get(accessKey); ok {
torrents.Remove(accessKey)
}
})
}
func (t *TorrentManager) Repair(accessKey string) {
@@ -567,7 +624,7 @@ func (t *TorrentManager) Repair(accessKey string) {
return
}
allTorrents, _ := t.DirectoryMap.Get(ALL_TORRENTS)
allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
torrent, _ := allTorrents.Get(accessKey)
if torrent == nil {
t.log.Warnf("Cannot find torrent %s anymore so we are not repairing it", accessKey)