feat: add some code to cancel

However, currently this never gets reached.

1 files changed, 10 insertions, 3 deletions

diff --git a/main.go b/main.go
index 30c9ebf..9c1a94d 100644
--- a/main.go
+++ b/main.go
@@ -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
+			}
 		}
 	})