From 05d2544fe85d306deef646a057f0a522bd0e096d Mon Sep 17 00:00:00 2001 From: Ben Sarmiento Date: Sat, 27 Jan 2024 13:17:45 +0100 Subject: [PATCH] Refactor workers --- internal/app.go | 1 + internal/torrent/hooks.go | 7 ++ internal/torrent/latestState.go | 87 ++++++---------------- internal/torrent/manager.go | 125 ++++++++++---------------------- internal/torrent/refresh.go | 71 ++++++++++++------ internal/torrent/repair.go | 7 +- 6 files changed, 127 insertions(+), 171 deletions(-) diff --git a/internal/app.go b/internal/app.go index 0409b3e..ead8a47 100644 --- a/internal/app.go +++ b/internal/app.go @@ -45,6 +45,7 @@ func MainApp(configPath string) { premium.MonitorPremiumStatus(rd, zurglog) + // extra 1 worker for the refresh job workerPool, err := ants.NewPool(config.GetNumOfWorkers() + 1) if err != nil { zurglog.Errorf("Failed to create worker pool: %v", err) diff --git a/internal/torrent/hooks.go b/internal/torrent/hooks.go index e7105df..12a43f7 100644 --- a/internal/torrent/hooks.go +++ b/internal/torrent/hooks.go @@ -41,6 +41,13 @@ func (se *ScriptExecutor) Execute() (string, error) { 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) { executor := &ScriptExecutor{ Script: config.GetOnLibraryUpdate(), diff --git a/internal/torrent/latestState.go b/internal/torrent/latestState.go index b7dead8..684bea9 100644 --- a/internal/torrent/latestState.go +++ b/internal/torrent/latestState.go @@ -37,74 +37,35 @@ func (t *TorrentManager) SetNewLatestState(checksum LibraryState) { 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 func (t *TorrentManager) getCurrentState() LibraryState { - torrentChan := make(chan torrentsResp, 1) - 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) + var state LibraryState - _ = t.workerPool.Submit(func() { - torrents, totalCount, err := t.Api.GetTorrents(1, false) - if err != nil { - errChan <- err - return - } - torrentChan <- torrentsResp{torrents: torrents, totalCount: totalCount} - }) + torrents, totalCount, err := t.Api.GetTorrents(1, false) + if err != nil { + t.log.Warnf("Checksum API Error (GetTorrents): %v", err) + return LibraryState{} + } + if len(torrents) > 0 { + state.FirstTorrent = &torrents[0] + } + state.TotalCount = totalCount - _ = t.workerPool.Submit(func() { - torrents, totalCount, err := t.Api.GetTorrents(1, true) - if err != nil { - errChan <- err - return - } - activeChan <- torrentsResp{torrents: torrents, totalCount: totalCount} - }) - - _ = 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{} - } + activeTorrents, _, err := t.Api.GetTorrents(1, true) + if err != nil { + t.log.Warnf("Checksum API Error (GetTorrents): %v", err) + return LibraryState{} + } + if len(activeTorrents) > 0 { + state.FirstActiveTorrent = &activeTorrents[0] } - return LibraryState{ - TotalCount: totalCount, - ActiveCount: count, - FirstTorrent: first, - FirstActiveTorrent: active, + count, err := t.Api.GetActiveTorrentCount() + if err != nil { + t.log.Warnf("Checksum API Error (GetActiveTorrentCount): %v", err) + return LibraryState{} } + state.ActiveCount = count.DownloadingCount + + return state } diff --git a/internal/torrent/manager.go b/internal/torrent/manager.go index 7342cdd..208f9db 100644 --- a/internal/torrent/manager.go +++ b/internal/torrent/manager.go @@ -3,13 +3,11 @@ package torrent import ( "io" "os" - "path/filepath" "strings" "github.com/debridmediamanager/zurg/internal/config" "github.com/debridmediamanager/zurg/pkg/logutil" "github.com/debridmediamanager/zurg/pkg/realdebrid" - "github.com/debridmediamanager/zurg/pkg/utils" mapset "github.com/deckarep/golang-set/v2" cmap "github.com/orcaman/concurrent-map/v2" "github.com/panjf2000/ants/v2" @@ -58,51 +56,12 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, w log: log, } - // 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 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.initializeDirectories() + t.mountDownloads() + t.refreshTorrents() t.SetNewLatestState(t.getCurrentState()) - - if t.Config.EnableRepair() { - 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() - }) + t.startRefreshJob() + t.startRepairJob() return t } @@ -127,45 +86,6 @@ func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download { 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 { if !t.Config.ShouldIgnoreRenames() && 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) } } + +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]()) + } +} diff --git a/internal/torrent/refresh.go b/internal/torrent/refresh.go index a594e14..b28291a 100644 --- a/internal/torrent/refresh.go +++ b/internal/torrent/refresh.go @@ -9,11 +9,12 @@ import ( "github.com/debridmediamanager/zurg/internal/config" "github.com/debridmediamanager/zurg/pkg/realdebrid" + "github.com/debridmediamanager/zurg/pkg/utils" mapset "github.com/deckarep/golang-set/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) if err != nil { t.log.Warnf("Cannot get torrents: %v", err) @@ -120,28 +121,24 @@ func (t *TorrentManager) RefreshTorrents() []string { // startRefreshJob periodically refreshes the torrents func (t *TorrentManager) startRefreshJob() { - t.log.Info("Starting periodic refresh") - for { - <-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second) + _ = t.workerPool.Submit(func() { + t.log.Info("Starting periodic refresh") + for { + <-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second) - checksum := t.getCurrentState() - if t.latestState.equal(checksum) { - continue + checksum := t.getCurrentState() + if t.latestState.equal(checksum) { + 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 @@ -313,3 +310,35 @@ func (t *TorrentManager) mergeToMain(existing, toMerge *Torrent) Torrent { 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 + } + } + } + } +} diff --git a/internal/torrent/repair.go b/internal/torrent/repair.go index 563c2c6..4c05acc 100644 --- a/internal/torrent/repair.go +++ b/internal/torrent/repair.go @@ -16,15 +16,18 @@ const ( 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 _ = t.repairPool.Submit(func() { - t.log.Info("Periodic repair invoked; searching for broken torrents") t.repairAll() }) } func (t *TorrentManager) repairAll() { + t.log.Info("Periodic repair invoked; searching for broken torrents") allTorrents, _ := t.DirectoryMap.Get(INT_ALL) // collect all torrents that need to be repaired