diff options
| -rw-r--r-- | main.go | 13 |
1 files changed, 10 insertions, 3 deletions
@@ -3,6 +3,7 @@ package main import ( + "context" "flag" "fmt" "log" @@ -50,9 +51,10 @@ func pool(startURL url.URL, maxConcurrency, maxURLs int) { var wg sync.WaitGroup urls := make(chan url.URL) outlets := make([]<-chan []url.URL, maxConcurrency) + ctx, cancel := context.WithCancel(context.Background()) for i := range maxConcurrency { - outlets[i] = worker(&wg, i+1, urls) + outlets[i] = worker(ctx, &wg, i+1, urls) } out := fanIn(outlets...) @@ -76,10 +78,11 @@ func pool(startURL url.URL, maxConcurrency, maxURLs int) { } } + cancel() wg.Wait() } -func worker(wg *sync.WaitGroup, id int, urls <-chan url.URL) <-chan []url.URL { +func worker(ctx context.Context, wg *sync.WaitGroup, id int, urls <-chan url.URL) <-chan []url.URL { out := make(chan []url.URL) wg.Go(func() { @@ -96,7 +99,11 @@ func worker(wg *sync.WaitGroup, id int, urls <-chan url.URL) <-chan []url.URL { batch := findURLs(u, doc) // fmt.Printf("(%d) finished %s\n", id, &u) - out <- batch + select { + case out <- batch: + case <-ctx.Done(): + return + } } }) |