From 7d8cbe1393c73db42697e18867276511c6dfb6f5 Mon Sep 17 00:00:00 2001
From: ElegantMonkey <ramonzitosdesorte@hotmail.com>
Date: Sun, 14 May 2017 11:35:03 -0300
Subject: [PATCH] Add Upload and Download limiters, rely on unique_index for
 avoiding duplicates

The FileList that is read with GetTorrents is empty, as it's only loaded with
GetTorrentById. So, always insert the new FileList on the database, and rely on
the unique_index to avoid duplications.
---
 config/metainfoFetcher.go                     | 12 +++--
 model/file.go                                 |  4 +-
 .../metainfoFetcher/metainfoFetcher.go        | 44 +++++++++++--------
 3 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/config/metainfoFetcher.go b/config/metainfoFetcher.go
index 792642f8..eceb37a1 100644
--- a/config/metainfoFetcher.go
+++ b/config/metainfoFetcher.go
@@ -5,12 +5,18 @@ type MetainfoFetcherConfig struct {
 	Timeout        int `json:"timeout"`
 	MaxDays        int `json:"max_days"`
 	WakeUpInterval int `json:"wake_up_interval"`
+
+	UploadRateLimiter   int `json:"upload_rate_limiter"`
+	DownloadRateLimiter int `json:"download_rate_limiter"`
 }
 
 var DefaultMetainfoFetcherConfig = MetainfoFetcherConfig{
-	QueueSize: 10,
-	Timeout: 120, // 2 min
-	MaxDays: 90,
+	QueueSize:      10,
+	Timeout:        120, // 2 min
+	MaxDays:        90,
 	WakeUpInterval: 300, // 5 min
+
+	UploadRateLimiter:   1024,
+	DownloadRateLimiter: 1024,
 }
 
diff --git a/model/file.go b/model/file.go
index 3f27e944..67b64168 100644
--- a/model/file.go
+++ b/model/file.go
@@ -2,8 +2,8 @@ package model
 
 type File struct {
 	ID        uint   `gorm:"column:file_id;primary_key"`
-	TorrentID uint   `gorm:"column:torrent_id"`
-	Path      string `gorm:"column:path"`
+	TorrentID uint   `gorm:"column:torrent_id;unique_index:idx_tid_path"`
+	Path      string `gorm:"column:path;unique_index:idx_tid_path"`
 	Filesize  int64  `gorm:"column:filesize"`
 }
 
diff --git a/service/torrent/metainfoFetcher/metainfoFetcher.go b/service/torrent/metainfoFetcher/metainfoFetcher.go
index 95e97103..e05263e5 100644
--- a/service/torrent/metainfoFetcher/metainfoFetcher.go
+++ b/service/torrent/metainfoFetcher/metainfoFetcher.go
@@ -9,6 +9,7 @@ import (
 	"github.com/ewhal/nyaa/util/log"
 	serviceBase "github.com/ewhal/nyaa/service"
 	torrentService "github.com/ewhal/nyaa/service/torrent"
+	"golang.org/x/time/rate"
 	"sync"
 	"time"
 )
@@ -28,7 +29,18 @@ type MetainfoFetcher struct {
 }
 
 func New(fetcherConfig *config.MetainfoFetcherConfig) (fetcher *MetainfoFetcher, err error) {
-	client, err := torrent.NewClient(nil)
+	clientConfig := torrent.Config{}
+	// Well, it seems this is the right way to convert speed -> rate.Limiter
+	// https://github.com/anacrolix/torrent/blob/master/cmd/torrent/main.go
+	if fetcherConfig.UploadRateLimiter != -1 {
+		clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(fetcherConfig.UploadRateLimiter * 1024), 256<<10)
+	}
+	if fetcherConfig.DownloadRateLimiter != -1 {
+		clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(fetcherConfig.DownloadRateLimiter * 1024), 1<<20)
+	}
+
+	client, err := torrent.NewClient(&clientConfig)
+
 	fetcher = &MetainfoFetcher{
 		torrentClient:    client,
 		results:          make(chan Result, fetcherConfig.QueueSize),
@@ -86,26 +98,19 @@ func updateFileList(dbEntry model.Torrent, info *metainfo.Info) error {
 	log.Infof("TID %d has %d files.", dbEntry.ID, len(torrentFiles))
 	for _, file := range torrentFiles {
 		path := file.DisplayPath(info)
-		fileExists := false
-		for _, existingFile := range dbEntry.FileList {
-			if existingFile.Path == path {
-				fileExists = true
-				break
-			}
+
+		// Can't read FileList from the GetTorrents output, rely on the unique_index
+		// to ensure no files are duplicated.
+		log.Infof("Adding file %s to filelist of TID %d", path, dbEntry.ID)
+		dbFile := model.File{
+			TorrentID: dbEntry.ID,
+			Path: path,
+			Filesize: file.Length,
 		}
 
-		if !fileExists {
-			log.Infof("Adding file %s to filelist of TID %d", path, dbEntry.ID)
-			dbFile := model.File{
-				TorrentID: dbEntry.ID,
-				Path: path,
-				Filesize: file.Length,
-			}
-
-			err := db.ORM.Create(&dbFile).Error
-			if err != nil {
-				return err
-			}
+		err := db.ORM.Create(&dbFile).Error
+		if err != nil {
+			return err
 		}
 	}
 
@@ -223,6 +228,7 @@ func (fetcher *MetainfoFetcher) Close() error {
 	}
 
 	fetcher.done <- 1
+	log.Infof("Send done signal to everyone, waiting...")
 	fetcher.wg.Wait()
 	return nil
 }