performance and edge case handling improvements

+ new, better and cuter logging

Author: nyancrimew

cmd/goop.go

@@ -1,10 +1,11 @@
 package cmd
 
 import (
-	"fmt"
+	"os"
+
 	"github.com/deletescape/goop/pkg/goop"
+	"github.com/phuslu/log"
 	"github.com/spf13/cobra"
-	"os"
 )
 
 var force bool
@@ -21,12 +22,12 @@ var rootCmd = &cobra.Command{
 		}
 		if list {
 			if err := goop.CloneList(args[0], dir, force, keep); err != nil {
-				fmt.Fprintln(os.Stderr, err)
+				log.Error().Err(err).Msg("exiting")
 				os.Exit(1)
 			}
 		} else {
 			if err := goop.Clone(args[0], dir, force, keep); err != nil {
-				fmt.Fprintln(os.Stderr, err)
+				log.Error().Err(err).Msg("exiting")
 				os.Exit(1)
 			}
 		}
@@ -41,7 +42,7 @@ func init() {
 
 func Execute() {
 	if err := rootCmd.Execute(); err != nil {
-		fmt.Fprintln(os.Stderr, err)
+		log.Error().Err(err).Msg("exiting")
 		os.Exit(1)
 	}
 }

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/PuerkitoBio/goquery v1.6.0
 	github.com/go-git/go-billy/v5 v5.0.0
 	github.com/go-git/go-git/v5 v5.2.0
+	github.com/phuslu/log v1.0.75
 	github.com/spf13/cobra v1.1.1
 	github.com/valyala/fasthttp v1.16.0
 )

go.sum

@@ -164,6 +164,8 @@ github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLA
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
+github.com/phuslu/log v1.0.75 h1:2Qcqgwo1sOsvj7QIuclIS92hmWxIISI2+XskYM1Nw2A=
+github.com/phuslu/log v1.0.75/go.mod h1:kzJN3LRifrepxThMjufQwS7S35yFAB+jAV1qgA7eBW4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

internal/jobtracker/jobtracker.go

@@ -0,0 +1,64 @@
+package jobtracker
+
+import (
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type JobTracker struct {
+	activeWorkers int32
+	queuedJobs    int32
+	didWork       bool
+	cond          *sync.Cond
+	Queue         chan string
+}
+
+func Nap() {
+	time.Sleep(40 * time.Millisecond)
+}
+
+func NewJobTracker() *JobTracker {
+	return &JobTracker{
+		cond:  sync.NewCond(&sync.Mutex{}),
+		Queue: make(chan string, 999999), // TODO: dont create oversized queues, we should try to save memory; maybe read the channel docs again
+	}
+}
+
+func (jt *JobTracker) AddJob(job string) {
+	// TODO: can we discard empty jobs here?
+	jt.cond.L.Lock()
+	atomic.AddInt32(&jt.queuedJobs, 1)
+	jt.Queue <- job
+	jt.cond.L.Unlock()
+}
+
+func (jt *JobTracker) StartWork() {
+	atomic.AddInt32(&jt.activeWorkers, 1)
+}
+
+func (jt *JobTracker) EndWork() {
+	jt.didWork = true
+	atomic.AddInt32(&jt.activeWorkers, -1)
+	atomic.AddInt32(&jt.queuedJobs, -1)
+}
+
+func (jt *JobTracker) HasWork() bool {
+	// TODO: didWork is a somewhat ugly workaround to ensure we dont exit before doing work at least once,
+	// this will however result in locking up if we create a JobTracker but never queue any jobs
+	hasWork := !jt.didWork || (atomic.LoadInt32(&jt.queuedJobs) > 0 && atomic.LoadInt32(&jt.activeWorkers) > 0)
+
+	if !hasWork {
+		jt.cond.Broadcast()
+	}
+	return hasWork
+}
+
+func (jt *JobTracker) Wait() {
+	defer close(jt.Queue)
+
+	jt.cond.L.Lock()
+	for jt.HasWork() {
+		jt.cond.Wait()
+	}
+}

internal/utils/html.go

@@ -2,9 +2,10 @@ package utils
 
 import (
 	"bytes"
-	"github.com/PuerkitoBio/goquery"
 	"net/url"
 	"strings"
+
+	"github.com/PuerkitoBio/goquery"
 )
 
 var htmlTag = []byte{'<', 'h', 't', 'm', 'l'}
@@ -36,5 +37,5 @@ func GetIndexedFiles(body []byte) ([]string, error) {
 		}
 		return true
 	})
-	return files, err
+	return files, exitErr
 }

internal/utils/urls.go

@@ -2,6 +2,7 @@ package utils
 
 import "strings"
 
