From 5cb62347e7d03b17518ff68913c48099902cb6e3 Mon Sep 17 00:00:00 2001
From: Tony Blyler <tony@blyler.cc>
Date: Sat, 23 Jul 2016 17:23:36 -0400
Subject: [PATCH] Add ability to prevent downloads based off of available disk
 space

---
 .gitmodules                             |  3 ++
 README.md                               |  6 ++++
 queue/queue.go                          | 47 +++++++++++++++++++++++++
 queue/vendor/github.com/shirou/gopsutil |  1 +
 4 files changed, 57 insertions(+)
 create mode 160000 queue/vendor/github.com/shirou/gopsutil

diff --git a/.gitmodules b/.gitmodules
index 29a3177..edeb42d 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -31,3 +31,6 @@
 [submodule "cmd/hoarder/vendor/github.com/go-yaml/yaml"]
 	path = cmd/hoarder/vendor/github.com/go-yaml/yaml
 	url = git://github.com/go-yaml/yaml
+[submodule "queue/vendor/github.com/shirou/gopsutil"]
+	path = queue/vendor/github.com/shirou/gopsutil
+	url = https://github.com/shirou/gopsutil
diff --git a/README.md b/README.md
index 15ce559..0904383 100644
--- a/README.md
+++ b/README.md
@@ -46,6 +46,12 @@ resume_downloads: true
 # Path to the unix socket file that hoarder uses for RPC
 rpc_socket_path: /tmp/hoarder.sock
 
+# Whether or not to see if there is enough disk space before starting a download
+check_disk_space: true
+
+# (must have check_disk_space set to true) Minimum disk space to have after completed downloads (measured in bytes, 0 to disable check)
+min_disk_space: 5368709120
+
 rtorrent:
   # The address to the rtorrent XMLRPC endpoint
   addr: https://mycoolrtorrentserver.com/XMLRPC
diff --git a/queue/queue.go b/queue/queue.go
index ad71b07..0adef27 100644
--- a/queue/queue.go
+++ b/queue/queue.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/fsnotify/fsnotify"
+	"github.com/shirou/gopsutil/disk"
 	"github.com/tblyler/easysftp"
 	"github.com/tblyler/go-rtorrent/rtorrent"
 	"github.com/tblyler/hoarder/metainfo"
@@ -45,6 +46,8 @@ type Config struct {
 	TorrentListUpdateInterval time.Duration     `json:"rtorrent_update_interval" yaml:"rtorrent_update_interval"`
 	ConcurrentDownloads       uint              `json:"download_jobs" yaml:"download_jobs"`
 	ResumeDownloads           bool              `json:"resume_downloads" yaml:"resume_downloads"`
+	CheckDiskSpace            bool              `json:"check_disk_space" yaml:"check_disk_space"`
+	MinDiskSpace              uint64            `json:"min_disk_space" yaml:"min_disk_space"`
 }
 
 // Queue watches the given folders for new .torrent files,
@@ -583,6 +586,50 @@ func (q *Queue) Run(stop <-chan bool) {
 
 				torrentFilePath := q.downloadQueue[torrent.Hash]
 
+				skip := false
+				if q.config.CheckDiskSpace {
+					diskSpacePaths := []string{q.config.WatchDownloadPaths[filepath.Dir(torrentFilePath)]}
+
+					if q.config.TempDownloadPath != "" {
+						diskSpacePaths = append(diskSpacePaths, q.config.TempDownloadPath)
+					}
+
+					downloadSizes := uint64(torrent.Size)
+					for _, dTorrent := range downloadsRunning {
+						downloadSizes += uint64(dTorrent.size)
+					}
+
+					for _, path := range diskSpacePaths {
+						fsStat, err := disk.Usage(path)
+						if err != nil {
+							q.logger.Printf("Failed to check disk space on '%s' for '%s' (%s): %s", path, torrent.Name, torrentFilePath, err)
+							continue
+						}
+
+						if q.config.MinDiskSpace == 0 {
+							if fsStat.Free > downloadSizes {
+								continue
+							}
+
+							q.logger.Printf("Not downloading '%s' (%s) not enough disk space, only %d bytes free on '%s'", torrent.Name, torrentFilePath, fsStat.Free, path)
+							skip = true
+							break
+						} else {
+							if fsStat.Free > downloadSizes && (fsStat.Free-downloadSizes) > q.config.MinDiskSpace {
+								continue
+							}
+
+							q.logger.Printf("Not downloading '%s' (%s) minimum disk space (%d) reached on '%s'", torrent.Name, torrentFilePath, q.config.MinDiskSpace, path)
+							skip = true
+							break
+						}
+					}
+				}
+
+				if skip {
+					continue
+				}
+
 				go func(torrent rtorrent.Torrent, torrentPath string, hashChan chan<- string) {
 					err := q.downloadTorrent(torrent, torrentPath)
 					if err != nil {
diff --git a/queue/vendor/github.com/shirou/gopsutil b/queue/vendor/github.com/shirou/gopsutil
new file mode 160000
index 0000000..ee66bc5
--- /dev/null
+++ b/queue/vendor/github.com/shirou/gopsutil
@@ -0,0 +1 @@
+Subproject commit ee66bc560c366dd33b9a4046ba0b644caba46bed