From dbd787aa81d604fbe8ef0310807f131d9d176190 Mon Sep 17 00:00:00 2001
From: Richard Patel <terorie@alphakevin.club>
Date: Sun, 3 Feb 2019 17:09:43 +0100
Subject: [PATCH] Fix WaitGroup crash

---
 main.go      |  8 +-------
 resume.go    |  2 ++
 scheduler.go | 11 +++++++----
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index c13cef1..b9b4c96 100644
--- a/main.go
+++ b/main.go
@@ -115,13 +115,7 @@ func cmdBase(_ *cobra.Command, _ []string) {
 			if urlErr, ok := err.(*fasturl.Error); ok && urlErr.Err == fasturl.ErrUnknownScheme {
 				// Not an error
 				err = nil
-
-				// Give back task
-				//err2 := CancelTask(t.WebsiteId)
-				//if err2 != nil {
-				//	logrus.Error(err2)
-				//}
-
+				// TODO FTP crawler
 				continue
 			} else if err != nil {
 				logrus.WithError(err).
diff --git a/resume.go b/resume.go
index e5da8bf..4092a4c 100644
--- a/resume.go
+++ b/resume.go
@@ -30,6 +30,8 @@ func LoadResumeTasks(inRemotes chan<- *OD) {
 	}
 
 	for _, remote := range resumed {
+		// TODO Cleanup globalWait management
+		globalWait.Add(1)
 		inRemotes <- remote
 	}
 }
diff --git a/scheduler.go b/scheduler.go
index c7fa4cb..dbab95e 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -141,7 +141,9 @@ func (o *OD) Watch(results chan File) {
 
 	// Block until all results are written
 	// (closes results channel)
-	o.handleCollect(results, f, collectErrC)
+	if !o.handleCollect(results, f, collectErrC) {
+		return
+	}
 
 	// Exit code of Collect()
 	err = <-collectErrC
@@ -161,7 +163,8 @@ func (o *OD) Watch(results chan File) {
 	}
 }
 
-func (o *OD) handleCollect(results chan File, f *os.File, collectErrC chan error) {
+// Returns if aborted naturally (results ready for upload)
+func (o *OD) handleCollect(results chan File, f *os.File, collectErrC chan error) bool {
 	// Begin collecting results
 	go o.Task.Collect(results, f, collectErrC)
 	defer close(results)
@@ -171,14 +174,14 @@ func (o *OD) handleCollect(results chan File, f *os.File, collectErrC chan error
 		// Natural finish
 		if atomic.LoadInt64(&o.InProgress) == 0 {
 			o.onTaskFinished()
-			return
+			return true
 		}
 		// Abort
 		if atomic.LoadInt32(&o.WCtx.aborted) != 0 {
 			// Wait for all workers to finish
 			o.WCtx.workers.Wait()
 			o.onTaskPaused()
-			return
+			return false
 		}
 
 		time.Sleep(500 * time.Millisecond)