77 lines
1.8 KiB
Go
77 lines
1.8 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"io/ioutil"
|
||
|
|
"path"
|
||
|
|
"strconv"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"github.com/dustin/go-humanize"
|
||
|
|
)
|
||
|
|
|
||
|
|
// BeginDownload is meant to be called as a goroutine and begins the post download process.
|
||
|
|
func BeginDownload(posts *[]Post, saveDirectory *string, maxConcurrents *int) {
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
var completed int
|
||
|
|
|
||
|
|
total := len(*posts)
|
||
|
|
|
||
|
|
// Distribute the posts based on the number of workers
|
||
|
|
ppw := len(*posts) / *maxConcurrents
|
||
|
|
mod := len(*posts) % *maxConcurrents
|
||
|
|
|
||
|
|
for i := 0; i < *maxConcurrents; i++ {
|
||
|
|
postsLower := i * ppw
|
||
|
|
postsUpper := i*ppw + ppw
|
||
|
|
|
||
|
|
if i == *maxConcurrents-1 {
|
||
|
|
// Give the last worker the remaining posts
|
||
|
|
// TODO: compensate it for labor
|
||
|
|
postsUpper += mod
|
||
|
|
}
|
||
|
|
|
||
|
|
wg.Add(1)
|
||
|
|
go work(i+1, (*posts)[postsLower:postsUpper], *saveDirectory, &completed, &total, &wg)
|
||
|
|
}
|
||
|
|
|
||
|
|
wg.Wait()
|
||
|
|
}
|
||
|
|
|
||
|
|
func work(wn int, posts []Post, directory string, completed *int, total *int, wg *sync.WaitGroup) {
|
||
|
|
defer wg.Done()
|
||
|
|
|
||
|
|
for _, post := range posts {
|
||
|
|
*completed++
|
||
|
|
fmt.Printf("[%d/%d] [w%d] Downloading post %d (%s)...\n", *completed, *total, wn, post.ID, humanize.Bytes(uint64(post.FileSize)))
|
||
|
|
downloadPost(&post, directory)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func downloadPost(post *Post, directory string) {
|
||
|
|
pathSliced := strings.Split(post.FileURL, ".")
|
||
|
|
extension := pathSliced[len(pathSliced)-1]
|
||
|
|
|
||
|
|
resp, err := HTTPGet(post.FileURL)
|
||
|
|
if err != nil {
|
||
|
|
fmt.Println("Unable to download, skipping...")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
body, err := ioutil.ReadAll(resp.Body)
|
||
|
|
if err != nil {
|
||
|
|
fmt.Println("Unable to read post response body, skipping...")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
savePath := path.Join(directory, strconv.Itoa(post.ID)+"."+extension)
|
||
|
|
|
||
|
|
err = ioutil.WriteFile(savePath, body, 0755)
|
||
|
|
if err != nil {
|
||
|
|
fmt.Printf("Error: could not write to file: %v\n", err)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|