diff options
| -rw-r--r-- | classic.go | 22 |
1 files changed, 19 insertions, 3 deletions
@@ -1,9 +1,11 @@ package main import ( + "context" "fmt" "log" "net/url" + "sync" ) func classic(startURL url.URL, maxConcurrency, maxURLs int) { @@ -19,6 +21,10 @@ func classic(startURL url.URL, maxConcurrency, maxURLs int) { seen := make(map[url.URL]bool) count := 1 + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + +loop: for ; numPendingSends > 0; numPendingSends-- { batch := <-worklist for _, u := range batch { @@ -27,14 +33,24 @@ func classic(startURL url.URL, maxConcurrency, maxURLs int) { count++ seen[u] = true + if len(seen) == maxURLs { + break loop + } numPendingSends++ - go func() { - worklist <- getBatch(u) - }() + wg.Go(func() { + select { + case <-ctx.Done(): + return + case worklist <- getBatch(u): + } + }) } } } + + cancel() + wg.Wait() } func getBatch(u url.URL) []url.URL { |