Bump to v1.1.0

config.go

@@ -20,12 +20,12 @@ var config struct {
 	Retries    int
 	Workers    int
 	UserAgent  string
-	Timeout    time.Duration
 	Tasks      int32
 	CrawlStats time.Duration
 	AllocStats time.Duration
 	Verbose    bool
 	PrintHTTP  bool
+	JobBufferSize int
 }
 
 const (
@@ -38,7 +38,9 @@ const (
 	ConfRetries    = "crawl.retries"
 	ConfWorkers    = "crawl.connections"
 	ConfUserAgent  = "crawl.user-agent"
+	ConfDialTimeout = "crawl.dial_timeout"
 	ConfTimeout    = "crawl.timeout"
+	ConfJobBufferSize = "crawl.job_buffer"
 	ConfCrawlStats = "output.crawl_stats"
 	ConfAllocStats = "output.resource_stats"
 	ConfVerbose    = "output.verbose"
@@ -51,7 +53,9 @@ func prepareConfig() {
 	viper.SetDefault(ConfWorkers, 2)
 	viper.SetDefault(ConfTasks, 3)
 	viper.SetDefault(ConfUserAgent, "")
-	viper.SetDefault(ConfTimeout, 10 * time.Second)
+	viper.SetDefault(ConfDialTimeout, 10 * time.Second)
+	viper.SetDefault(ConfTimeout, 30 * time.Second)
+	viper.SetDefault(ConfJobBufferSize, 5000)
 	viper.SetDefault(ConfCrawlStats, 3 * time.Second)
 	viper.SetDefault(ConfAllocStats, 0)
 	viper.SetDefault(ConfVerbose, false)
@@ -107,7 +111,11 @@ func readConfig() {
 
 	config.UserAgent = viper.GetString(ConfUserAgent)
 
-	config.Timeout = viper.GetDuration(ConfTimeout)
+	setDialTimeout(viper.GetDuration(ConfDialTimeout))
+
+	setTimeout(viper.GetDuration(ConfTimeout))
+
+	config.JobBufferSize = viper.GetInt(ConfJobBufferSize)
 
 	config.CrawlStats = viper.GetDuration(ConfCrawlStats)
 

config.yml

@@ -46,15 +46,32 @@ crawl:
   # Please be careful with this setting!
   # The crawler fires fast and more than
   # ten connections can overwhelm a server.
-  connections: 10
+  connections: 4
 
   # How often to retry getting data
   # from the site before giving up
   retries: 5
 
+  # Time before discarding a failed connection attempt
+  dial_timeout: 10s
+
   # Time before discarding a network request
-  timeout: 10s
+  timeout: 30s
 
   # Crawler User-Agent
   # If empty, no User-Agent header is sent.
   user-agent: "Mozilla/5.0 (X11; od-database-crawler) Gecko/20100101 Firefox/52.0"
+
+  # Job buffer size (per task)
+  # Higher values cause less disk writes
+  # but require more memory.
+  #
+  # The job queue contains all URLs
+  # that should be crawled next.
+  # As it grows very large over time,
+  # it's kept mainly on disk.
+  # This sets how many jobs are kept
+  # in memory.
+  # A negative value will cause all jobs
+  # to be stored in memory. (Don't do this)
+  job_buffer: 5000

crawl.go

@@ -8,6 +8,7 @@ import (
 	"github.com/valyala/fasthttp"
 	"golang.org/x/crypto/blake2b"
 	"golang.org/x/net/html"
+	"net"
 	"path"
 	"strconv"
 	"strings"
@@ -20,6 +21,17 @@ var client = fasthttp.Client {
 	},
 }
 
+func setDialTimeout(d time.Duration) {
+	client.Dial = func(addr string) (net.Conn, error) {
+		return fasthttp.DialTimeout(addr, d)
+	}
+}
+
+func setTimeout(d time.Duration) {
+	client.ReadTimeout = d
+	client.WriteTimeout = d / 2
+}
+
 func GetDir(j *Job, f *File) (links []fasturl.URL, err error) {
 	f.IsDir = true
 	f.Name = path.Base(j.Uri.Path)
@@ -33,7 +45,7 @@ func GetDir(j *Job, f *File) (links []fasturl.URL, err error) {
 	res := fasthttp.AcquireResponse()
 	defer fasthttp.ReleaseResponse(res)
 
-	err = client.DoTimeout(req, res, config.Timeout)
+	err = client.Do(req, res)
 	fasthttp.ReleaseRequest(req)
 
 	if err != nil {
@@ -132,7 +144,7 @@ func GetFile(u fasturl.URL, f *File) (err error) {
 	res.SkipBody = true
 	defer fasthttp.ReleaseResponse(res)
 
-	err = client.DoTimeout(req, res, config.Timeout)
+	err = client.Do(req, res)
 	fasthttp.ReleaseRequest(req)
 
 	if err != nil {

main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"github.com/sirupsen/logrus"
+	"github.com/spf13/viper"
 	"github.com/terorie/od-database-crawler/fasturl"
 	"github.com/urfave/cli"
 	"os"
@@ -11,10 +12,12 @@ import (
 	"time"
 )
 
+var configFile string
+
 var app = cli.App {
 	Name:         "od-database-crawler",
 	Usage:        "OD-Database Go crawler",
-	Version:      "1.0.2",
+	Version:      "1.1.0",
 	BashComplete: cli.DefaultAppComplete,
 	Writer:       os.Stdout,
 	Action:       cmdBase,
@@ -26,6 +29,19 @@ var app = cli.App {
 			Action:    cmdCrawler,
 		},
 	},
+	Flags: []cli.Flag {
+		cli.StringFlag {
+			Name: "config",
+			EnvVar: "CONFIG",
+			Destination: &configFile,
+		},
+	},
+	Before: func(i *cli.Context) error {
+		if configFile != "" {
+			viper.SetConfigFile(configFile)
+		}
+		return nil
+	},
 	After: func(i *cli.Context) error {
 		exitHooks.Execute()
 		return nil
@@ -39,10 +55,11 @@ func init() {
 }
 
 func main() {
-	err := os.MkdirAll("crawled", 0755)
+	if err := os.MkdirAll("crawled", 0755);
-	if err != nil {
+		err != nil { panic(err) }
-		panic(err)
+
-	}
+	if err := os.MkdirAll("queue", 0755);
+		err != nil { panic(err) }
 
 	readConfig()
 	app.Run(os.Args)

model.go

@@ -23,7 +23,6 @@ type TaskResult struct {
 }
 
 type Job struct {
-	OD        *OD
 	Uri       fasturl.URL
 	UriStr    string
 	Fails     int
@@ -57,3 +56,8 @@ func (o *OD) LoadOrStoreKey(k *redblackhash.Key) (exists bool) {
 	o.Scanned.Put(k)
 	return false
 }
+
+type errorString string
+func (e errorString) Error() string {
+	return string(e)
+}

queue.go

@@ -0,0 +1,129 @@
+package main
+
+import (
+	"github.com/beeker1121/goque"
+	"os"
+	"sync"
+	"sync/atomic"
+)
+
+type BufferedQueue struct {
+	dataDir string
+	q       *goque.Queue
+	buf     []Job
+	m       sync.Mutex
+}
+
+func OpenQueue(dataDir string) (bq *BufferedQueue, err error) {
+	bq = new(BufferedQueue)
+	if config.JobBufferSize < 0 {
+		return
+	}
+	bq.dataDir = dataDir
+	bq.q, err = goque.OpenQueue(dataDir)
+	if err != nil { return nil, err }
+	return
+}
+
+func (q *BufferedQueue) Enqueue(job *Job) error {
+	atomic.AddInt64(&totalQueued, 1)
+	if q.directEnqueue(job) {
+		return nil
+	}
+
+	var gob JobGob
+	gob.ToGob(job)
+	_, err := q.q.EnqueueObject(gob)
+	return err
+}
+
+func (q *BufferedQueue) Dequeue() (job Job, err error) {
+	if q.directDequeue(&job) {
+		atomic.AddInt64(&totalQueued, -1)
+		return job, nil
+	}
+
+	if config.JobBufferSize < 0 {
+		err = goque.ErrEmpty
+		return
+	}
+
+	var item *goque.Item
+	item, err = q.q.Dequeue()
+	if err != nil { return }
+
+	atomic.AddInt64(&totalQueued, -1)
+
+	var gob JobGob
+	err = item.ToObject(&gob)
+	if err != nil { return }
+	gob.FromGob(&job)
+
+	return
+}
+
+func (q *BufferedQueue) directEnqueue(job *Job) bool {
+	q.m.Lock()
+	defer q.m.Unlock()
+
+	bs := config.JobBufferSize
+	if len(q.buf) < bs || bs < 0 {
+		q.buf = append(q.buf, *job)
+		return true
+	} else {
+		return false
+	}
+}
+
+func (q *BufferedQueue) directDequeue(job *Job) bool {
+	q.m.Lock()
+	defer q.m.Unlock()
+
+	if len(q.buf) > 0 {
+		*job = q.buf[0]
+		q.buf = q.buf[1:]
+		return true
+	} else {
+		return false
+	}
+}
+
+// Always returns nil (But implements io.Closer)
+func (q *BufferedQueue) Close() error {
+	if config.JobBufferSize < 0 {
+		return nil
+	}
+
+	// Close ignoring errors
+	q.q.Close()
+
+	// Delete files
+	if err := os.RemoveAll(q.dataDir);
+		err != nil { panic(err) }
+
+	return nil
+}
+
+type JobGob struct {
+	Uri string
+	Fails int
+	LastError string
+}
+
+func (g *JobGob) ToGob(j *Job) {
+	g.Uri = j.UriStr
+	g.Fails = j.Fails
+	if j.LastError != nil {
+		g.LastError = j.LastError.Error()
+	}
+}
+
+func (g *JobGob) FromGob(j *Job) {
+	if err := j.Uri.Parse(g.Uri);
+		err != nil { panic(err) }
+	j.UriStr = g.Uri
+	j.Fails = g.Fails
+	if g.LastError != "" {
+		j.LastError = errorString(g.LastError)
+	}
+}

scheduler.go

@@ -16,7 +16,7 @@ import (
 var activeTasksLock sync.Mutex
 var activeTasks = make(map[uint64]bool)
 var numActiveTasks int32
-var totalBuffered int64
+var totalQueued int64
 
 func Schedule(c context.Context, remotes <-chan *OD) {
 	go Stats(c)
@@ -28,8 +28,21 @@ func Schedule(c context.Context, remotes <-chan *OD) {
 		// Collect results
 		results := make(chan File)
 
+		remote.WCtx.OD = remote
+
+		// Get queue path
+		queuePath := path.Join("queue", fmt.Sprintf("%d", remote.Task.WebsiteId))
+
+		// Delete existing queue
+		if err := os.RemoveAll(queuePath);
+			err != nil { panic(err) }
+
+		// Start new queue
+		var err error
+		remote.WCtx.Queue, err = OpenQueue(queuePath)
+		if err != nil { panic(err) }
+
 		// Spawn workers
-		remote.WCtx.in, remote.WCtx.out = makeJobBuffer(c)
 		for i := 0; i < config.Workers; i++ {
 			go remote.WCtx.Worker(results)
 		}
@@ -37,7 +50,6 @@ func Schedule(c context.Context, remotes <-chan *OD) {
 		// Enqueue initial job
 		atomic.AddInt32(&numActiveTasks, 1)
 		remote.WCtx.queueJob(Job{
-			OD:     remote,
 			Uri:    remote.BaseUri,
 			UriStr: remote.BaseUri.String(),
 			Fails:  0,
@@ -148,7 +160,11 @@ func (o *OD) handleCollect(results chan File, f *os.File, collectErrC chan error
 
 	// Wait for all jobs on remote to finish
 	o.Wait.Wait()
-	close(o.WCtx.in)
+
+	// Close queue
+	if err := o.WCtx.Queue.Close(); err != nil {
+		panic(err)
+	}
 	atomic.AddInt32(&numActiveTasks, -1)
 
 	// Log finish
@@ -198,51 +214,3 @@ func (t *Task) collect(results chan File, f *os.File) error {
 
 	return nil
 }
-
-func makeJobBuffer(c context.Context) (chan<- Job, <-chan Job) {
-	in  := make(chan Job)
-	out := make(chan Job)
-	go bufferJobs(c, in, out)
-	return in, out
-}
-
-func bufferJobs(c context.Context, in chan Job, out chan Job) {
-	defer close(out)
-	var inQueue []Job
-	outCh := func() chan Job {
-		if len(inQueue) == 0 {
-			return nil
-		}
-		return out
-	}
-	for len(inQueue) > 0 || in != nil {
-		if len(inQueue) == 0 {
-			select {
-			case v, ok := <-in:
-				if !ok {
-					in = nil
-				} else {
-					atomic.AddInt64(&totalBuffered, 1)
-					inQueue = append(inQueue, v)
-				}
-			case <-c.Done():
-				return
-			}
-		} else {
-			select {
-			case v, ok := <-in:
-				if !ok {
-					in = nil
-				} else {
-					atomic.AddInt64(&totalBuffered, 1)
-					inQueue = append(inQueue, v)
-				}
-			case outCh() <- inQueue[0]:
-				atomic.AddInt64(&totalBuffered, -1)
-				inQueue = inQueue[1:]
-			case <-c.Done():
-				return
-			}
-		}
-	}
-}

server.go

@@ -102,9 +102,8 @@ func uploadChunks(websiteId uint64, f *os.File) error {
 
 		multi.Close()
 
-		for retried := false; true; retried = true {
+		for retries := 0; retries < 10; retries++ {
-			err = nil
+			if retries > 0 {
-			if retried {
 				// Error occurred, retry upload
 				time.Sleep(5 * time.Second)
 			}

stats.go

@@ -57,7 +57,7 @@ func Stats(c context.Context) {
 			runtime.ReadMemStats(&mem)
 
 			logrus.WithFields(logrus.Fields{
-				"queue_count": atomic.LoadInt64(&totalBuffered),
+				"queue_count": atomic.LoadInt64(&totalQueued),
 				"heap": FormatByteCount(mem.Alloc),
 				"objects": mem.HeapObjects,
 				"num_gc": mem.NumGC,

worker.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"github.com/beeker1121/goque"
 	"github.com/sirupsen/logrus"
 	"github.com/valyala/fasthttp"
 	"math"
@@ -14,24 +15,38 @@ import (
 var globalWait sync.WaitGroup
 
 type WorkerContext struct {
-	in chan<- Job
+	OD *OD
-	out <-chan Job
+	Queue *BufferedQueue
 	lastRateLimit time.Time
 	numRateLimits int
 }
 
-func (w WorkerContext) Worker(results chan<- File) {
+func (w *WorkerContext) Worker(results chan<- File) {
-	for job := range w.out {
+	for {
+		job, err := w.Queue.Dequeue()
+		switch err {
+		case goque.ErrEmpty:
+			time.Sleep(500 * time.Millisecond)
+			continue
+
+		case goque.ErrDBClosed:
+			return
+
+		case nil:
 			w.step(results, job)
+
+		default:
+			panic(err)
+		}
 	}
 }
 
-func (w WorkerContext) step(results chan<- File, job Job) {
+func (w *WorkerContext) step(results chan<- File, job Job) {
 	defer w.finishJob(&job)
 
 	var f File
 
-	newJobs, err := DoJob(&job, &f)
+	newJobs, err := w.DoJob(&job, &f)
 	atomic.AddUint64(&totalStarted, 1)
 	if err == ErrKnown {
 		return
@@ -75,7 +90,7 @@ func (w WorkerContext) step(results chan<- File, job Job) {
 	}
 }
 
-func DoJob(job *Job, f *File) (newJobs []Job, err error) {
+func (w *WorkerContext) DoJob(job *Job, f *File) (newJobs []Job, err error) {
 	if len(job.Uri.Path) == 0 { return }
 	if job.Uri.Path[len(job.Uri.Path)-1] == '/' {
 		// Load directory
@@ -93,7 +108,7 @@ func DoJob(job *Job, f *File) (newJobs []Job, err error) {
 		hash := f.HashDir(links)
 
 		// Skip symlinked dirs
-		if job.OD.LoadOrStoreKey(&hash) {
+		if w.OD.LoadOrStoreKey(&hash) {
 			return nil, ErrKnown
 		}
 
@@ -114,7 +129,6 @@ func DoJob(job *Job, f *File) (newJobs []Job, err error) {
 			lastLink = uriStr
 
 			newJobs = append(newJobs, Job{
-				OD:     job.OD,
 				Uri:    link,
 				UriStr: uriStr,
 				Fails:  0,
@@ -139,13 +153,13 @@ func DoJob(job *Job, f *File) (newJobs []Job, err error) {
 			}
 			return nil, err
 		}
-		atomic.AddUint64(&job.OD.Result.FileCount, 1)
+		atomic.AddUint64(&w.OD.Result.FileCount, 1)
 	}
 	return
 }
 
-func (w WorkerContext) queueJob(job Job) {
+func (w *WorkerContext) queueJob(job Job) {
-	job.OD.Wait.Add(1)
+	w.OD.Wait.Add(1)
 
 	if w.numRateLimits > 0 {
 		if time.Since(w.lastRateLimit) > 5 * time.Second {
@@ -156,11 +170,13 @@ func (w WorkerContext) queueJob(job Job) {
 		}
 	}
 
-	w.in <- job
+	if err := w.Queue.Enqueue(&job); err != nil {
+		panic(err)
+	}
 }
 
-func (w WorkerContext) finishJob(job *Job) {
+func (w *WorkerContext) finishJob(job *Job) {
-	job.OD.Wait.Done()
+	w.OD.Wait.Done()
 }
 
 func isErrSilent(err error) bool {