tblyler/hoarder
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Handles SSH/SFTP connection closing gracefully. Adds auto restarting of processes config option. Makes dying a little more uniform. Still need to make dying graceful.

Browse source
This commit is contained in:
Tony Blyler 2015-01-24 11:45:28 -05:00
parent d36b7ff8d6
commit 758ce5cb8e
3 changed files with 79 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
hoarder.conf.ex
View file

@ -1,6 +1,8 @@
// Configuration for Hoarder
// All fields and values are necessary
// Should Hoarder automatically restart when one of its processes errors? (default true if not set)
restart_on_error: "true"
// Username for XMLRPC
xml_user: "testuser"
// Password for XMLRPC

53
hoarder.go
View file

@ -43,14 +43,14 @@ func loadConfig(configPath string) (map[string]string, error) {
continue
}
config[line[:sepPosition]] = line[sepPosition + 3:len(line) - 1]
config[line[:sepPosition]] = line[sepPosition+3 : len(line)-1]
}
return config, nil
}
// Checker routine to see if torrents are completed
func checker(config map[string]string, checkerChan <- chan map[string]string, com chan <- error) error {
func checker(config map[string]string, checkerChan <-chan map[string]string, com chan<- error) error {
for {
torrentInfo := <-checkerChan
@ -66,6 +66,7 @@ func checker(config map[string]string, checkerChan <- chan map[string]string, co
}
syncer, err := NewSync(config["threads"], config["ssh_user"], config["ssh_pass"], config["ssh_server"], config["ssh_port"])
defer syncer.Close()
if err != nil {
log.Println("Failed to create a new sync: " + err.Error())
com <- err
@ -142,6 +143,8 @@ func checker(config map[string]string, checkerChan <- chan map[string]string, co
} else {
log.Println(name + " is not completed, waiting for it to finish")
}
syncer.Close()
}
com <- nil
@ -149,7 +152,7 @@ func checker(config map[string]string, checkerChan <- chan map[string]string, co
}
// Scanner routine to see if there are new torrent_files
func scanner(config map[string]string, checkerChan chan <- map[string]string, com chan <- error) error {
func scanner(config map[string]string, checkerChan chan<- map[string]string, com chan<- error) error {
watchDirs := map[string]string{config["local_torrent_dir"]: config["local_download_dir"]}
dirContents, err := ioutil.ReadDir(config["local_torrent_dir"])
@ -186,6 +189,7 @@ func scanner(config map[string]string, checkerChan chan <- map[string]string, co
syncer, err := NewSync("1", config["ssh_user"], config["ssh_pass"], config["ssh_server"], config["ssh_port"])
if err != nil {
log.Println("Failed to create a new sync: " + err.Error())
syncer.Close()
continue
}
@ -193,6 +197,7 @@ func scanner(config map[string]string, checkerChan chan <- map[string]string, co
exists, err := syncer.Exists(destinationTorrent)
if err != nil {
log.Println("Failed to see if " + torrentPath + " already exists on the server: " + err.Error())
syncer.Close()
continue
}
@ -207,12 +212,15 @@ func scanner(config map[string]string, checkerChan chan <- map[string]string, co
log.Println("Failed to upload " + torrentPath + " to " + destinationTorrent + ": " + err.Error())
}
syncer.Close()
continue
}
syncer.Close()
}
downloadInfo := map[string]string{
"torrent_path": torrentPath,
"torrent_path": torrentPath,
"local_download_dir": downloadDir,
}
@ -239,10 +247,14 @@ func scanner(config map[string]string, checkerChan chan <- map[string]string, co
return nil
}
func die(exitCode int) {
log.Println("Quiting")
os.Exit(exitCode)
}
func main() {
sigint.ListenForSIGINT(func() {
log.Println("Quiting")
os.Exit(1)
die(1)
})
var configPath string
@ -252,14 +264,14 @@ func main() {
if configPath == "" {
log.Println("Missing argument for configuration file path")
flag.PrintDefaults()
os.Exit(1)
die(1)
}
log.Println("Reading configuration file")
config, err := loadConfig(configPath)
if err != nil {
log.Println(err)
os.Exit(1)
die(1)
}
log.Println("Successfully read configuration file")
@ -268,7 +280,7 @@ func main() {
if err != nil {
log.Println(err)
os.Exit(1)
die(1)
}
log.Println("Starting the scanner routine")
@ -279,17 +291,36 @@ func main() {
checkerCom := make(chan error)
go checker(config, checkerChan, checkerCom)
restartOnError := true
if config["restart_on_error"] != "" {
restartOnError = config["restart_on_error"] == "true"
}
for {
select {
case err := <-scannerCom:
if err != nil {
log.Println("Scanner failed: " + err.Error())
os.Exit(1)
if restartOnError {
log.Println("Restarting scanner")
go scanner(config, checkerChan, scannerCom)
} else {
log.Println("Quiting due to scanner error")
die(1)
}
}
case err := <-checkerCom:
if err != nil {
log.Println("Checker failed: " + err.Error())
os.Exit(1)
if restartOnError {
log.Println("Restarting checker")
go checker(config, checkerChan, checkerCom)
} else {
log.Println("Quiting due to checker error")
die(1)
}
}
default:
break

52
sync.go
View file

@ -22,8 +22,8 @@ const (
// Sync Keeps track of the connection information.
type Sync struct {
sshClient *ssh.Client
sftpClients []*sftp.Client
sshClient *ssh.Client
sftpClients []*sftp.Client
sftpClientCount int
}
@ -58,6 +58,24 @@ func NewSync(threads string, user string, pass string, server string, port strin
return &Sync{sshClient, sftpClients, clientCount}, nil
}
// Close Closes all of the ssh and sftp connections to the SSH server.
func (s Sync) Close() error {
var returnError error
for i := 0; i < s.sftpClientCount; i++ {
err := s.sftpClients[i].Close()
if err != nil {
returnError = err
}
}
err := s.sshClient.Close()
if err != nil {
return err
}
return returnError
}
// Create a new SSH client instance and confirm that we can make sessions
func newSSHClient(user string, pass string, server string, port string) (*ssh.Client, error) {
sshConfig := &ssh.ClientConfig{
@ -67,7 +85,7 @@ func newSSHClient(user string, pass string, server string, port string) (*ssh.Cl
},
}
sshClient, err := ssh.Dial("tcp", server + ":" + port, sshConfig)
sshClient, err := ssh.Dial("tcp", server+":"+port, sshConfig)
if err != nil {
return sshClient, err
@ -121,8 +139,8 @@ func SendFiles(sftpClient *sftp.Client, files map[string]string) error {
var currentPosition int64
for currentPosition < sourceSize {
// If the next iteration will be greater than the file size, reduce to the data size
if currentPosition + int64(len(data)) > sourceSize {
data = make([]byte, sourceSize - currentPosition)
if currentPosition+int64(len(data)) > sourceSize {
data = make([]byte, sourceSize-currentPosition)
}
// Read data from the source file
@ -176,7 +194,7 @@ func (s Sync) GetFile(sourceFile string, destinationFile string) error {
// Start the concurrent downloads
for i := 0; i < s.sftpClientCount; i++ {
go GetFile(s.sftpClients[i], sourceFile, destinationFile, i + 1, s.sftpClientCount, channels[i])
go GetFile(s.sftpClients[i], sourceFile, destinationFile, i+1, s.sftpClientCount, channels[i])
}
// Block until all downloads are completed or one errors
@ -189,7 +207,7 @@ func (s Sync) GetFile(sourceFile string, destinationFile string) error {
}
select {
case err := <- channel:
case err := <-channel:
if err != nil {
return err
}
@ -218,7 +236,7 @@ func (s Sync) GetFile(sourceFile string, destinationFile string) error {
// GetFile Get a file from the source_file path to be stored in destination_file path.
// worker_number and work_total are not zero indexed, but 1 indexed
func GetFile(sftpClient *sftp.Client, sourceFile string, destinationFile string, workerNumber int, workerTotal int, com chan <- error) error {
func GetFile(sftpClient *sftp.Client, sourceFile string, destinationFile string, workerNumber int, workerTotal int, com chan<- error) error {
// Open source_data for reading
sourceData, err := sftpClient.OpenFile(sourceFile, os.O_RDONLY)
if err != nil {
@ -247,7 +265,7 @@ func GetFile(sftpClient *sftp.Client, sourceFile string, destinationFile string,
if workerNumber == 1 {
start = 0
} else {
start = (statSize * int64(workerNumber - 1)) / int64(workerTotal)
start = (statSize * int64(workerNumber-1)) / int64(workerTotal)
}
}
@ -260,7 +278,7 @@ func GetFile(sftpClient *sftp.Client, sourceFile string, destinationFile string,
}
// Create the new file for writing
newFile, err := os.OpenFile(destinationFile, os.O_WRONLY | os.O_CREATE, 0777)
newFile, err := os.OpenFile(destinationFile, os.O_WRONLY|os.O_CREATE, 0777)
if err != nil {
com <- err
return err
@ -313,7 +331,7 @@ func GetFile(sftpClient *sftp.Client, sourceFile string, destinationFile string,
}
// Adjust the size of the buffer if the next iteration will be greater than what has yet to be read
if currentSize + dataSize > stop {
if currentSize+dataSize > stop {
dataSize = stop - currentSize
data = make([]byte, dataSize)
}
@ -369,8 +387,8 @@ func (s Sync) GetPath(sourcePath string, destinationPath string) error {
dirs, files, err := s.getChildren(sourcePath)
// Remove the trailing slash if it exists
if sourcePath[len(sourcePath) - 1] == '/' {
sourcePath = sourcePath[:len(sourcePath) - 1]
if sourcePath[len(sourcePath)-1] == '/' {
sourcePath = sourcePath[:len(sourcePath)-1]
}
// Get the parent path of source_path
@ -467,7 +485,7 @@ func getProgress(filePath string, workerNumber int) (int64, error) {
var progress int64
progressSize := int64(binary.Size(progress))
offset := progressSize * int64(workerNumber - 1)
offset := progressSize * int64(workerNumber-1)
realOffset, err := file.Seek(offset, os.SEEK_SET)
if err != nil {
@ -502,13 +520,13 @@ func getProgress(filePath string, workerNumber int) (int64, error) {
}
func updateProgress(filePath string, written int64, workerNumber int) error {
file, err := os.OpenFile(getProgressPath(filePath), os.O_WRONLY | os.O_CREATE, 0777)
file, err := os.OpenFile(getProgressPath(filePath), os.O_WRONLY|os.O_CREATE, 0777)
if err != nil {
return err
}
writtenSize := int64(binary.Size(written))
offset := writtenSize * int64(workerNumber - 1)
offset := writtenSize * int64(workerNumber-1)
realOffset, err := file.Seek(offset, os.SEEK_SET)
if err != nil {
@ -532,5 +550,5 @@ func destroyProgress(filePath string) error {
}
func getProgressPath(filePath string) string {
return filepath.Join(filepath.Dir(filePath), "." + filepath.Base(filePath) + ".progress")
return filepath.Join(filepath.Dir(filePath), "."+filepath.Base(filePath)+".progress")
}