+//TODO: replace all uses of this with the proper path utils
 func Url(base, path string) string {
 	return strings.TrimSuffix(base, "/") + "/" + strings.TrimPrefix(path, "/")
 }

internal/workers/consts.go

@@ -1,65 +1,73 @@
 package workers
 
 import (
-	"fmt"
-	"github.com/deletescape/goop/internal/utils"
-	"github.com/valyala/fasthttp"
 	"io/ioutil"
 	"os"
-	"sync"
-	"time"
+
+	"github.com/deletescape/goop/internal/jobtracker"
+	"github.com/deletescape/goop/internal/utils"
+	"github.com/phuslu/log"
+	"github.com/valyala/fasthttp"
 )
 
-func DownloadWorker(c *fasthttp.Client, queue chan string, baseUrl, baseDir string, wg *sync.WaitGroup, allowHtml bool) {
-	defer wg.Done()
-	var ctr int
+func DownloadWorker(c *fasthttp.Client, baseUrl, baseDir string, jt *jobtracker.JobTracker, allowHtml, allowEmpty bool) {
 	for {
 		select {
-		case file := <-queue:
-
-			checkRatelimted()
-			if file == "" {
-				continue
-			}
-			targetFile := utils.Url(baseDir, file)
-			if utils.Exists(targetFile) {
-				fmt.Printf("%s was downloaded already, skipping\n", targetFile)
-				continue
-			}
-			uri := utils.Url(baseUrl, file)
-			code, body, err := c.Get(nil, uri)
-			fmt.Printf("[-] Fetching %s [%d]\n", uri, code)
-			if err != nil {
-				fmt.Fprintf(os.Stderr, "error: %s\n", err)
-				continue
-			}
-			if code == 200 {
-				if !allowHtml && utils.IsHtml(body) {
-					fmt.Printf("warning: %s appears to be an html file, skipping\n", uri)
-					continue
-				}
-				if utils.IsEmptyBytes(body) {
-					fmt.Printf("warning: %s appears to be an empty file, skipping\n", uri)
-					continue
-				}
-				if err := utils.CreateParentFolders(targetFile); err != nil {
-					fmt.Fprintf(os.Stderr, "error: %s\n", err)
-					continue
-				}
-				if err := ioutil.WriteFile(targetFile, body, os.ModePerm); err != nil {
-					fmt.Fprintf(os.Stderr, "error: %s\n", err)
-				}
-			} else if code == 429 {
-				setRatelimited()
-				queue <- file
-			}
+		case file := <-jt.Queue:
+			downloadWork(c, baseUrl, baseDir, file, jt, allowHtml, allowEmpty)
 		default:
-			// TODO: get rid of dirty hack somehow
-			if ctr >= graceTimes {
+			if !jt.HasWork() {
 				return
 			}
-			ctr++
-			time.Sleep(gracePeriod)
+			jobtracker.Nap()
 		}
 	}
 }
+
+func downloadWork(c *fasthttp.Client, baseUrl, baseDir, file string, jt *jobtracker.JobTracker, allowHtml, allowEmpty bool) {
+	jt.StartWork()
+	defer jt.EndWork()
+
+	if file == "" {
+		return
+	}
+	checkRatelimted()
+
+	targetFile := utils.Url(baseDir, file)
+	if utils.Exists(targetFile) {
+		log.Info().Str("file", targetFile).Msg("already fetched, skipping redownload")
+		return
+	}
+	uri := utils.Url(baseUrl, file)
+	code, body, err := c.Get(nil, uri)
+	if err == nil && code != 200 {
+		if code == 429 {
+			setRatelimited()
+			jt.AddJob(file)
+			return
+		}
+		log.Warn().Str("uri", uri).Int("code", code).Msg("couldn't fetch file")
+		return
+	} else if err != nil {
+		log.Error().Str("uri", uri).Int("code", code).Err(err).Msg("couldn't fetch file")
+		return
+	}
+
+	if !allowHtml && utils.IsHtml(body) {
+		log.Warn().Str("uri", uri).Msg("file appears to be html, skipping")
+		return
+	}
+	if !allowEmpty && utils.IsEmptyBytes(body) {
+		log.Warn().Str("uri", uri).Msg("file appears to be empty, skipping")
+		return
+	}
+	if err := utils.CreateParentFolders(targetFile); err != nil {
+		log.Error().Str("uri", uri).Str("file", targetFile).Err(err).Msg("couldn't create parent directories")
+		return
+	}
+	if err := ioutil.WriteFile(targetFile, body, os.ModePerm); err != nil {
+		log.Error().Str("uri", uri).Str("file", targetFile).Err(err).Msg("clouldn't write file")
+		return
+	}
+	log.Info().Str("uri", uri).Str("file", file).Msg("fetched file")
+}