Refactor workers

2024-01-27 13:17:45 +01:00
parent e01622064d
commit 05d2544fe8
6 changed files with 127 additions and 171 deletions
--- 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)
--- 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(),
--- 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
 }
--- 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]())
+	}
+}
--- 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
+				}
+			}
+		}
+	}
+}
--- 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