Will continue where left off at when restarted

src/viddown.go

@@ -20,6 +20,7 @@
 package main
 
 import (
+	"encoding/gob"
 	"encoding/json"
 	"flag"
 	"io"
@@ -75,16 +76,15 @@ type ops struct {
 	Res chan string
 }
 
-// List of the operations currently being worked on
-var currentWork map[int]string
-
 type WorkItem struct {
 	Status string
 	URL    string
 }
 
+type CurrentWork map[int]WorkItem
+
 // CurrentWork
-var CurrentWork = make(map[int]WorkItem)
+var currentWork CurrentWork = make(map[int]WorkItem)
 
 // program configurations
 var config struct {
@@ -99,10 +99,50 @@ var config struct {
 // Global channel to send operation requests to
 var opsc = make(chan ops)
 
+// read the saved currentWork file
+func (cw *CurrentWork) read() (err error) {
+	dir := path.Join(config.saveLocation, ".viddown", "currentWork")
+
+	f, err := os.Open(dir)
+	defer f.Close()
+	if os.IsNotExist(err) {
+		return nil
+	} else if err != nil {
+		return
+	}
+
+	// Decode data
+	decoder := gob.NewDecoder(f)
+	err = decoder.Decode(cw)
+	return
+}
+
+// save the currentWork to the save file
+func (cw *CurrentWork) save() (err error) {
+	dir := path.Join(config.saveLocation, ".viddown", "currentWork")
+	f, err := os.Create(dir)
+
+	defer f.Close()
+	if err != nil {
+		return
+	}
+
+	encoder := gob.NewEncoder(f)
+	err = encoder.Encode(cw)
+	return
+}
+
+// apply an item then save it
+func (cw *CurrentWork) apply(i int, item WorkItem) (err error) {
+	(*cw)[i] = item
+	cw.save()
+	return
+}
+
 // handleRequest handles http requests
 func handleRequest(w http.ResponseWriter, r *http.Request) {
 	if r.Method == http.MethodGet && r.URL.Path == "/currentwork" {
-		list, err := json.Marshal(CurrentWork)
+		list, err := json.Marshal(currentWork)
 		if err != nil {
 			log.Panic(err)
 		}
@@ -124,7 +164,7 @@ func handleRequest(w http.ResponseWriter, r *http.Request) {
 
 	// Modify template and send it to the requester
 	var insert string
-	for _, workItem := range CurrentWork {
+	for _, workItem := range currentWork {
 		insert += "<br>" + workItem.URL
 	}
 	temp := strings.Replace(template, "currentWork", insert, 1)
@@ -138,6 +178,12 @@ func operationHandler() {
 		dir := path.Join(config.saveLocation, ".viddown", strconv.Itoa(i))
 		var err error
 		// Check if index exists
+		_, ok := currentWork[i]
+		if ok {
+			continue
+		}
+
+		// Check if directory already exists
 		_, err = os.Stat(dir)
 		if err == nil {
 			continue
@@ -147,58 +193,61 @@ func operationHandler() {
 		log.Print("Start Downloading: " + o.URL)
 
 		// Spin off a go routine to handle a download so the handler can continue to receive requests
-		go func() {
-			var item WorkItem
-			item.URL = o.URL
-			item.Status = "downloading"
-			CurrentWork[i] = item
-
-			// Create a unique directory to download the video in
-			os.MkdirAll(dir, 0755)
-
-			// Build the download command
-			cmd := exec.Command(config.youtubedll, o.URL)
-			cmd.Dir = dir + "/"
-			output, err := cmd.CombinedOutput()
-			log.Print("Finished: " + o.URL)
-			if err != nil {
-				item.Status = "failed"
-				CurrentWork[i] = item
-
-				log.Print("Failed to download")
-				log.Print(err)
-			}
-
-			item.Status = "finished"
-			CurrentWork[i] = item
-
-			log.Print(string(output))
-
-			// Clean up directory
-			files, err := ioutil.ReadDir(dir + "/")
-			if err != nil {
-				log.Print(err)
-			}
-
-			for _, file := range files {
-				log.Print("Moving: " + file.Name())
-				err := os.Rename(dir+"/"+file.Name(), file.Name())
-				if err != nil {
-					log.Print(err)
-					log.Print("Failed to move: " + dir + "/" + file.Name() + " to: " + file.Name())
-				}
-			}
-
-			err = os.RemoveAll(dir)
-			if err != nil {
-				log.Print("Failed to remove: " + dir)
-			}
-
-		}()
+		go download(i, o.URL)
 	}
 	log.Fatal("Exited operationsHandler")
 }
 
+func download(i int, url string) {
+	dir := path.Join(config.saveLocation, ".viddown", strconv.Itoa(i))
+	var item WorkItem
+	item.URL = url
+	item.Status = "downloading"
+	currentWork.apply(i, item)
+
+	// Create a unique directory to download the video in
+	os.MkdirAll(dir, 0755)
+
+	// Build the download command
+	cmd := exec.Command(config.youtubedll, url)
+	cmd.Dir = dir + "/"
+	output, err := cmd.CombinedOutput()
+	log.Print("Finished: " + url)
+	if err != nil {
+		item.Status = "failed"
+		currentWork[i] = item
+
+		log.Print("Failed to download")
+		log.Print(err)
+	}
+
+	item.Status = "finished"
+	currentWork.apply(i, item)
+
+	log.Print(string(output))
+
+	// Clean up directory
+	files, err := ioutil.ReadDir(dir + "/")
+	if err != nil {
+		log.Print(err)
+	}
+
+	for _, file := range files {
+		log.Print("Moving: " + file.Name())
+		err := os.Rename(dir+"/"+file.Name(), file.Name())
+		if err != nil {
+			log.Print(err)
+			log.Print("Failed to move: " + dir + "/" + file.Name() + " to: " + file.Name())
+		}
+	}
+
+	err = os.RemoveAll(dir)
+	if err != nil {
+		log.Print("Failed to remove: " + dir)
+	}
+
+}
+
 // handleConfigs configures the program at start, handling any flags or env variables
 func handleConfigs() {
 	flag.StringVar(&config.saveLocation, "o", "./", "The output location for videos")
@@ -211,8 +260,18 @@ func handleConfigs() {
 
 func main() {
 	handleConfigs()
-	// initialize currentWork with a blank map
-	currentWork = make(map[int]string)
+	// initialize currentWork
+	err := currentWork.read()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Continue old Downloads
+	for i, item := range currentWork {
+		if item.Status == "downloading" {
+			go download(i, item.URL)
+		}
+	}
 
 	// Create the work directory
 	os.MkdirAll(path.Join(config.saveLocation, ".viddown"), 0755)