Refactor workers

This commit is contained in:
Ben Sarmiento
2024-01-27 13:17:45 +01:00
parent e01622064d
commit 05d2544fe8
6 changed files with 127 additions and 171 deletions

View File

@@ -45,6 +45,7 @@ func MainApp(configPath string) {
premium.MonitorPremiumStatus(rd, zurglog) premium.MonitorPremiumStatus(rd, zurglog)
// extra 1 worker for the refresh job
workerPool, err := ants.NewPool(config.GetNumOfWorkers() + 1) workerPool, err := ants.NewPool(config.GetNumOfWorkers() + 1)
if err != nil { if err != nil {
zurglog.Errorf("Failed to create worker pool: %v", err) zurglog.Errorf("Failed to create worker pool: %v", err)

View File

@@ -41,6 +41,13 @@ func (se *ScriptExecutor) Execute() (string, error) {
return out.String(), nil return out.String(), nil
} }
func (t *TorrentManager) TriggerHookOnLibraryUpdate(updatedPaths []string) {
_ = t.workerPool.Submit(func() {
OnLibraryUpdateHook(updatedPaths, t.Config, t.log)
t.log.Debugf("Triggered hook on_library_update for %d path(s)", len(updatedPaths))
})
}
func OnLibraryUpdateHook(paths []string, config config.ConfigInterface, log *logutil.Logger) { func OnLibraryUpdateHook(paths []string, config config.ConfigInterface, log *logutil.Logger) {
executor := &ScriptExecutor{ executor := &ScriptExecutor{
Script: config.GetOnLibraryUpdate(), Script: config.GetOnLibraryUpdate(),

View File

@@ -37,74 +37,35 @@ func (t *TorrentManager) SetNewLatestState(checksum LibraryState) {
t.latestState.FirstActiveTorrent = checksum.FirstActiveTorrent t.latestState.FirstActiveTorrent = checksum.FirstActiveTorrent
} }
type torrentsResp struct {
torrents []realdebrid.Torrent
totalCount int
}
// generates a checksum based on the number of torrents, the first torrent id and the number of active torrents // generates a checksum based on the number of torrents, the first torrent id and the number of active torrents
func (t *TorrentManager) getCurrentState() LibraryState { func (t *TorrentManager) getCurrentState() LibraryState {
torrentChan := make(chan torrentsResp, 1) var state LibraryState
activeChan := make(chan torrentsResp, 1)
countChan := make(chan int, 1)
errChan := make(chan error, 3) // accommodate errors from all goroutines
defer close(torrentChan)
defer close(activeChan)
defer close(countChan)
defer close(errChan)
_ = t.workerPool.Submit(func() { torrents, totalCount, err := t.Api.GetTorrents(1, false)
torrents, totalCount, err := t.Api.GetTorrents(1, false) if err != nil {
if err != nil { t.log.Warnf("Checksum API Error (GetTorrents): %v", err)
errChan <- err return LibraryState{}
return }
} if len(torrents) > 0 {
torrentChan <- torrentsResp{torrents: torrents, totalCount: totalCount} state.FirstTorrent = &torrents[0]
}) }
state.TotalCount = totalCount
_ = t.workerPool.Submit(func() { activeTorrents, _, err := t.Api.GetTorrents(1, true)
torrents, totalCount, err := t.Api.GetTorrents(1, true) if err != nil {
if err != nil { t.log.Warnf("Checksum API Error (GetTorrents): %v", err)
errChan <- err return LibraryState{}
return }
} if len(activeTorrents) > 0 {
activeChan <- torrentsResp{torrents: torrents, totalCount: totalCount} state.FirstActiveTorrent = &activeTorrents[0]
})
_ = t.workerPool.Submit(func() {
count, err := t.Api.GetActiveTorrentCount()
if err != nil {
errChan <- err
return
}
countChan <- count.DownloadingCount
})
var first, active *realdebrid.Torrent
var totalCount, count int
for i := 0; i < 3; i++ {
select {
case firstResp := <-torrentChan:
if len(firstResp.torrents) > 0 {
first = &firstResp.torrents[0]
}
totalCount = firstResp.totalCount
case activeResp := <-activeChan:
if len(activeResp.torrents) > 0 {
active = &activeResp.torrents[0]
}
case count = <-countChan:
case err := <-errChan:
t.log.Warnf("Checksum API Error: %v", err)
return LibraryState{}
}
} }
return LibraryState{ count, err := t.Api.GetActiveTorrentCount()
TotalCount: totalCount, if err != nil {
ActiveCount: count, t.log.Warnf("Checksum API Error (GetActiveTorrentCount): %v", err)
FirstTorrent: first, return LibraryState{}
FirstActiveTorrent: active,
} }
state.ActiveCount = count.DownloadingCount
return state
} }

View File

@@ -3,13 +3,11 @@ package torrent
import ( import (
"io" "io"
"os" "os"
"path/filepath"
"strings" "strings"
"github.com/debridmediamanager/zurg/internal/config" "github.com/debridmediamanager/zurg/internal/config"
"github.com/debridmediamanager/zurg/pkg/logutil" "github.com/debridmediamanager/zurg/pkg/logutil"
"github.com/debridmediamanager/zurg/pkg/realdebrid" "github.com/debridmediamanager/zurg/pkg/realdebrid"
"github.com/debridmediamanager/zurg/pkg/utils"
mapset "github.com/deckarep/golang-set/v2" mapset "github.com/deckarep/golang-set/v2"
cmap "github.com/orcaman/concurrent-map/v2" cmap "github.com/orcaman/concurrent-map/v2"
"github.com/panjf2000/ants/v2" "github.com/panjf2000/ants/v2"
@@ -58,51 +56,12 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w
log: log, log: log,
} }
// create internal directories t.initializeDirectories()
t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is GetAccessKey() t.mountDownloads()
t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID t.refreshTorrents()
// create directory maps
for _, directory := range cfg.GetDirectories() {
t.DirectoryMap.Set(directory, cmap.New[*Torrent]())
}
// Fetch downloads
if t.Config.EnableDownloadMount() {
_ = t.workerPool.Submit(func() {
page := 1
offset := 0
for {
downloads, totalDownloads, err := t.Api.GetDownloads(page, offset)
if err != nil {
t.log.Fatalf("Cannot get downloads: %v", err)
}
for i := range downloads {
t.DownloadMap.Set(downloads[i].Filename, &downloads[i])
}
offset += len(downloads)
page++
if offset >= totalDownloads {
t.log.Infof("Compiled into %d downloads", t.DownloadCache.Count())
break
}
}
})
}
t.RefreshTorrents()
t.SetNewLatestState(t.getCurrentState()) t.SetNewLatestState(t.getCurrentState())
t.startRefreshJob()
if t.Config.EnableRepair() { t.startRepairJob()
t.RepairAll() // initial repair
} else {
t.log.Info("Repair is disabled, skipping repair check")
}
t.log.Info("Finished initializing torrent manager")
_ = t.workerPool.Submit(func() {
t.startRefreshJob()
})
return t return t
} }
@@ -127,45 +86,6 @@ func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download {
return ret return ret
} }
func (t *TorrentManager) TriggerHookOnLibraryUpdate(updatedPaths []string) {
t.log.Debugf("Triggering hook on_library_update for %d path(s)", len(updatedPaths))
_ = t.workerPool.Submit(func() {
OnLibraryUpdateHook(updatedPaths, t.Config, t.log)
})
}
func (t *TorrentManager) assignedDirectoryCb(tor *Torrent, cb func(string)) {
torrentIDs := tor.DownloadedIDs.Union(tor.InProgressIDs).ToSlice()
// get filenames needed for directory conditions
var filenames []string
var fileSizes []int64
unplayable := true
tor.SelectedFiles.IterCb(func(key string, file *File) {
filenames = append(filenames, filepath.Base(file.Path))
fileSizes = append(fileSizes, file.Bytes)
if unplayable && utils.IsStreamable(file.Path) {
unplayable = false
}
})
// Map torrents to directories
switch t.Config.GetVersion() {
case "v1":
configV1 := t.Config.(*config.ZurgConfigV1)
for _, directories := range configV1.GetGroupMap() {
for _, directory := range directories {
if unplayable {
cb(config.UNPLAYABLE_TORRENTS)
break
}
if t.Config.MeetsConditions(directory, t.GetKey(tor), tor.ComputeTotalSize(), torrentIDs, filenames, fileSizes) {
cb(directory)
break
}
}
}
}
}
func (t *TorrentManager) GetKey(torrent *Torrent) string { func (t *TorrentManager) GetKey(torrent *Torrent) string {
if !t.Config.ShouldIgnoreRenames() && torrent.Rename != "" { if !t.Config.ShouldIgnoreRenames() && torrent.Rename != "" {
return torrent.Rename return torrent.Rename
@@ -242,3 +162,38 @@ func (t *TorrentManager) deleteTorrentFile(torrentID string) {
t.log.Warnf("Cannot delete file %s: %v", filePath, err) t.log.Warnf("Cannot delete file %s: %v", filePath, err)
} }
} }
func (t *TorrentManager) mountDownloads() {
if !t.Config.EnableDownloadMount() {
return
}
_ = t.workerPool.Submit(func() {
page := 1
offset := 0
for {
downloads, totalDownloads, err := t.Api.GetDownloads(page, offset)
if err != nil {
t.log.Fatalf("Cannot get downloads: %v", err)
}
for i := range downloads {
t.DownloadMap.Set(downloads[i].Filename, &downloads[i])
}
offset += len(downloads)
page++
if offset >= totalDownloads {
break
}
}
t.log.Infof("Compiled into %d downloads", t.DownloadCache.Count())
})
}
func (t *TorrentManager) initializeDirectories() {
// create internal directories
t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is GetAccessKey()
t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID
// create directory maps
for _, directory := range t.Config.GetDirectories() {
t.DirectoryMap.Set(directory, cmap.New[*Torrent]())
}
}

View File

@@ -9,11 +9,12 @@ import (
"github.com/debridmediamanager/zurg/internal/config" "github.com/debridmediamanager/zurg/internal/config"
"github.com/debridmediamanager/zurg/pkg/realdebrid" "github.com/debridmediamanager/zurg/pkg/realdebrid"
"github.com/debridmediamanager/zurg/pkg/utils"
mapset "github.com/deckarep/golang-set/v2" mapset "github.com/deckarep/golang-set/v2"
cmap "github.com/orcaman/concurrent-map/v2" cmap "github.com/orcaman/concurrent-map/v2"
) )
func (t *TorrentManager) RefreshTorrents() []string { func (t *TorrentManager) refreshTorrents() []string {
instances, _, err := t.Api.GetTorrents(0, false) instances, _, err := t.Api.GetTorrents(0, false)
if err != nil { if err != nil {
t.log.Warnf("Cannot get torrents: %v", err) t.log.Warnf("Cannot get torrents: %v", err)
@@ -120,28 +121,24 @@ func (t *TorrentManager) RefreshTorrents() []string {
// startRefreshJob periodically refreshes the torrents // startRefreshJob periodically refreshes the torrents
func (t *TorrentManager) startRefreshJob() { func (t *TorrentManager) startRefreshJob() {
t.log.Info("Starting periodic refresh") _ = t.workerPool.Submit(func() {
for { t.log.Info("Starting periodic refresh")
<-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second) for {
<-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second)
checksum := t.getCurrentState() checksum := t.getCurrentState()
if t.latestState.equal(checksum) { if t.latestState.equal(checksum) {
continue continue
}
t.SetNewLatestState(checksum)
t.log.Infof("Detected changes! Refreshing %d torrents", checksum.TotalCount)
updatedPaths := t.refreshTorrents()
t.log.Info("Finished refreshing torrents")
t.TriggerHookOnLibraryUpdate(updatedPaths)
} }
t.SetNewLatestState(checksum) })
t.log.Infof("Detected changes! Refreshing %d torrents", checksum.TotalCount)
updatedPaths := t.RefreshTorrents()
t.log.Info("Finished refreshing torrents")
t.TriggerHookOnLibraryUpdate(updatedPaths)
if t.Config.EnableRepair() {
t.RepairAll()
} else {
t.log.Info("Repair is disabled, skipping repair check")
}
}
} }
// getMoreInfo gets original name, size and files for a torrent // getMoreInfo gets original name, size and files for a torrent
@@ -313,3 +310,35 @@ func (t *TorrentManager) mergeToMain(existing, toMerge *Torrent) Torrent {
return mainTorrent return mainTorrent
} }
func (t *TorrentManager) assignedDirectoryCb(tor *Torrent, cb func(string)) {
torrentIDs := tor.DownloadedIDs.Union(tor.InProgressIDs).ToSlice()
// get filenames needed for directory conditions
var filenames []string
var fileSizes []int64
unplayable := true
tor.SelectedFiles.IterCb(func(key string, file *File) {
filenames = append(filenames, filepath.Base(file.Path))
fileSizes = append(fileSizes, file.Bytes)
if unplayable && utils.IsStreamable(file.Path) {
unplayable = false
}
})
// Map torrents to directories
switch t.Config.GetVersion() {
case "v1":
configV1 := t.Config.(*config.ZurgConfigV1)
for _, directories := range configV1.GetGroupMap() {
for _, directory := range directories {
if unplayable {
cb(config.UNPLAYABLE_TORRENTS)
break
}
if t.Config.MeetsConditions(directory, t.GetKey(tor), tor.ComputeTotalSize(), torrentIDs, filenames, fileSizes) {
cb(directory)
break
}
}
}
}
}

View File

@@ -16,15 +16,18 @@ const (
EXPIRED_LINK_TOLERANCE_HOURS = 24 EXPIRED_LINK_TOLERANCE_HOURS = 24
) )
func (t *TorrentManager) RepairAll() { func (t *TorrentManager) startRepairJob() {
if !t.Config.EnableRepair() {
t.log.Info("Repair is disabled, skipping repair job")
}
// there is 1 repair worker, with max 1 blocking task // there is 1 repair worker, with max 1 blocking task
_ = t.repairPool.Submit(func() { _ = t.repairPool.Submit(func() {
t.log.Info("Periodic repair invoked; searching for broken torrents")
t.repairAll() t.repairAll()
}) })
} }
func (t *TorrentManager) repairAll() { func (t *TorrentManager) repairAll() {
t.log.Info("Periodic repair invoked; searching for broken torrents")
allTorrents, _ := t.DirectoryMap.Get(INT_ALL) allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
// collect all torrents that need to be repaired // collect all torrents that need to be repaired