Remember scanned URLs

scheduler.go

@@ -2,67 +2,68 @@ package main
 
 import (
 	"context"
-	"net/url"
 	"sync/atomic"
+	"time"
 )
 
-type Job struct {
-	Remote *RemoteDir
-	Uri url.URL
-	UriStr string
-	Fails int
-	LastError error
-}
-
 var activeTasks int32
 
-func Schedule(c context.Context, remotes <-chan *RemoteDir) {
-	in, out := makeJobBuffer()
-	wCtx := WorkerContext{
-		in: in,
-		out: out,
-	}
-	for i := 0; i < config.Workers; i++ {
-		go wCtx.Worker()
-	}
+func Schedule(c context.Context, remotes <-chan *OD) {
 	go Stats(c)
 
 	for {
 		select {
 		case <-c.Done():
-			close(in)
 			return
 
 		case remote := <-remotes:
+			for atomic.LoadInt32(&activeTasks) > config.Tasks {
+				select {
+				case <-time.After(time.Second):
+					break
+				case <-c.Done():
+					return
+				}
+			}
+
+			// Spawn workers
+			remote.WCtx.in, remote.WCtx.out = makeJobBuffer(c)
+			for i := 0; i < config.Workers; i++ {
+				go remote.WCtx.Worker()
+			}
+
 			// Enqueue initial job
 			atomic.AddInt32(&activeTasks, 1)
-			wCtx.queueJob(Job{
-				Remote: remote,
-				Uri: remote.BaseUri,
+			remote.WCtx.queueJob(Job{
+				OD:     remote,
+				Uri:    remote.BaseUri,
 				UriStr: remote.BaseUri.String(),
-				Fails: 0,
+				Fails:  0,
 			})
 			globalWait.Done()
+
 			// Upload result when ready
 			go remote.Watch()
 		}
 	}
 }
 
-func (r *RemoteDir) Watch() {
+func (r *OD) Watch() {
 	// Wait for all jobs on remote to finish
 	r.Wait.Wait()
+	close(r.WCtx.in)
 	atomic.AddInt32(&activeTasks, -1)
 }
 
-func makeJobBuffer() (chan<- Job, <-chan Job) {
+func makeJobBuffer(c context.Context) (chan<- Job, <-chan Job) {
 	in  := make(chan Job)
 	out := make(chan Job)
-	go bufferJobs(in, out)
+	go bufferJobs(c, in, out)
 	return in, out
 }
 
-func bufferJobs(in chan Job, out chan Job) {
+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 {
@@ -72,11 +73,15 @@ func bufferJobs(in chan Job, out chan Job) {
 	}
 	for len(inQueue) > 0 || in != nil {
 		if len(inQueue) == 0 {
-			v, ok := <-in
-			if !ok {
-				in = nil
-			} else {
-				inQueue = append(inQueue, v)
+			select {
+			case v, ok := <-in:
+				if !ok {
+					in = nil
+				} else {
+					inQueue = append(inQueue, v)
+				}
+			case <-c.Done():
+				return
 			}
 		} else {
 			select {
@@ -88,8 +93,9 @@ func bufferJobs(in chan Job, out chan Job) {
 				}
 			case outCh() <- inQueue[0]:
 				inQueue = inQueue[1:]
+			case <-c.Done():
+				return
 			}
 		}
 	}
-	close(out)
 }