hpz937/zurg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use worker pool extensively

Browse Source
This commit is contained in:
Ben Sarmiento
2023-11-30 00:40:26 +01:00
parent 6e54fa760b
commit 9e3760f275
5 changed files with 97 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
internal/torrent/manager.go
View File

@@ -25,7 +25,6 @@ import (
const (
INT_ALL = "int__all__"
INT_INFO_CACHE = "int__info__"
DATA_DIR = "data"
)
type TorrentManager struct {
@@ -34,10 +33,9 @@ type TorrentManager struct {
DirectoryMap cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
DownloadCache cmap.ConcurrentMap[string, *realdebrid.Download]
ResponseCache *ristretto.Cache
checksum string
latestAdded string
latestState *LibraryState
requiredVersion string
antsPool *ants.Pool
workerPool *ants.Pool
unrestrictPool *ants.Pool
log *zap.SugaredLogger
}
@@ -46,19 +44,23 @@ type TorrentManager struct {
// 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, cache *ristretto.Cache, log *zap.SugaredLogger) *TorrentManager {
initialSate := EmptyState()
t := &TorrentManager{
Config: cfg,
Api: api,
DirectoryMap: cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
ResponseCache: cache,
latestState: &initialSate,
requiredVersion: "18.11.2023",
antsPool: p,
workerPool: p,
log: log,
}
// create unrestrict pool
unrestrictPool, err := ants.NewPool(t.Config.GetUnrestrictWorkers())
if err != nil {
t.unrestrictPool = t.antsPool
t.unrestrictPool = t.workerPool
} else {
t.unrestrictPool = unrestrictPool
}
@@ -66,17 +68,16 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
// 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]())
}
var initWait sync.WaitGroup
initWait.Add(2)
// Fetch downloads
go func() {
initWait.Add(1)
_ = t.workerPool.Submit(func() {
defer initWait.Done()
downloads, _, err := t.Api.GetDownloads()
if err != nil {
@@ -88,17 +89,17 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
t.DownloadCache.Set(downloads[i].Link, &downloads[i])
}
}
}()
})
// Fetch torrents
var newTorrents []realdebrid.Torrent
go func() {
initWait.Add(1)
_ = t.workerPool.Submit(func() {
defer initWait.Done()
newTorrents, _, err = t.Api.GetTorrents(0)
if err != nil {
t.log.Fatalf("Cannot get torrents: %v\n", err)
}
}()
})
initWait.Wait()
@@ -108,12 +109,11 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
var wg sync.WaitGroup
for i := range newTorrents {
wg.Add(1)
go func(idx int) {
_ = t.antsPool.Submit(func() {
defer wg.Done()
torrentsChan <- t.getMoreInfo(newTorrents[idx])
})
}(i)
idx := i // capture the loop variable
_ = t.workerPool.Submit(func() {
defer wg.Done()
torrentsChan <- t.getMoreInfo(newTorrents[idx])
})
}
wg.Wait()
close(torrentsChan)
@@ -164,16 +164,18 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
t.log.Infof("Compiled into %d torrents, %d were missing info", allTorrents.Count(), noInfoCount)
t.SetChecksum(t.getChecksum())
t.SetNewLatestState(t.getCurrentState())
if t.Config.EnableRepair() {
t.log.Info("Checking for torrents to repair")
t.repairAll()
t.log.Info("Finished checking for torrents to repair")
}
go t.startRefreshJob()
t.latestAdded = newTorrents[0].Added // set the latest added to the first torrent's added
_ = t.workerPool.Submit(func() {
t.startRefreshJob()
})
t.log.Info("Finished initializing torrent manager")
return t
@@ -219,65 +221,67 @@ func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download {
// return t.api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
}
type torrentsResponse struct {
func (t *TorrentManager) SetNewLatestState(checksum LibraryState) {
t.latestState.DownloadingCount = checksum.DownloadingCount
t.latestState.FirstTorrent = checksum.FirstTorrent
t.latestState.TotalCount = checksum.TotalCount
}
type torrentsResp struct {
torrents []realdebrid.Torrent
totalCount int
}
func (t *TorrentManager) SetChecksum(checksum string) {
// t.mu.Lock()
t.checksum = checksum
// t.mu.Unlock()
}
// generates a checksum based on the number of torrents, the first torrent id and the number of active torrents
func (t *TorrentManager) getChecksum() string {
torrentsChan := make(chan torrentsResponse, 1)
func (t *TorrentManager) getCurrentState() LibraryState {
torrentsChan := make(chan torrentsResp, 1)
countChan := make(chan int, 1)
errChan := make(chan error, 2) // accommodate errors from both goroutines
// GetTorrents request
go func() {
_ = t.workerPool.Submit(func() {
torrents, totalCount, err := t.Api.GetTorrents(1)
if err != nil {
errChan <- err
return
}
torrentsChan <- torrentsResponse{torrents: torrents, totalCount: totalCount}
}()
torrentsChan <- torrentsResp{torrents: torrents, totalCount: totalCount}
})
// GetActiveTorrentCount request
go func() {
_ = t.workerPool.Submit(func() {
count, err := t.Api.GetActiveTorrentCount()
if err != nil {
errChan <- err
return
}
countChan <- count.DownloadingCount
}()
})
// Existing goroutines for GetTorrents and GetActiveTorrentCount
var torrents []realdebrid.Torrent
var totalCount, count int
for i := 0; i < 2; i++ {
select {
case torrentsResp := <-torrentsChan:
torrents = torrentsResp.torrents
totalCount = torrentsResp.totalCount
case resp := <-torrentsChan:
torrents = resp.torrents
totalCount = resp.totalCount
case count = <-countChan:
case err := <-errChan:
t.log.Warnf("Checksum API Error: %v\n", err)
return ""
return EmptyState()
}
}
if len(torrents) == 0 {
t.log.Error("Huh, no torrents returned")
return ""
return EmptyState()
}
checksum := fmt.Sprintf("%d%s%d", totalCount, torrents[0].ID, count)
return checksum
return LibraryState{
TotalCount: totalCount,
FirstTorrent: &torrents[0],
DownloadingCount: count,
}
}
// startRefreshJob periodically refreshes the torrents
@@ -286,8 +290,8 @@ func (t *TorrentManager) startRefreshJob() {
for {
<-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second)
checksum := t.getChecksum()
if checksum == t.checksum {
checksum := t.getCurrentState()
if t.latestState.equal(checksum) {
continue
}
@@ -319,12 +323,11 @@ func (t *TorrentManager) startRefreshJob() {
var wg sync.WaitGroup
for i := range newTorrents {
wg.Add(1)
go func(idx int) {
_ = t.antsPool.Submit(func() {
defer wg.Done()
torrentsChan <- t.getMoreInfo(newTorrents[idx])
})
}(i)
idx := i // capture the loop variable
_ = t.workerPool.Submit(func() {
defer wg.Done()
torrentsChan <- t.getMoreInfo(newTorrents[idx])
})
}
wg.Wait()
close(torrentsChan)
@@ -368,7 +371,7 @@ func (t *TorrentManager) startRefreshJob() {
if t.Config.MeetsConditions(directory, torrent.AccessKey, torrentIDs, filenames) {
torrents, _ := t.DirectoryMap.Get(directory)
torrents.Set(torrent.AccessKey, torrent)
if torrent.LatestAdded > t.latestAdded {
if torrent.LatestAdded > t.latestState.FirstTorrent.Added {
updatedPaths = append(updatedPaths, fmt.Sprintf("%s/%s", directory, torrent.AccessKey))
}
break
@@ -392,7 +395,7 @@ func (t *TorrentManager) startRefreshJob() {
t.log.Infof("Compiled into %d torrents, %d were missing info", oldTorrents.Count(), noInfoCount)
t.SetChecksum(t.getChecksum())
t.SetNewLatestState(t.getCurrentState())
if t.Config.EnableRepair() {
t.log.Info("Checking for torrents to repair")
@@ -401,9 +404,10 @@ func (t *TorrentManager) startRefreshJob() {
} else {
t.log.Info("Repair is disabled, skipping repair check")
}
go OnLibraryUpdateHook(updatedPaths, t.Config, t.log)
_ = t.workerPool.Submit(func() {
OnLibraryUpdateHook(updatedPaths, t.Config, t.log)
})
t.latestAdded = newTorrents[0].Added
t.log.Info("Finished refreshing torrents")
}
}
@@ -494,7 +498,7 @@ func (t *TorrentManager) getName(name, originalName string) string {
}
func (t *TorrentManager) writeTorrentToFile(torrent *realdebrid.TorrentInfo) error {
filePath := DATA_DIR + "/" + torrent.ID + ".bin"
filePath := "data/" + torrent.ID + ".bin"
file, err := os.Create(filePath)
if err != nil {
return fmt.Errorf("failed creating file: %w", err)
@@ -513,7 +517,7 @@ func (t *TorrentManager) writeTorrentToFile(torrent *realdebrid.TorrentInfo) err
}
func (t *TorrentManager) readTorrentFromFile(torrentID string) *realdebrid.TorrentInfo {
filePath := DATA_DIR + "/" + torrentID + ".bin"
filePath := "data/" + torrentID + ".bin"
file, err := os.Open(filePath)
if err != nil {
if os.IsNotExist(err) {