e6dl/concurrent/download.go

135 lines
3.0 KiB
Go
Raw Normal View History

2019-03-14 03:20:04 +08:00
package concurrent
2019-03-01 03:47:50 +08:00
import (
"fmt"
"io/ioutil"
"path"
"strconv"
2019-03-30 00:33:56 +08:00
"time"
2019-03-01 03:47:50 +08:00
"github.com/dustin/go-humanize"
2019-03-14 04:21:50 +08:00
"github.com/logrusorgru/aurora"
2019-03-14 03:20:04 +08:00
"github.com/tjhorner/e6dl/e621"
2019-03-01 03:47:50 +08:00
)
2019-03-01 08:06:58 +08:00
// BeginDownload takes a slice of posts, a directory to save them in, and a
// number of concurrent workers to make. It blocks until all the post have
// been processed. It returns the number of successes, failures, and the total
// amount of posts.
2019-03-14 03:20:04 +08:00
func BeginDownload(posts *[]e621.Post, saveDirectory *string, maxConcurrents *int) (*int, *int, *int) {
2019-03-30 00:17:46 +08:00
// Channel for worker goroutines to notify the main goroutine that it is done
// downloading a post
done := make(chan interface{})
// Channel for main goroutine to give workers a post when they are done downloading
// one
pc := make(chan *e621.Post)
2019-03-01 03:47:50 +08:00
var completed int
2019-03-01 08:06:58 +08:00
var successes int
var failures int
2019-03-30 00:33:56 +08:00
var current int
2019-03-01 03:47:50 +08:00
total := len(*posts)
2019-03-30 00:17:46 +08:00
// If we have more workers than posts, then we don't need all of them
if *maxConcurrents > total {
*maxConcurrents = total
}
2019-03-01 03:47:50 +08:00
for i := 0; i < *maxConcurrents; i++ {
2019-03-30 00:17:46 +08:00
// Create our workers
go work(i+1, *saveDirectory, &completed, &total, &successes, &failures, done, pc)
2019-03-01 03:47:50 +08:00
2019-03-30 00:17:46 +08:00
// Give them their initial posts
pc <- &(*posts)[current]
current++
2019-03-30 00:33:56 +08:00
time.Sleep(time.Millisecond * 50)
2019-03-30 00:17:46 +08:00
}
for {
// Wait for a worker to be done
<-done
// If we finished downloading all posts, break out of the loop
if successes+failures == total {
break
2019-03-01 03:47:50 +08:00
}
2019-03-30 00:17:46 +08:00
// If there's no more posts to give, stop the worker
2019-03-30 00:33:56 +08:00
if current >= total {
2019-03-30 00:17:46 +08:00
pc <- nil
continue
}
2019-03-01 03:47:50 +08:00
2019-03-30 00:17:46 +08:00
// Give the worker the next post in the array
pc <- &(*posts)[current]
2019-03-30 00:33:56 +08:00
current++
2019-03-30 00:17:46 +08:00
}
2019-03-01 08:06:58 +08:00
return &successes, &failures, &total
2019-03-01 03:47:50 +08:00
}
2019-03-30 00:17:46 +08:00
func work(wn int, directory string, completed *int, total *int, successes *int, failures *int, done chan interface{}, pc chan *e621.Post) {
for {
2019-03-01 03:47:50 +08:00
*completed++
2019-03-01 08:06:58 +08:00
2019-03-30 00:33:56 +08:00
// Wait for a post from main
2019-03-30 00:17:46 +08:00
post := <-pc
2019-03-30 00:33:56 +08:00
if post == nil { // nil means there aren't any more posts, so we're OK to break
2019-03-30 00:17:46 +08:00
return
}
2019-03-14 04:21:50 +08:00
progress := aurora.Sprintf(aurora.Green("[%d/%d]"), *completed, *total)
workerText := aurora.Sprintf(aurora.Cyan("[w%d]"), wn)
fmt.Println(aurora.Sprintf(
"%s %s Downloading post %d (%s) -> %s...",
progress,
workerText,
2019-03-01 03:57:42 +08:00
post.ID,
humanize.Bytes(uint64(post.FileSize)),
2019-03-30 00:17:46 +08:00
getSavePath(post, &directory),
2019-03-14 04:21:50 +08:00
))
2019-03-01 08:06:58 +08:00
2019-03-30 00:17:46 +08:00
err := downloadPost(post, directory)
2019-03-01 08:06:58 +08:00
if err != nil {
2019-03-01 08:14:09 +08:00
fmt.Printf("[w%d] Failed to download post %d: %v\n", wn, post.ID, err)
2019-03-01 08:06:58 +08:00
*failures++
} else {
*successes++
}
2019-03-30 00:17:46 +08:00
done <- nil
2019-03-01 03:47:50 +08:00
}
}
2019-03-14 03:20:04 +08:00
func getSavePath(post *e621.Post, directory *string) string {
2019-03-14 04:10:05 +08:00
savePath := path.Join(*directory, strconv.Itoa(post.ID)+"."+post.FileExt)
2019-03-01 03:57:42 +08:00
return savePath
}
2019-03-14 03:20:04 +08:00
func downloadPost(post *e621.Post, directory string) error {
2019-03-01 03:57:42 +08:00
savePath := getSavePath(post, &directory)
2019-03-14 03:20:04 +08:00
resp, err := e621.HTTPGet(post.FileURL)
2019-03-01 03:47:50 +08:00
if err != nil {
2019-03-01 08:06:58 +08:00
return err
2019-03-01 03:47:50 +08:00
}
2019-03-01 08:06:58 +08:00
2019-03-01 03:47:50 +08:00
defer resp.Body.Close()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
2019-03-01 08:06:58 +08:00
return err
2019-03-01 03:47:50 +08:00
}
err = ioutil.WriteFile(savePath, body, 0755)
if err != nil {
2019-03-01 08:06:58 +08:00
return err
2019-03-01 03:47:50 +08:00
}
2019-03-01 08:06:58 +08:00
return nil
2019-03-01 03:47:50 +08:00
}